diff --git a/.gitignore b/.gitignore index a5f1ea2..a95ea58 100644 --- a/.gitignore +++ b/.gitignore @@ -445,3 +445,7 @@ FodyWeavers.xsd *.msix *.msm *.msp + +# Local environment files with credentials (templates use .env.example and stay tracked). +# Must come after the '!python/packages/**' negation above so it wins for test .env files. +**/.env diff --git a/docs/features/durable-agents/README.md b/docs/features/durable-agents/README.md index 3a801be..d8dfe07 100644 --- a/docs/features/durable-agents/README.md +++ b/docs/features/durable-agents/README.md @@ -2,13 +2,13 @@ ## Overview -Durable agents extend the standard Microsoft Agent Framework with **durable state management** powered by the Durable Task framework. An ordinary Agent Framework agent runs in-process: its conversation history lives in memory and is lost when the process ends. A durable agent persists conversation history and execution state in external storage so that sessions survive process restarts, failures, and scale-out events. +Durable agents extend the standard Microsoft Agent Framework with **durable execution state** powered by the Durable Task framework. Ordinary agents can already use in-memory, external-provider or service-owned history. Durable hosting persists execution and session control so that compatible workers can continue sessions across process restarts and scale-out. | Capability | Ordinary agent | Durable agent | | --- | --- | --- | -| Conversation history | In-memory only | Durably persisted | -| Failure recovery | State lost on crash | Automatically resumed | -| Multi-instance scale-out | Not supported | Any worker can resume a session | +| Conversation history | Selected provider or model service | Selected owner, with durable-backed local history when configured | +| Failure recovery | Application-owned | Persisted orchestration and entity state; uncommitted external effects can repeat | +| Multi-instance scale-out | Application-owned coordination | Compatible workers serialize access to each entity | | Multi-agent orchestrations | Manual coordination | Deterministic, checkpointed workflows | | Human-in-the-loop | Must keep process alive | Can wait days/weeks with zero compute | | Hosting | Any process | Console app, Azure Functions, or any Durable Task–compatible host | @@ -18,12 +18,15 @@ Durable agents extend the standard Microsoft Agent Framework with **durable stat ## How durable agents work -Durable agents are implemented on top of [Durable Entities](https://learn.microsoft.com/azure/azure-functions/durable/durable-functions-entities) (also called "virtual actors"). Each **agent session** maps to one entity instance whose state contains the full conversation history. When you send a message to a durable agent, the following happens: +Durable agents are implemented on top of [Durable Entities](https://learn.microsoft.com/azure/azure-functions/durable/durable-functions-entities) (also called "virtual actors"). Each **agent session** maps to one entity instance. Transcript ownership and response storage depend on the runtime version and selected history provider. When you send a message to a durable agent, the following happens: 1. The message is dispatched to the entity identified by an `AgentSessionId` (a composite of the agent name and a unique session key). -2. The entity loads its persisted `DurableAgentState`, which includes the complete conversation history. -3. The entity invokes the underlying `AIAgent` with the full conversation history, collects the response, and appends both the request and the response to the state. -4. The updated state is persisted back to durable storage automatically. +2. The entity loads its persisted `DurableAgentState` and session control. +3. The underlying agent obtains context from its configured history path and executes the request. +4. Entity-local changes are persisted. External provider writes and tool effects are not part of a distributed transaction. + +> [!WARNING] +> The [Python prototype in PR #59](https://github.com/microsoft/agent-framework-durable-extension/pull/59) uses schema `2.0.0`, independent response/completion storage and an explicit `isolated_v2` deployment gate. It does not require a local mirror of external/service-owned history. Existing .NET readers do not support this layout. Do not mix these writers or replay old workflow histories through the new Python engine. These are provisional [prototype deployment constraints](../../../python/packages/durabletask/README.md#version-2-deployment-warning). Design review belongs in [ADR PR #88](https://github.com/microsoft/agent-framework-durable-extension/pull/88); the agreed implementation will follow in stacked PRs after ADR approval. Because the entity framework serializes access to each entity instance, concurrent messages to the same session are processed one at a time, eliminating race conditions. @@ -110,7 +113,7 @@ Alternatively, `ConfigureDurableOptions` configures both from a single delegate **Python example:** ```python -app = AgentFunctionApp(agents=[agent]) +app = AgentFunctionApp(agents=[agent], deployment_mode="isolated_v2") ``` ### Console apps / generic hosts @@ -134,7 +137,7 @@ IHost host = Host.CreateDefaultBuilder(args) **Python example:** ```python -worker = DurableAIAgentWorker(TaskHubGrpcWorker(host_address="localhost:4001")) +worker = DurableAIAgentWorker(TaskHubGrpcWorker(host_address="localhost:4001"), deployment_mode="isolated_v2") worker.add_agent(agent) worker.start() ``` diff --git a/python/conftest.py b/python/conftest.py new file mode 100644 index 0000000..32c2102 --- /dev/null +++ b/python/conftest.py @@ -0,0 +1,20 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Declare isolated deployment mode for the local test hubs. + +This pytest-only fixture explicitly acknowledges the tests' isolated mode. It does +not bypass the production default. Gate tests remove or replace the variable with +their function-scoped monkeypatch fixture to exercise missing and invalid modes. +""" + +from collections.abc import Iterator + +import pytest + + +@pytest.fixture(scope="session", autouse=True) +def isolated_test_deployment() -> Iterator[None]: + """Declare isolated mode for tests and restore the environment at session end.""" + with pytest.MonkeyPatch.context() as monkeypatch: + monkeypatch.setenv("DURABLE_AGENTS_DEPLOYMENT_MODE", "isolated_v2") + yield \ No newline at end of file diff --git a/python/packages/azurefunctions/README.md b/python/packages/azurefunctions/README.md index 91d3bb8..e217b23 100644 --- a/python/packages/azurefunctions/README.md +++ b/python/packages/azurefunctions/README.md @@ -8,21 +8,216 @@ Please install this package via pip: pip install agent-framework-azurefunctions --pre ``` +Requires Python 3.10+ and `agent-framework-core>=1.13.0,<2`. The Durable Task dependency requires +`pydantic>=2.11,<3`. Full unit runs passed on Python 3.13/core 1.16, Python 3.13/core 1.13 and +Python 3.10/core 1.16. Pydantic 2.11 runtime validation remains blocked by dependency artifact +downloads. The offline lock, lint, typing and both package builds passed for this follow-up. +See [prototype validation](../../samples/README.md#prototype-validation) +for recorded results and limitations. + +## Version 2 deployment warning + +The settings below describe the [PR #59 prototype](https://github.com/microsoft/agent-framework-durable-extension/pull/59), +not an approved design or the contents of a published package. Design review belongs in +[ADR PR #88](https://github.com/microsoft/agent-framework-durable-extension/pull/88). +The outcome, retention telemetry and validation follow-up builds on prototype baseline `9b4550d`. +It does not establish shared-schema acceptance or a released package contract. +After ADR approval, the agreed implementation will be submitted as stacked PRs rather than merged +from this prototype as-is. + +> **Breaking deployment and state contract.** `AgentFunctionApp` and standalone `create_agent_entity` +> require `deployment_mode="isolated_v2"`, or `DURABLE_AGENTS_DEPLOYMENT_MODE=isolated_v2` when the +> argument is omitted/`None`. This is operator acknowledgement, not a handshake, security boundary +> or proof of isolation. Use a separate hub/deployment with compatible workers and all clients. +> Keep old workers and workflow histories on the old engine. The current .NET reader rejects version 2. + +Only `schemaVersion="2.0.0"` is writable. Legacy `1.x.y` and supported later `2.x.y` state can be +read/round-tripped, but `run`, `reset` and `expire_responses` reject those layouts. No operation +silently upgrades legacy state. Rollback requires compatible version-2 workers, clients and workflow +protocol. Names are unchanged. Reusing an old `@name@key` on an empty new hub is not migration. + +A matching version label does not prove layout compatibility. The shared Python reader rejects +known alternate `data.terminalResults` or `data.completionReceipts` containers instead of treating +their completed requests as new work. Unrelated optional metadata remains opaque, including nested +uses of those names. This guard is not a general format detector or a schema conversion. + +Generated workflow start routes and internal child dispatch wrap new starts with workflow engine +version 2. Raw/legacy starts reject before revised actions execute. Native custom scheduling must use +public `wrap_workflow_input` for new instances. It does not authorize input or migrate old histories. + +Both hosts expose privileged backend `AgentEntity.migrate`, supported by the pure +`migrate_legacy_state` helper. The request requires `source`, `sourceDigest`, `sourceSessionId`, +`destinationSessionId`, `migrationId` and `ownershipTransferId`, with optional `deliveryEvidence` +and `requireKnownOutcomes`. +Use an empty, separately addressed destination after quiescing and authorizing transfer from the +old owner. Nonempty scalar `ingestedPositions` requires a complete accepted-message journal, +including evicted inputs. `complete=True` is an operator assertion. Digest/max-position checks do +not prove authority/completeness or justify inferring a delivered prefix. Without the journal, keep +the old session on the old engine. + +Only recorded responses receive legacy completion backfill and a delivery grace window. Surviving +transcript payloads may be partial, so absence of error content does not prove success. Existing +original mailbox records keep their payload and expiry. A missing matching receipt gets its +`completedAt` from the mailbox's `createdAt`, not migration time. `requireKnownOutcomes=True` on +the entity request, or `require_known_outcomes=True` on the helper, rejects imports without +trustworthy known outcomes. The default legacy-compatible path preserves unknown completion +evidence and duplicate suppression rather than inventing an outcome or rerunning completed work. +Whole-request digest idempotency prevents grace refresh after an exact retry, cold reload or +subsequent run. The original logical session ID is retained for external history. Migration does +not copy that store or move workflow action histories. +No generated HTTP/MCP migration endpoint is provided. These are prototype constraints, not an agreed +cross-runtime migration contract. See [ADR PR #88](https://github.com/microsoft/agent-framework-durable-extension/pull/88) +for the design discussion and [prototype validation](../../samples/README.md#prototype-validation) +for recorded checks and remaining gaps. + ## Durable Agent Extension The durable agent extension lets you host Microsoft Agent Framework agents on Azure Durable Functions so they can persist state, replay conversation history, and recover from failures automatically. ### Basic Usage Example -See the durable functions integration sample in the repository to learn how to: - ```python +from agent_framework import Agent +from agent_framework.openai import OpenAIChatCompletionClient from agent_framework_azurefunctions import AgentFunctionApp -_app = AgentFunctionApp() +assistant = Agent(client=OpenAIChatCompletionClient(), name="assistant") +# Configure the Functions task hub/deployment separately from the old worker +app = AgentFunctionApp(agents=[assistant], deployment_mode="isolated_v2") +``` + +Post messages using the generated `/api/agents/{agent_name}/run` endpoint. + +### History and retention settings + +`AgentFunctionApp` uses the same entity/history integration as the direct Durable Task worker. +Automatic durable history is appended after existing providers to match core's reverse after-hook +order. Only exact built-in `InMemoryHistoryProvider` instances are replaced, preserving `source_id`, +`skip_excluded`, storage flags and optional `after_run_once_per_turn` metadata. Core 1.13 does not +require that hint. Custom in-memory subclasses retain their hooks/session transcripts in the +protected floor, outside durable transcript eviction. Other custom durable-provider JSON state +persists except the transient message buffer and position index. + +Registration does not enable compaction. External primaries and store-only sinks retain their +policies, subject to the intentional service-branch restriction below. Multiple primaries or +duplicate `source_id` values are rejected. Providers append through core hooks, followed by a final +durable flush. Only agents without a context pipeline use direct entity transcript appends. + +Eager pruning and pressure eviction are independent. The matrix assumes no explicit provider +`prune_excluded` override. + +| `retention` | `max_state_bytes=None` | Positive integer byte budget | +| --- | --- | --- | +| `"keep_all"` (default) | No transcript deletion (default) | Evict eligible oldest groups only under pressure | +| `"follow_compaction"` | Prune eligible compaction exclusions only | Prune exclusions, then evict under pressure if needed | + +- `max_state_bytes` defaults to `None`. Azure Functions cannot resolve its backend's hard limit, + so `"backend_limit"` is rejected at registration. Use an explicit positive integer to enable + pressure eviction. A budget does not enable blob offload or raise a backend limit. `"auto"` is + no longer a retention mode. +- The direct Scheduler host's `"backend_limit"` remains a non-normative Python-only convenience, + not part of the portable `None` or positive-integer contract or an agreed shared API. +- Watermarks default to `high_watermark=0.85` and `low_watermark=0.70`, with + `0 < low_watermark < high_watermark <= 1`. The whole serialized entity counts, including mailbox, + completion, session and ingestion state. Protected data can prevent a commit even after pruning. +- `response_delivery_window_seconds` defaults to `60` and must be a positive integer. Delivery + expiry is independent of transcript retention. +- `add_agent()` overrides app defaults. Constructor `workflow_*` settings supply workflow defaults, + and `configure_workflow()` can override them for newly registered nodes, including nested workflows. + Omitted budgets or `INHERIT` inherit the enclosing default. Explicit `None` disables that budget. +- Explicit `prune_excluded=False` on `DurableHistoryProvider` disables eager pruning even with + `follow_compaction`. It does not disable pressure eviction or change an external store's policy. + +As an alternative to the default app above, configure an explicit budget for standalone agents and +disable it for an existing named `workflow`. The sample byte budget is an application choice, not +an inferred Functions backend limit. + +```python +from agent_framework_durabletask import INHERIT + +app = AgentFunctionApp(deployment_mode="isolated_v2", max_state_bytes=800_000, workflow_max_state_bytes=None) +app.add_agent(assistant, retention="follow_compaction", max_state_bytes=INHERIT) +app.configure_workflow(workflow) ``` -- Register agents with `AgentFunctionApp` -- Post messages using the generated `/api/agents/{agent_name}/run` endpoint +`follow_compaction` only prunes exclusions produced by configured compaction. Workflow `full`, +`last_agent` and `custom` projection runs before per-target delta transport. Custom filters execute +during orchestration replay and must be synchronous, deterministic and side-effect-free, but need +not select monotonically increasing positions. Parallel `contextMessageIds` carry occurrence hashes +without rewriting public message IDs. Private forwarding provenance stays in internal checkpoints, +not application metadata. The outgoing logical conversation includes the full selection and all +response messages, not just the delta. Typed/cache-only requests, agent approval/HITL and +output-designated agents use the same contract. + +Generated agent outputs and intermediate events use portable response snapshots. HTTP workflow +results retain structured `value`, including null and falsey values, and response metadata. +External clients do not need the worker's Pydantic class. Worker-side conditions and activities +still receive the locally declared model. Arbitrary activity outputs keep the existing checkpoint +codec and its importable-type requirements. Parent designations also gate direct child outputs. + +### Service ownership, delivery and reset + +Effective `store` follows run options, then agent defaults, then the client's `STORES_BY_DEFAULT`. +For example, `options={"store": False}` selects client-owned history even on a service-storing +client. Explicit `False` excludes saved or supplied service conversation IDs from that invocation +and its history hooks. A later service-owned run can reuse the saved service ID without importing +the intervening client-owned transcript. Switching branches does not migrate or merge history. +External and service-owned runs create no local request-message mirror. + +Durable deliberately suppresses **both load and store hooks** on the inactive external/custom +primary during service-owned runs, including per-service-call persistence. Core 1.16 can still save +to a configured primary on such runs. This branch-isolation restriction is not universal unchanged +hook semantics. Use a distinct store-only sink with its own `source_id` to audit both branches. +Its configured storage flags still apply. + +HTTP polling uses independent original response snapshots in `responseMailbox`, including +serializable metadata and structured `value`. Transcript pruning or reset cannot change those +results. New `completedCorrelations` receipts retain `completedAt` and `outcome` (`succeeded` or +`failed`) after payload expiry. Expired lookup returns `response_expired` with +`durable_status="already_completed"` and `durable_outcome` set to `succeeded`, `failed` or `unknown`. +Older timestamp-only receipts still suppress duplicates when the outcome is unknown. Cleanup can +backfill known outcomes from independent original mailboxes before removing them, even after the +delivery deadline. A possibly pruned legacy transcript without error content is not success evidence. + +For expired delivery, HTTP returns 410. JSON includes top-level `outcome` and +`agent_response.additional_properties.durable_outcome`. Plain text carries `x-ms-durable-outcome`, +and the MCP error includes the invocation outcome. These additions do not modify retained original +response payloads or the standalone SDK API. Acceptance alone cannot create a new completion +receipt. A fresh response with no known invocation outcome raises before either delivery map +changes, while legacy-compatible receipts and fire-and-forget acceptance remain supported. + +Expiry is a logical deadline, not an idle timer. New runs, duplicates and reset remove expired +payloads. Both hosts also expose backend `expire_responses` without model/tool/provider execution. +Idle physical cleanup needs an application-owned schedule or explicit backend signal/manual +operation. No public HTTP/MCP cleanup endpoint is generated. Completion receipts are never removed +by expiry, cleanup or reset. + +Local reset clears session and transcript context but preserves live mailbox payloads, completion +receipts and ingestion evidence. Normal delivery expiry still applies. Reset with a non-durable +custom/external primary raises `NotImplementedError` until provider-owned clearing is available. + +Entity-local state commits once per operation. Only structured `previous_response_not_found` on a +service-owned run permits bounded retries, and only before a stream update, function execution or +service-session advancement. Otherwise fail without restarting the conversation. Provider-hook side +effects are not guaranteed safe or identical on retry. There is no generic non-streaming retry after +runtime failure. Only matching unsupported-stream `TypeError` before consumption negotiates fallback. +Final callbacks receive deep copies preserving Pydantic fields. Opaque SDK `raw_representation` +detachment is best effort and that field is omitted if it cannot be copied. +Uncommitted model/tool effects and external appends can repeat after failure. Completion receipts +last until entity deletion and can exhaust capacity. A bounded receipt protocol and optional +retry-safe external-history adapters remain deferred, with no mandatory core API changes or +guarantee of a distributed transaction or exactly-once uncommitted effects. + +### Retention telemetry + +The shared runtime emits the [retention instruments and bounded attributes](../durabletask/README.md#retention-telemetry) +under scope `agent_framework.durabletask`. Only the OpenTelemetry API is a direct runtime dependency +for this instrumentation. The SDK remains a development dependency, with application-owned meter +providers and exporters. Metrics contain no payloads or session, request or message IDs. + +Removal counts describe staged changes, not confirmed deletion. Host `set_state` returns and +failures both leave commit status `unknown`. Separate persisted-state readback and subsequent model +input are needed to validate retention. Telemetry does not change warm-state rollback or protect +uncommitted external effects from repetition. For more details, review the Python [README](https://github.com/microsoft/agent-framework/tree/main/python/README.md) and the samples directory. diff --git a/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py b/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py index 88f409d..1ab0a6d 100644 --- a/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py +++ b/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py @@ -25,8 +25,14 @@ from agent_framework._telemetry import mark_feature_used from agent_framework_durabletask import ( DEFAULT_MAX_POLL_RETRIES, + DEFAULT_MAX_STATE_BYTES, DEFAULT_POLL_INTERVAL_SECONDS, + DEFAULT_RETENTION, + DELIVERY_WINDOW_SECONDS, + HIGH_WATERMARK, + INHERIT, LEGACY_THREAD_ID_FIELD, + LOW_WATERMARK, MIMETYPE_APPLICATION_JSON, MIMETYPE_TEXT_PLAIN, REQUEST_RESPONSE_FORMAT_JSON, @@ -35,15 +41,30 @@ SESSION_ID_HEADER, WAIT_FOR_RESPONSE_FIELD, WAIT_FOR_RESPONSE_HEADER, + AgentRegistrationSettings, AgentResponseCallbackProtocol, AgentSessionId, ApiResponseFields, DurableAgentState, DurableAIAgent, + RegistrationIdentity, + RetentionMode, RunRequest, + StateBudget, + StateBudgetOverride, deserialize_workflow_output, execute_workflow_activity, + is_terminal_agent_response, plan_workflow_registration, + resolve_state_budget, + resolve_state_budget_override, + serialize_agent_response, + unwrap_workflow_input, + validate_agent_configuration, + validate_response_delivery_window, + validate_retention, + validate_runtime_deployment, + wrap_workflow_input, ) from agent_framework_durabletask._workflows.naming import ( SUBWORKFLOW_REQUEST_SEPARATOR, @@ -56,6 +77,7 @@ ) from agent_framework_durabletask._workflows.registration import collect_hosted_workflows from agent_framework_durabletask._workflows.serialization import strip_pickle_markers, strip_subworkflow_markers +from azure.functions.decorators.function_app import Function from ._entities import create_agent_entity from ._errors import IncomingRequestError @@ -97,9 +119,13 @@ def _json_default(obj: Any) -> Any: A workflow's yielded outputs are reconstructed (see ``deserialize_workflow_output``) before they reach the HTTP response, so they may be framework models (e.g. ``AgentResponse``), dataclasses, or other non-JSON-native objects. - Prefer the type's own serialization so the response carries clean domain - JSON, falling back to ``str`` for anything without one. + Preserve agent response values with the public durable serializer; use the + type's own serialization for other objects, then fall back to ``str``. """ + from agent_framework import AgentResponse + + if isinstance(obj, AgentResponse): + return serialize_agent_response(cast("AgentResponse[Any]", obj)) to_dict = getattr(obj, "to_dict", None) if callable(to_dict): try: @@ -137,6 +163,8 @@ class AgentMetadata: class DFAppBase: def __init__(self, http_auth_level: func.AuthLevel = func.AuthLevel.FUNCTION) -> None: ... + def get_functions(self) -> list[Function]: ... + def function_name(self, name: str) -> Callable[[HandlerT], HandlerT]: ... def route(self, route: str, methods: list[str]) -> Callable[[HandlerT], HandlerT]: ... @@ -172,6 +200,11 @@ class AgentFunctionApp(DFAppBase): - Signal-based operation invocation - Better state management than orchestrations + Set ``deployment_mode="isolated_v2"`` or ``DURABLE_AGENTS_DEPLOYMENT_MODE=isolated_v2`` + to acknowledge an isolated schema 2 task hub/deployment with upgraded clients. + Old workflow histories must remain on the old engine. This acknowledgement is + not runtime proof of isolation and cannot detect peer workers. + Example: ------- @@ -193,11 +226,12 @@ class AgentFunctionApp(DFAppBase): tools=[calculate], ) + # Both options acknowledge an isolated schema 2 deployment. # Option 1: Pass list of agents during initialization - app = AgentFunctionApp(agents=[weather_agent, math_agent]) + app = AgentFunctionApp(agents=[weather_agent, math_agent], deployment_mode="isolated_v2") # Option 2: Add agents after initialization - app = AgentFunctionApp() + app = AgentFunctionApp(deployment_mode="isolated_v2") app.add_agent(weather_agent) app.add_agent(math_agent) @@ -246,6 +280,18 @@ def __init__( poll_interval_seconds: float = DEFAULT_POLL_INTERVAL_SECONDS, enable_mcp_tool_trigger: bool = False, default_callback: AgentResponseCallbackProtocol | None = None, + retention: RetentionMode = DEFAULT_RETENTION, + workflow_retention: RetentionMode | None = None, + max_state_bytes: StateBudget = DEFAULT_MAX_STATE_BYTES, + *, + deployment_mode: str | None = None, + high_watermark: float = HIGH_WATERMARK, + low_watermark: float = LOW_WATERMARK, + response_delivery_window_seconds: int = DELIVERY_WINDOW_SECONDS, + workflow_max_state_bytes: StateBudgetOverride = INHERIT, + workflow_high_watermark: float | None = None, + workflow_low_watermark: float | None = None, + workflow_response_delivery_window_seconds: int | None = None, ): """Initialize the AgentFunctionApp. @@ -265,14 +311,52 @@ def __init__( :param poll_interval_seconds: Delay in seconds between polling attempts. Defaults to ``DEFAULT_POLL_INTERVAL_SECONDS``. :param default_callback: Optional callback invoked for agents without specific callbacks. + :param deployment_mode: Exactly ``isolated_v2`` to acknowledge an isolated schema 2 + deployment with upgraded clients. None reads ``DURABLE_AGENTS_DEPLOYMENT_MODE``. + Old workflow histories stay on the old engine. This is not runtime proof of isolation. + :param retention: Eager pruning policy. ``keep_all`` does not prune compaction exclusions; + ``follow_compaction`` does. Pressure eviction is controlled separately by the budget. + :param max_state_bytes: Positive integer serialized-state budget, or None to disable pressure + eviction. ``backend_limit`` is unsupported because Functions cannot infer its backend limit. + :param workflow_retention: Retention for agent nodes inside hosted workflows. When None, + ``retention`` applies. + :param high_watermark: Budget fraction at which pressure eviction starts. + :param low_watermark: Target budget fraction after pressure eviction. + :param response_delivery_window_seconds: Positive integer response delivery window in seconds. + :param workflow_max_state_bytes: Workflow budget default. INHERIT uses the host budget; + None disables pressure eviction for workflow agents. + :param workflow_high_watermark: Workflow pressure trigger, or None to inherit. + :param workflow_low_watermark: Workflow pressure target, or None to inherit. + :param workflow_response_delivery_window_seconds: Workflow delivery window, or None to inherit. :note: If no agents are provided, they can be added later using :meth:`add_agent`. """ + validate_runtime_deployment(deployment_mode) + validate_retention(retention, high_watermark, low_watermark) + resolved_budget = resolve_state_budget(max_state_bytes) + validate_response_delivery_window(response_delivery_window_seconds) + resolved_workflow_retention = retention if workflow_retention is None else workflow_retention + resolved_workflow_budget = resolve_state_budget_override(workflow_max_state_bytes, resolved_budget) + resolved_workflow_high = high_watermark if workflow_high_watermark is None else workflow_high_watermark + resolved_workflow_low = low_watermark if workflow_low_watermark is None else workflow_low_watermark + resolved_workflow_window = ( + response_delivery_window_seconds + if workflow_response_delivery_window_seconds is None + else workflow_response_delivery_window_seconds + ) + validate_retention(resolved_workflow_retention, resolved_workflow_high, resolved_workflow_low) + validate_response_delivery_window(resolved_workflow_window) + + initial_workflows = self._collect_workflows(workflow, workflows) + logger.debug("[AgentFunctionApp] Initializing with Durable Entities...") # Initialize parent DFApp super().__init__(http_auth_level=http_auth_level) + # Validation accepts only this mode. Retain it rather than re-reading the environment. + self._deployment_mode = "isolated_v2" + # Initialize agent metadata dictionary self._agent_metadata = {} self._workflows: dict[str, Workflow] = {} @@ -281,10 +365,56 @@ def __init__( # so a shared sub-workflow is registered once while two different workflows # whose names collide (including case-only differences) are rejected. self._registered_orchestrations: dict[str, Workflow] = {} + self._registration_identities: dict[tuple[str, str], RegistrationIdentity] = {} + self._registration_failed = False self.enable_health_check = enable_health_check self.enable_http_endpoints = enable_http_endpoints self.enable_mcp_tool_trigger = enable_mcp_tool_trigger self.default_callback = default_callback + self._retention: RetentionMode = retention + self._workflow_retention: RetentionMode = resolved_workflow_retention + self._max_state_bytes = resolved_budget + self._high_watermark = high_watermark + self._low_watermark = low_watermark + self._response_delivery_window_seconds = response_delivery_window_seconds + self._workflow_max_state_bytes = resolved_workflow_budget + self._workflow_high_watermark = resolved_workflow_high + self._workflow_low_watermark = resolved_workflow_low + self._workflow_response_delivery_window_seconds = resolved_workflow_window + + # Validate the whole constructor composition, including later standalone agents, + # before installing any triggers. Never publish these temporary reservations. + agent_settings = AgentRegistrationSettings( + retention, + resolved_budget, + high_watermark, + low_watermark, + response_delivery_window_seconds, + default_callback, + ) + workflow_settings = AgentRegistrationSettings( + resolved_workflow_retention, + resolved_workflow_budget, + resolved_workflow_high, + resolved_workflow_low, + resolved_workflow_window, + default_callback, + ) + if enable_health_check: + RegistrationIdentity(self, self, "health", agent_settings, "health check").reserve( + self._registration_identities, "health_check", namespace="function-name" + ) + identities = dict(self._registration_identities) + for initial_workflow in initial_workflows: + self._preflight_workflow(initial_workflow, workflow_settings, identities) + for agent_instance in agents or []: + self._preflight_agent( + agent_instance, + getattr(agent_instance, "name", None), + agent_settings, + (enable_http_endpoints, enable_mcp_tool_trigger), + identities, + ) try: retries = int(max_poll_retries) @@ -300,7 +430,7 @@ def __init__( # Register each hosted workflow. ``workflow=`` is a convenience alias for a # single-element ``workflows``; both may be combined. - for wf in self._collect_workflows(workflow, workflows): + for wf in initial_workflows: self._register_workflow(wf) # Back-compat: expose the sole workflow as ``.workflow`` when exactly one is @@ -315,11 +445,105 @@ def __init__( # Setup health check if enabled if self.enable_health_check: - self._setup_health_route() + try: + self._setup_health_route() + except Exception: + self._registration_failed = True + raise mark_feature_used(FeatureIndex.AZUREFUNCTIONS) logger.debug("[AgentFunctionApp] Initialization complete") + def _ensure_registration_usable(self) -> None: + if self._registration_failed: + raise RuntimeError( + "Backend registration failed; this app may be partially registered. " + "Create a new app before registering or indexing functions." + ) + + def get_functions(self) -> list[Function]: + """Do not index an app whose backend registration failed.""" + self._ensure_registration_usable() + return super().get_functions() + + def _preflight_agent( + self, + agent: SupportsAgentRun, + name: str | None, + settings: AgentRegistrationSettings, + endpoints: tuple[bool, bool], + identities: dict[tuple[str, str], RegistrationIdentity], + *, + owner: Workflow | None = None, + ) -> None: + if not isinstance(name, str) or not name: + raise ValueError("Agent must have a name to be registered") + validate_agent_configuration(agent, retention=settings.retention) + label = f"workflow '{owner.name}' agent '{name}'" if owner is not None else f"agent '{name}'" + identity = RegistrationIdentity(agent if owner is None else owner, agent, "entity", settings, label, endpoints) + entity_name = AgentSessionId.to_entity_name(name) + identity.reserve(identities, entity_name, namespace="entity-name") + identity.reserve(identities, entity_name, namespace="function-name") + if endpoints[0]: + identity.reserve(identities, self._build_function_name(name, "http"), namespace="function-name") + if endpoints[1]: + identity.reserve(identities, self._build_function_name(name, "mcptool"), namespace="function-name") + + def _preflight_workflow( + self, + workflow: Workflow, + settings: AgentRegistrationSettings, + identities: dict[tuple[str, str], RegistrationIdentity], + ) -> list[Workflow]: + validate_workflow_name(workflow.name) + hosted_workflows = list(collect_hosted_workflows(workflow)) + for hosted in hosted_workflows: + validate_workflow_name(hosted.name) + for executor_id in hosted.executors: + validate_executor_id(executor_id) + label = f"workflow '{hosted.name}'" + identity = RegistrationIdentity(hosted, hosted, "orchestration", settings, label) + orchestrator_name = workflow_orchestrator_name(hosted.name) + identity.reserve(identities, orchestrator_name, namespace="orchestrator-name") + identity.reserve(identities, orchestrator_name, namespace="function-name") + plan = plan_workflow_registration(hosted) + for agent_executor in plan.agent_executors: + validate_executor_id(agent_executor.id) + self._preflight_agent( + agent_executor.agent, + workflow_scoped_executor_id(hosted.name, agent_executor.id), + settings, + (self.enable_http_endpoints, self.enable_mcp_tool_trigger), + identities, + owner=hosted, + ) + for executor in plan.activity_executors: + validate_executor_id(executor.id) + identity = RegistrationIdentity( + hosted, executor, "activity", settings, f"{label} executor '{executor.id}'" + ) + activity_name = workflow_executor_activity_name(hosted.name, executor.id) + identity.reserve(identities, activity_name, namespace="activity-name") + identity.reserve(identities, activity_name, namespace="function-name") + for suffix in ("start", "status", "respond"): + RegistrationIdentity(workflow, workflow, "route", settings, f"workflow '{workflow.name}' routes").reserve( + identities, self._workflow_route_function_name(workflow, suffix), namespace="function-name" + ) + return hosted_workflows + + @staticmethod + def _workflow_route_function_name(workflow: Workflow, suffix: str) -> str: + """Preserve legacy HTTP function names unless a workflow executor occupies one. + + Durable executor names and public URLs stay unchanged. The HTTP prefix keeps + an executor named start, status, or respond indexable alongside its route. + """ + name = f"{workflow_orchestrator_name(workflow.name)}-{suffix}" + plan = plan_workflow_registration(workflow) + if any(executor.id.casefold() == suffix for executor in (*plan.agent_executors, *plan.activity_executors)): + return f"http-{name}" + return name + def _collect_workflows( self, workflow: Workflow | None, @@ -345,7 +569,49 @@ def _collect_workflows( collected.extend(workflows) return collected - def _register_workflow(self, workflow: Workflow) -> None: + def configure_workflow( + self, + workflow: Workflow, + *, + retention: RetentionMode | None = None, + max_state_bytes: StateBudgetOverride = INHERIT, + high_watermark: float | None = None, + low_watermark: float | None = None, + response_delivery_window_seconds: int | None = None, + ) -> None: + """Register a workflow with overrides for its newly registered agent nodes. + + Args: + workflow: Named workflow to register, including its nested workflows. + retention: Eager pruning policy, or None to use the app's workflow default. + max_state_bytes: Workflow budget. INHERIT uses the workflow default; None disables it. + high_watermark: Pressure trigger override, or None to inherit the workflow default. + low_watermark: Pressure target override, or None to inherit the workflow default. + response_delivery_window_seconds: Delivery window override, or None to inherit. + + Raises: + ValueError: Workflow names, agent history providers, or retention settings are invalid. + """ + self._register_workflow( + workflow, + retention=retention, + max_state_bytes=max_state_bytes, + high_watermark=high_watermark, + low_watermark=low_watermark, + response_delivery_window_seconds=response_delivery_window_seconds, + ) + self.workflow = next(iter(self._workflows.values())) if len(self._workflows) == 1 else None + + def _register_workflow( + self, + workflow: Workflow, + *, + retention: RetentionMode | None = None, + max_state_bytes: StateBudgetOverride = INHERIT, + high_watermark: float | None = None, + low_watermark: float | None = None, + response_delivery_window_seconds: int | None = None, + ) -> None: """Register a top-level workflow's durable primitives and HTTP routes. The "what to register" decision (agent -> entity, non-agent -> activity, @@ -357,58 +623,70 @@ def _register_workflow(self, workflow: Workflow) -> None: Raises: ValueError: If the workflow (or a nested sub-workflow) name is - missing/invalid/auto-generated, or a top-level workflow with the - same name is already registered. + missing/invalid/auto-generated, a derived name has a different owner, + or a shared workflow has different settings. """ - validate_workflow_name(workflow.name) - if any(name.casefold() == workflow.name.casefold() for name in self._workflows): - raise ValueError( - f"Workflow '{workflow.name}' is already registered on this app " - "(workflow names are compared case-insensitively)." - ) - - # Validate the whole composition (top-level plus every nested sub-workflow) - # up front, so an invalid/auto-generated nested name (or an executor id that - # would break durable naming / nested-HITL addressing) fails before any - # registration side effects leave the app partially configured. - hosted_workflows = list(collect_hosted_workflows(workflow)) - for hosted in hosted_workflows: - validate_workflow_name(hosted.name) - for executor_id in hosted.executors: - validate_executor_id(executor_id) - - # Check every cross-call collision *before* mutating any state, so a clash - # between a nested sub-workflow and an already-registered orchestration cannot - # leave the app partially configured (e.g. the top-level name added to - # ``_workflows`` while a later child fails). Registration below is then a pure - # commit step. - for hosted in hosted_workflows: - existing = self._registered_orchestrations.get(hosted.name.casefold()) - if existing is not None and existing is not hosted: - raise ValueError( - f"A different workflow named '{hosted.name}' collides with already-registered " - f"'{existing.name}' on this app. A workflow name maps to a single durable " - f"orchestration ('dafx-{hosted.name}'), compared case-insensitively; rename one " - "of them." + self._ensure_registration_usable() + effective_retention = self._workflow_retention if retention is None else retention + effective_budget = resolve_state_budget_override(max_state_bytes, self._workflow_max_state_bytes) + effective_high = self._workflow_high_watermark if high_watermark is None else high_watermark + effective_low = self._workflow_low_watermark if low_watermark is None else low_watermark + effective_window = ( + self._workflow_response_delivery_window_seconds + if response_delivery_window_seconds is None + else response_delivery_window_seconds + ) + validate_retention(effective_retention, effective_high, effective_low) + validate_response_delivery_window(effective_window) + + settings = AgentRegistrationSettings( + effective_retention, + effective_budget, + effective_high, + effective_low, + effective_window, + self.default_callback, + ) + identities = dict(self._registration_identities) + hosted_workflows = self._preflight_workflow(workflow, settings, identities) + previous_metadata = dict(self._agent_metadata) + previous_identities = self._registration_identities + try: + for hosted in hosted_workflows: + if hosted.name.casefold() in self._registered_orchestrations: + continue + self._register_workflow_primitives( + hosted, + retention=effective_retention, + max_state_bytes=effective_budget, + high_watermark=effective_high, + low_watermark=effective_low, + response_delivery_window_seconds=effective_window, ) - + if workflow.name not in self._workflows: + self._register_workflow_routes(workflow) + except Exception: + # SDK trigger decorators have no public rollback. Keep metadata honest and fail closed. + self._registration_failed = True + self._agent_metadata = previous_metadata + self._registration_identities = previous_identities + raise + self._registration_identities = identities + self._registered_orchestrations.update({hosted.name.casefold(): hosted for hosted in hosted_workflows}) self._workflows[workflow.name] = workflow - # Commit: register orchestration primitives for the top-level workflow and every - # nested sub-workflow (deduped by name). - for hosted in hosted_workflows: - if hosted.name.casefold() in self._registered_orchestrations: - continue - self._register_workflow_primitives(hosted) - - # HTTP routes are only exposed for the top-level workflow; sub-workflows are - # driven by the parent via call_sub_orchestrator, not addressed directly. - self._register_workflow_routes(workflow) - - def _register_workflow_primitives(self, workflow: Workflow) -> None: + def _register_workflow_primitives( + self, + workflow: Workflow, + *, + retention: RetentionMode, + max_state_bytes: int | None, + high_watermark: float, + low_watermark: float, + response_delivery_window_seconds: int, + ) -> None: """Register one workflow's entities, activities, and orchestrator (no routes).""" validate_workflow_name(workflow.name) - self._registered_orchestrations[workflow.name.casefold()] = workflow logger.debug("[AgentFunctionApp] Registering workflow '%s'", workflow.name) plan = plan_workflow_registration(workflow) @@ -416,12 +694,16 @@ def _register_workflow_primitives(self, workflow: Workflow) -> None: # Register each workflow agent through the same surface as a standalone # agent (so it stays tracked in ``agents`` / ``get_agent``), under the # workflow-scoped entity id ``{workflow}-{executor}`` the orchestrator - # dispatches to. This keeps two co-hosted workflows that reuse an executor - # id from colliding on one global entity name. + # dispatches to. Preflight has already rejected ambiguous derived names. self.add_agent( agent_executor.agent, callback=self.default_callback, entity_id=workflow_scoped_executor_id(workflow.name, agent_executor.id), + retention=retention, + max_state_bytes=max_state_bytes, + high_watermark=high_watermark, + low_watermark=low_watermark, + response_delivery_window_seconds=response_delivery_window_seconds, ) for executor in plan.activity_executors: # Set up a Functions activity trigger for each non-agent executor, scoped @@ -484,9 +766,8 @@ def workflow_orchestrator(context: df.DurableOrchestrationContext) -> Any: """Generic orchestrator for running the configured workflow.""" input_data = context.get_input() - # Pass the deserialized client input straight to the shared engine, which - # reconstructs the start executor's declared type (see _coerce_initial_input). - initial_message = input_data + # Reject legacy recorded starts before entering the changed engine. + initial_message = unwrap_workflow_input(input_data) # Create local shared state dict for cross-executor state sharing shared_state: dict[str, Any] = {} @@ -508,7 +789,7 @@ def _register_workflow_routes(self, workflow: Workflow) -> None: workflow_name = workflow.name orchestrator_name = workflow_orchestrator_name(workflow_name) - @self.function_name(f"{orchestrator_name}-start") + @self.function_name(self._workflow_route_function_name(workflow, "start")) @self.route(route=f"workflow/{workflow_name}/run", methods=["POST"]) @self.durable_client_input(client_name="client") async def start_workflow_orchestration( @@ -548,7 +829,7 @@ async def start_workflow_orchestration( instance_id = await client.start_new( orchestrator_name, instance_id=requested_instance_id, - client_input=client_input, + client_input=wrap_workflow_input(client_input), ) if wait_for_response: @@ -568,7 +849,7 @@ async def start_workflow_orchestration( return self._build_workflow_accepted_response(req, workflow_name, instance_id) - @self.function_name(f"{orchestrator_name}-status") + @self.function_name(self._workflow_route_function_name(workflow, "status")) @self.route(route=f"workflow/{workflow_name}/status/{{instanceId}}", methods=["GET"]) @self.durable_client_input(client_name="client") async def get_workflow_status( @@ -636,7 +917,7 @@ async def get_workflow_status( mimetype="application/json", ) - @self.function_name(f"{orchestrator_name}-respond") + @self.function_name(self._workflow_route_function_name(workflow, "respond")) @self.route(route=f"workflow/{workflow_name}/respond/{{instanceId}}/{{requestId}}", methods=["POST"]) @self.durable_client_input(client_name="client") async def send_hitl_response(req: func.HttpRequest, client: df.DurableOrchestrationClient) -> func.HttpResponse: @@ -829,6 +1110,11 @@ def add_agent( enable_mcp_tool_trigger: bool | None = None, *, entity_id: str | None = None, + retention: RetentionMode | None = None, + max_state_bytes: StateBudgetOverride = INHERIT, + high_watermark: float | None = None, + low_watermark: float | None = None, + response_delivery_window_seconds: int | None = None, ) -> None: """Add an agent to the function app after initialization. @@ -845,25 +1131,41 @@ def add_agent( durable entity (and the ``agents`` / ``get_agent`` key) matches the identity the orchestrator dispatches to. Mirrors ``DurableAIAgentWorker.add_agent(entity_id=...)``. + retention: Per-agent retention override. When None, the app-level setting is used. + max_state_bytes: Per-agent budget. INHERIT uses the host default; None disables it. + Functions requires an explicit integer instead of ``backend_limit``. + high_watermark: Pressure trigger override, or None to inherit the host default. + low_watermark: Pressure target override, or None to inherit the host default. + response_delivery_window_seconds: Delivery window override, or None to inherit. Raises: - ValueError: If the agent doesn't have a 'name' attribute. + ValueError: If the agent has no name, retention settings or history providers are invalid, + or an existing registration has a different agent, owner, or configuration. """ + self._ensure_registration_usable() # Get agent name from the agent's name attribute name = getattr(agent, "name", None) - if name is None: + if name is None and not entity_id: raise ValueError("Agent does not have a 'name' attribute. All agents must have a 'name' attribute.") # The registration name keys the agent everywhere on this app (metadata, # routes, entity). It defaults to the agent name but can be overridden so a # workflow agent is keyed by its executor id. registration_name = entity_id or name - - if registration_name in self._agent_metadata: - logger.warning( - "[AgentFunctionApp] Agent '%s' is already registered, skipping duplicate.", registration_name - ) - return + if not isinstance(registration_name, str) or not registration_name.strip(): + raise ValueError("Agent registration requires a nonblank name or explicit entity_id.") + + effective_retention = self._retention if retention is None else retention + effective_budget = resolve_state_budget_override(max_state_bytes, self._max_state_bytes) + effective_high = self._high_watermark if high_watermark is None else high_watermark + effective_low = self._low_watermark if low_watermark is None else low_watermark + effective_window = ( + self._response_delivery_window_seconds + if response_delivery_window_seconds is None + else response_delivery_window_seconds + ) + validate_retention(effective_retention, effective_high, effective_low) + validate_response_delivery_window(effective_window) effective_enable_http_endpoint = ( self.enable_http_endpoints if enable_http_endpoint is None else self._coerce_to_bool(enable_http_endpoint) @@ -873,6 +1175,20 @@ def add_agent( if enable_mcp_tool_trigger is None else self._coerce_to_bool(enable_mcp_tool_trigger) ) + effective_callback = self.default_callback if callback is None else callback + settings = AgentRegistrationSettings( + effective_retention, effective_budget, effective_high, effective_low, effective_window, effective_callback + ) + identities = dict(self._registration_identities) + self._preflight_agent( + agent, + registration_name, + settings, + (effective_enable_http_endpoint, effective_enable_mcp_endpoint), + identities, + ) + if any(name.casefold() == registration_name.casefold() for name in self._agent_metadata): + return logger.debug(f"[AgentFunctionApp] Adding agent: {registration_name}") logger.debug(f"[AgentFunctionApp] Route: /api/agents/{registration_name}") @@ -885,18 +1201,29 @@ def add_agent( f"[AgentFunctionApp] MCP tool trigger: {'enabled' if effective_enable_mcp_endpoint else 'disabled'}" ) - # Store agent metadata + try: + self._setup_agent_functions( + agent, + registration_name, + effective_callback, + effective_enable_http_endpoint, + effective_enable_mcp_endpoint, + retention=effective_retention, + max_state_bytes=effective_budget, + high_watermark=effective_high, + low_watermark=effective_low, + response_delivery_window_seconds=effective_window, + ) + except Exception: + self._registration_failed = True + raise + self._agent_metadata[registration_name] = AgentMetadata( agent=agent, http_endpoint_enabled=effective_enable_http_endpoint, mcp_tool_enabled=effective_enable_mcp_endpoint, ) - - effective_callback = callback or self.default_callback - - self._setup_agent_functions( - agent, registration_name, effective_callback, effective_enable_http_endpoint, effective_enable_mcp_endpoint - ) + self._registration_identities = identities logger.debug(f"[AgentFunctionApp] Agent '{registration_name}' added successfully") @@ -940,6 +1267,12 @@ def _setup_agent_functions( callback: AgentResponseCallbackProtocol | None, enable_http_endpoint: bool, enable_mcp_tool_trigger: bool, + *, + retention: RetentionMode = DEFAULT_RETENTION, + max_state_bytes: int | None = None, + high_watermark: float = HIGH_WATERMARK, + low_watermark: float = LOW_WATERMARK, + response_delivery_window_seconds: int = DELIVERY_WINDOW_SECONDS, ) -> None: """Set up the HTTP trigger, entity, and MCP tool trigger for a specific agent. @@ -949,6 +1282,11 @@ def _setup_agent_functions( callback: Optional callback to receive response updates enable_http_endpoint: Whether to create HTTP endpoint enable_mcp_tool_trigger: Whether to create MCP tool trigger + retention: How much of the conversation durable state may discard. + max_state_bytes: Resolved pressure budget, or None to disable pressure eviction. + high_watermark: Budget fraction at which pressure eviction starts. + low_watermark: Target budget fraction after pressure eviction. + response_delivery_window_seconds: Response delivery window in seconds. """ logger.debug(f"[AgentFunctionApp] Setting up functions for agent '{agent_name}'...") @@ -959,7 +1297,16 @@ def _setup_agent_functions( "[AgentFunctionApp] HTTP run route disabled for agent '%s'", agent_name, ) - self._setup_agent_entity(agent, agent_name, callback) + self._setup_agent_entity( + agent, + agent_name, + callback, + retention=retention, + max_state_bytes=max_state_bytes, + high_watermark=high_watermark, + low_watermark=low_watermark, + response_delivery_window_seconds=response_delivery_window_seconds, + ) if enable_mcp_tool_trigger: agent_description = agent.description @@ -1049,9 +1396,15 @@ async def http_start(req: func.HttpRequest, client: df.DurableOrchestrationClien ) logger.debug(f"[HTTP Trigger] Result status: {result.get('status', 'unknown')}") + if result.get("status") == "success": + status_code = 200 + elif result.get("error_code") == "response_expired": + status_code = 410 + else: + status_code = 500 return self._create_http_response( payload=result, - status_code=200 if result.get("status") == "success" else 500, + status_code=status_code, request_response_format=request_response_format, session_id=session_id, ) @@ -1101,6 +1454,12 @@ def _setup_agent_entity( agent: SupportsAgentRun, agent_name: str, callback: AgentResponseCallbackProtocol | None, + *, + retention: RetentionMode = DEFAULT_RETENTION, + max_state_bytes: int | None = None, + high_watermark: float = HIGH_WATERMARK, + low_watermark: float = LOW_WATERMARK, + response_delivery_window_seconds: int = DELIVERY_WINDOW_SECONDS, ) -> None: """Register the durable entity responsible for agent state. @@ -1108,9 +1467,24 @@ def _setup_agent_entity( agent: The agent instance agent_name: The agent name (used for both entity identification and function naming) callback: Optional callback for response updates + retention: How much of the conversation durable state may discard. + max_state_bytes: Resolved pressure budget, or None to disable pressure eviction. + high_watermark: Budget fraction at which pressure eviction starts. + low_watermark: Target budget fraction after pressure eviction. + response_delivery_window_seconds: Response delivery window in seconds. """ # Use the prefixed entity name for both registration and function naming entity_name_with_prefix = AgentSessionId.to_entity_name(agent_name) + entity_handler = create_agent_entity( + agent, + callback, + deployment_mode=self._deployment_mode, + retention=retention, + max_state_bytes=max_state_bytes, + high_watermark=high_watermark, + low_watermark=low_watermark, + response_delivery_window_seconds=response_delivery_window_seconds, + ) def entity_function(context: df.DurableEntityContext) -> None: """Durable entity that manages agent execution and conversation state. @@ -1118,9 +1492,8 @@ def entity_function(context: df.DurableEntityContext) -> None: Operations: - run: Execute the agent with a message - run_agent: (Deprecated) Execute the agent with a message - - reset: Clear conversation history + - reset: Delegate reset to AgentEntity """ - entity_handler = create_agent_entity(agent, callback) entity_handler(context) # Set function name for Azure Functions (used in function.json generation) @@ -1286,6 +1659,8 @@ async def _handle_mcp_tool_invocation( logger.info("[MCP Tool] Agent '%s' responded successfully", agent_name) return response_text error_msg = result.get("error", "Unknown error") + if result.get("status") == "already_completed": + error_msg = f"{error_msg} Invocation outcome: {result.get('outcome', 'unknown')}." logger.error("[MCP Tool] Agent '%s' execution failed: %s", agent_name, error_msg) raise RuntimeError(f"Agent execution failed: {error_msg}") @@ -1407,14 +1782,53 @@ async def _poll_entity_for_response( return None agent_response = state.try_get_agent_response(correlation_id) - if agent_response: - result = self._build_success_result( - response_message=agent_response.text, - message=message, - session_id=session_id, - correlation_id=correlation_id, - state=state, + if agent_response is not None: + snapshot = serialize_agent_response(agent_response) + errors = [ + content + for response_message in agent_response.messages + if response_message.role != "tool" + for content in response_message.contents + if content.type == "error" + ] + expired_error = next((error for error in errors if error.error_code == "response_expired"), None) + expired = ( + expired_error is not None + or agent_response.additional_properties.get("durable_status") == "already_completed" ) + if is_terminal_agent_response(agent_response): + error = expired_error or (errors[0] if errors else None) + error_message = error.message if error is not None else None + error_code = "response_expired" if expired else (error.error_code if error is not None else None) + if not error_message: + error_message = agent_response.text or ( + "This request completed, but its response delivery window has expired." + if expired + else "Agent execution failed." + ) + result = self._build_response_payload( + response=None, + message=message, + session_id=session_id, + status="already_completed" if expired else "error", + correlation_id=correlation_id, + extra_fields={ + "error": error_message, + "error_code": error_code, + ApiResponseFields.MESSAGE_COUNT: state.message_count, + }, + ) + else: + result = self._build_success_result( + response_message=agent_response.text, + message=message, + session_id=session_id, + correlation_id=correlation_id, + state=state, + ) + result["agent_response"] = snapshot + if expired: + result["outcome"] = agent_response.additional_properties.get("durable_outcome", "unknown") logger.debug(f"[HTTP Trigger] Found response for correlation ID: {correlation_id}") except Exception as exc: @@ -1574,6 +1988,8 @@ def _build_plain_text_response( """Return a plain-text response with optional session identifier header.""" body_text = payload if isinstance(payload, str) else self._convert_payload_to_text(payload) headers = {SESSION_ID_HEADER: session_id} if session_id is not None else None + if isinstance(payload, dict) and payload.get("status") == "already_completed": + headers = {**(headers or {}), "x-ms-durable-outcome": str(payload.get("outcome", "unknown"))} return func.HttpResponse(body_text, status_code=status_code, mimetype=MIMETYPE_TEXT_PLAIN, headers=headers) def _build_json_response(self, payload: dict[str, Any] | str, status_code: int) -> func.HttpResponse: diff --git a/python/packages/azurefunctions/agent_framework_azurefunctions/_entities.py b/python/packages/azurefunctions/agent_framework_azurefunctions/_entities.py index 83ad50a..1e0d176 100644 --- a/python/packages/azurefunctions/agent_framework_azurefunctions/_entities.py +++ b/python/packages/azurefunctions/agent_framework_azurefunctions/_entities.py @@ -16,10 +16,23 @@ import azure.durable_functions as df from agent_framework import SupportsAgentRun from agent_framework_durabletask import ( + DEFAULT_MAX_STATE_BYTES, + DEFAULT_RETENTION, + DELIVERY_WINDOW_SECONDS, + HIGH_WATERMARK, + LOW_WATERMARK, AgentEntity, AgentEntityStateProviderMixin, AgentResponseCallbackProtocol, + RetentionMode, + StateBudget, + resolve_state_budget, run_agent_coroutine, + serialize_agent_response, + validate_agent_configuration, + validate_response_delivery_window, + validate_retention, + validate_runtime_deployment, ) logger = logging.getLogger("agent_framework.azurefunctions") @@ -37,8 +50,10 @@ def __init__(self, context: df.DurableEntityContext) -> None: def _get_state_dict(self) -> dict[str, Any]: raw_state = self._context.get_state(lambda: {}) - if not isinstance(raw_state, dict): + if raw_state is None: return {} + if not isinstance(raw_state, dict): + raise ValueError("Existing durable entity state must be a dictionary; refusing to replace malformed state.") return cast(dict[str, Any], raw_state) def _set_state_dict(self, state: dict[str, Any]) -> None: @@ -47,10 +62,20 @@ def _set_state_dict(self, state: dict[str, Any]) -> None: def _get_session_id_from_entity(self) -> str: return str(self._context.entity_key) + def _get_entity_name_from_entity(self) -> str: + return str(self._context.entity_name) + def create_agent_entity( agent: SupportsAgentRun, callback: AgentResponseCallbackProtocol | None = None, + *, + deployment_mode: str | None = None, + retention: RetentionMode = DEFAULT_RETENTION, + max_state_bytes: StateBudget = DEFAULT_MAX_STATE_BYTES, + high_watermark: float = HIGH_WATERMARK, + low_watermark: float = LOW_WATERMARK, + response_delivery_window_seconds: int = DELIVERY_WINDOW_SECONDS, ) -> Callable[[df.DurableEntityContext], None]: """Factory function to create an agent entity class. @@ -58,9 +83,29 @@ def create_agent_entity( agent: The Microsoft Agent Framework agent instance (must implement SupportsAgentRun) callback: Optional callback invoked during streaming and final responses + Keyword Args: + deployment_mode: Exactly ``isolated_v2`` to acknowledge an isolated schema 2 + deployment with upgraded clients. None reads ``DURABLE_AGENTS_DEPLOYMENT_MODE``. + Old workflow histories stay on the old engine. This acknowledgement is not runtime + proof of isolation and cannot detect peer workers. + retention: Eager pruning policy, independent of pressure eviction. + max_state_bytes: Positive integer pressure budget, or None to disable it. Functions cannot + resolve ``backend_limit`` because the storage backend is configured outside Python. + high_watermark: Budget fraction at which pressure eviction starts. + low_watermark: Target budget fraction after pressure eviction. + response_delivery_window_seconds: Positive integer response delivery window in seconds. + Returns: Entity function configured with the agent + + Raises: + ValueError: Deployment mode, retention settings, or the agent's history providers are invalid. """ + validate_runtime_deployment(deployment_mode) + validate_retention(retention, high_watermark, low_watermark) + resolved_budget = resolve_state_budget(max_state_bytes) + validate_response_delivery_window(response_delivery_window_seconds) + validate_agent_configuration(agent, retention=retention) async def _entity_coroutine(context: df.DurableEntityContext) -> None: """Async handler that executes the entity operations.""" @@ -69,7 +114,16 @@ async def _entity_coroutine(context: df.DurableEntityContext) -> None: logger.debug("[entity_function] Operation: %s", context.operation_name) state_provider = AzureFunctionEntityStateProvider(context) - entity = AgentEntity(agent, callback, state_provider=state_provider) + entity = AgentEntity( + agent, + callback, + state_provider=state_provider, + retention=retention, + max_state_bytes=resolved_budget, + high_watermark=high_watermark, + low_watermark=low_watermark, + response_delivery_window_seconds=response_delivery_window_seconds, + ) operation = context.operation_name @@ -84,12 +138,18 @@ async def _entity_coroutine(context: df.DurableEntityContext) -> None: request = "" if input_data is None else str(cast(object, input_data)) result = await entity.run(request) - context.set_result(result.to_dict()) + context.set_result(serialize_agent_response(result)) elif operation == "reset": entity.reset() context.set_result({"status": "reset"}) + elif operation == "expire_responses": + context.set_result({"expired": entity.expire_responses()}) + + elif operation == "migrate": + context.set_result(entity.migrate(context.get_input())) + else: logger.error("[entity_function] Unknown operation: %s", operation) context.set_result({"error": f"Unknown operation: {operation}"}) diff --git a/python/packages/azurefunctions/agent_framework_azurefunctions/_orchestration.py b/python/packages/azurefunctions/agent_framework_azurefunctions/_orchestration.py index cbbd134..98fa06e 100644 --- a/python/packages/azurefunctions/agent_framework_azurefunctions/_orchestration.py +++ b/python/packages/azurefunctions/agent_framework_azurefunctions/_orchestration.py @@ -149,28 +149,8 @@ def __init__(self, context: AgentOrchestrationContextType): def generate_unique_id(self) -> str: return str(self.context.new_uuid()) - def get_run_request( - self, - message: str, - *, - options: dict[str, Any] | None = None, - ) -> RunRequest: - """Get the current run request from the orchestration context. - - Args: - message: The message to send to the agent - options: Optional options dictionary. Supported keys include - ``response_format``, ``enable_tool_calls``, and ``wait_for_response``. - Additional keys are forwarded to the agent execution. - - Returns: - RunRequest: The current run request - """ - # Create a copy to avoid modifying the caller's dict - - request = super().get_run_request(message, options=options) - request.orchestration_id = self.context.instance_id - return request + def _orchestration_id(self) -> str | None: + return self.context.instance_id def run_durable_agent( self, diff --git a/python/packages/azurefunctions/agent_framework_azurefunctions/_workflow.py b/python/packages/azurefunctions/agent_framework_azurefunctions/_workflow.py index e15ae94..5353e50 100644 --- a/python/packages/azurefunctions/agent_framework_azurefunctions/_workflow.py +++ b/python/packages/azurefunctions/agent_framework_azurefunctions/_workflow.py @@ -25,7 +25,6 @@ PendingHITLRequest, TaskMetadata, TaskType, - _extract_message_content, # pyright: ignore[reportPrivateUsage] build_agent_executor_response, execute_hitl_response_handler, route_message_through_edge_groups, @@ -48,7 +47,6 @@ "PendingHITLRequest", "TaskMetadata", "TaskType", - "_extract_message_content", "build_agent_executor_response", "execute_hitl_response_handler", "route_message_through_edge_groups", diff --git a/python/packages/azurefunctions/agent_framework_azurefunctions/_workflow_af_context.py b/python/packages/azurefunctions/agent_framework_azurefunctions/_workflow_af_context.py index eaf99a5..2bf84a9 100644 --- a/python/packages/azurefunctions/agent_framework_azurefunctions/_workflow_af_context.py +++ b/python/packages/azurefunctions/agent_framework_azurefunctions/_workflow_af_context.py @@ -12,7 +12,7 @@ from datetime import datetime from typing import Any -from agent_framework_durabletask import AgentSessionId, DurableAgentSession, DurableAIAgent +from agent_framework_durabletask import WorkflowOrchestrationContext, build_agent_task from azure.durable_functions import DurableOrchestrationContext from ._orchestration import AzureFunctionsAgentExecutor @@ -56,12 +56,22 @@ def current_utc_datetime(self) -> datetime: # -- Agent / Activity dispatch -------------------------------------------- - def prepare_agent_task(self, executor_id: str, message: str, orchestration_instance_id: str) -> Any: - session_id = AgentSessionId(name=executor_id, key=orchestration_instance_id) - session = DurableAgentSession(durable_session_id=session_id) - az_executor = AzureFunctionsAgentExecutor(self._context) - agent = DurableAIAgent(az_executor, executor_id) - return agent.run(message, session=session) + def prepare_agent_task( + self, + executor_id: str, + message: str, + orchestration_instance_id: str, + context_messages: list[dict[str, Any]] | None = None, + context_message_ids: list[str] | None = None, + ) -> Any: + return build_agent_task( + AzureFunctionsAgentExecutor(self._context), + executor_id, + message, + orchestration_instance_id, + context_messages, + context_message_ids, + ) def prepare_activity_task(self, activity_name: str, input_json: str) -> Any: orchestration_context: Any = self._context @@ -103,3 +113,8 @@ def cancel_task(self, task: Any) -> None: def get_task_result(self, task: Any) -> Any: return getattr(task, "result", None) + + +# Ensure the adapter satisfies the protocol. Validated statically by the type checker, +# so a signature change on the protocol is caught here rather than at a distant call site. +_protocol_check: type[WorkflowOrchestrationContext] = AzureFunctionsWorkflowContext diff --git a/python/packages/azurefunctions/tests/integration_tests/conftest.py b/python/packages/azurefunctions/tests/integration_tests/conftest.py index c2e6bb7..7a9ba77 100644 --- a/python/packages/azurefunctions/tests/integration_tests/conftest.py +++ b/python/packages/azurefunctions/tests/integration_tests/conftest.py @@ -344,6 +344,8 @@ def _start_function_app(sample_path: Path, port: int) -> subprocess.Popen[Any]: # This prevents conflicts between parallel or repeated test runs, as Durable Functions # use the task hub name to separate orchestration state. env["TASKHUB_NAME"] = f"test{uuid.uuid4().hex[:8]}" + # Opt in only for the subprocess using this isolated test hub. + env["DURABLE_AGENTS_DEPLOYMENT_MODE"] = "isolated_v2" # On Windows, use CREATE_NEW_PROCESS_GROUP to allow proper termination # shell=True only on Windows to handle PATH resolution diff --git a/python/packages/azurefunctions/tests/integration_tests/live_media_app/function_app.py b/python/packages/azurefunctions/tests/integration_tests/live_media_app/function_app.py new file mode 100644 index 0000000..4bedb38 --- /dev/null +++ b/python/packages/azurefunctions/tests/integration_tests/live_media_app/function_app.py @@ -0,0 +1,186 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Copied into a temporary app by test_16, never deployed as a sample. + +Only model I/O is substituted. AgentFunctionApp registers the production entity +handler and the Functions worker supplies its DurableEntityContext. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import uuid +from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence +from pathlib import Path +from typing import Any + +import agent_framework_durabletask +import azure.durable_functions as df +import azure.functions as func +from agent_framework import Agent, BaseChatClient, ChatResponse, ChatResponseUpdate, Content, Message, ResponseStream +from agent_framework_durabletask import AgentSessionId, DurableHistoryProvider, RunRequest +from agent_framework_durabletask._history_provider import current_durable_history_binding +from opentelemetry.metrics import get_meter_provider, set_meter_provider +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import InMemoryMetricReader, Sum + +import agent_framework_azurefunctions +from agent_framework_azurefunctions import AgentFunctionApp +from agent_framework_azurefunctions._entities import AzureFunctionEntityStateProvider + +ROOT = Path(__file__).resolve().parent +CONFIG = json.loads((ROOT / "live_config.json").read_text(encoding="utf-8")) +BOOT_ID = uuid.uuid4().hex +ENTITY = df.EntityId(AgentSessionId.to_entity_name(CONFIG["agent"]), CONFIG["session"]) + +for package, expected in ( + (agent_framework_azurefunctions, CONFIG["azurefunctions_source"]), + (agent_framework_durabletask, CONFIG["durabletask_source"]), +): + assert package.__file__ is not None + if Path(package.__file__).resolve().parent != Path(expected).resolve(): + raise RuntimeError("Live Functions test imported an extension from a different checkout") + +READER = InMemoryMetricReader() +METERS = MeterProvider(metric_readers=[READER], shutdown_on_exit=False) +set_meter_provider(METERS) +if get_meter_provider() is not METERS: + raise RuntimeError("Live test requires its own in-memory metric reader in the Functions worker") + + +class RecordingModel(BaseChatClient): + def __init__(self) -> None: + super().__init__() + self.calls = 0 + + def _inner_get_response( + self, + *, + messages: Sequence[Message], + stream: bool, + options: Mapping[str, Any], + **kwargs: Any, + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + async def update() -> ChatResponseUpdate: + await self._validate_options(options) + binding = current_durable_history_binding() + if binding is None or not isinstance(binding.state_provider, AzureFunctionEntityStateProvider): + raise RuntimeError("Expected the production Functions state provider on the async bridge") + context = binding.state_provider._context + if not isinstance(context, df.DurableEntityContext): + raise RuntimeError("Expected a real Functions DurableEntityContext") + if context.entity_name != ENTITY.name or context.entity_key != ENTITY.key: + raise RuntimeError("Functions context addressed the wrong test entity") + current_id = next(message.message_id for message in reversed(messages) if message.role == "user") + self.calls += 1 + await asyncio.to_thread( + (ROOT / f"model-{BOOT_ID}-{self.calls}.json").write_text, + json.dumps({ + "boot": BOOT_ID, + "calls": self.calls, + "current_id": current_id, + "messages": [message.to_dict() for message in messages], + "context": { + "provider": type(binding.state_provider).__name__, + "type": type(context).__name__, + "entity_name": context.entity_name, + "entity_key": context.entity_key, + "operation": context.operation_name, + }, + }), + encoding="utf-8", + ) + return ChatResponseUpdate( + role="assistant", + author_name="retention-model", + contents=[Content.from_text(f"answer:{current_id}")], + message_id=f"{current_id}-answer", + response_id=f"response:{current_id}", + finish_reason="stop", + ) + + async def updates() -> AsyncIterable[ChatResponseUpdate]: + yield await update() + + async def response() -> ChatResponse: + return ChatResponse.from_updates([await update()]) + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) if stream else response() + + +MODEL = RecordingModel() +app = AgentFunctionApp( + agents=[ + Agent( + client=MODEL, + name=CONFIG["agent"], + id=CONFIG["agent"], + default_options={"store": False}, + context_providers=[DurableHistoryProvider()], + ) + ], + http_auth_level=func.AuthLevel.ANONYMOUS, + enable_http_endpoints=False, + deployment_mode="isolated_v2", + retention="keep_all", + max_state_bytes=CONFIG["max_state_bytes"], + response_delivery_window_seconds=CONFIG["delivery_window_seconds"], +) + + +def _json(value: Any, status: int = 200) -> func.HttpResponse: + return func.HttpResponse(json.dumps(value), status_code=status, mimetype="application/json") + + +@app.route(route="retention/run", methods=["POST"]) +@app.durable_client_input(client_name="client") +async def run(req: func.HttpRequest, client: df.DurableOrchestrationClient) -> func.HttpResponse: + # The destination is fixed by the test, not a caller-supplied entity or path. + if len(req.get_body()) > 100_000: + return _json({"error": "oversized test request"}, 413) + try: + payload = req.get_json() + if not isinstance(payload, dict): + raise ValueError("Object required") + request = RunRequest.from_dict(payload) + if not request.context_messages: + raise ValueError("Projected messages required") + except (KeyError, TypeError, ValueError): + return _json({"error": "invalid test request"}, 400) + await client.signal_entity(ENTITY, "run", request.to_dict()) + return _json({"correlation": request.correlation_id, "session": ENTITY.key}, 202) + + +@app.route(route="retention/state", methods=["GET"]) +@app.durable_client_input(client_name="client") +async def state(req: func.HttpRequest, client: df.DurableOrchestrationClient) -> func.HttpResponse: + result = await client.read_entity_state(ENTITY) + # No typed-state reserialization here. Return the backend JSON unchanged. + return _json(result.entity_state) if result.entity_exists else _json(None, 404) + + +@app.route(route="retention/capture", methods=["GET"]) +def capture(req: func.HttpRequest) -> func.HttpResponse: + # Only the latest synthetic artifact is readable. A restarted worker must not + # mistake the old process's on-disk capture for a new model invocation. + path = ROOT / f"model-{BOOT_ID}-{MODEL.calls}.json" + record = json.loads(path.read_text(encoding="utf-8")) if MODEL.calls else None + return _json({"boot": BOOT_ID, "pid": os.getpid(), "calls": MODEL.calls, "capture": record}) + + +@app.route(route="retention/metrics", methods=["GET"]) +def metrics(req: func.HttpRequest) -> func.HttpResponse: + rows: list[dict[str, Any]] = [] + data = READER.get_metrics_data() + if data is not None: + for resource in data.resource_metrics: + for scope in resource.scope_metrics: + for metric in scope.metrics: + if metric.name.startswith("durable.retention.") and isinstance(metric.data, Sum): + rows.extend( + {"name": metric.name, "value": point.value, "attributes": dict(point.attributes or {})} + for point in metric.data.data_points + ) + return _json({"boot": BOOT_ID, "rows": rows}) diff --git a/python/packages/azurefunctions/tests/integration_tests/test_01_single_agent.py b/python/packages/azurefunctions/tests/integration_tests/test_01_single_agent.py index ff0e425..b57922f 100644 --- a/python/packages/azurefunctions/tests/integration_tests/test_01_single_agent.py +++ b/python/packages/azurefunctions/tests/integration_tests/test_01_single_agent.py @@ -14,8 +14,11 @@ uv run pytest packages/azurefunctions/tests/integration_tests/test_01_single_agent.py -v """ +import json + import pytest -from agent_framework_durabletask import SESSION_ID_HEADER +from agent_framework import AgentResponse +from agent_framework_durabletask import SESSION_ID_HEADER, serialize_agent_response # Module-level markers - applied to all tests in this file pytestmark = [ @@ -35,6 +38,29 @@ def _setup(self, base_url: str, sample_helper) -> None: self.base_url = f"{base_url}/api/agents/Joker" self.helper = sample_helper + def _assert_success_response(self, response, session_id: str) -> dict: + """Check actual response delivery independently of local transcript storage.""" + assert response.status_code == 200, response.text + data = response.json() + assert data["status"] == "success", data + assert data["session_id"] == session_id + assert data["correlation_id"] + assert data["response"].strip() + + # The legacy field counts local transcript entries, not completed executions. + # Sample 01 uses Foundry's default service-managed history. + assert data["message_count"] == 0 + + snapshot = data["agent_response"] + assert snapshot["type"] == "agent_response" + assert snapshot["created_at"], "the Foundry result timestamp was lost" + delivered = AgentResponse.from_dict(snapshot) + assert delivered.messages + assert delivered.text == data["response"] + assert all(content.type != "error" for message in delivered.messages for content in message.contents) + assert json.loads(json.dumps(serialize_agent_response(delivered))) == snapshot + return data + def test_health_check(self, base_url: str, sample_helper) -> None: """Test health check endpoint.""" response = sample_helper.get(f"{base_url}/api/health") @@ -48,23 +74,12 @@ def test_simple_message_json(self) -> None: f"{self.base_url}/run", {"message": "Tell me a short joke about cloud computing.", "session_id": "test-simple-json"}, ) - # Agent can return 200 (immediate) or 202 (async with wait_for_response=false) - assert response.status_code in [200, 202] - data = response.json() - - if response.status_code == 200: - # Synchronous response - check result directly - assert data["status"] == "success" - assert "response" in data - assert data["message_count"] >= 1 - else: - # Async response - check we got correlation info - assert "correlation_id" in data or "session_id" in data + self._assert_success_response(response, "test-simple-json") def test_simple_message_plain_text(self) -> None: """Test sending a message with plain text payload.""" response = self.helper.post_text(f"{self.base_url}/run", "Tell me a short joke about networking.") - assert response.status_code in [200, 202] + assert response.status_code == 200, response.text # Agent responded with plain text when the request body was text/plain. assert response.text.strip() @@ -75,7 +90,7 @@ def test_session_id_in_query(self) -> None: response = self.helper.post_text( f"{self.base_url}/run?session_id=test-query-session", "Tell me a short joke about weather in Texas." ) - assert response.status_code in [200, 202] + assert response.status_code == 200, response.text assert response.text.strip() assert response.headers.get(SESSION_ID_HEADER) == "test-query-session" @@ -85,7 +100,7 @@ def test_legacy_thread_id_in_query_still_accepted(self) -> None: response = self.helper.post_text( f"{self.base_url}/run?thread_id=test-legacy-query", "Tell me a short joke about weather in Texas." ) - assert response.status_code in [200, 202] + assert response.status_code == 200, response.text assert response.text.strip() assert response.headers.get(SESSION_ID_HEADER) == "test-legacy-query" @@ -93,34 +108,26 @@ def test_legacy_thread_id_in_query_still_accepted(self) -> None: assert response.headers.get("x-ms-thread-id") is None def test_conversation_continuity(self) -> None: - """Test conversation context is maintained across requests.""" + """Service-managed context must reach the model without a local transcript.""" session_id = "test-continuity" - # First message + # First message establishes a fact that exists nowhere else. response1 = self.helper.post_json( f"{self.base_url}/run", - {"message": "Tell me a short joke about weather in Seattle.", "session_id": session_id}, + {"message": "My favorite animal is the axolotl. Tell me a short joke about it.", "session_id": session_id}, + ) + data1 = self._assert_success_response(response1, session_id) + + # The follow-up needs the same session's service-managed context. + response2 = self.helper.post_json( + f"{self.base_url}/run", + {"message": "What is my favorite animal? Reply with just the animal name.", "session_id": session_id}, + ) + data2 = self._assert_success_response(response2, session_id) + assert data2["correlation_id"] != data1["correlation_id"] + assert "axolotl" in data2["response"].lower(), ( + f"Agent lost conversation context across turns. Got: {data2['response']!r}" ) - assert response1.status_code in [200, 202] - - if response1.status_code == 200: - data1 = response1.json() - assert data1["message_count"] == 2 # Initial + reply - - # Second message in same session - response2 = self.helper.post_json( - f"{self.base_url}/run", {"message": "What about San Francisco?", "session_id": session_id} - ) - assert response2.status_code == 200 - data2 = response2.json() - assert data2["message_count"] == 4 - else: - # In async mode, we can't easily test message count - # Just verify we can make multiple calls - response2 = self.helper.post_json( - f"{self.base_url}/run", {"message": "What about Texas?", "session_id": session_id} - ) - assert response2.status_code == 202 if __name__ == "__main__": diff --git a/python/packages/azurefunctions/tests/integration_tests/test_14_conversation_compaction.py b/python/packages/azurefunctions/tests/integration_tests/test_14_conversation_compaction.py new file mode 100644 index 0000000..fc6813a --- /dev/null +++ b/python/packages/azurefunctions/tests/integration_tests/test_14_conversation_compaction.py @@ -0,0 +1,126 @@ +# Copyright (c) Microsoft. All rights reserved. +""" +Integration Tests for the Conversation Compaction Sample + +Verifies that an agent configured the ordinary core way - an in-memory history provider plus a +compaction provider - runs durably under the Azure Functions host with no durable-specific +agent configuration. The sample explicitly uses ``store=False``, ``retention="keep_all"`` and +``max_state_bytes=None``. Transcript counts below apply only to its local input/output provider, +not to external or service-managed history. Completed HTTP results carry the original response +payload separately from the local transcript count. + +The function app is automatically started by the test fixture. + +Prerequisites: +- FOUNDRY_PROJECT_ENDPOINT and FOUNDRY_MODEL configured, with Azure CLI authentication +- Azure Functions Core Tools, Durable Task Scheduler, and Azurite or Azure Storage configured + +Usage: + uv run pytest packages/azurefunctions/tests/integration_tests/test_14_conversation_compaction.py -v +""" + +import json +import uuid + +import pytest +from agent_framework import AgentResponse +from agent_framework_durabletask import serialize_agent_response + +# Matches function_app.py: only the most recent groups stay in the model's context. +KEEP_LAST_GROUPS = 4 + +# Module-level markers - applied to all tests in this file +pytestmark = [ + pytest.mark.flaky, + pytest.mark.integration, + pytest.mark.sample("14_conversation_compaction"), + pytest.mark.usefixtures("function_app_for_test"), +] + + +class TestSampleConversationCompaction: + """Tests for 14_conversation_compaction sample.""" + + @pytest.fixture(autouse=True) + def _setup(self, base_url: str, sample_helper) -> None: + """Provide agent-specific base URL and helper for the tests.""" + self.base_url = f"{base_url}/api/agents/Historian" + self.helper = sample_helper + + def _run(self, message: str, session_id: str) -> dict: + """Send one turn to the agent and return the parsed response. + + Args: + message: The user message for this turn. + session_id: The session id tying the turns into one conversation. + + Returns: + The parsed JSON response body. + """ + response = self.helper.post_json( + f"{self.base_url}/run", + {"message": message, "session_id": session_id, "wait_for_response": True}, + ) + assert response.status_code == 200, response.text + result = response.json() + assert result["status"] == "success", result + assert result["session_id"] == session_id + assert result["correlation_id"] + + # A 202 acceptance is not an agent result. Successful polling returns the mailbox + # snapshot, including original result metadata, not a transcript reconstruction. + snapshot = result["agent_response"] + assert snapshot["type"] == "agent_response" + assert snapshot["created_at"], "the Foundry result timestamp was lost" + delivered = AgentResponse.from_dict(snapshot) + assert delivered.text == result["response"] + assert all(content.type != "error" for message in delivered.messages for content in message.contents) + assert json.loads(json.dumps(serialize_agent_response(delivered))) == snapshot + return result + + def test_health_check(self, base_url: str, sample_helper) -> None: + """Test health check endpoint.""" + response = sample_helper.get(f"{base_url}/api/health") + assert response.status_code == 200 + assert response.json()["status"] == "healthy" + + def test_recent_context_survives_compaction(self) -> None: + """A fact inside the retained window is still answerable after the window fills.""" + session_id = f"compaction-recent-{uuid.uuid4().hex[:8]}" + + for index in range(KEEP_LAST_GROUPS): + self._run(f"Name animal number {index + 1}.", session_id) + + self._run("My project codename is BLUEHERON.", session_id) + answer = self._run("What is my project codename? Reply with just the codename.", session_id) + + assert "blueheron" in str(answer["response"]).lower() + + def test_conversation_continues_across_turns(self) -> None: + """Durable history reaches the model, so the agent recalls an earlier turn.""" + session_id = f"compaction-continuity-{uuid.uuid4().hex[:8]}" + + self._run("My favorite animal is the axolotl.", session_id) + answer = self._run("What is my favorite animal? Reply with just the animal name.", session_id) + + assert "axolotl" in str(answer["response"]).lower() + + def test_local_transcript_count_is_separate_from_the_original_response_payload(self) -> None: + """Only this store=False input/output provider has two transcript entries per turn.""" + session_id = f"compaction-delivery-{uuid.uuid4().hex[:8]}" + correlations: set[str] = set() + + for turn, prompt in enumerate(("Name a river.", "Name an ocean."), start=1): + result = self._run(prompt, session_id) + assert result["correlation_id"] not in correlations + correlations.add(result["correlation_id"]) + assert result["message_count"] == turn * 2 + + # The original payload's message count and date must round-trip on their own. + # They are not synthesized from the echoed request or message_count above. + snapshot = result["agent_response"] + delivered = AgentResponse.from_dict(snapshot) + round_tripped = json.loads(json.dumps(serialize_agent_response(delivered))) + assert len(delivered.messages) == len(snapshot["messages"]) + assert delivered.messages + assert round_tripped["created_at"] == snapshot["created_at"] diff --git a/python/packages/azurefunctions/tests/integration_tests/test_16_live_media_retention.py b/python/packages/azurefunctions/tests/integration_tests/test_16_live_media_retention.py new file mode 100644 index 0000000..7605d00 --- /dev/null +++ b/python/packages/azurefunctions/tests/integration_tests/test_16_live_media_retention.py @@ -0,0 +1,406 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Live Functions/Azure Storage retention with a deterministic core model, not Foundry. + +Requires func v4, Azurite on 10000/10001/10002 (with --skipApiVersionCheck), +DTS on 8080, and the selected venv's test dependencies including psutil and +opentelemetry-sdk. The test generates its app/settings under tmp_path and supplies +local emulator defaults. It never uses the sample-starting fixture. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import shutil +import site +import struct +import subprocess +import sys +import time +import uuid +import zlib +from collections.abc import Callable, Iterator +from contextlib import contextmanager, suppress +from copy import deepcopy +from datetime import datetime +from pathlib import Path +from threading import Event +from typing import Any + +import agent_framework_durabletask +import psutil +import pytest +import requests +from agent_framework import Content, Message +from agent_framework_durabletask import AgentSessionId, DurableAgentState, DurableHistoryProvider + +import agent_framework_azurefunctions + +pytestmark = [ + pytest.mark.integration, + pytest.mark.orchestration, + pytest.mark.timeout(170), + # Collection-only reuse of the existing no-LLM category. Without this marker + # conftest requires Foundry even for generated apps. No sample is launched. + pytest.mark.sample("13_subworkflow_hitl"), +] +PYTHON_ROOT = Path(__file__).resolve().parents[4] +TEMPLATE = Path(__file__).with_name("live_media_app") / "function_app.py" +AGENT = "live-media-retention" +MAX_STATE_BYTES = 50_000 +DELIVERY_WINDOW_SECONDS = 3600 +TURNS = 8 + + +def _equal(actual: Any, expected: Any, label: str) -> None: + # Equality covers full JSON, not just text/counts. Only hashes reach failures. + if actual != expected: + + def digest(value: Any) -> str: + return hashlib.sha256(json.dumps(value, sort_keys=True).encode()).hexdigest() + + pytest.fail(f"{label}: full JSON mismatch ({digest(actual)} != {digest(expected)})", pytrace=False) + + +def _stored(raw: dict[str, Any]) -> list[dict[str, Any]]: + state = DurableAgentState.from_json(json.dumps(raw)) + return [ + message.to_chat_message().to_dict() for entry in state.data.conversation_history for message in entry.messages + ] + + +def _model_history(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + result = deepcopy(messages) + for message in result: + message.setdefault("additional_properties", {})["_attribution"] = { + "source_id": DurableHistoryProvider.DEFAULT_SOURCE_ID, + "source_type": "DurableHistoryProvider", + } + return result + + +def _inputs(kind: str, turn: str) -> list[dict[str, Any]]: + # Valid PNG scanlines with incompressible pixels make binary payload bytes a + # substantial part of pressure. The text is not the dominant storage cost. + def chunk(tag: bytes, value: bytes) -> bytes: + return struct.pack(">I", len(value)) + tag + value + struct.pack(">I", zlib.crc32(tag + value)) + + width, height = 128, 64 + pixels = hashlib.shake_256(b"durable-media-pressure").digest(width * height) + rows = b"".join(b"\x00" + pixels[row * width : (row + 1) * width] for row in range(height)) + png = ( + b"\x89PNG\r\n\x1a\n" + + chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 0, 0, 0, 0)) + + chunk(b"IDAT", zlib.compress(rows)) + + chunk(b"IEND", b"") + ) + if kind == "inline-png": + media = Content.from_data(png, "image/png") + elif kind == "inline-file": + media = Content.from_data((f"{turn}: inline document 界\n" * 256).encode(), "text/plain") + else: + raise ValueError(f"Unknown media case: {kind}") + properties = {"application": {"type": "text", "values": [turn, "界", 0, False, None]}} + return [ + Message( + "user", + [Content.from_text(f"{turn}: " + "context " * 100), media], + message_id=f"{turn}-input", + author_name="media-user", + additional_properties=deepcopy(properties), + ).to_dict(), + Message( + "assistant", + [Content.from_function_call(f"{turn}-call", "lookup", arguments={"query": turn})], + message_id=f"{turn}-call-message", + author_name="planner", + additional_properties=deepcopy(properties), + ).to_dict(), + Message( + "tool", + [Content.from_function_result(f"{turn}-call", result={"records": [turn, "界", False]})], + message_id=f"{turn}-result-message", + author_name="lookup", + additional_properties=deepcopy(properties), + ).to_dict(), + ] + + +def _answer(current_id: str) -> dict[str, Any]: + return Message( + "assistant", [f"answer:{current_id}"], message_id=f"{current_id}-answer", author_name="retention-model" + ).to_dict() + + +class _Host: + def __init__(self, app: Path, port: int, env: dict[str, str], deadline: float, epoch: str) -> None: + self.url = f"http://127.0.0.1:{port}/api" + self.deadline = deadline + self.log = (app / f"{epoch}-host.log").open("w", encoding="utf-8") + creationflags = 0 + if sys.platform == "win32": + creationflags = subprocess.CREATE_NEW_PROCESS_GROUP + try: + self.process = subprocess.Popen( + ["func", "start", "--port", str(port)], + cwd=app, + env=env, + stdin=subprocess.DEVNULL, + stdout=self.log, + stderr=subprocess.STDOUT, + shell=sys.platform == "win32", # Core Tools can be a .cmd shim. + creationflags=creationflags, + start_new_session=sys.platform != "win32", + ) + except BaseException: + self.log.close() + raise + + def wait(self, probe: Callable[[], Any], description: str, seconds: int = 30) -> Any: + end = min(self.deadline, time.monotonic() + seconds) + pause = Event() + while time.monotonic() < end: + assert self.process.poll() is None, "Functions host exited. Inspect the temporary host log" + try: + if result := probe(): + return result + except (requests.ConnectionError, requests.Timeout): + pass + # Pacing is not proof of completion. Only the HTTP/backend predicate is. + pause.wait(min(0.1, max(0, end - time.monotonic()))) + pytest.fail(f"Timed out waiting for {description}. Inspect the temporary host log", pytrace=False) + + def get(self, path: str, missing_ok: bool = False) -> Any: + response = requests.get(f"{self.url}/{path}", timeout=3) + if missing_ok and response.status_code in (404, 503): + return None + assert response.status_code == 200, f"GET {path}: HTTP {response.status_code}" + return response.json() + + def turn(self, correlation: str, messages: list[dict[str, Any]], session: str) -> dict[str, Any]: + response = requests.post( + f"{self.url}/retention/run", + json={"message": "synthetic projected input", "correlationId": correlation, "contextMessages": messages}, + timeout=5, + ) + assert response.status_code == 202, f"Signal returned HTTP {response.status_code}" + _equal(response.json(), {"correlation": correlation, "session": session}, "signal acknowledgement") + + def committed() -> dict[str, Any] | None: + raw = self.get("retention/state", missing_ok=True) + if raw and correlation in raw.get("data", {}).get("completedCorrelations", {}): + receipt = raw["data"]["completedCorrelations"][correlation] + assert receipt["outcome"] == "succeeded", "The real Functions agent operation failed" + return raw + return None + + return self.wait(committed, f"persisted completion receipt for {correlation}") + + +@contextmanager +def _host(app: Path, env: dict[str, str], deadline: float, epoch: str, harness: Any) -> Iterator[_Host]: + host = _Host(app, harness._find_available_port(), env, deadline, epoch) + try: + host.wait(lambda: host.get("health", missing_ok=True), "Functions health", seconds=50) + yield host + finally: + # Require psutil (imported above) so the existing cleanup also kills workers. + descendants: list[psutil.Process] = [] + try: + with suppress(psutil.NoSuchProcess): + descendants = psutil.Process(host.process.pid).children(recursive=True) + finally: + try: + harness._cleanup_function_app(host.process) + host.process.wait(timeout=5) + assert not any(child.is_running() for child in descendants), "Functions worker survived host cleanup" + finally: + host.log.close() + + +def _prepare(app: Path, session: str, hub: str) -> dict[str, str]: + app.mkdir() + source_paths = [PYTHON_ROOT / "packages" / name for name in ("azurefunctions", "durabletask")] + for package, path in zip((agent_framework_azurefunctions, agent_framework_durabletask), source_paths): + assert package.__file__ is not None + assert Path(package.__file__).resolve().parent == path / package.__name__, ( + "Run using packages from this exact pr59 worktree" + ) + shutil.copyfile(TEMPLATE, app / "function_app.py") + config = { + "agent": AGENT, + "session": session, + "max_state_bytes": MAX_STATE_BYTES, + "delivery_window_seconds": DELIVERY_WINDOW_SECONDS, + "azurefunctions_source": str(source_paths[0] / "agent_framework_azurefunctions"), + "durabletask_source": str(source_paths[1] / "agent_framework_durabletask"), + } + (app / "live_config.json").write_text(json.dumps(config), encoding="utf-8") + (app / "host.json").write_text( + json.dumps({ + "version": "2.0", + "extensionBundle": {"id": "Microsoft.Azure.Functions.ExtensionBundle", "version": "[4.*, 5.0.0)"}, + "extensions": {"durableTask": {"hubName": hub}}, + "logging": {"logLevel": {"default": "Warning"}}, + }), + encoding="utf-8", + ) + settings = { + "FUNCTIONS_WORKER_RUNTIME": "python", + "FUNCTIONS_WORKER_PROCESS_COUNT": "1", + "AzureWebJobsStorage": "UseDevelopmentStorage=true", + "DURABLE_TASK_SCHEDULER_CONNECTION_STRING": "Endpoint=http://localhost:8080;Authentication=None", + "TASKHUB_NAME": hub, + "AzureFunctionsJobHost__extensions__durableTask__hubName": hub, + "DURABLE_AGENTS_DEPLOYMENT_MODE": "isolated_v2", + "languageWorkers__python__defaultExecutablePath": sys.executable, + } + (app / "local.settings.json").write_text(json.dumps({"IsEncrypted": False, "Values": settings}), encoding="utf-8") + return { + **os.environ, + **settings, + "VIRTUAL_ENV": sys.prefix, + "PATH": str(Path(sys.executable).parent) + os.pathsep + os.environ.get("PATH", ""), + "PYTHONDONTWRITEBYTECODE": "1", + # Keep parent worker/grpc workarounds, but prefer both exact source roots. + "PYTHONPATH": os.pathsep.join([*map(str, source_paths), os.getenv("PYTHONPATH", ""), *site.getsitepackages()]), + } + + +def _capture(host: _Host, boot: str, calls: int, current_id: str, session: str, expected: Any) -> None: + record = host.get("retention/capture") + assert record["boot"] == boot and record["calls"] == calls + capture = record["capture"] + assert capture["boot"] == boot and capture["calls"] == calls and capture["current_id"] == current_id + _equal(capture["messages"], expected, "full next-model input") + _equal( + capture["context"], + { + "provider": "AzureFunctionEntityStateProvider", + "type": "DurableEntityContext", + "entity_name": AgentSessionId.to_entity_name(AGENT), + "entity_key": session, + "operation": "run", + }, + "real Functions context on the async bridge", + ) + + +def _retained(raw: dict[str, Any], originals: dict[str, dict[str, Any]]) -> tuple[list[dict[str, Any]], int]: + retained = _stored(raw) + ids = {message["message_id"] for message in retained} + _equal(retained, [message for key, message in originals.items() if key in ids], "persisted media/metadata/order") + for turn in range(TURNS): + pair = {f"turn-{turn}-call-message", f"turn-{turn}-result-message"} + assert pair <= ids or pair.isdisjoint(ids), "Pressure split an atomic tool pair" + removed = len(originals) - len(retained) + assert (raw["data"].get("truncation") or {}).get("evictedMessageCount", 0) == removed + assert len(json.dumps(raw)) < int(MAX_STATE_BYTES * 0.85) + return retained, removed + + +def _metrics(host: _Host, boot: str, removed: int, calls: int) -> None: + record = host.get("retention/metrics") + assert record["boot"] == boot + rows = record["rows"] + deletions = [row for row in rows if row["name"] == "durable.retention.removed_messages"] + for row in deletions: + _equal( + row["attributes"], + {"mechanism": "pressure", "outcome": "staged", "commit_status": "not_attempted"}, + "bounded deletion metric labels", + ) + assert sum(row["value"] for row in deletions) == removed + for metric, extra in ( + ("write_attempts", {"stage": "set_state"}), + ("operations", {}), + ): + observations = [row for row in rows if row["name"] == f"durable.retention.{metric}"] + assert sum(row["value"] for row in observations) == calls + for row in observations: + attributes = row["attributes"] + assert isinstance(attributes["deletion_staged"], bool) + _equal( + attributes, + { + **extra, + "outcome": "returned", + "commit_status": "unknown", + "deletion_staged": attributes["deletion_staged"], + }, + "host write observations are not commit proof", + ) + + +@pytest.mark.parametrize("kind", ["inline-png", "inline-file"]) +def test_live_functions_media_pressure_cold_json_and_exact_next_model( + kind: str, tmp_path: Path, request: pytest.FixtureRequest +) -> None: + # Resolve the already-loaded local conftest, not an identically named DTS module. + harness_path = Path(__file__).with_name("conftest.py").resolve() + harness = next( + plugin + for plugin in request.config.pluginmanager.get_plugins() + if getattr(plugin, "__file__", None) and Path(plugin.__file__).resolve() == harness_path + ) + deadline = time.monotonic() + 150 + session = f"media-{kind}-{uuid.uuid4().hex[:12]}" + app = tmp_path / "app" + env = _prepare(app, session, f"media{uuid.uuid4().hex[:16]}") + originals: dict[str, dict[str, Any]] = {} + previous: list[dict[str, Any]] = [] + raw: dict[str, Any] = {} + removed = 0 + + with _host(app, env, deadline, "warm", harness) as warm: + boot = warm.get("retention/capture")["boot"] + for index in range(TURNS): + correlation = f"turn-{index}" + inputs = _inputs(kind, correlation) + raw = warm.turn(correlation, inputs, session) + _capture(warm, boot, index + 1, f"{correlation}-input", session, [*_model_history(previous), *inputs]) + current = [*inputs, _answer(f"{correlation}-input")] + originals.update({message["message_id"]: message for message in current}) + previous, removed = _retained(raw, originals) + assert {message["message_id"] for message in current} <= {message["message_id"] for message in previous} + _metrics(warm, boot, removed, index + 1) + (tmp_path / f"warm-{index}-state.json").write_text(json.dumps(raw), encoding="utf-8") + assert removed >= 4, "Must actually evict messages, not merely round-trip media" + assert sum(message["role"] == "user" for message in previous) >= 2, "Keep older media for cold replay" + assert len(raw["data"]["completedCorrelations"]) == len(raw["data"]["responseMailbox"]) == TURNS + + with _host(app, env, deadline, "cold", harness) as cold: + initial = cold.get("retention/capture") + cold_boot = initial["boot"] + assert cold_boot != boot and initial["calls"] == 0 and initial["capture"] is None + cold_read = cold.get("retention/state") + (tmp_path / "cold-read-state.json").write_text(json.dumps(cold_read), encoding="utf-8") + _equal(cold_read, raw, "exact backend JSON after killing and restarting the host") + _metrics(cold, cold_boot, 0, 0) + evicted = set(originals) - {message["message_id"] for message in previous} + assert "turn-0-input" in evicted + next_input = Message("user", ["next turn"], message_id="cold-input").to_dict() + final = cold.turn("cold", [*_inputs(kind, "turn-0"), next_input], session) + _capture(cold, cold_boot, 1, "cold-input", session, [*_model_history(previous), next_input]) + originals.update({"cold-input": next_input, "cold-input-answer": _answer("cold-input")}) + retained, total_removed = _retained(final, originals) + ids = {message["message_id"] for message in retained} + assert evicted.isdisjoint(ids), "Cold projected replay resurrected evicted media" + assert {"cold-input", "cold-input-answer"} <= ids + _metrics(cold, cold_boot, total_removed - removed, 1) + _equal( + {key: final["data"]["ingestedMessages"][key] for key in raw["data"]["ingestedMessages"]}, + raw["data"]["ingestedMessages"], + "cold replay preserves ingestion receipts", + ) + assert set(final["data"]["ingestedMessages"]) == {*raw["data"]["ingestedMessages"], "cold-input"} + for field in ("completedCorrelations", "responseMailbox"): + _equal({key: final["data"][field][key] for key in raw["data"][field]}, raw["data"][field], field) + assert set(final["data"][field]) == {*raw["data"][field], "cold"} + for mailbox in final["data"]["responseMailbox"].values(): + assert ( + datetime.fromisoformat(mailbox["expiresAt"]) - datetime.fromisoformat(mailbox["createdAt"]) + ).total_seconds() == DELIVERY_WINDOW_SECONDS + (tmp_path / "cold-final-state.json").write_text(json.dumps(final), encoding="utf-8") diff --git a/python/packages/azurefunctions/tests/test_app.py b/python/packages/azurefunctions/tests/test_app.py index 5751d3a..ca86ce3 100644 --- a/python/packages/azurefunctions/tests/test_app.py +++ b/python/packages/azurefunctions/tests/test_app.py @@ -14,6 +14,7 @@ import pytest from agent_framework import AgentResponse, Message from agent_framework_durabletask import ( + DEFAULT_MAX_STATE_BYTES, MIMETYPE_APPLICATION_JSON, MIMETYPE_TEXT_PLAIN, SESSION_ID_HEADER, @@ -23,6 +24,7 @@ AgentEntityStateProviderMixin, DurableAgentState, workflow_orchestrator_name, + wrap_workflow_input, ) from agent_framework_azurefunctions import AgentFunctionApp @@ -273,7 +275,16 @@ def test_agent_override_enables_http_route_when_app_disabled(self) -> None: app.add_agent(mock_agent, enable_http_endpoint=True) http_route_mock.assert_called_once_with("OverrideAgent") - agent_entity_mock.assert_called_once_with(mock_agent, "OverrideAgent", ANY) + agent_entity_mock.assert_called_once_with( + mock_agent, + "OverrideAgent", + None, + retention="keep_all", + max_state_bytes=DEFAULT_MAX_STATE_BYTES, + high_watermark=0.85, + low_watermark=0.70, + response_delivery_window_seconds=60, + ) assert app._agent_metadata["OverrideAgent"].http_endpoint_enabled is True def test_agent_override_disables_http_route_when_app_enabled(self) -> None: @@ -290,9 +301,52 @@ def test_agent_override_disables_http_route_when_app_enabled(self) -> None: app.add_agent(mock_agent, enable_http_endpoint=False) http_route_mock.assert_not_called() - agent_entity_mock.assert_called_once_with(mock_agent, "DisabledOverride", ANY) + agent_entity_mock.assert_called_once_with( + mock_agent, + "DisabledOverride", + None, + retention="keep_all", + max_state_bytes=DEFAULT_MAX_STATE_BYTES, + high_watermark=0.85, + low_watermark=0.70, + response_delivery_window_seconds=60, + ) assert app._agent_metadata["DisabledOverride"].http_endpoint_enabled is False + def test_configured_state_budget_reaches_the_entity(self) -> None: + """A budget set on the app has to bound the entity, not just sit on the app. + + Asserting that ``_setup_agent_entity`` was called is not enough, because the value can + still be dropped below that point and the agent would silently keep the default budget. + So the registered entity function is invoked and the factory call is inspected. + """ + mock_agent = Mock() + mock_agent.name = "BudgetAgent" + registered: list[Callable[[Any], None]] = [] + + def _capture_entity_trigger(**kwargs: Any) -> Callable[[FuncT], FuncT]: + def decorator(entity_function: FuncT) -> FuncT: + registered.append(entity_function) + return entity_function + + return decorator + + with ( + patch.object(AgentFunctionApp, "entity_trigger", side_effect=_capture_entity_trigger), + patch("agent_framework_azurefunctions._app.create_agent_entity") as create_entity_mock, + ): + app = AgentFunctionApp( + enable_health_check=False, + enable_http_endpoints=False, + max_state_bytes=4096, + ) + app.add_agent(mock_agent) + + assert registered, "no entity function was registered" + registered[0](Mock()) + + assert create_entity_mock.call_args.kwargs["max_state_bytes"] == 4096 + def test_multiple_apps_independent(self) -> None: """Test that multiple AgentFunctionApp instances are independent.""" agent1 = Mock() @@ -585,11 +639,11 @@ def test_entity_function_handles_reset_operation(self) -> None: mock_agent = Mock() entity_function = create_agent_entity(mock_agent) - # Mock context + # Reset an admitted v2 target, not a legacy session. mock_context = Mock() mock_context.operation_name = "reset" mock_context.get_state.return_value = { - "schemaVersion": "1.0.0", + "schemaVersion": DurableAgentState.SCHEMA_VERSION, "data": { "conversationHistory": [ { @@ -1042,6 +1096,7 @@ def decorator(func: FuncT) -> FuncT: workflow = Mock() workflow.name = workflow_name + workflow.executors = {} app = AgentFunctionApp(enable_health_check=False) with ( @@ -1089,7 +1144,7 @@ async def test_wait_for_response_query_waits_with_timeout(self) -> None: client.start_new.assert_awaited_once_with( "dafx-test_workflow", instance_id="custom-run", - client_input={"message": "hello"}, + client_input=wrap_workflow_input({"message": "hello"}), ) async def test_wait_for_response_header_waits_with_default_timeout(self) -> None: @@ -2236,9 +2291,8 @@ def test_different_subworkflow_sharing_a_name_is_rejected(self) -> None: def test_cross_registration_nested_collision_is_atomic(self) -> None: """A later top-level workflow whose nested child collides aborts before committing it. - Hosting ``[first, second]`` where ``second``'s nested sub-workflow reuses - ``first``'s child name must raise *before* ``second`` registers any primitives, - so the app is never left with ``second`` half-configured. + Configuring ``second`` after ``first`` must preserve the first registration + when the second workflow's nested child has a conflicting identity. """ shared_a, _ = self._inner_agent_wf("shared", "agent_node") shared_b, _ = self._inner_agent_wf("shared", "other_node") # different instance, same name @@ -2248,15 +2302,43 @@ def test_cross_registration_nested_collision_is_atomic(self) -> None: with ( patch.object(AgentFunctionApp, "_setup_executor_activity"), patch.object(AgentFunctionApp, "_setup_workflow_orchestration") as setup_orch, - pytest.raises(ValueError, match="collides"), ): - AgentFunctionApp(workflows=[first, second]) + app = AgentFunctionApp(workflow=first) + identities = dict(app._registration_identities) + agents = app.agents + with pytest.raises(ValueError, match="collides"): + app.configure_workflow(second) # Only 'first' and its child 'shared' committed primitives; the collision aborted # before 'second' (or its colliding child) registered anything. registered = {call.args[0].name for call in setup_orch.call_args_list} assert registered == {"first", "shared"} - assert "second" not in registered + assert setup_orch.call_count == 2 + assert app._registered_orchestrations == {"first": first, "shared": shared_a} + assert app._registration_identities == identities + assert app.agents == agents + assert app.workflows == {"first": first} + assert app.workflow is first + + def test_constructor_nested_collision_is_preflighted_before_any_setup(self) -> None: + """Constructor validation covers every workflow before installing any triggers.""" + shared_a, _ = self._inner_agent_wf("shared", "agent_node") + shared_b, _ = self._inner_agent_wf("shared", "other_node") + first = self._outer_wf("first", shared_a) + second = self._outer_wf("second", shared_b) + + with ( + patch.object(AgentFunctionApp, "_setup_agent_functions") as setup_agent, + patch.object(AgentFunctionApp, "_setup_executor_activity") as setup_activity, + patch.object(AgentFunctionApp, "_setup_workflow_orchestration") as setup_orch, + patch.object(AgentFunctionApp, "_register_workflow_routes") as setup_routes, + patch.object(AgentFunctionApp, "_setup_health_route") as setup_health, + pytest.raises(ValueError, match="collides"), + ): + AgentFunctionApp(workflows=[first, second]) + + for setup in (setup_agent, setup_activity, setup_orch, setup_routes, setup_health): + setup.assert_not_called() def test_executor_id_with_reserved_separator_is_rejected(self) -> None: """An executor id containing the nested-HITL separator is rejected at registration.""" diff --git a/python/packages/azurefunctions/tests/test_azurefunctions_workflow_initial_input.py b/python/packages/azurefunctions/tests/test_azurefunctions_workflow_initial_input.py index b1de0bc..e82dd0d 100644 --- a/python/packages/azurefunctions/tests/test_azurefunctions_workflow_initial_input.py +++ b/python/packages/azurefunctions/tests/test_azurefunctions_workflow_initial_input.py @@ -9,6 +9,7 @@ from unittest.mock import AsyncMock, Mock, patch from agent_framework import Executor, Workflow, WorkflowBuilder, WorkflowContext, handler +from agent_framework_durabletask import unwrap_workflow_input, wrap_workflow_input from agent_framework_azurefunctions import AgentFunctionApp @@ -71,4 +72,7 @@ async def test_workflow_run_route_neutralizes_reserved_marker_shaped_input() -> await handler(request, client) - assert client.start_new.await_args.kwargs["client_input"] is None + client.start_new.assert_awaited_once_with( + "dafx-input_boundary", instance_id=None, client_input=wrap_workflow_input(None) + ) + assert unwrap_workflow_input(client.start_new.await_args.kwargs["client_input"]) is None diff --git a/python/packages/azurefunctions/tests/test_delivery_consumers_af.py b/python/packages/azurefunctions/tests/test_delivery_consumers_af.py new file mode 100644 index 0000000..5606710 --- /dev/null +++ b/python/packages/azurefunctions/tests/test_delivery_consumers_af.py @@ -0,0 +1,573 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Azure Functions delivery through real state readers, HTTP handlers, and tasks.""" + +import json +from collections.abc import Awaitable, Callable +from copy import deepcopy +from datetime import date, datetime, timezone +from types import SimpleNamespace +from typing import Any, cast +from unittest.mock import AsyncMock, Mock, patch + +import azure.durable_functions as df +import azure.functions as func +import pytest +from agent_framework import AgentResponse, Content, ContinuationToken, Message +from agent_framework_durabletask import ( + MIMETYPE_APPLICATION_JSON, + MIMETYPE_TEXT_PLAIN, + SESSION_ID_HEADER, + WAIT_FOR_RESPONSE_HEADER, + DurableAgentState, + DurableAgentStateErrorResponse, + DurableAgentStateResponse, + RunRequest, + serialize_agent_response, +) +from azure.durable_functions.models.actions.NoOpAction import NoOpAction +from azure.durable_functions.models.Task import AtomicTask, TaskState +from pydantic import BaseModel + +from agent_framework_azurefunctions import AgentFunctionApp +from agent_framework_azurefunctions._entities import create_agent_entity +from agent_framework_azurefunctions._orchestration import AgentTask + +CORRELATION_ID = "consumer-correlation" +SESSION_ID = "consumer-session" +AGENT_NAME = "consumer" +HISTORICAL_TIME = datetime(2024, 1, 1, tzinfo=timezone.utc) +EXPIRED_MESSAGE = "This request completed, but its response delivery window has expired." +HttpHandler = Callable[[func.HttpRequest, Any], Awaitable[func.HttpResponse]] + + +class Answer(BaseModel): + answer: int + + +def _response(*, value: Any = None, text: str = "Readable answer") -> AgentResponse[Any]: + return AgentResponse( + messages=[ + Message( + "assistant", + [ + Content.from_text( + text, + annotations=[{"type": "citation", "title": "Source", "url": "https://example.test/source"}], + additional_properties={"provider": {"labels": ["content"]}}, + raw_representation=object(), + ) + ], + author_name="writer", + message_id="answer-message", + additional_properties={"provider": {"labels": ["message"]}}, + raw_representation=object(), + ) + ], + response_id="response-1", + agent_id="agent-1", + created_at=HISTORICAL_TIME.isoformat(), + finish_reason="stop", + usage_details={"input_token_count": 3, "output_token_count": 2, "total_token_count": 5}, + continuation_token=cast(ContinuationToken, {"cursor": {"pages": [1, 2]}}), + additional_properties={"provider": {"labels": ["response"]}}, + raw_representation=object(), + value=value, + ) + + +def _runtime_error(*, include_text: bool = True) -> AgentResponse[Any]: + response = _response() + contents = [ + Content.from_error( + message="Model endpoint unavailable", + error_code="ProviderUnavailable", + error_details="provider details", + additional_properties={"retryable": False}, + ) + ] + if include_text: + contents.append(Content.from_text("ProviderUnavailable: Model endpoint unavailable")) + response.messages.append(Message("system", contents, author_name="runtime", message_id="error-message")) + return response + + +def _mailbox_state( + response: AgentResponse[Any], *, expired: bool = False, cleanup: bool = False, legacy: bool = False +) -> dict[str, Any]: + state = DurableAgentState() + state.data.conversation_history.append(DurableAgentStateResponse.from_run_response(CORRELATION_ID, response)) + state.record_response( + CORRELATION_ID, + response, + delivery_window_seconds=3600, + now=HISTORICAL_TIME if expired else None, + legacy=legacy, + ) + if not expired: + state.data.conversation_history.clear() + if cleanup: + state.expire_responses() + return json.loads(state.to_json()) + + +def _legacy_state(response: AgentResponse[Any], version: str, *, failed: bool = False) -> dict[str, Any]: + state = DurableAgentState(schema_version=version) + entry_type = DurableAgentStateErrorResponse if failed else DurableAgentStateResponse + state.data.conversation_history.append(entry_type.from_run_response(CORRELATION_ID, response)) + return json.loads(state.to_json()) + + +def _client(payload: dict[str, Any] | None) -> Mock: + client = Mock(spec=df.DurableOrchestrationClient) + client.signal_entity = AsyncMock() + client.read_entity_state = AsyncMock( + return_value=SimpleNamespace(entity_exists=payload is not None, entity_state=deepcopy(payload)) + ) + return client + + +def _request(*, plain_text: bool = False, wait: bool = True) -> func.HttpRequest: + content_type = MIMETYPE_TEXT_PLAIN if plain_text else MIMETYPE_APPLICATION_JSON + body = b"question" if plain_text else json.dumps({"message": "question", "session_id": SESSION_ID}).encode() + return func.HttpRequest( + method="POST", + url=f"https://example.test/api/agents/{AGENT_NAME}/run", + headers={"Content-Type": content_type, "Accept": content_type, WAIT_FOR_RESPONSE_HEADER: str(wait).lower()}, + params={"session_id": SESSION_ID}, + body=body, + ) + + +def _entity_context(payload: dict[str, Any] | None, operation: str = "run") -> Mock: + context = Mock(spec=df.DurableEntityContext) + context.operation_name = operation + context.entity_name = f"dafx-{AGENT_NAME}" + context.entity_key = SESSION_ID + context.get_state.return_value = deepcopy(payload) + context.get_input.return_value = RunRequest(message="question", correlation_id=CORRELATION_ID).to_dict() + return context + + +def _task(payload: dict[str, Any], response_format: type[BaseModel] | None, *, precompleted: bool = False) -> AgentTask: + child = AtomicTask(1, NoOpAction()) + if precompleted: + child.set_value(is_error=False, value=payload) + task = AgentTask(child, response_format, CORRELATION_ID) + if not precompleted: + assert not task.is_completed + child.set_value(is_error=False, value=payload) + return task + + +@pytest.fixture +def sleep(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: + mocked = AsyncMock() + monkeypatch.setattr("agent_framework_azurefunctions._app.asyncio.sleep", mocked) + return mocked + + +@pytest.fixture +def app(monkeypatch: pytest.MonkeyPatch) -> AgentFunctionApp: + result = AgentFunctionApp( + enable_health_check=False, enable_http_endpoints=False, max_poll_retries=3, poll_interval_seconds=0.01 + ) + monkeypatch.setattr(result, "_generate_unique_id", Mock(return_value=CORRELATION_ID)) + return result + + +@pytest.fixture +def http_handler(app: AgentFunctionApp, monkeypatch: pytest.MonkeyPatch) -> HttpHandler: + handlers: list[HttpHandler] = [] + + def identity(*args: Any, **kwargs: Any) -> Callable[[HttpHandler], HttpHandler]: + return lambda handler: handler + + def route(*args: Any, **kwargs: Any) -> Callable[[HttpHandler], HttpHandler]: + def capture(handler: HttpHandler) -> HttpHandler: + handlers.append(handler) + return handler + + return capture + + monkeypatch.setattr(app, "function_name", identity) + monkeypatch.setattr(app, "route", route) + monkeypatch.setattr(app, "durable_client_input", identity) + app._setup_http_run_route(AGENT_NAME) + return handlers[0] + + +@pytest.mark.parametrize("value", [None, 0, False, "", [], {}, {"answer": 42}]) +async def test_http_success_keeps_text_and_adds_full_mailbox_snapshot( + value: Any, http_handler: HttpHandler, sleep: AsyncMock +) -> None: + original = _response(value=deepcopy(value)) + state = _mailbox_state(original) + assert state["data"]["conversationHistory"] == [] + client = _client(state) + + response = await http_handler(_request(), client) + + assert response.status_code == 200 + assert response.mimetype == MIMETYPE_APPLICATION_JSON + payload = json.loads(response.get_body()) + assert payload == { + "response": original.text, + "message": "question", + "session_id": SESSION_ID, + "status": "success", + "correlation_id": CORRELATION_ID, + "message_count": 0, + "agent_response": state["data"]["responseMailbox"][CORRELATION_ID]["response"], + } + delivered = AgentResponse.from_dict(payload["agent_response"]) + assert delivered.to_dict() == original.to_dict() + assert delivered.value == value + assert type(delivered.value) is type(value) + assert delivered.messages[0].author_name == "writer" + assert delivered.messages[0].message_id == "answer-message" + client.signal_entity.assert_awaited_once() + entity_id = client.signal_entity.call_args.args[0] + assert entity_id.name == f"dafx-{AGENT_NAME}" + assert entity_id.key == SESSION_ID + client.read_entity_state.assert_awaited_once_with(entity_id) + sleep.assert_awaited_once_with(0.01) + + +@pytest.mark.parametrize("version", ["1.0.0", "1.1.0"]) +async def test_http_keeps_legacy_transcript_lookup(version: str, http_handler: HttpHandler, sleep: AsyncMock) -> None: + original = _response() + state = _legacy_state(original, version) + client = _client(state) + + response = await http_handler(_request(), client) + + assert response.status_code == 200 + payload = json.loads(response.get_body()) + assert payload["status"] == "success" + assert payload["response"] == original.text + assert payload["message_count"] == 1 + delivered = AgentResponse.from_dict(payload["agent_response"]) + assert delivered.messages[0].author_name == "writer" + assert delivered.messages[0].message_id == "answer-message" + assert delivered.usage_details == original.usage_details + assert state["schemaVersion"] == version + client.read_entity_state.assert_awaited_once() + sleep.assert_awaited_once_with(0.01) + + +@pytest.mark.parametrize("cleanup", [False, True]) +@pytest.mark.parametrize("plain_text", [False, True]) +@pytest.mark.parametrize("failed", [False, True]) +async def test_http_expired_delivery_returns_410_without_waiting_for_more_polls( + cleanup: bool, plain_text: bool, failed: bool, http_handler: HttpHandler, sleep: AsyncMock +) -> None: + original = _runtime_error() if failed else _response(value={"answer": 42}) + client = _client(_mailbox_state(original, expired=True, cleanup=cleanup)) + + response = await http_handler(_request(plain_text=plain_text), client) + + assert response.status_code == 410 + if plain_text: + assert response.mimetype == MIMETYPE_TEXT_PLAIN + assert response.get_body().decode() == EXPIRED_MESSAGE + assert response.headers[SESSION_ID_HEADER] == SESSION_ID + assert response.headers["x-ms-durable-outcome"] == ("failed" if failed else "succeeded") + else: + payload = json.loads(response.get_body()) + assert payload["status"] == "already_completed" + assert payload["error_code"] == "response_expired" + assert payload["error"] == EXPIRED_MESSAGE + assert payload["response"] is None + assert payload["message_count"] == 1 + assert payload["outcome"] == ("failed" if failed else "succeeded") + assert payload["agent_response"]["additional_properties"] == { + "durable_status": "already_completed", + "correlation_id": CORRELATION_ID, + "durable_outcome": "failed" if failed else "succeeded", + } + error = payload["agent_response"]["messages"][0]["contents"][0] + assert error["error_code"] == "response_expired" + assert error["message"] == EXPIRED_MESSAGE + client.read_entity_state.assert_awaited_once() + sleep.assert_awaited_once_with(0.01) + + +@pytest.mark.parametrize("expiry_marker", ["error_code", "durable_status"]) +async def test_http_accepts_either_terminal_expiry_marker( + expiry_marker: str, http_handler: HttpHandler, sleep: AsyncMock +) -> None: + original = AgentResponse(messages=[]) + if expiry_marker == "error_code": + original.messages = [ + Message("system", [Content.from_error(message=EXPIRED_MESSAGE, error_code="response_expired")]) + ] + else: + original.additional_properties["durable_status"] = "already_completed" + # An older result can itself be unavailable. Revised writers require a known + # invocation outcome, but readers must keep suppressing that old completion. + client = _client(_mailbox_state(original, legacy=True)) + + response = await http_handler(_request(), client) + + assert response.status_code == 410 + payload = json.loads(response.get_body()) + assert payload["status"] == "already_completed" + assert payload["error_code"] == "response_expired" + assert payload["error"] == EXPIRED_MESSAGE + assert payload["agent_response"] == original.to_dict() + client.read_entity_state.assert_awaited_once() + sleep.assert_awaited_once_with(0.01) + + +async def test_http_error_without_code_or_message_is_still_a_failure( + http_handler: HttpHandler, sleep: AsyncMock +) -> None: + original = AgentResponse(messages=[Message("system", [Content.from_error()])]) + client = _client(_mailbox_state(original)) + + response = await http_handler(_request(), client) + + assert response.status_code == 500 + payload = json.loads(response.get_body()) + assert payload["status"] == "error" + assert payload["error_code"] is None + assert payload["error"] == "Agent execution failed." + assert payload["agent_response"] == original.to_dict() + client.read_entity_state.assert_awaited_once() + sleep.assert_awaited_once_with(0.01) + + +@pytest.mark.parametrize("version", ["1.0.0", "1.1.0", "2.0.0"]) +@pytest.mark.parametrize("include_text", [False, True]) +async def test_http_runtime_error_is_500_not_success( + version: str, include_text: bool, http_handler: HttpHandler, sleep: AsyncMock +) -> None: + original = _runtime_error(include_text=include_text) + state = _mailbox_state(original) if version == "2.0.0" else _legacy_state(original, version, failed=True) + client = _client(state) + + response = await http_handler(_request(), client) + + assert response.status_code == 500 + payload = json.loads(response.get_body()) + assert payload["status"] == "error" + assert payload["error_code"] == "ProviderUnavailable" + assert payload["error"] == "Model endpoint unavailable" + assert payload["response"] is None + assert payload["message"] == "question" + assert payload["session_id"] == SESSION_ID + assert payload["correlation_id"] == CORRELATION_ID + delivered = AgentResponse.from_dict(payload["agent_response"]) + assert delivered.messages[1].contents[0].error_details == "provider details" + if version == "2.0.0": + assert delivered.to_dict() == original.to_dict() + assert payload["message_count"] == 0 + client.read_entity_state.assert_awaited_once() + sleep.assert_awaited_once_with(0.01) + + +async def test_plain_text_runtime_error_returns_the_error_not_partial_success( + http_handler: HttpHandler, sleep: AsyncMock +) -> None: + client = _client(_mailbox_state(_runtime_error())) + + response = await http_handler(_request(plain_text=True), client) + + assert response.status_code == 500 + assert response.get_body().decode() == "Model endpoint unavailable" + client.read_entity_state.assert_awaited_once() + sleep.assert_awaited_once_with(0.01) + + +async def test_http_missing_response_keeps_the_existing_timeout( + http_handler: HttpHandler, app: AgentFunctionApp, sleep: AsyncMock +) -> None: + client = _client(None) + + response = await http_handler(_request(), client) + + assert response.status_code == 500 + assert json.loads(response.get_body()) == { + "response": "Agent is still processing or timed out...", + "message": "question", + "session_id": SESSION_ID, + "status": "timeout", + "correlation_id": CORRELATION_ID, + } + client.signal_entity.assert_awaited_once() + assert client.read_entity_state.await_count == app.max_poll_retries + assert sleep.await_count == app.max_poll_retries + + +async def test_http_signal_without_waiting_still_returns_202(http_handler: HttpHandler, sleep: AsyncMock) -> None: + client = _client(_mailbox_state(_response(), expired=True)) + + response = await http_handler(_request(wait=False), client) + + assert response.status_code == 202 + assert json.loads(response.get_body()) == { + "response": "Agent request accepted", + "message": "question", + "session_id": SESSION_ID, + "status": "accepted", + "correlation_id": CORRELATION_ID, + } + client.signal_entity.assert_awaited_once() + client.read_entity_state.assert_not_awaited() + sleep.assert_not_awaited() + + +@pytest.mark.parametrize("expired", [False, True]) +async def test_mcp_raises_for_expired_and_failed_delivery( + expired: bool, app: AgentFunctionApp, sleep: AsyncMock +) -> None: + client = _client(_mailbox_state(_runtime_error(), expired=expired, cleanup=expired)) + expected = EXPIRED_MESSAGE if expired else "Model endpoint unavailable" + + with pytest.raises(RuntimeError, match=expected) as error: + await app._handle_mcp_tool_invocation( + AGENT_NAME, json.dumps({"arguments": {"query": "question", "sessionId": SESSION_ID}}), client + ) + if expired: + assert "Invocation outcome: failed." in str(error.value) + + client.signal_entity.assert_awaited_once() + client.read_entity_state.assert_awaited_once() + sleep.assert_awaited_once_with(0.01) + + +def test_success_helper_keeps_its_existing_signature_and_payload(app: AgentFunctionApp) -> None: + state = DurableAgentState() + + assert app._build_success_result("answer", "question", SESSION_ID, CORRELATION_ID, state) == { + "response": "answer", + "message": "question", + "session_id": SESSION_ID, + "status": "success", + "correlation_id": CORRELATION_ID, + "message_count": 0, + } + + +@pytest.mark.parametrize("value", [None, 0, False, "", [], {}, {"answer": 42}]) +@pytest.mark.parametrize("operation", ["run", "run_agent"]) +def test_entity_factory_delivers_cold_mailbox_without_rerunning_agent(value: Any, operation: str) -> None: + original = _response(value=deepcopy(value)) + state = _mailbox_state(original) + agent = Mock(context_providers=None) + agent.run = AsyncMock() + context = _entity_context(state, operation) + + create_agent_entity(agent)(context) + + context.set_result.assert_called_once() + payload = json.loads(json.dumps(context.set_result.call_args.args[0], allow_nan=False)) + assert payload == state["data"]["responseMailbox"][CORRELATION_ID]["response"] + delivered = AgentResponse.from_dict(payload) + assert delivered.value == value + assert type(delivered.value) is type(value) + agent.run.assert_not_called() + context.set_state.assert_not_called() + + +def test_entity_factory_serializes_live_pydantic_value_in_json_mode() -> None: + class DatedAnswer(BaseModel): + answer: int + day: date + + original = _response(value=DatedAnswer(answer=42, day=date(2026, 9, 8))) + context = _entity_context(None) + with patch("agent_framework_azurefunctions._entities.AgentEntity") as entity: + entity.return_value.run = AsyncMock(return_value=original) + create_agent_entity(Mock(context_providers=None))(context) + entity.return_value.run.assert_awaited_once_with(context.get_input.return_value) + + context.set_result.assert_called_once() + payload = json.loads(json.dumps(context.set_result.call_args.args[0], allow_nan=False)) + assert payload == {**original.to_dict(), "value": {"answer": 42, "day": "2026-09-08"}} + + +@pytest.mark.parametrize("cleanup", [False, True]) +def test_entity_factory_and_task_keep_expired_delivery_terminal(cleanup: bool) -> None: + agent = Mock(context_providers=None) + agent.run = AsyncMock() + state = _mailbox_state(_response(), expired=True, cleanup=cleanup) + before = deepcopy(state) + context = _entity_context(state) + + create_agent_entity(agent)(context) + + context.set_result.assert_called_once() + payload = json.loads(json.dumps(context.set_result.call_args.args[0])) + task = _task(payload, Answer) + assert task.state == TaskState.SUCCEEDED + assert isinstance(task.result, AgentResponse) + assert task.result.additional_properties == { + "durable_status": "already_completed", + "correlation_id": CORRELATION_ID, + "durable_outcome": "succeeded", + } + assert task.result.messages[0].contents[0].error_code == "response_expired" + assert task.result.messages[0].contents[0].message == EXPIRED_MESSAGE + assert task.result.value is None + agent.run.assert_not_called() + if cleanup: + context.set_state.assert_not_called() + else: + context.set_state.assert_called_once() + persisted = context.set_state.call_args.args[0] + expected = deepcopy(state) + del expected["data"]["responseMailbox"] + assert persisted == expected + assert persisted["data"]["completedCorrelations"] == state["data"]["completedCorrelations"] + assert CORRELATION_ID in persisted["data"]["completedCorrelations"] + assert state == before + + +@pytest.mark.parametrize("response_format", [None, Answer]) +@pytest.mark.parametrize("precompleted", [False, True]) +def test_functions_task_keeps_snapshot_metadata_and_structured_value( + response_format: type[BaseModel] | None, precompleted: bool +) -> None: + original = _response(value={"answer": 42}) + payload = _mailbox_state(original)["data"]["responseMailbox"][CORRELATION_ID]["response"] + + task = _task(payload, response_format, precompleted=precompleted) + + assert task.state == TaskState.SUCCEEDED + assert isinstance(task.result, AgentResponse) + assert task.result.to_dict() == original.to_dict() + if response_format is None: + assert task.result.value == {"answer": 42} + else: + assert isinstance(task.result.value, Answer) + assert task.result.value.answer == 42 + + +@pytest.mark.parametrize("terminal_kind", ["error", "already_completed"]) +@pytest.mark.parametrize("text", ["not JSON", '{"answer":0}']) +@pytest.mark.parametrize("precompleted", [False, True]) +def test_functions_task_does_not_parse_error_or_status_only_responses( + terminal_kind: str, text: str, precompleted: bool +) -> None: + original = _response(text=text) + if terminal_kind == "error": + original.messages.append(Message("system", [Content.from_error(message="Failure", error_code="RuntimeError")])) + else: + original.additional_properties["durable_status"] = "already_completed" + + payload = json.loads(json.dumps(serialize_agent_response(original))) + task = _task(payload, Answer, precompleted=precompleted) + + assert task.state == TaskState.SUCCEEDED + assert isinstance(task.result, AgentResponse) + assert task.result.to_dict() == original.to_dict() + assert task.result.value is None + + +def test_functions_task_still_rejects_invalid_success_schema() -> None: + task = _task(_response(text='{"wrong":42}').to_dict(), Answer) + + assert task.state == TaskState.FAILED + assert isinstance(task.result, ValueError) diff --git a/python/packages/azurefunctions/tests/test_deployment_gate_review_af.py b/python/packages/azurefunctions/tests/test_deployment_gate_review_af.py new file mode 100644 index 0000000..d297a97 --- /dev/null +++ b/python/packages/azurefunctions/tests/test_deployment_gate_review_af.py @@ -0,0 +1,116 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Deployment acknowledgement validation for Functions hosts and entity factories.""" + +from typing import Any +from unittest.mock import Mock + +import pytest + +from agent_framework_azurefunctions import AgentFunctionApp +from agent_framework_azurefunctions import _app as app_module +from agent_framework_azurefunctions import _entities as entities_module +from agent_framework_azurefunctions._entities import create_agent_entity + +_ENVIRONMENT_VARIABLE = "DURABLE_AGENTS_DEPLOYMENT_MODE" +_INVALID_MODES = ("", "isolated_v1", "mixed", "ISOLATED_V2", " isolated_v2", "isolated_v2 ", "isolated_v2\n") + + +@pytest.fixture +def agent() -> Mock: + instance = Mock(context_providers=None) + instance.name = "assistant" + return instance + + +@pytest.mark.parametrize("surface", ["app", "factory"]) +def test_missing_mode_fails_before_native_constructor_or_agent_configuration( + monkeypatch: pytest.MonkeyPatch, agent: Mock, surface: str +) -> None: + monkeypatch.delenv(_ENVIRONMENT_VARIABLE, raising=False) + native_init = Mock(side_effect=AssertionError("Native constructor ran before deployment validation")) + configuration = Mock(side_effect=AssertionError("Agent configuration ran before deployment validation")) + reserve = Mock(side_effect=AssertionError("Registry reservation ran before deployment validation")) + monkeypatch.setattr(app_module.DFAppBase, "__init__", native_init) + monkeypatch.setattr(app_module, "validate_agent_configuration", configuration) + monkeypatch.setattr(entities_module, "validate_agent_configuration", configuration) + monkeypatch.setattr(app_module.RegistrationIdentity, "reserve", reserve) + + with pytest.raises(ValueError, match="isolated_v2"): + if surface == "app": + AgentFunctionApp(agents=[agent]) + else: + create_agent_entity(agent) + + native_init.assert_not_called() + configuration.assert_not_called() + reserve.assert_not_called() + + +@pytest.mark.parametrize("surface", ["app", "factory"]) +@pytest.mark.parametrize("deployment_mode", _INVALID_MODES) +def test_invalid_explicit_mode_is_not_overridden_by_valid_environment( + monkeypatch: pytest.MonkeyPatch, agent: Mock, surface: str, deployment_mode: str +) -> None: + monkeypatch.setenv(_ENVIRONMENT_VARIABLE, "isolated_v2") + with pytest.raises(ValueError, match="isolated_v2"): + if surface == "app": + AgentFunctionApp(agents=[agent], deployment_mode=deployment_mode) + else: + create_agent_entity(agent, deployment_mode=deployment_mode) + + +@pytest.mark.parametrize("surface", ["app", "factory"]) +@pytest.mark.parametrize("deployment_mode", _INVALID_MODES) +def test_invalid_environment_mode_is_rejected( + monkeypatch: pytest.MonkeyPatch, agent: Mock, surface: str, deployment_mode: str +) -> None: + monkeypatch.setenv(_ENVIRONMENT_VARIABLE, deployment_mode) + with pytest.raises(ValueError, match="isolated_v2"): + if surface == "app": + AgentFunctionApp(agents=[agent]) + else: + create_agent_entity(agent) + + +@pytest.mark.parametrize("source", ["explicit", "environment"]) +def test_direct_factory_accepts_isolated_mode(monkeypatch: pytest.MonkeyPatch, agent: Mock, source: str) -> None: + if source == "explicit": + monkeypatch.delenv(_ENVIRONMENT_VARIABLE, raising=False) + handler = create_agent_entity(agent, deployment_mode="isolated_v2") + else: + monkeypatch.setenv(_ENVIRONMENT_VARIABLE, "isolated_v2") + handler = create_agent_entity(agent) + assert callable(handler) + + +@pytest.mark.parametrize("source", ["explicit", "environment"]) +def test_app_passes_effective_mode_to_initial_and_later_factories( + monkeypatch: pytest.MonkeyPatch, agent: Mock, source: str +) -> None: + kwargs: dict[str, Any] = {} + if source == "explicit": + monkeypatch.delenv(_ENVIRONMENT_VARIABLE, raising=False) + kwargs["deployment_mode"] = "isolated_v2" + else: + monkeypatch.setenv(_ENVIRONMENT_VARIABLE, "isolated_v2") + factory = Mock(wraps=entities_module.create_agent_entity) + monkeypatch.setattr(app_module, "create_agent_entity", factory) + app = AgentFunctionApp(agents=[agent], enable_health_check=False, enable_http_endpoints=False, **kwargs) + assert app._deployment_mode == "isolated_v2" + + # Later registrations retain the acknowledged mode instead of reading the environment again. + monkeypatch.delenv(_ENVIRONMENT_VARIABLE, raising=False) + later = Mock(context_providers=None) + later.name = "later" + app.add_agent(later) + + assert factory.call_count == 2 + assert all(call.kwargs["deployment_mode"] == "isolated_v2" for call in factory.call_args_list) + assert app.agents == {"assistant": agent, "later": later} + names: list[str] = [] + for function in app.get_functions(): + name = function.get_function_name() + assert name is not None + names.append(name) + assert sorted(names) == ["dafx-assistant", "dafx-later"] diff --git a/python/packages/azurefunctions/tests/test_entities.py b/python/packages/azurefunctions/tests/test_entities.py index cc2fc75..2c6b297 100644 --- a/python/packages/azurefunctions/tests/test_entities.py +++ b/python/packages/azurefunctions/tests/test_entities.py @@ -11,6 +11,7 @@ import pytest from agent_framework import AgentResponse, Message +from agent_framework_durabletask import DurableAgentState, migrate_legacy_state, state_snapshot_digest from agent_framework_azurefunctions._entities import create_agent_entity @@ -66,11 +67,11 @@ def test_entity_function_handles_reset(self) -> None: entity_function = create_agent_entity(mock_agent) - # Mock context with existing state + # Reset an admitted v2 target, not a legacy session. mock_context = Mock() mock_context.operation_name = "reset" mock_context.get_state.return_value = { - "schemaVersion": "1.0.0", + "schemaVersion": DurableAgentState.SCHEMA_VERSION, "data": { "conversationHistory": [ { @@ -147,7 +148,7 @@ def test_entity_function_restores_existing_state(self) -> None: entity_function = create_agent_entity(mock_agent) - existing_state = { + existing_state: dict[str, Any] = { "schemaVersion": "1.0.0", "data": { "conversationHistory": [ @@ -188,17 +189,33 @@ def test_entity_function_restores_existing_state(self) -> None: } mock_context = Mock() + mock_context.entity_name = "dafx-restore" + mock_context.entity_key = "destination" mock_context.operation_name = "reset" - mock_context.get_state.return_value = existing_state + # Import the legacy history explicitly before the normal reset operation. + migrated = migrate_legacy_state( + existing_state, + source_digest=state_snapshot_digest(existing_state), + source_session_id="@dafx-restore@legacy-source", + migration_id="restore-migration-1", + ownership_transfer_id="restore-transfer-1", + delivery_window_seconds=3600, + ).to_dict() + mock_context.get_state.return_value = migrated entity_function(mock_context) assert mock_context.set_result.called + assert mock_context.set_result.call_args[0][0] == {"status": "reset"} # Reset should clear history and persist via set_state assert mock_context.set_state.called persisted_state = mock_context.set_state.call_args[0][0] assert persisted_state["data"]["conversationHistory"] == [] + assert persisted_state["data"]["completedCorrelations"] == migrated["data"]["completedCorrelations"] + assert persisted_state["data"]["responseMailbox"] == migrated["data"]["responseMailbox"] + assert persisted_state["data"]["migration"] == migrated["data"]["migration"] + assert existing_state["schemaVersion"] == "1.0.0" def test_entity_function_handles_string_input(self) -> None: """Test that the entity function handles non-dict input by converting to string.""" diff --git a/python/packages/azurefunctions/tests/test_failure_boundary_consumers.py b/python/packages/azurefunctions/tests/test_failure_boundary_consumers.py new file mode 100644 index 0000000..ba42f46 --- /dev/null +++ b/python/packages/azurefunctions/tests/test_failure_boundary_consumers.py @@ -0,0 +1,170 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""AF polling at deterministic execution boundaries, without a live Functions host.""" + +import asyncio +import importlib +import json +from copy import deepcopy +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import azure.durable_functions as df +import pytest +from agent_framework_durabletask import AgentEntity, DurableAgentState + +from agent_framework_azurefunctions import AgentFunctionApp +from agent_framework_azurefunctions import _app as app_module + + +@pytest.fixture +def boundaries(monkeypatch: pytest.MonkeyPatch) -> Any: + # Allow this file to run alone, without requiring DT test collection first. + # The directory is derived from this worktree, never from another checkout. + tests = Path(__file__).resolve().parents[2] / "durabletask" / "tests" + monkeypatch.syspath_prepend(str(tests)) + module = importlib.import_module("test_cancellation_boundaries") + assert module.__file__ is not None + assert Path(module.__file__).resolve().parent == tests + return module + + +@pytest.fixture +def app(monkeypatch: pytest.MonkeyPatch) -> AgentFunctionApp: + async def immediate_poll_interval(interval: float) -> None: + assert interval == 0.01 + + # Replace only this module's scheduling seam, not asyncio.sleep process-wide. + monkeypatch.setattr(app_module, "asyncio", SimpleNamespace(sleep=immediate_poll_interval)) + return AgentFunctionApp( + enable_health_check=False, + enable_http_endpoints=False, + max_poll_retries=3, + poll_interval_seconds=0.01, + ) + + +class _JsonAFBackend: + """Each backend read returns a fresh real JSON snapshot, never a synthetic response.""" + + def __init__(self, provider: Any) -> None: + self.provider = provider + self.reads = 0 + self.entered = asyncio.Event() + self.release = asyncio.Event() + self.pause = False + self.observed: list[dict[str, Any]] = [] + + async def read_entity_state(self, entity_id: df.EntityId) -> Any: + self.reads += 1 + raw = json.loads(json.dumps(self.provider.raw)) + self.observed.append(raw) + if self.pause: + self.entered.set() + await self.release.wait() + return SimpleNamespace(entity_exists=True, entity_state=raw) + + +async def _poll(app: AgentFunctionApp, backend: Any, correlation: str) -> dict[str, Any]: + return await app._get_response_from_entity( + client=backend, + entity_instance_id=df.EntityId("dafx-boundary", "revision-session"), + correlation_id=correlation, + message="boundary request", + session_id="revision-session", + ) + + +@pytest.mark.parametrize("phase", ["load", "store"]) +async def test_external_failure_and_rejected_error_write_timeout_until_a_real_commit( + phase: str, boundaries: Any, app: AgentFunctionApp +) -> None: + entity, provider, external, client = boundaries.failure_boundary(phase) + before = deepcopy(provider.raw) + request = boundaries.projected_request("provider-failed") + with pytest.raises(OSError, match="entity storage write rejected"): + await entity.run(request) + boundaries.assert_staged_not_committed(provider, before, phase) + assert entity.state.to_dict() == before + assert external.loads == 1 and len(external.saved) == int(phase == "store") + backend = _JsonAFBackend(provider) + + result = await _poll(app, backend, "provider-failed") + + assert result["status"] == "timeout" + assert result["correlation_id"] == "provider-failed" + assert "agent_response" not in result + assert backend.reads == 3 and backend.observed == [before] * 3 + assert provider.writes == 0 and len(provider.attempts) == 1 + assert external.loads == 1 and len(client.effects) == int(phase == "store") + + provider.reject = False + failed = await entity.run(request) + assert provider.writes == 1 + assert failed.additional_properties["durable_status"] == "error" + calls = (external.loads, len(external.saved), len(client.effects)) + delivered = await _poll(app, backend, "provider-failed") + assert delivered["status"] == "error" and delivered["error_code"] == "OSError" + assert delivered["agent_response"] == json.loads(json.dumps(failed.to_dict())) + assert backend.reads == 4 + external.phase = None + cold_provider = boundaries.JsonStateProvider(provider.raw) + cold = AgentEntity(entity.agent, state_provider=cold_provider) + assert (await cold.run(request)).to_dict() == failed.to_dict() + assert (external.loads, len(external.saved), len(client.effects)) == calls + assert cold_provider.writes == 0 + + +async def test_cancelling_caller_poll_does_not_cancel_concurrent_entity_or_allow_duplicate_run( + boundaries: Any, app: AgentFunctionApp +) -> None: + barrier = boundaries.PhaseBarrier() + barrier.phase = "model" + client = boundaries.BarrierClient(barrier) + agent = boundaries.NonStreamingAgent(client=client, name="boundary") + provider = boundaries.JsonStateProvider() + entity = AgentEntity(agent, state_provider=provider) + request = boundaries.projected_request("caller-cancelled") + backend = _JsonAFBackend(provider) + backend.pause = True + execution = asyncio.create_task(entity.run(request)) + polling: asyncio.Task[dict[str, Any]] | None = None + try: + await boundaries.await_boundary(execution, barrier.entered) + polling = asyncio.create_task(_poll(app, backend, "caller-cancelled")) + await boundaries.await_boundary(polling, backend.entered) + assert not execution.done() and not polling.done() + assert backend.observed == [{}] and provider.writes == 0 + assert len(client.effects) == 1 + + polling.cancel() + with pytest.raises(asyncio.CancelledError): + await polling + assert polling.cancelled() + assert not execution.done() and not execution.cancelled() + assert provider.writes == 0 and provider.raw == {} + + barrier.release.set() + response = await execution + assert response.text == "boundary answer" and provider.writes == 1 + assert len(client.effects) == 1 + raw = json.loads(json.dumps(provider.raw)) + assert DurableAgentState.from_json(json.dumps(raw)).try_get_agent_response("caller-cancelled") is not None + backend.pause = False + delivered = await _poll(app, backend, "caller-cancelled") + assert delivered["status"] == "success" + assert delivered["agent_response"] == response.to_dict() + assert backend.reads == 2 + + cold_provider = boundaries.JsonStateProvider(raw) + cold_agent = boundaries.NonStreamingAgent(client=client, name="boundary") + cold = AgentEntity(cold_agent, state_provider=cold_provider) + assert (await cold.run(request)).to_dict() == response.to_dict() + assert len(client.effects) == 1 and cold_provider.writes == 0 + finally: + tasks = [execution, *([polling] if polling is not None else [])] + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) diff --git a/python/packages/azurefunctions/tests/test_hosting_review_af.py b/python/packages/azurefunctions/tests/test_hosting_review_af.py new file mode 100644 index 0000000..a3c7dd0 --- /dev/null +++ b/python/packages/azurefunctions/tests/test_hosting_review_af.py @@ -0,0 +1,512 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Functions hosting preflight, fail-closed registration, and response classification.""" + +import json +from collections.abc import Awaitable, Callable +from copy import deepcopy +from dataclasses import fields +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, Mock + +import azure.durable_functions as df +import azure.functions as func +import pytest +from agent_framework import ( + Agent, + AgentExecutor, + AgentResponse, + Content, + Executor, + InMemoryHistoryProvider, + Message, + WorkflowExecutor, +) +from agent_framework_durabletask import DurableAgentState, ensure_response_format, load_agent_response +from agent_framework_durabletask._configuration import AgentRegistrationSettings +from pydantic import BaseModel + +from agent_framework_azurefunctions import AgentFunctionApp +from agent_framework_azurefunctions._entities import AzureFunctionEntityStateProvider, create_agent_entity + + +class RecordingApp(AgentFunctionApp): + def __init__(self, *, calls: list[tuple[str, str]] | None = None, **kwargs: Any) -> None: + self.calls = calls if calls is not None else [] + self.fail_at: int | None = None + super().__init__(enable_health_check=False, **kwargs) + + def _record(self, kind: str, name: str) -> None: + self.calls.append((kind, name)) + if len(self.calls) == self.fail_at: + raise RuntimeError("injected trigger registration failure") + + def _setup_agent_functions( + self, + agent: Any, + agent_name: str, + callback: Any, + enable_http_endpoint: bool, + enable_mcp_tool_trigger: bool, + **kwargs: Any, + ) -> None: + create_agent_entity(agent, callback, **kwargs) + self._record("entity", f"dafx-{agent_name}") + + def _setup_executor_activity(self, workflow: Any, executor_id: str) -> None: + self._record("activity", f"dafx-{workflow.name}-{executor_id}") + + def _setup_workflow_orchestration(self, workflow: Any) -> None: + self._record("orchestration", f"dafx-{workflow.name}") + + def _register_workflow_routes(self, workflow: Any) -> None: + self._record("routes", workflow.name) + + +def _agent(name: str = "assistant") -> Agent: + client: Any = Mock(additional_properties={}, STORES_BY_DEFAULT=False) + return Agent(client=client, name=name, context_providers=[InMemoryHistoryProvider("history")]) + + +def _workflow(name: str, executor_id: str = "node", *, agent: Any = None, children: tuple[Any, ...] = ()) -> Any: + executor = Mock(spec=Executor if agent is None else AgentExecutor) + executor.id = executor_id + if agent is not None: + executor.agent = agent + executors = {executor_id: executor} + for index, child in enumerate(children): + nested = Mock(spec=WorkflowExecutor) + nested.id = f"child{index}" + nested.workflow = child + executors[nested.id] = nested + workflow = Mock() + workflow.name = name + workflow.executors = executors + return workflow + + +@pytest.mark.parametrize("surface", ["constructor", "nested", "later"]) +@pytest.mark.parametrize("kinds", [(False, False), (False, True), (True, False), (True, True)]) +def test_ambiguous_derived_names_are_preflighted_for_the_entire_composition( + surface: str, kinds: tuple[bool, bool] +) -> None: + left = _workflow("alpha-beta", "gamma", agent=_agent("left") if kinds[0] else None) + right = _workflow("alpha", "beta-gamma", agent=_agent("right") if kinds[1] else None) + calls: list[tuple[str, str]] = [] + if surface == "constructor": + with pytest.raises(ValueError, match="Derived name.*collides"): + RecordingApp(calls=calls, workflows=[left, right]) + assert calls == [] + return + app = RecordingApp(calls=calls) + if surface == "later": + app.configure_workflow(left) + candidate = right + else: + candidate = _workflow("root", children=(left, right)) + before = list(calls) + agents, workflows = app.agents, app.workflows + with pytest.raises(ValueError, match="Derived name.*collides"): + app.configure_workflow(candidate) + assert calls == before + assert app.agents == agents and app.workflows == workflows + app.configure_workflow(_workflow("corrected")) + + +@pytest.mark.parametrize("surface", ["constructor", "standalone_first", "workflow_first"]) +@pytest.mark.parametrize("same_agent", [False, True]) +def test_standalone_agent_cannot_occupy_a_workflow_owned_identity(surface: str, same_agent: bool) -> None: + standalone = _agent("flow-node") + workflow = _workflow("flow", agent=standalone if same_agent else _agent()) + calls: list[tuple[str, str]] = [] + if surface == "constructor": + with pytest.raises(ValueError, match="collides"): + RecordingApp(calls=calls, agents=[standalone], workflow=workflow) + assert calls == [] + return + app = RecordingApp(calls=calls) + if surface == "standalone_first": + app.add_agent(standalone) + else: + app.configure_workflow(workflow) + before = list(calls) + with pytest.raises(ValueError, match="collides"): + if surface == "standalone_first": + app.configure_workflow(workflow) + else: + app.add_agent(standalone) + assert calls == before + + +def test_different_agent_with_same_name_is_not_silently_skipped() -> None: + first = _agent() + app = RecordingApp(agents=[first]) + calls = list(app.calls) + with pytest.raises(ValueError, match="collides"): + app.add_agent(_agent()) + app.add_agent(first) + assert app.calls == calls + assert app.agents == {"assistant": first} + + +def test_constructor_rejects_different_agents_with_duplicate_names_before_any_setup() -> None: + calls: list[tuple[str, str]] = [] + with pytest.raises(ValueError, match="collides"): + RecordingApp(calls=calls, agents=[_agent(), _agent()]) + assert calls == [] + + +def test_case_only_agent_names_fail_before_setup() -> None: + app = RecordingApp(agents=[_agent("Assistant")]) + calls = list(app.calls) + with pytest.raises(ValueError, match="case-insensitively"): + app.add_agent(_agent("assistant")) + assert app.calls == calls + + +_CHANGED_SETTINGS: dict[str, Any] = { + "retention": "follow_compaction", + "max_state_bytes": 8192, + "high_watermark": 0.99, + "low_watermark": 0.1, + "response_delivery_window_seconds": 17, + "callback": Mock(), +} + + +def test_changed_settings_cover_every_shared_configuration_field() -> None: + assert set(_CHANGED_SETTINGS) == {field.name for field in fields(AgentRegistrationSettings)} + + +@pytest.mark.parametrize("setting", [*_CHANGED_SETTINGS, "enable_http_endpoint", "enable_mcp_tool_trigger"]) +def test_same_agent_with_different_configuration_is_rejected(setting: str) -> None: + agent = _agent() + app = RecordingApp(agents=[agent]) + changed = {**_CHANGED_SETTINGS, "enable_http_endpoint": False, "enable_mcp_tool_trigger": True} + calls = list(app.calls) + with pytest.raises(ValueError, match="different settings"): + app.add_agent(agent, **{setting: changed[setting]}) + assert app.calls == calls and app.agents["assistant"] is agent + + +@pytest.mark.parametrize("setting", _CHANGED_SETTINGS) +def test_shared_child_requires_identical_workflow_configuration(setting: str) -> None: + child = _workflow("shared", agent=_agent()) + app = RecordingApp(workflow=_workflow("first", children=(child,))) + calls = list(app.calls) + if setting == "callback": + app.default_callback = _CHANGED_SETTINGS[setting] + overrides: dict[str, Any] = {} + else: + overrides = {setting: _CHANGED_SETTINGS[setting]} + with pytest.raises(ValueError, match="different settings"): + app.configure_workflow(_workflow("second", children=(child,)), **overrides) + assert app.calls == calls and list(app.workflows) == ["first"] + + +def test_shared_workflow_and_explicit_original_agent_remain_benign() -> None: + agent = _agent() + providers = agent.context_providers + child = _workflow("shared", agent=agent) + first = _workflow("first", children=(child,)) + app = RecordingApp(agents=[agent, agent], workflows=[first, first, _workflow("second", children=(child,))]) + assert app.calls.count(("entity", "dafx-shared-node")) == 1 + assert app.calls.count(("routes", "first")) == 1 + assert app.agents["assistant"] is agent and app.agents["shared-node"] is agent + assert agent.context_providers is providers + assert isinstance(providers[0], InMemoryHistoryProvider) + + +@pytest.mark.parametrize("endpoint", ["http", "mcp"]) +def test_sanitized_endpoint_names_are_also_preflighted(endpoint: str) -> None: + calls: list[tuple[str, str]] = [] + with pytest.raises(ValueError, match="Derived name.*collides"): + RecordingApp( + calls=calls, + agents=[_agent("alpha-beta"), _agent("alpha_beta")], + enable_http_endpoints=endpoint == "http", + enable_mcp_tool_trigger=endpoint == "mcp", + ) + assert calls == [] + + +@pytest.mark.parametrize("suffix", ["start", "status", "respond"]) +@pytest.mark.parametrize("agent_executor", [False, True]) +@pytest.mark.parametrize("uppercase", [False, True]) +@pytest.mark.parametrize("surface", ["constructor", "later"]) +def test_workflow_route_suffixes_remain_valid_executor_ids_with_real_triggers( + suffix: str, agent_executor: bool, uppercase: bool, surface: str +) -> None: + executor_id = suffix.upper() if uppercase else suffix + workflow = _workflow("input_boundary", executor_id, agent=_agent() if agent_executor else None) + app = AgentFunctionApp( + workflow=workflow if surface == "constructor" else None, + enable_health_check=False, + enable_http_endpoints=False, + ) + if surface == "later": + app.configure_workflow(workflow) + functions = app.get_functions() + names = [function.get_function_name() for function in functions] + durable_name = f"dafx-input_boundary-{executor_id}" + assert len(names) == len(set(names)) == 5 + assert durable_name in names + assert f"http-dafx-input_boundary-{suffix}" in names + assert "dafx-input_boundary" in names + routes = {} + for function in functions: + trigger = function.get_trigger() + assert trigger is not None + binding = trigger.get_dict_repr() + if binding["type"] == "httpTrigger": + routes[binding["route"]] = function.get_function_name() + if function.get_function_name() == durable_name: + assert binding["type"] == ("entityTrigger" if agent_executor else "activityTrigger") + assert set(routes) == { + "workflow/input_boundary/run", + "workflow/input_boundary/status/{instanceId}", + "workflow/input_boundary/respond/{instanceId}/{requestId}", + } + assert set(routes.values()) == { + f"{'http-' if route_suffix == suffix else ''}dafx-input_boundary-{route_suffix}" + for route_suffix in ("start", "status", "respond") + } + namespace = "entity-name" if agent_executor else "activity-name" + assert (namespace, durable_name.casefold()) in app._registration_identities + assert ("function-name", durable_name.casefold()) in app._registration_identities + + +@pytest.mark.parametrize("agent_executor", [False, True]) +@pytest.mark.parametrize("workflow_first", [False, True]) +def test_real_native_function_collisions_still_fail_before_setup(agent_executor: bool, workflow_first: bool) -> None: + app = AgentFunctionApp(enable_health_check=False, enable_http_endpoints=False) + first = _workflow("alpha", "beta", agent=_agent() if agent_executor else None) + second = _workflow("alpha-beta") + if workflow_first: + first, second = second, first + app.configure_workflow(first) + identities = dict(app._registration_identities) + agents = app.agents + with pytest.raises(ValueError, match="Derived name.*collides"): + app.configure_workflow(second) + assert app._registration_identities == identities + assert app.agents == agents and app.workflows == {first.name: first} + assert app._registered_orchestrations == {first.name.casefold(): first} + assert len(app.get_functions()) == 5 + + +def test_logical_agent_name_can_match_an_http_function_name() -> None: + app = AgentFunctionApp(agents=[_agent("Assistant"), _agent("http-Assistant")], enable_health_check=False) + names = [function.get_function_name() for function in app.get_functions()] + assert set(names) == {"dafx-Assistant", "http-Assistant", "dafx-http-Assistant", "http-http_Assistant"} + assert len(names) == 4 + + +def test_cross_workflow_route_function_collision_is_still_preflighted() -> None: + calls: list[tuple[str, str]] = [] + with pytest.raises(ValueError, match="collides"): + RecordingApp(calls=calls, workflows=[_workflow("flow"), _workflow("flow-start")]) + assert calls == [] + + +def test_real_trigger_names_keep_deployment_compatibility() -> None: + app = AgentFunctionApp( + enable_health_check=False, + agents=[_agent("Assistant")], + workflow=_workflow("Orders", "review", agent=_agent("reviewer")), + ) + names = {function.get_function_name() for function in app.get_functions()} + assert names == { + "dafx-Assistant", + "http-Assistant", + "dafx-Orders-review", + "http-Orders_review", + "dafx-Orders", + "dafx-Orders-start", + "dafx-Orders-status", + "dafx-Orders-respond", + } + + +class UncopyableAgent: + name = "uncopyable" + context_providers = [InMemoryHistoryProvider("history")] + + def __copy__(self) -> Any: + raise TypeError("cannot copy") + + +class ReadOnlyProviders: + name = "readonly" + + @property + def context_providers(self) -> list[Any]: + return [InMemoryHistoryProvider("history")] + + +@pytest.mark.parametrize("factory", [UncopyableAgent, ReadOnlyProviders]) +@pytest.mark.parametrize("surface", ["constructor", "agent", "workflow", "factory"]) +def test_adapter_copy_and_attachment_fail_before_any_triggers(factory: Any, surface: str) -> None: + agent = factory() + calls: list[tuple[str, str]] = [] + app = RecordingApp(calls=calls) + with pytest.raises(ValueError, match="attach durable history"): + if surface == "constructor": + RecordingApp(calls=calls, agents=[_agent(), agent]) + elif surface == "agent": + app.add_agent(agent) + elif surface == "workflow": + app.configure_workflow(_workflow("outer", agent=_agent(), children=(_workflow("inner", agent=agent),))) + else: + create_agent_entity(agent) + assert calls == [] and app.agents == {} and app.workflows == {} + + +@pytest.mark.parametrize("fail_at", [1, 2, 3, 4, 5]) +def test_backend_failure_prevents_retry_and_function_indexing(fail_at: int) -> None: + app = RecordingApp(agents=[_agent("existing")]) + app.fail_at = len(app.calls) + fail_at + workflow = _workflow("root", agent=_agent(), children=(_workflow("child"),)) + with pytest.raises(RuntimeError, match="injected trigger"): + app.configure_workflow(workflow) + calls = list(app.calls) + assert list(app.agents) == ["existing"] + assert app.workflows == {} and app._registered_orchestrations == {} + assert app.workflow is None + for action in (lambda: app.add_agent(_agent("retry")), lambda: app.configure_workflow(workflow), app.get_functions): + with pytest.raises(RuntimeError, match="partially registered"): + action() + assert app.calls == calls + + +def test_standalone_setup_failure_also_blocks_indexing_and_retry() -> None: + app = RecordingApp() + app.fail_at = 1 + with pytest.raises(RuntimeError, match="injected trigger"): + app.add_agent(_agent()) + assert app.agents == {} + for action in (lambda: app.add_agent(_agent()), app.get_functions): + with pytest.raises(RuntimeError, match="partially registered"): + action() + assert len(app.calls) == 1 + + +def _provider(raw_state: Any) -> tuple[AzureFunctionEntityStateProvider, Mock]: + context = Mock(spec=df.DurableEntityContext) + context.get_state.return_value = raw_state + return AzureFunctionEntityStateProvider(context), context + + +@pytest.mark.parametrize("raw", [[], ["state"], "state", 0, False, 2.5]) +def test_non_dictionary_existing_state_is_rejected_not_replaced(raw: Any) -> None: + provider, context = _provider(raw) + with pytest.raises(ValueError, match="Existing durable entity state"): + _ = provider.state + context.set_state.assert_not_called() + + +@pytest.mark.parametrize("raw", [None, {}]) +def test_absent_state_still_initializes(raw: Any) -> None: + provider, context = _provider(raw) + assert provider.state.message_count == 0 + context.set_state.assert_not_called() + + +def test_future_state_fields_survive_adapter_read_and_write() -> None: + raw: dict[str, Any] = { + "schemaVersion": "2.0.0", + "futureEnvelope": {"opaque": [1, {"nested": True}]}, + "data": {"conversationHistory": [], "futureSidecar": {"records": [{"version": 9}]}}, + } + before = deepcopy(raw) + provider, context = _provider(raw) + assert provider._get_state_dict() is raw + _ = provider.state + provider.persist_state() + saved = context.set_state.call_args.args[0] + assert saved["futureEnvelope"] == raw["futureEnvelope"] + assert saved["data"]["futureSidecar"] == raw["data"]["futureSidecar"] + assert raw == before + + +class Answer(BaseModel): + answer: int + + +HttpHandler = Callable[[func.HttpRequest, Any], Awaitable[func.HttpResponse]] + + +def _http_handler(app: AgentFunctionApp, monkeypatch: pytest.MonkeyPatch) -> HttpHandler: + handlers: list[HttpHandler] = [] + + def identity(*args: Any, **kwargs: Any) -> Callable[[HttpHandler], HttpHandler]: + return lambda handler: handler + + def route(*args: Any, **kwargs: Any) -> Callable[[HttpHandler], HttpHandler]: + def capture(handler: HttpHandler) -> HttpHandler: + handlers.append(handler) + return handler + + return capture + + monkeypatch.setattr(app, "function_name", identity) + monkeypatch.setattr(app, "route", route) + monkeypatch.setattr(app, "durable_client_input", identity) + monkeypatch.setattr(app, "_generate_unique_id", lambda: "correlation") + monkeypatch.setattr("agent_framework_azurefunctions._app.asyncio.sleep", AsyncMock()) + app._setup_http_run_route("assistant") + return handlers[0] + + +@pytest.mark.parametrize("kind,expected", [("recovered_tool", 200), ("explicit_error", 500), ("direct_error", 500)]) +async def test_http_uses_shared_terminal_classification_and_canonical_delivery( + monkeypatch: pytest.MonkeyPatch, kind: str, expected: int +) -> None: + app = AgentFunctionApp(enable_health_check=False, enable_http_endpoints=False, max_poll_retries=1) + handler = _http_handler(app, monkeypatch) + messages = [ + Message("tool", [Content.from_error(message="recovered tool failure", error_code="response_expired")]), + Message("assistant", [Content.from_text('{"answer":42}')]), + ] + properties = {"durable_status": "error"} if kind == "explicit_error" else {} + if kind == "direct_error": + messages.append(Message("system", [Content.from_error(message="runtime failure", error_code="runtime")])) + original: AgentResponse[Any] = AgentResponse( + messages=messages, value=Answer(answer=42), additional_properties=properties + ) + state = DurableAgentState() + state.record_response("correlation", original, delivery_window_seconds=3600) + stored = json.loads(state.to_json()) + before = deepcopy(stored) + client = Mock(spec=df.DurableOrchestrationClient) + client.signal_entity = AsyncMock() + client.read_entity_state = AsyncMock(return_value=SimpleNamespace(entity_exists=True, entity_state=stored)) + request = func.HttpRequest( + method="POST", + url="https://example.test/api/agents/assistant/run", + headers={"Content-Type": "application/json"}, + body=b'{"message":"question","session_id":"session"}', + ) + + response = await handler(request, client) + + assert response.status_code == expected + result = json.loads(response.get_body()) + assert result["status"] == ("success" if expected == 200 else "error") + assert result["agent_response"]["type"] == "agent_response" + assert result["agent_response"] == stored["data"]["responseMailbox"]["correlation"]["response"] + assert stored == before + delivered = load_agent_response(result["agent_response"]) + assert type(delivered) is AgentResponse + assert delivered.to_dict() == original.to_dict() + if expected == 200: + ensure_response_format(Answer, "correlation", delivered) + assert delivered.value == Answer(answer=42) + assert result["response"] == original.text + assert result["message_count"] == 0 and result["message"] == "question" + assert result["session_id"] == "session" and result["correlation_id"] == "correlation" + else: + assert result["response"] is None + assert result["error_code"] != "response_expired" + client.read_entity_state.assert_awaited_once() diff --git a/python/packages/azurefunctions/tests/test_integration_environment_af.py b/python/packages/azurefunctions/tests/test_integration_environment_af.py new file mode 100644 index 0000000..6009708 --- /dev/null +++ b/python/packages/azurefunctions/tests/test_integration_environment_af.py @@ -0,0 +1,125 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Regression tests for the integration function app's subprocess environment.""" + +import os +import subprocess +import sys +import uuid +from importlib.util import module_from_spec, spec_from_file_location +from pathlib import Path +from types import ModuleType, SimpleNamespace +from typing import Any +from unittest.mock import Mock, call + +import pytest + + +@pytest.fixture +def _af_harness(monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest) -> ModuleType: + path = Path(__file__).parent / "integration_tests" / "conftest.py" + spec = spec_from_file_location(f"_af_integration_environment_{uuid.uuid4().hex}", path) + assert spec is not None and spec.loader is not None + module = module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + # Keep the integration hooks out of unit-test discovery and avoid local dotenv inputs. + assert not request.config.pluginmanager.is_registered(module) + monkeypatch.setattr(module, "_load_env_file_if_present", Mock()) + return module + + +@pytest.mark.parametrize("parent_mode", [None, "legacy"], ids=["missing-mode", "invalid-mode"]) +@pytest.mark.parametrize("platform", ["win32", "linux"], ids=["windows", "unix"]) +@pytest.mark.parametrize("startup_failures", [0, 2], ids=["first-start", "third-start"]) +def test_function_app_subprocess_opts_into_isolated_mode_on_every_start( + _af_harness: ModuleType, + monkeypatch: pytest.MonkeyPatch, + parent_mode: str | None, + platform: str, + startup_failures: int, +) -> None: + harness = _af_harness + # Set the case after importing the harness so the root fixture cannot mask it. + if parent_mode is None: + monkeypatch.delenv("DURABLE_AGENTS_DEPLOYMENT_MODE", raising=False) + else: + monkeypatch.setenv("DURABLE_AGENTS_DEPLOYMENT_MODE", parent_mode) + monkeypatch.setenv("TASKHUB_NAME", "parent-hub") + monkeypatch.setenv("AzureWebJobsStorage", "UseDevelopmentStorage=true") + monkeypatch.setenv( + "DURABLE_TASK_SCHEDULER_CONNECTION_STRING", + "Endpoint=http://localhost:8080;TaskHub=parent-hub;Authentication=None", + ) + monkeypatch.setenv("FUNCTIONS_WORKER_RUNTIME", "python") + parent_env = dict(os.environ) + + processes = [Mock(spec=subprocess.Popen) for _ in range(startup_failures + 1)] + pending_processes = iter(processes) + + def start_app(*_args: object, **kwargs: Any) -> Mock: + assert dict(os.environ) == parent_env + assert kwargs["env"].get("DURABLE_AGENTS_DEPLOYMENT_MODE") == "isolated_v2" + return next(pending_processes) + + popen = Mock(side_effect=start_app) + monkeypatch.setattr(harness, "sys", SimpleNamespace(platform=platform)) + monkeypatch.setattr(harness, "subprocess", SimpleNamespace(Popen=popen, CREATE_NEW_PROCESS_GROUP=512)) + monkeypatch.setattr( + harness, + "time", + SimpleNamespace(monotonic=Mock(return_value=0), sleep=Mock(side_effect=AssertionError("Unexpected sleep"))), + ) + ports = list(range(17071, 17071 + len(processes))) + find_port = Mock(side_effect=ports) + monkeypatch.setattr(harness, "_find_available_port", find_port) + readiness = Mock( + side_effect=[harness.FunctionAppStartupError("Retry this test startup") for _ in range(startup_failures)] + + [None] + ) + monkeypatch.setattr(harness, "_wait_for_function_app_ready", readiness) + cleanup = Mock() + monkeypatch.setattr(harness, "_cleanup_function_app", cleanup) + for probe in ("_check_func_cli_available", "_check_azurite_available", "_check_dts_emulator_available"): + monkeypatch.setattr(harness, probe, Mock(side_effect=AssertionError("Unexpected infrastructure probe"))) + + assert harness.__file__ is not None + python_root = Path(harness.__file__).resolve().parents[4] + monkeypatch.setattr(harness, "_resolve_repo_root", Mock(return_value=python_root)) + sample_name = "13_subworkflow_hitl" + sample_path = python_root / "samples" / "azure_functions" / sample_name + request_stub = Mock(spec=pytest.FixtureRequest) + request_stub.node.get_closest_marker.return_value = SimpleNamespace(args=(sample_name,)) + lifecycle = harness.function_app_for_test.__wrapped__(request=request_stub) + try: + app_info = next(lifecycle) + assert popen.call_count == len(processes) + hubs: set[str] = set() + for invocation, port in zip(popen.call_args_list, ports, strict=True): + child_env = invocation.kwargs["env"] + assert child_env is not os.environ + hub = child_env["TASKHUB_NAME"] + assert hub.startswith("test") and hub != parent_env["TASKHUB_NAME"] + hubs.add(hub) + expected_options: dict[str, Any] = { + "cwd": str(sample_path), + "env": {**parent_env, "TASKHUB_NAME": hub, "DURABLE_AGENTS_DEPLOYMENT_MODE": "isolated_v2"}, + } + if platform == "win32": + expected_options.update(creationflags=512, shell=True) + else: + expected_options["start_new_session"] = True + assert invocation == call(["func", "start", "--port", str(port)], **expected_options) + assert len(hubs) == len(processes) + assert app_info == {"base_url": f"http://localhost:{ports[-1]}", "port": ports[-1]} + assert find_port.call_count == len(processes) + assert readiness.call_args_list == [ + call(process, port, max_wait=60) for process, port in zip(processes, ports, strict=True) + ] + request_stub.node.get_closest_marker.assert_called_once_with("sample") + assert dict(os.environ) == parent_env + finally: + lifecycle.close() + + assert cleanup.call_args_list == [call(process) for process in processes] + assert dict(os.environ) == parent_env diff --git a/python/packages/azurefunctions/tests/test_maintenance_review_af.py b/python/packages/azurefunctions/tests/test_maintenance_review_af.py new file mode 100644 index 0000000..fb57178 --- /dev/null +++ b/python/packages/azurefunctions/tests/test_maintenance_review_af.py @@ -0,0 +1,547 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Maintenance operations through the synchronous Functions wrapper and real AgentEntity.""" + +import hashlib +import json +from collections.abc import AsyncIterable, Sequence +from copy import deepcopy +from datetime import datetime, timedelta, timezone +from typing import Any +from unittest.mock import Mock + +import azure.durable_functions as df +import pytest +from agent_framework import ( + Agent, + AgentResponse, + ChatResponse, + ChatResponseUpdate, + Content, + ContextProvider, + HistoryProvider, + Message, + ResponseStream, +) +from agent_framework_durabletask import DurableAgentState, serialize_agent_response +from agent_framework_durabletask import _durable_agent_state as state_module +from agent_framework_durabletask import _entities as entities_module +from agent_framework_durabletask import _retention as retention_module +from agent_framework_durabletask import _state_migration as migration_module +from agent_framework_durabletask._message_identity import message_identity +from typing_extensions import Self + +from agent_framework_azurefunctions import _entities as af_entities +from agent_framework_azurefunctions._entities import AzureFunctionEntityStateProvider, create_agent_entity + +NOW = datetime(2040, 1, 1, 12, tzinfo=timezone.utc) +SOURCE_ID = "@dafx-maintenance@legacy-source" +DESTINATION_ID = "@dafx-maintenance@destination" + + +class _ClockType(type): + def __instancecheck__(cls, instance: Any) -> bool: + # Parsed timestamps remain real datetime objects, not instances of the test subclass. + return isinstance(instance, datetime) + + +class Clock(datetime, metaclass=_ClockType): + current = NOW + + @classmethod + def now(cls, tz: Any = None) -> Self: + return cls.fromtimestamp(cls.current.timestamp(), tz=tz) + + +@pytest.fixture +def clock(monkeypatch: pytest.MonkeyPatch) -> type[Clock]: + monkeypatch.setattr(Clock, "current", NOW) + for module in (state_module, entities_module, retention_module, migration_module): + monkeypatch.setattr(module, "datetime", Clock) + return Clock + + +def _wire(value: Any) -> Any: + return json.loads(json.dumps(value, allow_nan=False)) + + +def _digest(raw: dict[str, Any]) -> str: + text = json.dumps(raw, sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False) + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def _size(raw: dict[str, Any]) -> int: + return len(json.dumps(raw, allow_nan=False)) + + +class Model: + def __init__(self) -> None: + self.additional_properties: dict[str, Any] = {} + self.calls: list[list[Message]] = [] + + def get_response(self, messages: Any, *, stream: bool = False, **kwargs: Any) -> Any: + self.calls.append(list(messages)) + text = f"reply-{len(self.calls)}" + + async def complete() -> ChatResponse: + return ChatResponse(messages=[Message("assistant", [text])]) + + async def updates() -> AsyncIterable[ChatResponseUpdate]: + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text(text)]) + + def finalize(items: Sequence[ChatResponseUpdate]) -> ChatResponse: + return ChatResponse.from_updates(items) + + return ResponseStream(updates(), finalizer=finalize) if stream else complete() + + +class Hooks(ContextProvider): + def __init__(self) -> None: + super().__init__("maintenance-probe") + self.calls: list[str] = [] + + async def before_run(self, **kwargs: Any) -> None: + self.calls.append("before") + + async def after_run(self, **kwargs: Any) -> None: + self.calls.append("after") + + +class ExternalHistory(HistoryProvider): + def __init__(self) -> None: + super().__init__("external") + self.calls: list[tuple[str, str | None]] = [] + self.rows: dict[str | None, list[Message]] = { + SOURCE_ID: [Message("user", ["preexisting external input"], message_id="external-old")] + } + + async def get_messages(self, session_id: str | None, **kwargs: Any) -> list[Message]: + self.calls.append(("get", session_id)) + return deepcopy(self.rows.get(session_id, [])) + + async def save_messages(self, session_id: str | None, messages: Any, **kwargs: Any) -> None: + self.calls.append(("save", session_id)) + self.rows.setdefault(session_id, []).extend(deepcopy(list(messages))) + + +class Host: + def __init__( + self, + raw: dict[str, Any] | None = None, + *, + external: ExternalHistory | None = None, + **settings: Any, + ) -> None: + self.raw = _wire(raw or {}) + self.writes = 0 + self.fail_writes = False + self.model = Model() + self.hooks = Hooks() + self.callback = Mock(spec=["on_streaming_response_update", "on_agent_response"]) + client: Any = self.model + providers: list[Any] = [self.hooks] if external is None else [external, self.hooks] + agent = Agent(client=client, name="maintenance", context_providers=providers) + self.handler = create_agent_entity(agent, callback=self.callback, deployment_mode="isolated_v2", **settings) + self.contexts: list[Mock] = [] + + def _write(self, raw: dict[str, Any]) -> None: + if self.fail_writes: + raise OSError("injected commit failure") + self.raw = _wire(raw) + self.writes += 1 + + def invoke(self, operation: str, request: Any = None) -> Any: + # Fresh context/provider on every invocation, but the actual wrapper owns dispatch. + context = Mock(spec=df.DurableEntityContext) + context.entity_name = "dafx-maintenance" + context.entity_key = "destination" + context.operation_name = operation + context.get_input.return_value = request + context.get_state.side_effect = lambda *args, **kwargs: _wire(self.raw) + context.set_state.side_effect = self._write + self.contexts.append(context) + self.handler(context) + context.set_result.assert_called_once() + return context.set_result.call_args.args[0] + + def assert_quiet(self) -> None: + assert self.model.calls == [] and self.hooks.calls == [] and self.callback.mock_calls == [] + + +def _source() -> dict[str, Any]: + return { + "schemaVersion": "1.1.0", + "futureRoot": {"keep": ["雪", None]}, + "data": { + "conversationHistory": [ + { + "$type": kind, + "correlationId": "legacy-done", + "createdAt": "2024-01-01T00:00:00Z", + "messages": [ + { + "role": role, + "contents": [{"$type": "text", "text": text}], + "messageId": f"legacy-{kind}", + } + ], + } + for kind, role, text in ( + ("request", "user", "retained legacy input"), + ("response", "assistant", "retained legacy answer"), + ) + ], + "session": {"session_id": SOURCE_ID, "state": {"opaque": {"keep": [1, 3]}}}, + "futureData": {"keep": [False, 0]}, + }, + } + + +def _request(*, evidence: bool = False) -> dict[str, Any]: + source = _source() + if evidence: + source["data"]["ingestedPositions"] = {"upstream": 3} + source["data"]["conversationHistory"][0]["messages"][0]["messageId"] = "wf_upstream_3" + digest = _digest(source) + request: dict[str, Any] = { + "source": source, + "sourceDigest": digest, + "sourceSessionId": SOURCE_ID, + "destinationSessionId": DESTINATION_ID, + "migrationId": "migration-1", + "ownershipTransferId": "operator-transfer-1", + } + if evidence: + request["deliveryEvidence"] = { + "sourceDigest": digest, + "evidenceId": "operator-journal-1", + "complete": True, + "messages": [ + Message("user", [f"accepted {position}"], message_id=f"wf_upstream_{position}").to_dict() + for position in (1, 3) + ], + } + return request + + +def _mailboxes() -> dict[str, Any]: + raw = _source() + raw["schemaVersion"] = "2.0.0" + state = DurableAgentState.from_dict(raw) + state.data.ingested_messages = {"old-input": ["a" * 64]} + state.data.completed_correlations["long-gone"] = {"completedAt": "2020-01-01T00:00:00+00:00"} + for correlation, error in (("expired-success", False), ("expired-error", True), ("live", False)): + response = AgentResponse[Any]( + messages=[ + Message( + "assistant", + [Content.from_error(message="original failure", error_code="OriginalError")] + if error + else [Content.from_text("original result", additional_properties={"keep": ["雪"]})], + message_id=f"answer-{correlation}", + ) + ], + response_id=f"response-{correlation}", + additional_properties={"durable_status": "error" if error else "success", "nested": {"keep": [1]}}, + value=None if error else {"original": [1, 2]}, + ) + state.record_response( + correlation, + response, + delivery_window_seconds=60, + now=NOW if correlation == "live" else NOW - timedelta(minutes=2), + ) + assert state.data.response_mailbox[correlation]["response"] == serialize_agent_response(response) + return _wire(state.to_dict()) + + +def _without_expired(raw: dict[str, Any]) -> dict[str, Any]: + result = deepcopy(raw) + for correlation in ("expired-success", "expired-error"): + del result["data"]["responseMailbox"][correlation] + return result + + +@pytest.mark.parametrize( + ("operation", "correlation"), + [("run", "new"), ("run", "legacy-done"), ("run_agent", "legacy-done"), ("reset", None), ("expire_responses", None)], +) +def test_af_legacy_lookup_allowed_but_actual_writer_rejected_before_model_hooks_or_write( + operation: str, correlation: str | None +) -> None: + raw = _source() + lookup = DurableAgentState.from_dict(raw).try_get_agent_response("legacy-done") + assert lookup is not None and lookup.text == "retained legacy answer" + host = Host(raw) + result = host.invoke(operation, {"message": "blocked", "correlationId": correlation}) + assert result["status"] == "error" and "read-only" in result["error"] and "Legacy" in result["error"] + assert host.raw == raw and host.writes == 0 + host.contexts[-1].set_state.assert_not_called() + host.assert_quiet() + + +@pytest.mark.parametrize("evidence", [False, True]) +def test_af_migrate_uses_actual_destination_raw_digest_and_optional_evidence( + clock: type[Clock], evidence: bool +) -> None: + request = _request(evidence=evidence) + before = deepcopy(request) + assert request["sourceDigest"] != _digest(DurableAgentState.from_dict(request["source"]).to_dict()) + host = Host(response_delivery_window_seconds=17) + result = host.invoke("migrate", request) + assert result == {"status": "migrated", "migrationId": "migration-1", "sessionId": DESTINATION_ID} + assert host.writes == 1 and request == before + state = DurableAgentState.from_dict(host.raw) + assert state.schema_version == "2.0.0" + assert state.data.session == request["source"]["data"]["session"] + expected_history = DurableAgentState.from_dict(request["source"]).to_dict()["data"]["conversationHistory"] + assert host.raw["data"]["conversationHistory"] == expected_history + metadata = state.data.unknown_fields["migration"] + assert metadata == { + "id": "migration-1", + "sourceDigest": request["sourceDigest"], + "sourceSessionId": SOURCE_ID, + "destinationSessionId": DESTINATION_ID, + "requestDigest": _digest(request), + "ownershipTransferId": "operator-transfer-1", + "createdAt": NOW.isoformat(), + **({"evidenceId": "operator-journal-1"} if evidence else {}), + } + assert state.data.response_mailbox["legacy-done"]["expiresAt"] == (NOW + timedelta(seconds=17)).isoformat() + assert state.data.completed_correlations["legacy-done"]["legacy"] is True + if evidence: + assert state.data.ingested_messages == { + message["message_id"]: [message_identity(Message.from_dict(deepcopy(message)))] + for message in request["deliveryEvidence"]["messages"] + } + assert "wf_upstream_2" not in state.data.ingested_messages + assert host.raw["futureRoot"] == request["source"]["futureRoot"] + assert host.raw["data"]["futureData"] == request["source"]["data"]["futureData"] + host.assert_quiet() + + +def test_af_exact_cold_retry_after_v2_run_does_not_rewrite_or_refresh_expiry(clock: type[Clock]) -> None: + request = _request() + before_request = deepcopy(request) + host = Host(response_delivery_window_seconds=17) + result = host.invoke("migrate", request) + assert result["status"] == "migrated" + clock.current = NOW + timedelta(seconds=5) + assert host.invoke("run", {"message": "new turn", "correlationId": "v2-done"})["type"] == "agent_response" + assert len(host.model.calls) == 1 + before = deepcopy(host.raw) + hooks, callbacks = list(host.hooks.calls), list(host.callback.mock_calls) + clock.current = NOW + timedelta(days=1) + assert host.invoke("migrate", deepcopy(request)) == result + host.contexts[-1].set_state.assert_not_called() + assert host.raw == before and host.writes == 2 + assert host.invoke("migrate", deepcopy(request)) == result + assert host.raw == before and host.writes == 2 + assert len(host.model.calls) == 1 and host.hooks.calls == hooks and host.callback.mock_calls == callbacks + assert host.invoke("expire_responses") == {"expired": 2} + expected = deepcopy(before) + del expected["data"]["responseMailbox"] + assert host.raw == expected and host.writes == 3 + assert request == before_request + + +@pytest.mark.parametrize( + "invalid", ["sourceDigest", "destinationSessionId", "sourceSessionId", "migrationId", "ownershipTransferId"] +) +def test_af_migration_invalid_identity_or_digest_returns_error_without_write(invalid: str) -> None: + request = _request() + if invalid == "sourceDigest": + request[invalid] = _digest(DurableAgentState.from_dict(request["source"]).to_dict()) + elif invalid == "destinationSessionId": + request[invalid] = "@dafx-other@destination" + elif invalid == "sourceSessionId": + request[invalid] = DESTINATION_ID + else: + request[invalid] = " \t" + before = deepcopy(request) + host = Host() + result = host.invoke("migrate", request) + assert result["status"] == "error" and result["error"] + assert host.raw == {} and host.writes == 0 and request == before + host.contexts[-1].set_state.assert_not_called() + host.assert_quiet() + + +@pytest.mark.parametrize("change", ["source", "ownershipTransferId", "migrationId"]) +def test_af_mismatched_retry_never_replaces_committed_migration(clock: type[Clock], change: str) -> None: + request = _request() + host = Host() + assert host.invoke("migrate", request)["status"] == "migrated" + before = deepcopy(host.raw) + changed = deepcopy(request) + if change == "source": + changed["source"]["futureRoot"]["keep"].append("changed source") + changed["sourceDigest"] = _digest(changed["source"]) + else: + changed[change] += "-different" + result = host.invoke("migrate", changed) + assert result["status"] == "error" and "empty" in result["error"] + assert host.raw == before and host.writes == 1 + host.contexts[-1].set_state.assert_not_called() + host.assert_quiet() + + +def test_af_nonempty_destination_without_migration_is_not_overwritten(clock: type[Clock]) -> None: + raw = _mailboxes() + host = Host(raw) + result = host.invoke("migrate", _request()) + assert result["status"] == "error" and "empty" in result["error"] + assert host.raw == raw and host.writes == 0 + host.contexts[-1].set_state.assert_not_called() + host.assert_quiet() + + +@pytest.mark.parametrize("correlation", ["expired-success", "expired-error"]) +def test_af_expired_duplicate_removes_physical_mailbox_without_reexecution( + clock: type[Clock], correlation: str +) -> None: + raw = _mailboxes() + host = Host(raw) + result = host.invoke("run", {"message": "duplicate", "correlationId": correlation}) + assert result["type"] == "agent_response" + assert result["additional_properties"] == { + "durable_status": "already_completed", + "correlation_id": correlation, + "durable_outcome": "failed" if correlation == "expired-error" else "succeeded", + } + assert result["messages"][0]["contents"][0]["error_code"] == "response_expired" + assert raw["data"]["responseMailbox"][correlation]["response"]["response_id"] == f"response-{correlation}" + assert host.raw == _without_expired(raw) and host.writes == 1 + host.assert_quiet() + + +def test_af_idle_expiry_preserves_live_original_history_and_forever_completion_receipts(clock: type[Clock]) -> None: + raw = _mailboxes() + host = Host(raw) + assert host.invoke("expire_responses") == {"expired": 2} + assert host.raw == _without_expired(raw) and host.writes == 1 + assert host.invoke("expire_responses") == {"expired": 0} + host.contexts[-1].set_state.assert_not_called() + assert host.writes == 1 + live = host.invoke("run", {"message": "duplicate", "correlationId": "live"}) + assert live == raw["data"]["responseMailbox"]["live"]["response"] + assert host.writes == 1 + clock.current = NOW + timedelta(days=36500) + assert host.invoke("expire_responses") == {"expired": 1} + expected = _without_expired(raw) + del expected["data"]["responseMailbox"] + assert host.raw == expected and host.writes == 2 + for correlation in raw["data"]["completedCorrelations"]: + result = host.invoke("run", {"message": "old duplicate", "correlationId": correlation}) + assert result["additional_properties"]["durable_status"] == "already_completed" + assert host.raw == expected and host.writes == 2 + host.assert_quiet() + + +@pytest.mark.parametrize("operation", ["migrate", "expire_responses"]) +def test_af_failed_maintenance_commit_restores_real_provider_cache_and_retries( + clock: type[Clock], monkeypatch: pytest.MonkeyPatch, operation: str +) -> None: + providers: list[AzureFunctionEntityStateProvider] = [] + originals: list[DurableAgentState] = [] + + def capture(context: Any) -> AzureFunctionEntityStateProvider: + provider = AzureFunctionEntityStateProvider(context) + providers.append(provider) + originals.append(provider.state) + return provider + + # Observe the actual state provider, never replace AgentEntity or its maintenance methods. + monkeypatch.setattr(af_entities, "AzureFunctionEntityStateProvider", capture) + raw = {} if operation == "migrate" else _mailboxes() + host = Host(raw) + request = _request(evidence=True) + before_request = deepcopy(request) + host.fail_writes = True + result = host.invoke(operation, request) + assert result == {"status": "error", "error": "injected commit failure"} + assert providers[-1].state is originals[-1] + assert providers[-1].state.to_dict() == (raw or DurableAgentState().to_dict()) + assert host.raw == raw and host.writes == 0 + host.contexts[-1].set_state.assert_called_once() + host.fail_writes = False + result = host.invoke(operation, request) + expected = ( + {"status": "migrated", "migrationId": "migration-1", "sessionId": DESTINATION_ID} + if operation == "migrate" + else {"expired": 2} + ) + assert result == expected and host.writes == 1 and request == before_request + host.assert_quiet() + + +@pytest.mark.parametrize("operation", ["migrate", "expire_responses"]) +def test_af_strict_maintenance_budget_includes_full_retained_floor_and_metadata( + clock: type[Clock], operation: str +) -> None: + raw = {} if operation == "migrate" else _mailboxes() + request = _request() + request["source"]["data"]["conversationHistory"][0]["messages"][0]["contents"][0]["text"] = "雪" * 500 + request["sourceDigest"] = _digest(request["source"]) + before_request = deepcopy(request) + sizing = Host(raw) + expected_result = sizing.invoke(operation, request) + if operation == "migrate": + assert expected_result["status"] == "migrated" + else: + assert expected_result == {"expired": 2} + expected = deepcopy(sizing.raw) + full_size = _size(expected) + if operation == "migrate": + assert full_size > _size(request["source"]) + normalized = DurableAgentState.from_dict(request["source"]).to_dict() + assert expected["data"]["conversationHistory"] == normalized["data"]["conversationHistory"] + else: + assert expected == _without_expired(raw) + assert full_size > _size(expected["data"]["responseMailbox"]) + rejected = Host(raw, max_state_bytes=full_size - 1) + result = rejected.invoke(operation, request) + assert result["status"] == "error" and "max_state_bytes" in result["error"] + assert rejected.raw == raw and rejected.writes == 0 + rejected.contexts[-1].set_state.assert_not_called() + accepted = Host(raw, max_state_bytes=full_size) + assert accepted.invoke(operation, request) == expected_result + assert accepted.raw == expected and accepted.writes == 1 + assert request == before_request and "truncation" not in accepted.raw["data"] + rejected.assert_quiet() + accepted.assert_quiet() + + +def test_af_external_get_and_save_use_logical_source_identity_on_cold_destination_after_retry( + clock: type[Clock], +) -> None: + external = ExternalHistory() + host = Host(external=external) + request = _request() + before_request = deepcopy(request) + rows = {key: [message.to_dict() for message in messages] for key, messages in external.rows.items()} + assert host.invoke("migrate", request)["status"] == "migrated" + # Input IDs do not authorize or prove external transfer; the operator owns that boundary. + assert external.calls == [] + assert {key: [message.to_dict() for message in messages] for key, messages in external.rows.items()} == rows + for index in range(2): + if index: + clock.current = NOW + timedelta(seconds=10) + before = deepcopy(host.raw) + assert host.invoke("migrate", request)["status"] == "migrated" + assert host.raw == before + host.contexts[-1].set_state.assert_not_called() + response = host.invoke("run", {"message": f"new turn {index}", "correlationId": f"new-{index}"}) + assert response["type"] == "agent_response" + assert host.raw["data"]["session"]["session_id"] == SOURCE_ID + assert external.calls == [(phase, SOURCE_ID) for _ in range(2) for phase in ("get", "save")] + assert set(external.rows) == {SOURCE_ID} + assert [message.text for message in external.rows[SOURCE_ID]] == [ + "preexisting external input", + "new turn 0", + "reply-1", + "new turn 1", + "reply-2", + ] + assert "retained legacy input" not in [message.text for batch in host.model.calls for message in batch] + assert request == before_request diff --git a/python/packages/azurefunctions/tests/test_multi_agent.py b/python/packages/azurefunctions/tests/test_multi_agent.py index c03e00d..e9a468d 100644 --- a/python/packages/azurefunctions/tests/test_multi_agent.py +++ b/python/packages/azurefunctions/tests/test_multi_agent.py @@ -2,7 +2,7 @@ """Unit tests for multi-agent support in AgentFunctionApp.""" -from unittest.mock import Mock +from unittest.mock import Mock, patch import pytest @@ -40,17 +40,18 @@ def test_init_with_no_agents(self) -> None: assert len(app.agents) == 0 def test_init_with_duplicate_agent_names(self) -> None: - """Test initialization with duplicate agent names deduplicates with warning.""" + """Different agents with the same name fail before any registration.""" agent1 = Mock() agent1.name = "TestAgent" agent2 = Mock() agent2.name = "TestAgent" - app = AgentFunctionApp(agents=[agent1, agent2]) - - # Duplicate is skipped, only the first agent is registered - assert len(app.agents) == 1 - assert "TestAgent" in app.agents + with ( + patch.object(AgentFunctionApp, "_setup_agent_functions") as setup, + pytest.raises(ValueError, match="collides"), + ): + AgentFunctionApp(agents=[agent1, agent2]) + setup.assert_not_called() def test_init_with_agent_without_name(self) -> None: """Test initialization with agent missing name attribute raises error.""" @@ -58,7 +59,7 @@ def test_init_with_agent_without_name(self) -> None: agent1.name = "Agent1" agent2 = Mock(spec=[]) # Mock without name attribute - with pytest.raises(ValueError, match="does not have a 'name' attribute"): + with pytest.raises(ValueError, match="Agent must have a name"): AgentFunctionApp(agents=[agent1, agent2]) @@ -94,8 +95,8 @@ def test_add_multiple_agents(self) -> None: assert "Agent1" in app.agents assert "Agent2" in app.agents - def test_add_agent_with_duplicate_name_skips(self) -> None: - """Test that adding agent with duplicate name logs warning and skips.""" + def test_add_agent_with_duplicate_name_raises(self) -> None: + """A different agent cannot replace or reuse an existing registration.""" agent1 = Mock() agent1.name = "MyAgent" agent2 = Mock() @@ -103,11 +104,14 @@ def test_add_agent_with_duplicate_name_skips(self) -> None: app = AgentFunctionApp(agents=[agent1]) - # Duplicate is silently skipped with a warning - app.add_agent(agent2) + with ( + patch.object(app, "_setup_agent_functions") as setup, + pytest.raises(ValueError, match="collides"), + ): + app.add_agent(agent2) - # Only the original agent remains - assert len(app.agents) == 1 + setup.assert_not_called() + assert app.agents == {"MyAgent": agent1} def test_add_agent_to_app_with_existing_agents(self) -> None: """Test adding agent to app that already has agents.""" diff --git a/python/packages/azurefunctions/tests/test_orchestration.py b/python/packages/azurefunctions/tests/test_orchestration.py index ec387f6..d3d659b 100644 --- a/python/packages/azurefunctions/tests/test_orchestration.py +++ b/python/packages/azurefunctions/tests/test_orchestration.py @@ -129,14 +129,15 @@ def executor_with_context(mock_context_with_uuid: tuple[Mock, str]) -> tuple[Any class TestAgentResponseHelpers: """Tests for response handling through public AgentTask API.""" - def test_try_set_value_exception_handling(self) -> None: + @pytest.mark.parametrize("invalid_result", [{"invalid": "format"}, {}, {"messages": None}, {"messages": ""}]) + def test_try_set_value_exception_handling(self, invalid_result: dict[str, Any]) -> None: """Test try_set_value handles exceptions raised when converting a successful task result to AgentResponse.""" entity_task = _create_entity_task() task = AgentTask(entity_task, None, "correlation-id") # Simulate successful entity task with invalid result that causes exception entity_task.state = TaskState.SUCCEEDED - entity_task.result = {"invalid": "format"} # Missing required fields for AgentResponse + entity_task.result = invalid_result # Clear pending_tasks to simulate that parent has processed the child task.pending_tasks.clear() @@ -146,9 +147,10 @@ def test_try_set_value_exception_handling(self) -> None: # Verify task failed due to conversion exception assert task.state == TaskState.FAILED - assert isinstance(task.result, Exception) + assert isinstance(task.result, (TypeError, ValueError)) - def test_try_set_value_success(self) -> None: + @pytest.mark.parametrize("include_type", [False, True]) + def test_try_set_value_success(self, include_type: bool) -> None: """Test try_set_value correctly processes successful task completion.""" entity_task = _create_entity_task() task = AgentTask(entity_task, None, "correlation-id") @@ -156,6 +158,8 @@ def test_try_set_value_success(self) -> None: # Simulate successful entity task completion entity_task.state = TaskState.SUCCEEDED entity_task.result = AgentResponse(messages=[Message(role="assistant", contents=["Test response"])]).to_dict() + if not include_type: + entity_task.result.pop("type") # Clear pending_tasks to simulate that parent has processed the child task.pending_tasks.clear() diff --git a/python/packages/azurefunctions/tests/test_retention_registration_af.py b/python/packages/azurefunctions/tests/test_retention_registration_af.py new file mode 100644 index 0000000..93d88a7 --- /dev/null +++ b/python/packages/azurefunctions/tests/test_retention_registration_af.py @@ -0,0 +1,419 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Functions registration validation and settings forwarded to the AgentEntity consumer.""" + +from collections.abc import Callable, Iterator +from inspect import signature +from typing import Any, get_args +from unittest.mock import Mock, patch + +import azure.durable_functions as df +import pytest +from agent_framework import Agent, AgentExecutor, Executor, InMemoryHistoryProvider, WorkflowExecutor +from agent_framework_durabletask import ( + DEFAULT_MAX_STATE_BYTES, + DEFAULT_RETENTION, + DELIVERY_WINDOW_SECONDS, + HIGH_WATERMARK, + INHERIT, + LOW_WATERMARK, + RetentionMode, +) + +from agent_framework_azurefunctions import AgentFunctionApp +from agent_framework_azurefunctions._entities import AzureFunctionEntityStateProvider, create_agent_entity + +EntityHandler = Callable[[df.DurableEntityContext], None] + + +def _agent(name: str = "assistant", *, ambiguous_history: bool = False) -> Agent: + client: Any = Mock(additional_properties={}, STORES_BY_DEFAULT=False) + providers = ( + [InMemoryHistoryProvider(source_id="first"), InMemoryHistoryProvider(source_id="second")] + if ambiguous_history + else [InMemoryHistoryProvider(source_id="primary")] + ) + return Agent(client=client, name=name, context_providers=providers) + + +def _workflow(name: str, *agents: Agent, child: Mock | None = None) -> Mock: + executors: dict[str, Mock] = {} + for index, agent in enumerate(agents): + node = Mock(spec=AgentExecutor) + node.id = f"node{index}" + node.agent = agent + executors[node.id] = node + if child is not None: + nested = Mock(spec=WorkflowExecutor) + nested.id = "child" + nested.workflow = child + executors[nested.id] = nested + activity = Mock(spec=Executor) + activity.id = "activity" + executors[activity.id] = activity + workflow = Mock() + workflow.name = name + workflow.executors = executors + return workflow + + +@pytest.fixture +def registered_entities() -> Iterator[dict[str, EntityHandler]]: + registered: dict[str, EntityHandler] = {} + + def capture(*, context_name: str, entity_name: str) -> Callable[[EntityHandler], EntityHandler]: + assert context_name == "context" + + def decorate(handler: EntityHandler) -> EntityHandler: + registered[entity_name] = handler + return handler + + return decorate + + with patch.object(AgentFunctionApp, "entity_trigger", side_effect=capture): + yield registered + + +def _consumer_settings(handler: EntityHandler) -> dict[str, Any]: + context = Mock() + context.operation_name = "reset" + with patch("agent_framework_azurefunctions._entities.AgentEntity") as consumer: + handler(context) + consumer.assert_called_once() + consumer.return_value.reset.assert_called_once_with() + context.set_result.assert_called_once_with({"status": "reset"}) + kwargs = consumer.call_args.kwargs + assert isinstance(kwargs["state_provider"], AzureFunctionEntityStateProvider) + return dict(kwargs) + + +def _assert_settings(actual: dict[str, Any], **expected: Any) -> None: + assert {key: actual[key] for key in expected} == expected + + +def _app(**kwargs: Any) -> AgentFunctionApp: + return AgentFunctionApp(enable_health_check=False, enable_http_endpoints=False, **kwargs) + + +def test_functions_uses_the_public_inheritance_sentinel() -> None: + assert signature(AgentFunctionApp).parameters["workflow_max_state_bytes"].default is INHERIT + assert signature(AgentFunctionApp.add_agent).parameters["max_state_bytes"].default is INHERIT + assert signature(AgentFunctionApp.configure_workflow).parameters["max_state_bytes"].default is INHERIT + + +def test_app_defaults_reach_the_entity_consumer(registered_entities: dict[str, EntityHandler]) -> None: + _app(agents=[_agent()]) + + assert DEFAULT_RETENTION == "keep_all" + assert DEFAULT_MAX_STATE_BYTES is None + _assert_settings( + _consumer_settings(registered_entities["dafx-assistant"]), + retention=DEFAULT_RETENTION, + max_state_bytes=None, + high_watermark=HIGH_WATERMARK, + low_watermark=LOW_WATERMARK, + response_delivery_window_seconds=DELIVERY_WINDOW_SECONDS, + ) + + +@pytest.mark.parametrize("retention", get_args(RetentionMode)) +@pytest.mark.parametrize("budget", [None, 8192]) +def test_pressure_budget_is_independent_of_retention( + registered_entities: dict[str, EntityHandler], retention: RetentionMode, budget: int | None +) -> None: + _app( + agents=[_agent()], + retention=retention, + max_state_bytes=budget, + high_watermark=0.95, + low_watermark=0.8, + response_delivery_window_seconds=120, + ) + _assert_settings( + _consumer_settings(registered_entities["dafx-assistant"]), + retention=retention, + max_state_bytes=budget, + high_watermark=0.95, + low_watermark=0.8, + response_delivery_window_seconds=120, + ) + + +@pytest.mark.parametrize("surface", ["agent", "workflow", "workflow_default"]) +@pytest.mark.parametrize( + "overrides,expected", + [ + ({}, 8192), + ({"max_state_bytes": INHERIT}, 8192), + ({"max_state_bytes": None}, None), + ({"max_state_bytes": 4096}, 4096), + ], +) +def test_budget_override_distinguishes_omitted_and_disabled( + registered_entities: dict[str, EntityHandler], surface: str, overrides: dict[str, Any], expected: int | None +) -> None: + if surface == "workflow_default": + _app( + workflow=_workflow("flow", _agent()), + max_state_bytes=8192, + **{f"workflow_{key}": value for key, value in overrides.items()}, + ) + else: + app = _app(max_state_bytes=8192) + if surface == "agent": + app.add_agent(_agent(), **overrides) + else: + app.configure_workflow(_workflow("flow", _agent()), **overrides) + + assert len(registered_entities) == 1 + handler = next(iter(registered_entities.values())) + assert _consumer_settings(handler)["max_state_bytes"] == expected + + +def test_per_agent_overrides_leave_host_defaults_unchanged(registered_entities: dict[str, EntityHandler]) -> None: + app = _app( + retention="follow_compaction", + max_state_bytes=8192, + high_watermark=0.95, + low_watermark=0.8, + response_delivery_window_seconds=120, + ) + app.add_agent( + _agent("override"), + retention="keep_all", + max_state_bytes=None, + high_watermark=0.9, + low_watermark=0.6, + response_delivery_window_seconds=15, + ) + app.add_agent(_agent("inherited")) + + _assert_settings( + _consumer_settings(registered_entities["dafx-override"]), + retention="keep_all", + max_state_bytes=None, + high_watermark=0.9, + low_watermark=0.6, + response_delivery_window_seconds=15, + ) + _assert_settings( + _consumer_settings(registered_entities["dafx-inherited"]), + retention="follow_compaction", + max_state_bytes=8192, + high_watermark=0.95, + low_watermark=0.8, + response_delivery_window_seconds=120, + ) + + +@pytest.mark.parametrize("retention", get_args(RetentionMode)) +def test_workflow_defaults_apply_to_nested_agents_not_standalone_agents( + registered_entities: dict[str, EntityHandler], retention: RetentionMode +) -> None: + inner = _workflow("inner", _agent("inneragent")) + outer = _workflow("outer", _agent("outeragent"), child=inner) + _app( + agents=[_agent("standalone")], + workflow=outer, + max_state_bytes=8192, + workflow_retention=retention, + workflow_max_state_bytes=None, + workflow_high_watermark=0.9, + workflow_low_watermark=0.6, + workflow_response_delivery_window_seconds=20, + ) + + for name in ("dafx-outer-node0", "dafx-inner-node0"): + _assert_settings( + _consumer_settings(registered_entities[name]), + retention=retention, + max_state_bytes=None, + high_watermark=0.9, + low_watermark=0.6, + response_delivery_window_seconds=20, + ) + _assert_settings( + _consumer_settings(registered_entities["dafx-standalone"]), + retention=DEFAULT_RETENTION, + max_state_bytes=8192, + high_watermark=HIGH_WATERMARK, + low_watermark=LOW_WATERMARK, + response_delivery_window_seconds=DELIVERY_WINDOW_SECONDS, + ) + + +def test_per_workflow_overrides_apply_to_all_new_nested_agents( + registered_entities: dict[str, EntityHandler], +) -> None: + app = _app( + workflow_retention="follow_compaction", + workflow_max_state_bytes=8192, + workflow_high_watermark=0.95, + workflow_low_watermark=0.8, + workflow_response_delivery_window_seconds=120, + ) + inner = _workflow("inner", _agent("inneragent")) + outer = _workflow("outer", _agent("outeragent"), child=inner) + app.configure_workflow( + outer, + retention="keep_all", + max_state_bytes=None, + high_watermark=0.9, + low_watermark=0.6, + response_delivery_window_seconds=15, + ) + assert app.workflow is outer + app.configure_workflow(_workflow("inherited", _agent("other"))) + assert app.workflow is None + + for name in ("dafx-outer-node0", "dafx-inner-node0"): + _assert_settings( + _consumer_settings(registered_entities[name]), + retention="keep_all", + max_state_bytes=None, + high_watermark=0.9, + low_watermark=0.6, + response_delivery_window_seconds=15, + ) + _assert_settings( + _consumer_settings(registered_entities["dafx-inherited-node0"]), + retention="follow_compaction", + max_state_bytes=8192, + high_watermark=0.95, + low_watermark=0.8, + response_delivery_window_seconds=120, + ) + + +_INVALID_SETTINGS: list[dict[str, Any]] = [ + {"retention": "auto"}, + {"retention": "invalid"}, + *({"max_state_bytes": value} for value in [0, -1, True, False, 1.5, "8192", "inherit", "backend_limit"]), + *({"high_watermark": value} for value in [0, 1.1, True, float("nan"), float("inf")]), + *({"low_watermark": value} for value in [0, -0.1, True, float("nan"), float("inf")]), + {"high_watermark": 0.7, "low_watermark": 0.7}, + {"high_watermark": 0.6, "low_watermark": 0.7}, + *( + {"response_delivery_window_seconds": value} + for value in [0, -1, True, False, 1.5, "60", float("nan"), float("inf")] + ), +] + + +@pytest.mark.parametrize("settings", _INVALID_SETTINGS) +@pytest.mark.parametrize("surface", ["host", "workflow_default", "agent", "workflow", "factory"]) +def test_invalid_settings_fail_before_registration( + registered_entities: dict[str, EntityHandler], surface: str, settings: dict[str, Any] +) -> None: + with ( + patch.object(AgentFunctionApp, "_setup_http_run_route") as http, + patch.object(AgentFunctionApp, "_setup_mcp_tool_trigger") as mcp, + patch.object(AgentFunctionApp, "_setup_executor_activity") as activity, + patch.object(AgentFunctionApp, "_setup_workflow_orchestration") as orchestration, + patch.object(AgentFunctionApp, "_register_workflow_routes") as routes, + ): + if surface in ("host", "workflow_default", "factory"): + with pytest.raises(ValueError): + if surface == "host": + _app(agents=[_agent()], **settings) + elif surface == "workflow_default": + _app( + workflow=_workflow("flow", _agent()), + **{f"workflow_{key}": value for key, value in settings.items()}, + ) + else: + create_agent_entity(_agent(), **settings) + else: + app = _app() + with pytest.raises(ValueError): + if surface == "agent": + app.add_agent(_agent(), **settings) + else: + app.configure_workflow(_workflow("flow", _agent()), **settings) + assert app.agents == {} + assert app.workflows == {} + assert app._registered_orchestrations == {} + for registration in (http, mcp, activity, orchestration, routes): + registration.assert_not_called() + assert registered_entities == {} + + +@pytest.mark.parametrize("surface", ["agent", "workflow", "nested_workflow"]) +def test_ambiguous_history_fails_before_any_registration( + registered_entities: dict[str, EntityHandler], surface: str +) -> None: + app = _app() + agent = _agent(ambiguous_history=True) + original_providers = agent.context_providers + with pytest.raises(ValueError, match="primary"): + if surface == "agent": + app.add_agent(agent) + elif surface == "workflow": + app.configure_workflow(_workflow("flow", _agent("good"), agent)) + else: + app.configure_workflow(_workflow("outer", _agent("good"), child=_workflow("inner", agent))) + + assert agent.context_providers is original_providers + assert all(isinstance(provider, InMemoryHistoryProvider) for provider in original_providers) + assert app.agents == {} + assert app.workflows == {} + assert app._registered_orchestrations == {} + assert registered_entities == {} + + +@pytest.mark.parametrize("surface", ["agents", "workflow", "workflows"]) +def test_constructor_preflights_all_initial_agents_and_workflows( + registered_entities: dict[str, EntityHandler], surface: str +) -> None: + good, bad = _agent("good"), _agent("bad", ambiguous_history=True) + with ( + patch.object(AgentFunctionApp, "_setup_agent_functions") as setup_agent, + patch.object(AgentFunctionApp, "_register_workflow_primitives") as setup_workflow, + pytest.raises(ValueError, match="primary"), + ): + if surface == "agents": + _app(agents=[good, bad]) + elif surface == "workflow": + _app(workflow=_workflow("outer", good, child=_workflow("inner", bad))) + else: + _app(workflows=[_workflow("first", good), _workflow("second", bad)]) + setup_agent.assert_not_called() + setup_workflow.assert_not_called() + assert registered_entities == {} + + +def test_registration_and_factory_validation_do_not_replace_history( + registered_entities: dict[str, EntityHandler], +) -> None: + agent = _agent() + original_providers = agent.context_providers + with patch("agent_framework_azurefunctions._entities.AgentEntity") as consumer: + _app(agents=[agent], retention="follow_compaction") + consumer.assert_not_called() + assert "dafx-assistant" in registered_entities + assert agent.context_providers is original_providers + assert isinstance(agent.context_providers[0], InMemoryHistoryProvider) + + +def test_functions_backend_limit_error_is_raised_before_invocation() -> None: + with pytest.raises(ValueError, match="max_state_bytes.*backend_limit"): + create_agent_entity(_agent(), max_state_bytes="backend_limit") + + +@pytest.mark.parametrize("invalid_name", [None, "", "invalid name"]) +def test_constructor_keeps_name_validation_before_workflow_traversal( + registered_entities: dict[str, EntityHandler], invalid_name: Any +) -> None: + with pytest.raises(ValueError, match="Workflow name"): + _app(workflows=[_workflow("valid", _agent()), _workflow(invalid_name, _agent("invalid"))]) + assert registered_entities == {} + + +def test_function_setup_failure_does_not_record_agent_metadata() -> None: + app = _app() + with ( + patch.object(app, "_setup_agent_functions", side_effect=RuntimeError("registration failed")), + pytest.raises(RuntimeError, match="registration failed"), + ): + app.add_agent(_agent()) + assert app.agents == {} diff --git a/python/packages/azurefunctions/tests/test_workflow.py b/python/packages/azurefunctions/tests/test_workflow.py index f68af16..2f45815 100644 --- a/python/packages/azurefunctions/tests/test_workflow.py +++ b/python/packages/azurefunctions/tests/test_workflow.py @@ -7,7 +7,6 @@ from typing import Any from agent_framework import ( - AgentExecutorRequest, AgentExecutorResponse, AgentResponse, Message, @@ -22,7 +21,6 @@ ) from agent_framework_azurefunctions._workflow import ( - _extract_message_content, build_agent_executor_response, route_message_through_edge_groups, ) @@ -198,76 +196,6 @@ def test_conversation_extends_previous_agent_executor_response(self) -> None: assert response.full_conversation[2].text == "Current response" -class TestExtractMessageContent: - """Test suite for _extract_message_content function.""" - - def test_extract_from_string(self) -> None: - """Test extracting content from plain string.""" - result = _extract_message_content("Hello, world!") - - assert result == "Hello, world!" - - def test_extract_from_agent_executor_response_with_text(self) -> None: - """Test extracting from AgentExecutorResponse with text.""" - response = AgentExecutorResponse( - executor_id="exec", - agent_response=AgentResponse(messages=[Message(role="assistant", contents=["Response text"])]), - full_conversation=[Message(role="assistant", contents=["Response text"])], - ) - - result = _extract_message_content(response) - - assert result == "Response text" - - def test_extract_from_agent_executor_response_with_messages(self) -> None: - """Test extracting from AgentExecutorResponse with messages.""" - response = AgentExecutorResponse( - executor_id="exec", - agent_response=AgentResponse( - messages=[ - Message(role="user", contents=["First"]), - Message(role="assistant", contents=["Last message"]), - ] - ), - full_conversation=[ - Message(role="user", contents=["First"]), - Message(role="assistant", contents=["Last message"]), - ], - ) - - result = _extract_message_content(response) - - # AgentResponse.text concatenates all message texts - assert result == "FirstLast message" - - def test_extract_from_agent_executor_request(self) -> None: - """Test extracting from AgentExecutorRequest.""" - request = AgentExecutorRequest( - messages=[ - Message(role="user", contents=["First"]), - Message(role="user", contents=["Last request"]), - ] - ) - - result = _extract_message_content(request) - - assert result == "Last request" - - def test_extract_from_dict_returns_empty(self) -> None: - """Test that dict messages return empty string (unexpected input).""" - msg_dict = {"messages": [{"text": "Hello"}]} - - result = _extract_message_content(msg_dict) - - assert result == "" - - def test_extract_returns_empty_for_unknown_type(self) -> None: - """Test that unknown types return empty string.""" - result = _extract_message_content(12345) - - assert result == "" - - class TestEdgeGroupIntegration: """Integration tests for edge group routing with realistic scenarios.""" diff --git a/python/packages/azurefunctions/tests/test_workflow_dispatch_revision_af.py b/python/packages/azurefunctions/tests/test_workflow_dispatch_revision_af.py new file mode 100644 index 0000000..6ea77c6 --- /dev/null +++ b/python/packages/azurefunctions/tests/test_workflow_dispatch_revision_af.py @@ -0,0 +1,238 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Workflow dispatch through the real Azure Functions adapter and shared shim.""" + +from __future__ import annotations + +import json +from copy import deepcopy +from datetime import datetime, timezone +from typing import Any +from unittest.mock import Mock +from uuid import UUID + +import azure.durable_functions as df +from agent_framework import AgentExecutor, AgentExecutorResponse, AgentResponse, AgentSession, Content, Message +from agent_framework_durabletask import DurableAgentStateRequest, RunRequest +from agent_framework_durabletask._workflows.orchestrator import _prepare_agent_task, _WorkflowDeliveryLedger +from azure.durable_functions.models.actions.NoOpAction import NoOpAction +from azure.durable_functions.models.Task import AtomicTask, TaskState + +from agent_framework_azurefunctions._orchestration import AgentTask +from agent_framework_azurefunctions._workflow_af_context import AzureFunctionsWorkflowContext + + +class _StubAgent: + name = "stub" + id = "stub" + description = None + + def create_session(self, **kwargs: Any) -> AgentSession: + return AgentSession(**kwargs) + + async def run(self, messages: Any = None, **kwargs: Any) -> AgentResponse: + raise AssertionError("Dispatch must schedule an entity, not invoke a model") + + +def _agent(**kwargs: Any) -> AgentExecutor: + stub: Any = _StubAgent() + return AgentExecutor(stub, id="target", **kwargs) + + +def _upstream(messages: list[Message]) -> AgentExecutorResponse: + return AgentExecutorResponse( + executor_id="source", + agent_response=AgentResponse(messages=messages[-1:]), + full_conversation=list(messages), + ) + + +def _context() -> tuple[AzureFunctionsWorkflowContext, Mock, list[AtomicTask]]: + host = Mock(spec=df.DurableOrchestrationContext) + host.instance_id = "dispatch-revision-run" + host.is_replaying = False + host.current_utc_datetime = datetime(2026, 1, 1, tzinfo=timezone.utc) + host.new_uuid.side_effect = [str(UUID(int=index + 1)) for index in range(2)] + children = [AtomicTask(index + 1, NoOpAction()) for index in range(2)] + host.call_entity.side_effect = children + return AzureFunctionsWorkflowContext(host), host, children + + +def _dispatch( + context: AzureFunctionsWorkflowContext, + host: Mock, + executor: AgentExecutor, + message: Any, + ledger: _WorkflowDeliveryLedger, +) -> tuple[AgentTask, dict[str, Any]]: + task = _prepare_agent_task(context, executor, executor.id, message, "dispatch-revision", ledger) + assert isinstance(task, AgentTask) + assert not task.is_completed + entity_id, operation, payload = host.call_entity.call_args.args + assert entity_id.name == "dafx-dispatch-revision-target" + assert entity_id.key == context.instance_id + assert operation == "run" + # Capture the real executor's serialized RunRequest after build_agent_task and the shim. + wire = json.loads(json.dumps(payload, allow_nan=False)) + assert wire["orchestrationId"] == context.instance_id + assert wire["correlationId"] == str(UUID(int=host.call_entity.call_count)) + assert host.new_uuid.call_count == host.call_entity.call_count + if "contextMessages" in wire: + assert len(wire["contextMessageIds"]) == len(wire["contextMessages"]) + assert all(isinstance(identity, str) and identity for identity in wire["contextMessageIds"]) + else: + assert "contextMessageIds" not in wire + host.signal_entity.assert_not_called() + return task, wire + + +def test_custom_empty_projection_reaches_the_af_entity_as_an_empty_list() -> None: + context, host, _ = _context() + executor = _agent(context_mode="custom", context_filter=lambda messages: []) + excluded = Message("assistant", ["unselected secret" * 1000], message_id="wf_source_0") + ledger = _WorkflowDeliveryLedger() + + _, wire = _dispatch(context, host, executor, _upstream([excluded]), ledger) + + assert wire["message"] == "" + assert wire["contextMessages"] == [] + assert wire["contextMessageIds"] == [] + assert "unselected secret" not in json.dumps(wire) + assert DurableAgentStateRequest.from_run_request(RunRequest.from_dict(wire)).messages == [] + assert ledger.sent == {} + assert ledger.handoffs == {"target": 1} + host.call_entity.assert_called_once() + + +def test_fully_duplicate_projection_reaches_the_af_entity_on_the_second_call() -> None: + context, host, _ = _context() + executor = _agent() + messages = [ + Message("user", ["question"], message_id="wf_source_0"), + Message("assistant", ["answer"], message_id="wf_source_1"), + ] + upstream = _upstream(messages) + expected = [message.to_dict() for message in messages] + ledger = _WorkflowDeliveryLedger() + + _, first = _dispatch(context, host, executor, upstream, ledger) + assert first["contextMessages"] == expected + assert len(set(first["contextMessageIds"])) == 2 + assert set(first["contextMessageIds"]).isdisjoint(message.message_id for message in messages) + _, repeated = _dispatch(context, host, executor, upstream, ledger) + + assert repeated["contextMessages"] == [] + assert repeated["contextMessageIds"] == [] + assert repeated["message"] == "" + assert first["correlationId"] != repeated["correlationId"] + assert DurableAgentStateRequest.from_run_request(RunRequest.from_dict(repeated)).messages == [] + assert len(ledger.sent["target"]) == 2 + assert ledger.handoffs == {"target": 2} + assert [message.to_dict() for message in messages] == expected + assert host.call_entity.call_count == 2 + + +def test_tool_only_projection_survives_af_dispatch_and_request_parsing() -> None: + context, host, children = _context() + result = {"type": "lookup_result", "items": [{"answer": 0, "label": "世界"}], "flags": [False, None]} + message = Message( + "tool", + [Content.from_function_result("lookup-1", result=result)], + message_id="wf_source_0", + author_name="lookup", + additional_properties={"provider": {"type": "context", "labels": []}}, + ) + expected = message.to_dict() + ledger = _WorkflowDeliveryLedger() + + task, wire = _dispatch(context, host, _agent(), _upstream([message]), ledger) + + assert wire["message"] == "" + assert wire["contextMessages"] == [expected] + request = RunRequest.from_json(json.dumps(wire)) + assert request.context_message_ids == wire["contextMessageIds"] + assert request.context_message_ids != [message.message_id] + entry = DurableAgentStateRequest.from_run_request(request) + assert len(entry.messages) == 1 + forwarded = entry.messages[0].to_chat_message() + assert isinstance(forwarded, Message) + assert forwarded.role == "tool" + assert forwarded.message_id == message.message_id + assert forwarded.text == "" + assert len(forwarded.contents) == 1 + assert forwarded.contents[0].type == "function_result" + assert forwarded.contents[0].call_id == "lookup-1" + assert forwarded.contents[0].result == message.contents[0].result + assert json.loads(forwarded.contents[0].result) == result + assert message.to_dict() == expected + + assert not children[0].is_completed + children[0].set_value(is_error=False, value=AgentResponse(messages=[Message("assistant", ["received"])]).to_dict()) + assert task.state == TaskState.SUCCEEDED + assert context.get_task_result(task).text == "received" + + +def test_af_adapter_does_not_preprocess_or_drop_raw_context_type_fields() -> None: + context, host, _ = _context() + context_messages = [ + { + "type": "message", + "role": "tool", + "message_id": "wf_source_0", + "contents": [ + { + "type": "function_result", + "call_id": "lookup-1", + "result": {"type": "application_payload", "items": [0, False, None, "世界"]}, + "future_content_field": {"type": "opaque", "items": []}, + }, + ], + "future_message_field": {"type": "opaque", "items": []}, + }, + ] + before = deepcopy(context_messages) + + task = context.prepare_agent_task( + "dispatch-revision-target", "", context.instance_id, context_messages, context_message_ids=["occurrence-0"] + ) + + assert isinstance(task, AgentTask) + assert not task.is_completed + host.call_entity.assert_called_once() + wire = json.loads(json.dumps(host.call_entity.call_args.args[2], allow_nan=False)) + assert wire["message"] == "" + assert wire["contextMessages"] == before + assert wire["contextMessageIds"] == ["occurrence-0"] + assert RunRequest.from_dict(wire).context_messages == before + assert RunRequest.from_dict(wire).context_message_ids == ["occurrence-0"] + assert context_messages == before + + +def test_standalone_af_input_is_not_truncated_or_deduplicated() -> None: + context, host, _ = _context() + executor = _agent() + ledger = _WorkflowDeliveryLedger() + prompt = "standalone input " * 1000 + + for _ in range(2): + _, wire = _dispatch(context, host, executor, prompt, ledger) + assert wire["message"] == prompt + assert "contextMessages" not in wire + assert RunRequest.from_dict(wire).context_messages is None + + assert ledger.sent == {} + assert host.call_entity.call_count == 2 + + +def test_empty_workflow_input_schedules_an_explicit_empty_user_message() -> None: + context, host, _ = _context() + ledger = _WorkflowDeliveryLedger() + + _, wire = _dispatch(context, host, _agent(), "", ledger) + + assert wire["message"] == "" + assert wire["contextMessages"] == [Message("user", [""]).to_dict()] + assert len(wire["contextMessageIds"]) == 1 + assert RunRequest.from_dict(wire).context_message_ids == wire["contextMessageIds"] + host.call_entity.assert_called_once() + host.new_uuid.assert_called_once() diff --git a/python/packages/azurefunctions/tests/test_workflow_output_boundaries_review_af.py b/python/packages/azurefunctions/tests/test_workflow_output_boundaries_review_af.py new file mode 100644 index 0000000..1857aa3 --- /dev/null +++ b/python/packages/azurefunctions/tests/test_workflow_output_boundaries_review_af.py @@ -0,0 +1,220 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Registered AF output/status boundaries preserve generated response JSON values.""" + +from __future__ import annotations + +import json +from collections.abc import Callable +from copy import deepcopy +from datetime import date, datetime, timezone +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, Mock, patch +from uuid import UUID + +import azure.durable_functions as df +import azure.functions as func +import pytest +from agent_framework import AgentExecutor, AgentResponse, AgentSession, Message, WorkflowBuilder, WorkflowExecutor +from agent_framework._workflows import _checkpoint_encoding +from agent_framework_durabletask import load_agent_response, serialize_agent_response +from agent_framework_durabletask._workflows.protocol import wrap_workflow_input +from azure.durable_functions.models.actions.NoOpAction import NoOpAction +from azure.durable_functions.models.ReplaySchema import ReplaySchema +from azure.durable_functions.models.Task import AtomicTask, WhenAllTask +from pydantic import BaseModel, Field + +from agent_framework_azurefunctions import AgentFunctionApp +from agent_framework_azurefunctions._app import _json_default + + +class _Agent: + id = name = "A" + description = None + + def __init__(self, response_format: type[BaseModel] | None) -> None: + self.default_options = {"response_format": response_format} + + def create_session(self, **kwargs: Any) -> AgentSession: + return AgentSession(**kwargs) + + async def run(self, *args: Any, **kwargs: Any) -> AgentResponse: + raise AssertionError("Entity completion is supplied at the native task boundary") + + +def _response_case(kind: str) -> tuple[AgentResponse, type[BaseModel] | None, Any, bool]: + class AliasAnswer(BaseModel): + answer: bool = Field(alias="wireAnswer") + day: date + + class ByNameAnswer(BaseModel): + answer: bool = Field(validation_alias="inputAnswer", serialization_alias="outputAnswer") + day: date + + model: type[BaseModel] | None = None + expected: Any = False if kind == "false" else None + value = expected + if kind == "alias": + model = AliasAnswer + value = AliasAnswer(wireAnswer=False, day=date(2026, 9, 9)) + expected = {"wireAnswer": False, "day": "2026-09-09"} + elif kind == "by_name": + model = ByNameAnswer + value = ByNameAnswer(inputAnswer=False, day=date(2026, 9, 9)) + expected = {"answer": False, "day": "2026-09-09"} + response = AgentResponse( + messages=[Message("assistant", ["not structured text"], message_id="message-id")], + response_id="response-id", + value=value, + additional_properties={"flag": False, "nullable": None}, + ) + if kind == "null": + response = load_agent_response({**response.to_dict(), "value": None}) + return response, model, expected, kind == "by_name" + + +def _register(model: type[BaseModel] | None, nested: bool) -> tuple[dict[str, Any], dict[str, Any]]: + agent: Any = _Agent(model) + workflow = WorkflowBuilder(name="inner" if nested else "portable", start_executor=AgentExecutor(agent)).build() + if nested: + child = WorkflowExecutor(workflow, id="child", allow_direct_output=True) + workflow = WorkflowBuilder(name="portable", start_executor=child, output_from=[child]).build() + app = AgentFunctionApp(workflow=workflow, enable_health_check=False, deployment_mode="isolated_v2") + orchestrators: dict[str, Any] = {} + routes: dict[str, Any] = {} + for function in app.get_functions(): + trigger = function.get_trigger() + assert trigger is not None + binding = trigger.get_dict_repr() + user_function: Any = function.get_user_function() + if binding["type"] == "orchestrationTrigger": + name = function.get_function_name() + assert name is not None + orchestrators[name] = user_function.orchestrator_function + elif binding["type"] == "httpTrigger": + routes[binding["route"]] = user_function.client_function + return orchestrators, routes + + +def _run_registered(orchestrators: dict[str, Any], response: AgentResponse, nested: bool) -> tuple[Any, list[Any]]: + statuses: list[Any] = [] + + def complete(value: Any) -> AtomicTask: + task = AtomicTask(0, NoOpAction()) + task.set_value(is_error=False, value=json.loads(json.dumps(value, allow_nan=False))) + return task + + def run(name: str, input_data: Any, instance_id: str) -> Any: + host = Mock(spec=df.DurableOrchestrationContext) + host.instance_id = instance_id + host.is_replaying = False + host.current_utc_datetime = datetime(2026, 9, 9, tzinfo=timezone.utc) + host.new_uuid.side_effect = [str(UUID(int=i)) for i in range(1, 10)] + host.get_input.return_value = input_data + host.call_entity.side_effect = lambda *args: complete(serialize_agent_response(response)) + host.call_sub_orchestrator.side_effect = lambda name, *, input_, instance_id: complete( + run(name, input_, instance_id) + ) + host.task_all.side_effect = lambda tasks: WhenAllTask(tasks, ReplaySchema.V1) + host.set_custom_status.side_effect = lambda status: statuses.append(deepcopy(status)) + generator = orchestrators[name](host) + value = None + while True: + try: + task = generator.send(value) + except StopIteration as completed: + host.call_activity.assert_not_called() + if name == "dafx-portable" and nested: + host.call_sub_orchestrator.assert_called_once() + else: + host.call_entity.assert_called_once() + return json.loads(json.dumps(completed.value, allow_nan=False)) + assert task.is_completed + value = task.result + + return run("dafx-portable", wrap_workflow_input("question"), "output-run"), statuses + + +@pytest.mark.parametrize("endpoint", ["status", "wait"]) +@pytest.mark.parametrize("nested", [False, True]) +@pytest.mark.parametrize("kind", ["false", "null", "alias", "by_name"]) +async def test_registered_status_and_terminal_run_return_generated_response_value( + endpoint: str, nested: bool, kind: str +) -> None: + response, model, expected, by_name = _response_case(kind) + orchestrators, routes = _register(model, nested) + with patch.object(_checkpoint_encoding, "_pickle_to_base64", side_effect=AssertionError("No worker pickle")): + raw, statuses = _run_registered(orchestrators, response, nested) + assert all("events" not in status for status in statuses) + assert len(raw) == 1 and raw[0]["_durable_agent_response"] == 1 + assert "__pickled__" not in json.dumps(raw) + client = AsyncMock(spec=df.DurableOrchestrationClient) + client.start_new.return_value = "output-run" + client.wait_for_completion_or_create_check_status_response.return_value = func.HttpResponse(status_code=200) + client.get_status.return_value = SimpleNamespace( + name="dafx-portable", + instance_id="output-run", + runtime_status=df.OrchestrationRuntimeStatus.Completed, + output=raw, + custom_status=statuses[-1], + created_time=None, + last_updated_time=None, + ) + route = "workflow/portable/status/{instanceId}" if endpoint == "status" else "workflow/portable/run" + request = func.HttpRequest( + method="GET" if endpoint == "status" else "POST", + url="https://example.test/api/" + route.replace("{instanceId}", "output-run"), + headers={"Content-Type": "application/json"}, + params={} if endpoint == "status" else {"waitForResponse": "true", "runId": "output-run"}, + route_params={"instanceId": "output-run"} if endpoint == "status" else {}, + body=b'"question"', + ) + handler: Callable[..., Any] = routes[route] + with ( + patch("importlib.import_module", side_effect=AssertionError("HTTP readers do not import worker models")), + patch.object(_checkpoint_encoding, "_base64_to_unpickle", side_effect=AssertionError("No response pickle")), + ): + http_response = await handler(request, client) + assert http_response.status_code == 200 + body = json.loads(http_response.get_body()) + assert body["runtimeStatus"] == "Completed" + assert len(body["output"]) == 1 + delivered = body["output"][0] + assert delivered == raw[0]["response"] + assert "value" in delivered and delivered["value"] == expected + assert type(delivered["value"]) is type(expected) + assert delivered.get("_durable_value_by_name", False) is by_name + assert delivered["response_id"] == "response-id" + assert delivered["additional_properties"] == {"flag": False, "nullable": None} + assert delivered["messages"][0]["message_id"] == "message-id" + client.get_status.assert_awaited_once_with("output-run") + if endpoint == "wait": + client.wait_for_completion_or_create_check_status_response.assert_awaited_once() + else: + client.start_new.assert_not_awaited() + + +@pytest.mark.parametrize("kind", ["false", "null", "alias", "by_name"]) +def test_json_default_uses_base_response_serializer_before_overridden_to_dict(kind: str) -> None: + response, _, expected, by_name = _response_case(kind) + + class WorkerResponse(AgentResponse): + def to_dict(self, **kwargs: Any) -> dict[str, Any]: + raise AssertionError("Provider override must not replace the durable response contract") + + subclass = WorkerResponse( + messages=response.messages, + value=load_agent_response(serialize_agent_response(response)).value, + response_id=response.response_id, + additional_properties=response.additional_properties, + ) + if kind == "null": + subclass._value_parsed = True + if by_name: + typed_snapshot: Any = subclass + typed_snapshot._durable_value_by_name = True + encoded = _json_default(subclass) + assert "value" in encoded and encoded["value"] == expected + assert encoded.get("_durable_value_by_name", False) is by_name + assert encoded["additional_properties"] == {"flag": False, "nullable": None} diff --git a/python/packages/azurefunctions/tests/test_workflow_protocol_review_af.py b/python/packages/azurefunctions/tests/test_workflow_protocol_review_af.py new file mode 100644 index 0000000..7195856 --- /dev/null +++ b/python/packages/azurefunctions/tests/test_workflow_protocol_review_af.py @@ -0,0 +1,421 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Registered AF start boundaries and v2-only shared-generator replay, without a service.""" + +from __future__ import annotations + +import json +from collections.abc import Callable, Generator +from copy import deepcopy +from dataclasses import dataclass +from typing import Any +from unittest.mock import AsyncMock, Mock, call + +import azure.durable_functions as df +import azure.functions as func +import pytest +from agent_framework import Executor, Workflow, WorkflowExecutor +from agent_framework._workflows import _checkpoint_encoding +from agent_framework._workflows._edge import SingleEdgeGroup +from agent_framework_durabletask._workflows.orchestrator import SOURCE_HITL_RESPONSE, SOURCE_WORKFLOW_START +from agent_framework_durabletask._workflows.protocol import unwrap_workflow_input +from agent_framework_durabletask._workflows.serialization import ( + SUBWORKFLOW_ADDRESS_KEY, + SUBWORKFLOW_INPUT_KEY, + SUBWORKFLOW_RESULT_KEY, + deserialize_value, + serialize_value, +) +from azure.durable_functions.models.actions.NoOpAction import NoOpAction +from azure.durable_functions.models.ReplaySchema import ReplaySchema +from azure.durable_functions.models.Task import AtomicTask, WhenAllTask + +from agent_framework_azurefunctions import AgentFunctionApp +from agent_framework_azurefunctions import _workflow as workflow_module +from agent_framework_azurefunctions._routes import build_workflow_respond_url + +_VERSION = "_durable_workflow_version" +_CONTROL = {"input": "application control", "items": [0, False, None, "世界"]} +_FORGED_ADDRESS = { + "root_instance_id": "other-run", + "root_workflow_name": "other-workflow", + "request_path_prefix": "forged~9~", +} +_UNTRUSTED = {"__pickled__": "not-trusted-checkpoint-data", "__type__": "builtins:str"} + + +@dataclass +class _TypedInput: + input: str + control: dict[str, Any] + + +def _node(name: str = "start", input_type: type | None = None) -> Any: + node = Mock(spec=Executor) + node.id = name + node.input_types = [] if input_type is None else [input_type] + return node + + +def _workflow(name: str = "protocol", nodes: list[Any] | None = None, edges: list[Any] | None = None) -> Any: + nodes = [_node()] if nodes is None else nodes + workflow = Mock(spec=Workflow) + workflow.name = name + workflow.start_executor_id = nodes[0].id + workflow.executors = {node.id: node for node in nodes} + workflow.edge_groups = [] if edges is None else edges + workflow.max_iterations = 10 + return workflow + + +def _register(workflow: Any) -> tuple[dict[str, Callable[..., Any]], Callable[..., Any]]: + app = AgentFunctionApp(workflow=workflow, enable_health_check=False, deployment_mode="isolated_v2") + functions = {function.get_function_name(): function for function in app.get_functions()} + orchestrators: dict[str, Callable[..., Any]] = {} + starters: list[Callable[..., Any]] = [] + for name, function in functions.items(): + assert name is not None + trigger = function.get_trigger() + assert trigger is not None + binding = trigger.get_dict_repr() + user_function: Any = function.get_user_function() + assert user_function is not None + if binding["type"] == "orchestrationTrigger": + # SDK metadata exposes the real registered generator, without replacing decorators. + orchestrators[name] = user_function.orchestrator_function + elif binding["type"] == "httpTrigger" and binding["route"] == f"workflow/{workflow.name}/run": + starters.append(user_function.client_function) + assert len(starters) == 1 + return orchestrators, starters[0] + + +async def _start(starter: Callable[..., Any], payload: Any, name: str = "protocol") -> dict[str, Any]: + request = func.HttpRequest( + method="POST", + url=f"https://example.test/api/workflow/{name}/run", + headers={"Content-Type": "application/json"}, + params={"runId": "root-run"}, + body=json.dumps(payload, allow_nan=False).encode("utf-8"), + ) + client = AsyncMock(spec=df.DurableOrchestrationClient) + client.start_new.return_value = "root-run" + response = await starter(request, client) + assert response.status_code == 202 + client.start_new.assert_awaited_once() + invocation = client.start_new.await_args + assert invocation is not None + assert invocation.args == (f"dafx-{name}",) + assert invocation.kwargs["instance_id"] == "root-run" + return json.loads(json.dumps(invocation.kwargs["client_input"], allow_nan=False)) + + +def _complete(value: Any) -> AtomicTask: + task = AtomicTask(0, NoOpAction()) + task.set_value(is_error=False, value=value) + return task + + +def _drain(generator: Generator[Any, Any, Any], value: Any = None) -> Any: + while True: + try: + task = generator.send(value) + except StopIteration as completed: + return completed.value + assert task.is_completed, "Use explicit event completion for a paused generator" + value = task.result + + +def _host( + wire: Any, + calls: list[dict[str, Any]], + result: Callable[[str, dict[str, Any]], dict[str, Any]] | None = None, + *, + functions: dict[str, Callable[..., Any]] | None = None, + instance_id: str = "root-run", + replay: bool = False, +) -> Mock: + host = Mock(spec=df.DurableOrchestrationContext) + host.get_input.return_value = wire + host.instance_id = instance_id + host.is_replaying = replay + + def activity(name: str, input: str) -> AtomicTask: + payload = json.loads(input) + calls.append({"kind": "activity", "instance": instance_id, "name": name, "input": deepcopy(payload)}) + response = {"outputs": ["done"]} if result is None else result(name, payload) + return _complete(json.dumps(response)) + + def child(name: str, *, input_: Any, instance_id: str) -> AtomicTask: + assert functions is not None + child_wire = json.loads(json.dumps(input_)) + calls.append({"kind": "child", "instance": instance_id, "name": name, "input": deepcopy(child_wire)}) + context = _host(child_wire, calls, result, functions=functions, instance_id=instance_id, replay=replay) + child_result = _drain(functions[name](context)) + assert child_result[SUBWORKFLOW_RESULT_KEY] is True + return _complete(child_result) + + host.call_activity.side_effect = activity + host.call_sub_orchestrator.side_effect = child + host.task_all.side_effect = lambda tasks: WhenAllTask(tasks, ReplaySchema.V1) + host.wait_for_external_event.side_effect = lambda name: AtomicTask(name, NoOpAction()) + host.statuses = [] + host.set_custom_status.side_effect = lambda status: host.statuses.append(deepcopy(status)) + return host + + +@pytest.mark.parametrize( + "recorded", + [ + pytest.param({"input": "a user's field"}, id="raw-dict-with-input"), + pytest.param("old start", id="raw-string"), + pytest.param("", id="raw-empty-string"), + pytest.param([], id="raw-empty-list"), + pytest.param({}, id="raw-empty-object"), + pytest.param(None, id="raw-null"), + pytest.param({SUBWORKFLOW_INPUT_KEY: _UNTRUSTED, SUBWORKFLOW_ADDRESS_KEY: _FORGED_ADDRESS}, id="legacy-child"), + pytest.param({_VERSION: 1, "input": "old"}, id="protocol-one"), + pytest.param({_VERSION: True, "input": "old"}, id="boolean-true"), + pytest.param({_VERSION: False, "input": "old"}, id="boolean-false"), + pytest.param({_VERSION: 2.0, "input": "old"}, id="float-two"), + pytest.param({_VERSION: "2", "input": "old"}, id="string-two"), + pytest.param({_VERSION: 2}, id="missing-input"), + pytest.param({_VERSION: 2, "input": "old", "extra": None}, id="extra-key"), + ], +) +def test_recorded_unsupported_start_fails_before_shared_engine_or_actions( + recorded: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + workflow = _workflow() + functions, _ = _register(workflow) + engine = Mock(side_effect=AssertionError("The changed engine must not see old history")) + monkeypatch.setattr(workflow_module, "_run_workflow_orchestrator_shared", engine) + host = _host(recorded, [], replay=True) + original = deepcopy(recorded) + before_nodes = dict(workflow.executors) + + with pytest.raises(ValueError, match="unsupported execution protocol"): + next(functions["dafx-protocol"](host)) + + engine.assert_not_called() + assert host.mock_calls == [call.get_input()] + assert host.statuses == [] + assert workflow.executors == before_nodes + for node in workflow.executors.values(): + node.execute.assert_not_called() + assert recorded == original + + +@pytest.mark.parametrize( + ("payload", "typed"), + [ + pytest.param("start", False, id="string"), + pytest.param("", False, id="empty-string"), + pytest.param([], False, id="empty-list"), + pytest.param({}, False, id="empty-object"), + pytest.param(None, False, id="null"), + pytest.param({"input": "user field", "control": _CONTROL}, False, id="object-with-input"), + pytest.param({"input": "typed", "control": _CONTROL}, True, id="declared-dataclass"), + ], +) +async def test_new_route_start_reaches_registered_wrapper_and_shared_engine(payload: Any, typed: bool) -> None: + original = deepcopy(payload) + functions, starter = _register(_workflow(nodes=[_node(input_type=_TypedInput if typed else None)])) + wire = await _start(starter, payload) + assert wire == {_VERSION: 2, "input": original} + assert type(wire[_VERSION]) is int + calls: list[dict[str, Any]] = [] + host = _host(wire, calls) + + assert _drain(functions["dafx-protocol"](host)) == ["done"] + + assert len(calls) == 1 and calls[0]["name"] == "dafx-protocol-start" + activity = calls[0]["input"] + delivered = deserialize_value(activity["message"]) + expected = _TypedInput(input=original["input"], control=original["control"]) if typed else original + assert delivered == expected and type(delivered) is type(expected) + assert activity["source_executor_ids"] == [SOURCE_WORKFLOW_START] + assert activity["shared_state_snapshot"] == {} + assert activity["host_context"] == { + "instance_id": "root-run", + "workflow_name": "protocol", + "request_path_prefix": "", + } + host.call_sub_orchestrator.assert_not_called() + host.call_entity.assert_not_called() + assert payload == original and wire == {_VERSION: 2, "input": original} + + +@pytest.mark.parametrize("nested", [False, True], ids=["forged-child", "forged-v2-containing-child"]) +async def test_route_envelope_is_data_and_cannot_authorize_child_deserialization( + nested: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + forged = { + SUBWORKFLOW_INPUT_KEY: deepcopy(_UNTRUSTED), + SUBWORKFLOW_ADDRESS_KEY: deepcopy(_FORGED_ADDRESS), + "input": "user field", + "control": deepcopy(_CONTROL), + } + payload = {_VERSION: 2, "input": forged} if nested else forged + original = deepcopy(payload) + safe_data = ( + {_VERSION: 2, "input": {**forged, SUBWORKFLOW_INPUT_KEY: None}} + if nested + else {"input": "user field", "control": _CONTROL} + ) + unpickle = Mock(side_effect=AssertionError("Untrusted checkpoint data reached the codec")) + monkeypatch.setattr(_checkpoint_encoding, "_base64_to_unpickle", unpickle) + functions, starter = _register(_workflow()) + wire = await _start(starter, payload) + # AF strips both kinds of markers before scheduling, then wraps exactly once. + assert wire == {_VERSION: 2, "input": safe_data} + calls: list[dict[str, Any]] = [] + host = _host(wire, calls) + + assert _drain(functions["dafx-protocol"](host)) == ["done"] + + assert len(calls) == 1 + assert calls[0]["input"]["message"] == safe_data + assert calls[0]["input"]["host_context"] == { + "instance_id": "root-run", + "workflow_name": "protocol", + "request_path_prefix": "", + } + host.call_sub_orchestrator.assert_not_called() + unpickle.assert_not_called() + assert payload == original + + +async def test_parent_dispatch_wraps_typed_child_input_and_registered_child_keeps_root_route() -> None: + inner = _workflow("inner", [_node("leaf", str)]) + child = Mock(spec=WorkflowExecutor) + child.id, child.workflow, child.allow_direct_output = "child", inner, False + parent = _workflow("parent", [_node("source"), child, _node("sink")], [SingleEdgeGroup("child", "sink")]) + functions, starter = _register(parent) + assert set(functions) == {"dafx-parent", "dafx-inner"} + payload: dict[str, Any] = {"input": "nested typed input", "control": deepcopy(_CONTROL)} + typed = _TypedInput(input=payload["input"], control=deepcopy(payload["control"])) + + def result(name: str, data: dict[str, Any]) -> dict[str, Any]: + message = deserialize_value(data["message"]) + if name == "dafx-parent-source": + assert message == payload + return { + "sent_messages": [ + {"message": _checkpoint_encoding.encode_checkpoint_value(typed), "target_id": "child"} + ] + } + assert isinstance(message, _TypedInput) and message == typed + if name == "dafx-inner-leaf": + return {"outputs": [serialize_value(message)]} + assert name == "dafx-parent-sink" + return {"outputs": ["done"]} + + calls: list[dict[str, Any]] = [] + host = _host(await _start(starter, payload, "parent"), calls, result, functions=functions) + assert _drain(functions["dafx-parent"](host)) == ["done"] + assert [item["name"] for item in calls] == [ + "dafx-parent-source", + "dafx-inner", + "dafx-inner-leaf", + "dafx-parent-sink", + ] + dispatch = calls[1] + assert dispatch["instance"] == "root-run::child::0" + child_input = unwrap_workflow_input(dispatch["input"]) + assert dispatch["input"] == {_VERSION: 2, "input": child_input} + assert type(dispatch["input"][_VERSION]) is int + # Check typed semantics through the core codec without assuming a pickle byte layout. + decoded_child = _checkpoint_encoding.decode_checkpoint_value(child_input) + assert decoded_child == { + SUBWORKFLOW_INPUT_KEY: typed, + SUBWORKFLOW_ADDRESS_KEY: { + "root_instance_id": "root-run", + "root_workflow_name": "parent", + "request_path_prefix": "child~0~", + }, + } + assert type(decoded_child[SUBWORKFLOW_INPUT_KEY]) is _TypedInput + assert type(deserialize_value(calls[2]["input"]["message"])) is _TypedInput + metadata = calls[2]["input"]["host_context"] + assert metadata == {"instance_id": "root-run", "workflow_name": "parent", "request_path_prefix": "child~0~"} + assert ( + build_workflow_respond_url( + "https://example.test", + metadata["workflow_name"], + metadata["instance_id"], + metadata["request_path_prefix"] + "approval", + prefix="api", + ) + == "https://example.test/api/workflow/parent/respond/root-run/child~0~approval" + ) + assert calls[2]["input"]["source_executor_ids"] == [SOURCE_WORKFLOW_START] + assert calls[3]["input"]["source_executor_ids"] == ["child"] + assert calls[0]["input"]["message"] == payload + + +async def test_v2_paused_hitl_replays_full_shared_generator_with_identical_dispatch_and_state() -> None: + """Cold generator replay of v2 only, not SDK history execution or old-history compatibility.""" + payload = {"input": "start", "control": deepcopy(_CONTROL)} + answer = {"input": "approved", "control": deepcopy(_CONTROL)} + + def result(name: str, data: dict[str, Any]) -> dict[str, Any]: + if data["source_executor_ids"] == [SOURCE_WORKFLOW_START]: + return { + "shared_state_updates": {"pending": payload}, + "pending_request_info_events": [ + { + "request_id": "approval", + "source_executor_id": "gate", + "data": payload, + "request_type": "builtins:dict", + "response_type": "builtins:dict", + } + ], + } + if name == "dafx-protocol-gate": + assert data["shared_state_snapshot"] == {"pending": payload} + assert deserialize_value(data["message"]) == { + "request_id": "approval", + "original_request": payload, + "response": answer, + "response_type": "builtins:dict", + } + return { + "shared_state_deletes": ["pending"], + "shared_state_updates": {"decision": answer}, + "sent_messages": [{"message": answer, "target_id": "sink"}], + } + assert name == "dafx-protocol-sink" + assert data["shared_state_snapshot"] == {"decision": answer} + assert data["message"] == answer + return {"outputs": ["done"]} + + _, starter = _register(_workflow(nodes=[_node("gate"), _node("sink")])) + wire = await _start(starter, payload) + executions = [] + for replay in (False, True): + functions, _ = _register(_workflow(nodes=[_node("gate"), _node("sink")])) + calls: list[dict[str, Any]] = [] + host = _host(deepcopy(wire), calls, result, replay=replay) + generator = functions["dafx-protocol"](host) + batch = next(generator) + assert batch.is_completed + waiting = generator.send(batch.result) + assert not waiting.is_completed and len(calls) == 1 + if not replay: + assert host.statuses[-1]["state"] == "waiting_for_human_input" + assert host.statuses[-1]["pending_requests"]["approval"]["data"] == payload + waiting.set_value(is_error=False, value=deepcopy(_UNTRUSTED)) + waiting_again = generator.send(waiting.result) + assert not waiting_again.is_completed and len(calls) == 1 + waiting_again.set_value(is_error=False, value=deepcopy(answer)) + assert _drain(generator, waiting_again.result) == ["done"] + assert [item.args[0] for item in host.wait_for_external_event.call_args_list] == ["approval", "approval"] + assert len(calls) == 3 + assert calls[1]["input"]["source_executor_ids"] == [f"{SOURCE_HITL_RESPONSE}_approval"] + assert calls[2]["input"]["source_executor_ids"] == ["gate"] + if replay: + host.set_custom_status.assert_not_called() + executions.append(calls) + assert executions[0] == executions[1] + assert wire == {_VERSION: 2, "input": payload} diff --git a/python/packages/durabletask/README.md b/python/packages/durabletask/README.md index 11885a3..902e665 100644 --- a/python/packages/durabletask/README.md +++ b/python/packages/durabletask/README.md @@ -8,6 +8,67 @@ Please install this package via pip: pip install agent-framework-durabletask --pre ``` +Requires Python 3.10+, `agent-framework-core>=1.13.0,<2` and `pydantic>=2.11,<3`. +The full unit suite passed on Python 3.13/core 1.16, Python 3.13/core 1.13 and Python 3.10/core 1.16. +Pydantic 2.11 runtime validation remains blocked by dependency artifact downloads. +The offline lock, lint, typing and both package builds passed for this follow-up. +See [prototype validation](../../samples/README.md#prototype-validation) for recorded results and limitations. + +## Version 2 deployment warning + +The settings below describe the [PR #59 prototype](https://github.com/microsoft/agent-framework-durable-extension/pull/59), +not an approved design or the contents of a published package. Design review belongs in +[ADR PR #88](https://github.com/microsoft/agent-framework-durable-extension/pull/88). +The outcome, retention telemetry and validation follow-up builds on prototype baseline `9b4550d`. +It does not establish shared-schema acceptance or a released package contract. +After ADR approval, the agreed implementation will be submitted as stacked PRs rather than merged +from this prototype as-is. + +> **Breaking deployment and state contract.** `DurableAIAgentWorker` requires +> `deployment_mode="isolated_v2"`, or `DURABLE_AGENTS_DEPLOYMENT_MODE=isolated_v2` when the argument +> is omitted/`None`. This is operator acknowledgement, not a handshake, security boundary or proof +> of isolation. Use a separate hub/deployment with compatible workers and all clients. Keep old +> workers and workflow histories on the old engine. The current .NET reader rejects version 2. + +Only `schemaVersion="2.0.0"` is writable. Legacy `1.x.y` and supported later `2.x.y` state can be +read/round-tripped, but `run`, `reset` and `expire_responses` reject those layouts. No operation +silently upgrades legacy state. Rollback requires compatible version-2 workers, clients and workflow +protocol. Names are unchanged. Reusing an old `@name@key` on an empty new hub is not migration. + +A matching version label does not prove layout compatibility. The reader rejects known alternate +`data.terminalResults` or `data.completionReceipts` containers instead of treating their completed +requests as new work. Unrelated optional metadata remains opaque, including nested uses of those +names. This guard is not a general format detector or a conversion between proposed schemas. + +`DurableWorkflowClient` and internal child dispatch wrap new starts with workflow engine version 2. +Raw/legacy starts reject before revised actions execute. Native custom scheduling must use public +`wrap_workflow_input` for new instances. It does not authorize input or migrate old action histories. + +Both hosts expose privileged backend `AgentEntity.migrate`, supported by the pure +`migrate_legacy_state` helper. The request requires `source`, `sourceDigest`, `sourceSessionId`, +`destinationSessionId`, `migrationId` and `ownershipTransferId`, with optional `deliveryEvidence` +and `requireKnownOutcomes`. +Use an empty, separately addressed destination after quiescing and authorizing transfer from the +old owner. Nonempty scalar `ingestedPositions` requires a complete accepted-message journal, +including evicted inputs. `complete=True` is an operator assertion. Digest/max-position validation +does not prove authority or completeness, and no delivered prefix is inferred. Without that journal, +keep the old session on the old engine. + +Only recorded responses receive legacy completion backfill and a delivery grace window. Surviving +transcript payloads may be partial, so absence of error content does not prove success. Existing +original mailbox records keep their payload and expiry. If a matching receipt is missing, its +`completedAt` comes from the mailbox's `createdAt`, not migration time. `requireKnownOutcomes=True` +on the entity request, or `require_known_outcomes=True` on the helper, rejects imports without +trustworthy known outcomes. The default legacy-compatible path preserves unknown completion +evidence and duplicate suppression rather than inventing an outcome or rerunning completed work. +Whole-request digest idempotency prevents grace refresh after an exact retry, cold reload or +subsequent run. Migration retains the original logical session ID for external history and does +not copy that store or migrate workflow histories. No generated HTTP/MCP migration endpoint is +provided. These are prototype constraints, not an agreed +cross-runtime migration contract. See [ADR PR #88](https://github.com/microsoft/agent-framework-durable-extension/pull/88) +for the design discussion and [prototype validation](../../samples/README.md#prototype-validation) +for recorded checks and remaining gaps. + ## Durable Task Integration The durable task integration lets you host Microsoft Agent Framework agents using the [Durable Task](https://github.com/microsoft/durabletask-python) framework so they can persist state, replay conversation history, and recover from failures automatically. @@ -20,13 +81,172 @@ from agent_framework.openai import OpenAIChatCompletionClient from agent_framework_durabletask import DurableAIAgentWorker from durabletask.worker import TaskHubGrpcWorker -# Create the worker +# Connect only to the separately configured version-2 deployment worker = TaskHubGrpcWorker(host_address="localhost:4001") -agent_worker = DurableAIAgentWorker(worker) +agent_worker = DurableAIAgentWorker(worker, deployment_mode="isolated_v2") chat_client = OpenAIChatCompletionClient() my_agent = Agent(client=chat_client, name="assistant") agent_worker.add_agent(my_agent) ``` +### History and retention settings + +Registration appends durable history when no load-enabled primary exists, matching core's automatic +injection and reverse after-hook order. Only the exact built-in `InMemoryHistoryProvider` is replaced, +preserving `source_id`, `skip_excluded`, storage flags and `after_run_once_per_turn` when available. +Core 1.13 does not require that optional hint. Custom in-memory subclasses keep their hooks and +session transcripts. Their state is a protected floor, not managed by durable transcript eviction. +Custom durable-provider JSON state persists except the transient message buffer and position index. + +External primaries and store-only sinks keep their storage policies, subject to the intentional +service-branch restriction below. Multiple load-enabled primaries or duplicate `source_id` values +are rejected. Registration does not enable compaction. The provider owns appends through core hooks, +with a final durable flush after all after-run callbacks. Only agents without a context pipeline use +direct entity transcript appends. + +Eager pruning and pressure eviction are independent. The matrix assumes no explicit provider +`prune_excluded` override. + +| `retention` | `max_state_bytes=None` | Positive byte budget or `"backend_limit"` | +| --- | --- | --- | +| `"keep_all"` (default) | No transcript deletion (default) | Evict eligible oldest groups only under pressure | +| `"follow_compaction"` | Prune eligible compaction exclusions only | Prune exclusions, then evict under pressure if needed | + +- `max_state_bytes` defaults to `None`. `"backend_limit"` resolves to 1,048,576 bytes (1 MiB) only + for `DurableTaskSchedulerWorker`, not a generic `TaskHubGrpcWorker`. An unresolved limit is + rejected. A positive integer sets an application budget, not a larger backend limit. `"auto"` + is no longer a retention mode. +- `"backend_limit"` remains a non-normative Python-only convenience, outside the portable `None` + or positive-integer contract. It does not account for transport overhead or imply shared-review + agreement. +- Watermarks default to `high_watermark=0.85` and `low_watermark=0.70`, with + `0 < low_watermark < high_watermark <= 1`. The whole serialized entity counts, including mailbox, + completion, session and ingestion state. Protected data can prevent a commit even after pruning. +- `response_delivery_window_seconds` defaults to `60` and must be a positive integer. Delivery + expiry is independent of transcript retention. +- `add_agent()` and `configure_workflow()` accept overrides. Omitted budgets or `INHERIT` use the + worker default. Explicit `None` disables that inherited budget. Workflow settings apply to newly + registered agent nodes, including nested workflows. +- Explicit `prune_excluded=False` on `DurableHistoryProvider` disables eager pruning even with + `follow_compaction`. It does not disable pressure eviction. Neither retention control configures + an external store's retention policy. + +As an alternative to the default registration above, use an unregistered worker, `my_agent` and an +existing named `workflow` to set an explicit byte budget while disabling it for workflow nodes. + +```python +from agent_framework_durabletask import INHERIT + +agent_worker = DurableAIAgentWorker(worker, deployment_mode="isolated_v2", max_state_bytes=800_000) +agent_worker.add_agent(my_agent, retention="follow_compaction", max_state_bytes=INHERIT) +agent_worker.configure_workflow(workflow, max_state_bytes=None) +``` + +`follow_compaction` only prunes exclusions produced by configured compaction. Without a strategy, +there are no exclusions to prune. Workflow `full`, `last_agent` and `custom` projection runs before +per-target delta transport. Custom filters must be synchronous, deterministic and side-effect-free, +but need not select monotonically increasing positions. Parallel `contextMessageIds` carry occurrence +hashes without rewriting public message IDs. Private forwarding provenance stays in internal +checkpoints, not application metadata. Outgoing context retains the full selected conversation plus +all response messages, not only the delta. Typed/cache-only requests, agent approval/HITL and +output-designated agents use the same workflow contract. + +Generated agent outputs and intermediate events use portable response snapshots. External clients +receive base `AgentResponse` objects with JSON structured values, without importing worker-local +response models. Worker-side conditions and activities still receive the locally declared model. +Arbitrary custom activity outputs retain the existing checkpoint codec and its importable-type +requirements. Parent output designations also apply to direct child-workflow outputs. + +### Service ownership, delivery and reset + +Effective `store` follows run options, then agent defaults, then the client's `STORES_BY_DEFAULT`. +For example, `options={"store": False}` selects client-owned history even on a service-storing +client. Explicit `False` excludes saved or supplied service conversation IDs from that invocation +and its history hooks. A later service-owned run can reuse the saved service ID without importing +the intervening client-owned transcript. Switching branches does not migrate or merge history. +External and service-owned runs create no local request-message mirror. + +On service-owned runs, durable deliberately suppresses **both load and store hooks** on the inactive +external/custom primary, including per-service-call persistence. Core 1.16 can still save to a +configured primary on such runs. This restriction avoids mixing service/client branches, but is not +universal unchanged-hook parity. Use a distinct store-only sink with its own `source_id` to audit +both branches. Its configured storage flags still apply. + +`responseMailbox` holds independent original serializable response snapshots, including metadata +and structured `value`, rather than rebuilding results from the mutable transcript. New +`completedCorrelations` receipts retain `completedAt` and the invocation `outcome`, either +`succeeded` or `failed`. After delivery expiry, lookup returns `durable_status="already_completed"` +with `response_expired` and `additional_properties["durable_outcome"]` set to `succeeded`, `failed` +or `unknown`. An older timestamp-only receipt with no trustworthy outcome still prevents +reinvocation. Cleanup can backfill a known outcome from an independent original mailbox before +removing its payload, even if the delivery deadline has passed. Version-2 lookup never uses the +possibly pruned transcript to infer success or reconstruct a result. + +The standalone SDK API is unchanged. Retained original responses are returned unmodified rather +than having receipt metadata injected into their payloads. Acceptance alone is not completion. +A fresh `record_response()` with no known invocation outcome raises before changing either delivery +map. Legacy-compatible receipts and fire-and-forget acceptance behavior remain supported. An +approval response can complete its invocation without proving that the guarded action executed. + +Expiry is a logical deadline, not an idle timer. New runs, duplicates and reset remove expired +payloads. Both hosts also expose backend `expire_responses` without model/tool/provider execution. +Idle physical cleanup needs an application-owned schedule or explicit backend signal/manual +operation. No public HTTP/MCP cleanup endpoint is generated. Completion receipts are never removed +by expiry, cleanup or reset. + +Local reset clears session and transcript context but preserves live mailbox payloads, completion +receipts and ingestion evidence. Normal delivery expiry still applies. Reset with a non-durable +custom/external primary raises `NotImplementedError` until provider-owned clearing is available. + +Entity-local state commits once per operation. Only structured `previous_response_not_found` on a +service-owned run permits bounded retries, and only before a stream update, function execution or +service-session advancement. Otherwise fail without restarting the conversation. Provider-hook side +effects are not guaranteed safe or identical on retry. There is no generic non-streaming retry after +runtime failure. Only matching unsupported-stream `TypeError` before consumption negotiates fallback. +Final callbacks receive deep copies preserving Pydantic fields. Opaque SDK `raw_representation` +detachment is best effort and that field is omitted if it cannot be copied. +Uncommitted model/tool effects and external appends can repeat after failure. Completion receipts +last until entity deletion and can exhaust capacity. A bounded receipt protocol and optional +retry-safe external-history adapters remain deferred, with no mandatory core API changes or +guarantee of a distributed transaction or exactly-once uncommitted effects. + +### Retention telemetry + +Both Python hosts use OpenTelemetry scope `agent_framework.durabletask`. The package directly +depends only on `opentelemetry-api` for these instruments. The SDK is a development dependency, +and applications own their meter provider, readers and exporters. The runtime configures none. + +| Instrument | Kind | Unit | +| --- | --- | --- | +| `durable.retention.evaluations` | Counter | `{evaluation}` | +| `durable.retention.budget` | Histogram | `By` | +| `durable.retention.state.size` | Histogram | `By` | +| `durable.retention.removed_messages` | Counter | `{message}` | +| `durable.retention.removed_entries` | Counter | `{entry}` | +| `durable.retention.reclaimed_bytes` | Counter | `By` | +| `durable.retention.capacity_failures` | Counter | `{failure}` | +| `durable.retention.write_attempts` | Counter | `{attempt}` | +| `durable.retention.operations` | Counter | `{operation}` | + +Attributes are bounded and apply only to the relevant instruments. + +| Attribute | Values | +| --- | --- | +| `mechanism` | `eager`, `pressure` | +| `outcome` | Retention uses `below_threshold`, `staged`, `protected_floor`, `unreachable_target`, `protected`. Write/operation observations use `returned`, `failed`. | +| `commit_status` | `not_attempted`, `unknown` | +| `phase` | `before`, `after` | +| `stage` | `serialization`, `set_state` | +| `deletion_staged` | `true`, `false` | + +No payloads or session, request or message IDs are recorded in these metrics. The budget is the +resolved whole-entity budget, and sizes describe serialized JSON at the retention boundary. +Removal counts and nonnegative reclaimed bytes describe staged changes, not detached trial plans +or committed deletion. Serialization failure leaves commit status `not_attempted`. A host +`set_state` return or failure leaves it `unknown`, since either can follow a staged write without +confirming persistence. Separate authoritative persisted-state readback is needed, paired with +subsequent model input when validating retention. Metrics do not change warm-state rollback or +make external effects transactional. + For more details, review the standalone [Durable Task samples](https://github.com/microsoft/agent-framework-durable-extension/tree/main/python/samples) and the full [Agent Framework Python documentation](https://github.com/microsoft/agent-framework/tree/main/python). diff --git a/python/packages/durabletask/agent_framework_durabletask/__init__.py b/python/packages/durabletask/agent_framework_durabletask/__init__.py index a3e2727..952643b 100644 --- a/python/packages/durabletask/agent_framework_durabletask/__init__.py +++ b/python/packages/durabletask/agent_framework_durabletask/__init__.py @@ -10,6 +10,17 @@ from ._async_bridge import run_agent_coroutine from ._callbacks import AgentCallbackContext, AgentResponseCallbackProtocol from ._client import DurableAIAgentClient +from ._configuration import ( + INHERIT, + AgentRegistrationSettings, + Inherit, + RegistrationIdentity, + StateBudgetOverride, + resolve_state_budget_override, + validate_agent_configuration, + validate_response_delivery_window, + validate_runtime_deployment, +) from ._constants import ( DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS, @@ -28,12 +39,14 @@ ) from ._durable_agent_state import ( DurableAgentState, + DurableAgentStateCompaction, DurableAgentStateContent, DurableAgentStateData, DurableAgentStateDataContent, DurableAgentStateEntry, DurableAgentStateEntryJsonType, DurableAgentStateErrorContent, + DurableAgentStateErrorResponse, DurableAgentStateFunctionCallContent, DurableAgentStateFunctionResultContent, DurableAgentStateHostedFileContent, @@ -50,10 +63,30 @@ ) from ._entities import AgentEntity, AgentEntityStateProviderMixin from ._executors import DurableAgentExecutor +from ._history_provider import DurableHistoryBinding, DurableHistoryProvider, validate_history_providers from ._models import AgentSessionId, DurableAgentSession, RunRequest from ._orchestration_context import DurableAIAgentOrchestrationContext -from ._response_utils import ensure_response_format, load_agent_response -from ._shim import DurableAIAgent +from ._response_utils import ( + ensure_response_format, + is_terminal_agent_response, + load_agent_response, + serialize_agent_response, +) +from ._retention import ( + DEFAULT_MAX_STATE_BYTES, + DEFAULT_RETENTION, + DELIVERY_WINDOW_SECONDS, + DTS_MAX_STATE_BYTES, + HIGH_WATERMARK, + LOW_WATERMARK, + RetentionMode, + StateBudget, + StateCapacityError, + resolve_state_budget, + validate_retention, +) +from ._shim import DurableAIAgent, build_agent_task +from ._state_migration import migrate_legacy_state, state_snapshot_digest from ._worker import DurableAIAgentWorker from ._workflows.activity import execute_workflow_activity from ._workflows.client import DurableWorkflowClient @@ -68,6 +101,7 @@ workflow_orchestrator_name, ) from ._workflows.orchestrator import run_workflow_orchestrator +from ._workflows.protocol import WORKFLOW_ENGINE_VERSION, unwrap_workflow_input, wrap_workflow_input from ._workflows.registration import WorkflowRegistrationPlan, collect_hosted_workflows, plan_workflow_registration from ._workflows.runner_context import CapturingRunnerContext from ._workflows.serialization import deserialize_workflow_output @@ -112,9 +146,16 @@ def __dir__() -> list[str]: __all__ = [ "DEFAULT_MAX_POLL_RETRIES", + "DEFAULT_MAX_STATE_BYTES", "DEFAULT_POLL_INTERVAL_SECONDS", + "DEFAULT_RETENTION", + "DELIVERY_WINDOW_SECONDS", + "DTS_MAX_STATE_BYTES", "DURABLE_NAME_PREFIX", + "HIGH_WATERMARK", + "INHERIT", "LEGACY_THREAD_ID_FIELD", + "LOW_WATERMARK", "MIMETYPE_APPLICATION_JSON", "MIMETYPE_TEXT_PLAIN", "REQUEST_RESPONSE_FORMAT_JSON", @@ -125,9 +166,11 @@ def __dir__() -> list[str]: "THREAD_ID_HEADER", "WAIT_FOR_RESPONSE_FIELD", "WAIT_FOR_RESPONSE_HEADER", + "WORKFLOW_ENGINE_VERSION", "AgentCallbackContext", "AgentEntity", "AgentEntityStateProviderMixin", + "AgentRegistrationSettings", "AgentResponseCallbackProtocol", "AgentSessionId", "ApiResponseFields", @@ -140,12 +183,14 @@ def __dir__() -> list[str]: "DurableAgentExecutor", "DurableAgentSession", "DurableAgentState", + "DurableAgentStateCompaction", "DurableAgentStateContent", "DurableAgentStateData", "DurableAgentStateDataContent", "DurableAgentStateEntry", "DurableAgentStateEntryJsonType", "DurableAgentStateErrorContent", + "DurableAgentStateErrorResponse", "DurableAgentStateFunctionCallContent", "DurableAgentStateFunctionResultContent", "DurableAgentStateHostedFileContent", @@ -159,24 +204,46 @@ def __dir__() -> list[str]: "DurableAgentStateUriContent", "DurableAgentStateUsage", "DurableAgentStateUsageContent", + "DurableHistoryBinding", + "DurableHistoryProvider", "DurableStateFields", "DurableTaskWorkflowContext", "DurableWorkflowClient", + "Inherit", + "RegistrationIdentity", + "RetentionMode", "RunRequest", + "StateBudget", + "StateBudgetOverride", + "StateCapacityError", "WorkflowOrchestrationContext", "WorkflowRegistrationPlan", "__version__", + "build_agent_task", "collect_hosted_workflows", "deserialize_workflow_output", "ensure_response_format", "execute_workflow_activity", "is_auto_generated_workflow_name", + "is_terminal_agent_response", "load_agent_response", + "migrate_legacy_state", "plan_workflow_registration", + "resolve_state_budget", + "resolve_state_budget_override", "run_agent_coroutine", "run_workflow_orchestrator", + "serialize_agent_response", + "state_snapshot_digest", + "unwrap_workflow_input", + "validate_agent_configuration", "validate_executor_id", + "validate_history_providers", + "validate_response_delivery_window", + "validate_retention", + "validate_runtime_deployment", "validate_workflow_name", "workflow_name_from_orchestrator", "workflow_orchestrator_name", + "wrap_workflow_input", ] diff --git a/python/packages/durabletask/agent_framework_durabletask/_configuration.py b/python/packages/durabletask/agent_framework_durabletask/_configuration.py new file mode 100644 index 0000000..f5279ad --- /dev/null +++ b/python/packages/durabletask/agent_framework_durabletask/_configuration.py @@ -0,0 +1,162 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Shared, typed overrides for durable agent registration.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from enum import Enum +from typing import Final, Literal, TypeAlias + +from agent_framework import SupportsAgentRun + +from ._callbacks import AgentResponseCallbackProtocol +from ._history_provider import ensure_durable_history +from ._retention import DEFAULT_RETENTION, RetentionMode, StateBudget, resolve_state_budget, validate_retention + +__all__ = [ + "INHERIT", + "AgentRegistrationSettings", + "Inherit", + "RegistrationIdentity", + "StateBudgetOverride", + "resolve_state_budget_override", + "validate_agent_configuration", + "validate_response_delivery_window", + "validate_runtime_deployment", +] + + +def validate_runtime_deployment(deployment_mode: str | None = None) -> None: + """Require explicit acknowledgement of an isolated schema 2 deployment. + + Schema 2 requires an isolated task hub/deployment with upgraded clients. + Old workflow histories must remain on the old engine. This is an operator + acknowledgement, not runtime proof of isolation, and cannot detect peer workers. + + Args: + deployment_mode: Exactly ``isolated_v2``. Only when None, read + ``DURABLE_AGENTS_DEPLOYMENT_MODE`` instead. + + Raises: + ValueError: The deployment mode is missing or is not exactly ``isolated_v2``. + """ + effective_mode = os.getenv("DURABLE_AGENTS_DEPLOYMENT_MODE") if deployment_mode is None else deployment_mode + if not isinstance(effective_mode, str) or effective_mode != "isolated_v2": + raise ValueError( + "Schema 2 requires an isolated task hub/deployment with upgraded clients. " + "Old workflow histories must remain on the old engine. " + "Set deployment_mode='isolated_v2' or DURABLE_AGENTS_DEPLOYMENT_MODE='isolated_v2'; " + "no other deployment mode is accepted. This is an explicit operator acknowledgement, " + "not runtime proof of isolation, and cannot detect peer workers." + ) + + +class Inherit(Enum): + """Use the enclosing host's setting instead of an explicit override.""" + + INHERIT = "inherit" + + +INHERIT: Final[Inherit] = Inherit.INHERIT +"""Inherit the configured budget; unlike None, this does not disable pressure eviction.""" + +StateBudgetOverride: TypeAlias = StateBudget | Inherit + + +@dataclass(frozen=True) +class AgentRegistrationSettings: + """Resolved settings used to check whether a hosted registration can be reused.""" + + retention: RetentionMode + max_state_bytes: int | None + high_watermark: float + low_watermark: float + response_delivery_window_seconds: int + callback: AgentResponseCallbackProtocol | None = field(default=None, compare=False) + + def matches(self, other: AgentRegistrationSettings) -> bool: + """Compare values, but require the same callback instance.""" + return self == other and self.callback is other.callback + + +@dataclass(frozen=True) +class RegistrationIdentity: + """Ownership of one derived host name, independent of backend registration APIs.""" + + owner: object + target: object + kind: str + settings: AgentRegistrationSettings + label: str + endpoints: tuple[bool, bool] = (False, False) + + def reserve( + self, + registrations: dict[tuple[str, str], RegistrationIdentity], + name: str, + *, + namespace: Literal["entity-name", "activity-name", "orchestrator-name", "function-name"], + ) -> None: + """Reserve a case-insensitive name in its backend artifact namespace. + + Call on a temporary mapping during preflight. Publishing that mapping is the + host's responsibility, after all backend registrations have succeeded. + """ + key = (namespace, name.casefold()) + existing = registrations.get(key) + if existing is not None: + if ( + existing.owner is not self.owner + or existing.target is not self.target + or existing.kind != self.kind + or existing.label != self.label + ): + raise ValueError( + f"Derived name '{name}' for {self.label} collides with already registered " + f"{existing.label}. Names are compared case-insensitively; " + "different registrations must not share a durable identity." + ) + if not existing.settings.matches(self.settings) or existing.endpoints != self.endpoints: + raise ValueError( + f"'{name}' is already registered with different settings for {existing.label}; " + "shared registrations require identical configuration." + ) + return + registrations[key] = self + + +def validate_agent_configuration(agent: SupportsAgentRun, *, retention: RetentionMode = DEFAULT_RETENTION) -> None: + """Dry-prepare durable history, including copy/attachment validation. + + Discard the prepared view so the host's agent registry retains the caller's + original instance. Entity construction prepares its own view at invocation. + """ + validate_retention(retention) + try: + ensure_durable_history(agent, prune_excluded=retention == "follow_compaction") + except ValueError: + raise + except Exception as exc: + raise ValueError("Could not prepare the agent's durable history configuration.") from exc + + +def resolve_state_budget_override( + value: StateBudgetOverride, + default: int | None, + *, + backend_limit: int | None = None, +) -> int | None: + """Resolve an inherited or explicit budget without conflating None with omission.""" + return resolve_state_budget(default if isinstance(value, Inherit) else value, backend_limit=backend_limit) + + +def validate_response_delivery_window(response_delivery_window_seconds: int) -> None: + """Require a positive integer delivery window, excluding booleans and non-finite floats.""" + if ( + isinstance(response_delivery_window_seconds, bool) + or not isinstance(response_delivery_window_seconds, int) + or response_delivery_window_seconds <= 0 + ): + raise ValueError("response_delivery_window_seconds must be a positive integer, not a boolean or another type.") diff --git a/python/packages/durabletask/agent_framework_durabletask/_constants.py b/python/packages/durabletask/agent_framework_durabletask/_constants.py index 9e48b51..945b037 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_constants.py +++ b/python/packages/durabletask/agent_framework_durabletask/_constants.py @@ -131,6 +131,31 @@ class DurableStateFields: # History field CONVERSATION_HISTORY: Final[str] = "conversationHistory" + # Stable per-message identity (used for compaction reconciliation and idempotency) + MESSAGE_ID: Final[str] = "messageId" + + # Serialized AgentSession: the provider state bag plus any service-issued conversation id + SESSION: Final[str] = "session" + + # Legacy scalar cursors are read for migration, never inferred to be exact receipts. + INGESTED_POSITIONS: Final[str] = "ingestedPositions" + INGESTED_MESSAGES: Final[str] = "ingestedMessages" + + # Result delivery is independent from the model transcript. + RESPONSE_MAILBOX: Final[str] = "responseMailbox" + COMPLETED_CORRELATIONS: Final[str] = "completedCorrelations" + RESPONSE: Final[str] = "response" + EXPIRES_AT: Final[str] = "expiresAt" + COMPLETED_AT: Final[str] = "completedAt" + OUTCOME: Final[str] = "outcome" + + # What retention has removed from this conversation. Present only once something has been + # evicted, so its absence means the record is complete. + TRUNCATION: Final[str] = "truncation" + EVICTED_MESSAGE_COUNT: Final[str] = "evictedMessageCount" + FIRST_EVICTED_AT: Final[str] = "firstEvictedAt" + LAST_EVICTED_AT: Final[str] = "lastEvictedAt" + class ContentTypes: """Content type discriminator values for the $type field. diff --git a/python/packages/durabletask/agent_framework_durabletask/_durable_agent_state.py b/python/packages/durabletask/agent_framework_durabletask/_durable_agent_state.py index f1fb577..84ae2f0 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_durable_agent_state.py +++ b/python/packages/durabletask/agent_framework_durabletask/_durable_agent_state.py @@ -31,8 +31,10 @@ import json import logging +import re from collections.abc import MutableMapping -from datetime import datetime, timezone +from copy import deepcopy +from datetime import datetime, timedelta, timezone from enum import Enum from typing import Any, ClassVar, cast @@ -45,19 +47,129 @@ from dateutil import parser as date_parser from ._constants import ContentTypes, DurableStateFields +from ._message_identity import message_identity from ._models import RunRequest, serialize_response_format +from ._response_utils import invocation_outcome, load_agent_response, serialize_agent_response logger = logging.getLogger("agent_framework.durabletask") +def _validate_delivery_layout(data: dict[str, Any]) -> None: + """Reject known alternate completion authorities, even with the same version label. + + These top-level data fields describe a different proposed delivery contract. + Preserving them as extensions while treating their requests as incomplete would + permit duplicate execution. This is rejection, not migration or schema agreement. + Unrelated metadata, including nested occurrences of these names, stays opaque. + """ + if "terminalResults" in data or "completionReceipts" in data: + raise ValueError( + "The durable agent state contains an incompatible delivery layout. " + "This prototype requires responseMailbox/completedCorrelations semantics; " + "a matching schemaVersion does not authorize interpreting another completion format." + ) + + +def _validate_completion_outcomes(records: dict[str, dict[str, Any]]) -> None: + """Absent outcomes are old completion evidence, not permission to invent success.""" + for record in records.values(): + if not isinstance(record, dict): + raise ValueError("completedCorrelations must contain objects keyed by correlation ID.") + if DurableStateFields.OUTCOME in record and record[DurableStateFields.OUTCOME] not in ("succeeded", "failed"): + raise ValueError("completedCorrelations.outcome must be 'succeeded' or 'failed' when present.") + + +def _validate_json(value: Any) -> None: + """Reject non-JSON values before the encoder can normalize them or collide keys.""" + if isinstance(value, dict): + for key, item in cast(dict[Any, Any], value).items(): + if not isinstance(key, str): + raise ValueError("JSON object keys must be strings.") + _validate_json(item) + elif isinstance(value, list): + for item in cast(list[Any], value): + _validate_json(item) + elif value is not None and not isinstance(value, (str, bool, int, float)): + raise ValueError("Values must contain only JSON objects, arrays and primitives.") + + +def _json_snapshot(value: Any) -> Any: + """Detach strict JSON without normalizing non-string keys or non-JSON containers.""" + try: + _validate_json(value) + return json.loads(json.dumps(value, allow_nan=False)) + except (TypeError, ValueError, RecursionError) as exc: + raise ValueError("State must be strict JSON with string keys and finite numbers.") from exc + + +def _parse_delivery_timestamp(value: Any) -> datetime: + """Parse offset-bearing RFC 3339 timestamps, including Z on Python 3.10.""" + if not isinstance(value, str) or not re.fullmatch( + r"[0-9]{4}-[0-9]{2}-[0-9]{2}[Tt](?:[01][0-9]|2[0-3]):[0-5][0-9]:[0-5][0-9]" + r"(?:\.[0-9]+)?(?:[Zz]|[+-](?:[01][0-9]|2[0-3]):[0-5][0-9])", + value, + ): + raise ValueError("Delivery timestamps must be RFC 3339 strings with an explicit offset.") + return datetime.fromisoformat(value[:-1] + "+00:00" if value[-1:] in ("Z", "z") else value) + + +def _array_field(data: dict[str, Any], name: str) -> list[Any]: + value = data.get(name, []) + if not isinstance(value, list): + raise ValueError(f"{name} must be an array.") + return cast(list[Any], value) + + +def _validate_core_message(data: Any) -> None: + if not isinstance(data, dict): + raise ValueError("Core messages must be objects.") + for content in _array_field(cast(dict[str, Any], data), "contents"): + if not isinstance(content, dict): + raise ValueError("Core contents must be objects with a non-empty type.") + content_type = cast(dict[str, Any], content).get("type") + if not isinstance(content_type, str) or not content_type: + raise ValueError("Core contents must be objects with a non-empty type.") + + +def _entry_unknown_fields(entry: DurableAgentStateEntry, data: dict[str, Any]) -> dict[str, Any]: + known = { + DurableStateFields.TYPE_DISCRIMINATOR, + DurableStateFields.JSON_TYPE, + DurableStateFields.CORRELATION_ID, + DurableStateFields.CREATED_AT, + DurableStateFields.MESSAGES, + DurableStateFields.EXTENSION_DATA, + } + if isinstance(entry, DurableAgentStateRequest): + known.update(( + DurableStateFields.ORCHESTRATION_ID, + DurableStateFields.RESPONSE_TYPE, + DurableStateFields.RESPONSE_SCHEMA, + )) + elif isinstance(entry, DurableAgentStateResponse): + known.add(DurableStateFields.USAGE) + return {key: deepcopy(value) for key, value in data.items() if key not in known} + + class DurableAgentStateEntryJsonType(str, Enum): """Enum for conversation history entry types. Discriminator values for the $type field in DurableAgentStateEntry objects. + + The type is what decides who may read an entry, rather than a flag alongside it. A flag has to + survive serialization to mean anything, and one that did not was how a failed turn came back as + ordinary assistant context after a cold start. + + ``errorResponse`` and ``compaction`` are opposites. A failed turn is worth returning to the + caller that is waiting for it but must never be replayed to the model. A compaction summary is + the reverse: it belongs in the model's transcript and must never be handed back as something + the agent said. """ REQUEST = "request" RESPONSE = "response" + ERROR_RESPONSE = "errorResponse" + COMPACTION = "compaction" def _parse_created_at(value: Any) -> datetime: @@ -90,12 +202,14 @@ def _parse_messages(data: dict[str, Any]) -> list[DurableAgentStateMessage]: List of DurableAgentStateMessage objects """ messages: list[DurableAgentStateMessage] = [] - raw_messages: list[Any] = data.get(DurableStateFields.MESSAGES, []) + raw_messages = _array_field(data, DurableStateFields.MESSAGES) for raw_msg in raw_messages: if isinstance(raw_msg, dict): messages.append(DurableAgentStateMessage.from_dict(cast(dict[str, Any], raw_msg))) elif isinstance(raw_msg, DurableAgentStateMessage): messages.append(raw_msg) + else: + raise ValueError("messages must contain message objects.") return messages @@ -108,7 +222,7 @@ def _parse_history_entries(data_dict: dict[str, Any]) -> list[DurableAgentStateE Returns: List of DurableAgentStateEntry objects (requests and responses) """ - history_data: list[Any] = data_dict.get(DurableStateFields.CONVERSATION_HISTORY, []) + history_data = _array_field(data_dict, DurableStateFields.CONVERSATION_HISTORY) deserialized_history: list[DurableAgentStateEntry] = [] for raw_entry in history_data: if isinstance(raw_entry, dict): @@ -116,14 +230,24 @@ def _parse_history_entries(data_dict: dict[str, Any]) -> list[DurableAgentStateE entry_type = entry_dict.get(DurableStateFields.TYPE_DISCRIMINATOR) or entry_dict.get( DurableStateFields.JSON_TYPE ) + if not isinstance(entry_type, str) or not entry_type: + raise ValueError("Conversation entries require a non-empty type discriminator.") if entry_type == DurableAgentStateEntryJsonType.RESPONSE: deserialized_history.append(DurableAgentStateResponse.from_dict(entry_dict)) + elif entry_type == DurableAgentStateEntryJsonType.ERROR_RESPONSE: + deserialized_history.append(DurableAgentStateErrorResponse.from_dict(entry_dict)) + elif entry_type == DurableAgentStateEntryJsonType.COMPACTION: + deserialized_history.append(DurableAgentStateCompaction.from_dict(entry_dict)) elif entry_type == DurableAgentStateEntryJsonType.REQUEST: deserialized_history.append(DurableAgentStateRequest.from_dict(entry_dict)) else: - deserialized_history.append(DurableAgentStateEntry.from_dict(entry_dict)) + deserialized_history.append(DurableAgentStateUnknownEntry(entry_dict)) + entry = deserialized_history[-1] + entry.unknown_fields = _entry_unknown_fields(entry, entry_dict) elif isinstance(raw_entry, DurableAgentStateEntry): deserialized_history.append(raw_entry) + else: + raise ValueError("conversationHistory must contain entry objects.") return deserialized_history @@ -137,7 +261,7 @@ def _parse_contents(data: dict[str, Any]) -> list[DurableAgentStateContent]: List of DurableAgentStateContent objects """ contents: list[DurableAgentStateContent] = [] - raw_contents: list[Any] = data.get(DurableStateFields.CONTENTS, []) + raw_contents = _array_field(data, DurableStateFields.CONTENTS) for raw_content in raw_contents: if isinstance(raw_content, DurableAgentStateContent): contents.append(raw_content) @@ -172,7 +296,7 @@ def _parse_contents(data: dict[str, Any]) -> list[DurableAgentStateContent]: DurableAgentStateFunctionCallContent( call_id=str(content_dict.get(DurableStateFields.CALL_ID, "")), name=str(content_dict.get(DurableStateFields.NAME, "")), - arguments=content_dict.get(DurableStateFields.ARGUMENTS, {}), + arguments=content_dict.get(DurableStateFields.ARGUMENTS), ) ) @@ -207,24 +331,36 @@ def _parse_contents(data: dict[str, Any]) -> list[DurableAgentStateContent]: contents.append( DurableAgentStateUriContent( uri=str(content_dict.get(DurableStateFields.URI, "")), - media_type=str(content_dict.get(DurableStateFields.MEDIA_TYPE, "")), + media_type=content_dict.get(DurableStateFields.MEDIA_TYPE), ) ) case ContentTypes.USAGE: usage_data = content_dict.get(DurableStateFields.USAGE) - if usage_data and isinstance(usage_data, dict): + if isinstance(usage_data, dict): contents.append( DurableAgentStateUsageContent( usage=DurableAgentStateUsage.from_dict(cast(dict[str, Any], usage_data)) ) ) + else: + raise ValueError("Usage content requires a usage object.") - case ContentTypes.UNKNOWN | _: - # Handle UNKNOWN type or any unexpected content types (including None) + case ContentTypes.UNKNOWN: contents.append( DurableAgentStateUnknownContent(content=content_dict.get(DurableStateFields.CONTENT, {})) ) + case _: + if not isinstance(content_type, str) or not content_type: + raise ValueError("Content requires a non-empty $type discriminator.") + contents.append(DurableAgentStateRawContent(content_dict)) + + content = contents[-1] + known = content.to_dict().keys() | {DurableStateFields.EXTENSION_DATA} + content.unknown_fields = {key: deepcopy(value) for key, value in content_dict.items() if key not in known} + content.extensionData = deepcopy(content_dict.get(DurableStateFields.EXTENSION_DATA)) + else: + raise ValueError("contents must contain content objects.") return contents @@ -241,12 +377,58 @@ class DurableAgentStateContent: between the durable state representation and the agent framework's content objects. Attributes: - extensionData: Optional additional metadata (not serialized per schema) + extensionData: Optional metadata, including unmapped canonical core fields. """ extensionData: dict[str, Any] | None = None + unknown_fields: dict[str, Any] | None = None type: str = "" + _NULLABLE_FIELDS: ClassVar[frozenset[str]] = frozenset() + + def to_persisted_dict(self) -> dict[str, Any]: + """Merge opaque fields without replacing mutable, known transcript fields.""" + result = { + **(self.unknown_fields or {}), + **{ + key: value for key, value in self.to_dict().items() if value is not None or key in self._NULLABLE_FIELDS + }, + } + if self.extensionData is not None: + result[DurableStateFields.EXTENSION_DATA] = self.extensionData + return _json_snapshot(result) + + def core_projection(self) -> dict[str, Any]: + """Map this subtype's durable fields to canonical core content fields. + + Returns: + Core field names and current values, including the content type. + """ + # Only map fields owned by this subtype, not a global union of content fields. + aliases = {"details": "error_details", "usage": "usage_details"} + fields = { + aliases.get(key, re.sub(r"(? Content: + """Restore canonical fields through the delivery loader, without dynamic type lookup.""" + extra = (self.extensionData or {}).get("coreContent") + if not isinstance(extra, dict): + return self.to_ai_content() + # The overlay contains extras only. Current text/result/arguments always win. + payload = {**deepcopy(cast(dict[str, Any], extra)), **self.core_projection()} + return load_agent_response({"messages": [{"role": "assistant", "contents": [payload]}]}).messages[0].contents[0] + def to_dict(self) -> dict[str, Any]: """Serialize this content to a dictionary for JSON storage. @@ -271,6 +453,24 @@ def to_ai_content(self) -> Any: @staticmethod def from_ai_content(content: Any) -> DurableAgentStateContent: + """Keep typed durable fields and persist only core fields they cannot represent. + + Args: + content: Core content to convert, or an unknown value to wrap as opaque content. + + Returns: + Durable content with canonical fields not owned by its subtype stored as metadata. + """ + stored = DurableAgentStateContent._from_ai_content(content) + if isinstance(content, Content) and not isinstance(stored, DurableAgentStateUnknownContent): + payload = _json_snapshot(content.to_dict()) + mapped = stored.core_projection() + # An empty overlay still identifies canonical rather than legacy conversion. + stored.extensionData = {"coreContent": {key: value for key, value in payload.items() if key not in mapped}} + return stored + + @staticmethod + def _from_ai_content(content: Any) -> DurableAgentStateContent: """Create a durable state content object from an agent framework content object. This factory method maps agent framework content types to their corresponding durable state representations. @@ -301,7 +501,7 @@ def from_ai_content(content: Any) -> DurableAgentStateContent: return DurableAgentStateHostedVectorStoreContent.from_hosted_vector_store_content(content) case "text": return DurableAgentStateTextContent.from_text_content(content) - case "reasoning": + case "reasoning" | "text_reasoning": return DurableAgentStateTextReasoningContent.from_text_reasoning_content(content) case "uri": return DurableAgentStateUriContent.from_uri_content(content) @@ -311,6 +511,27 @@ def from_ai_content(content: Any) -> DurableAgentStateContent: return DurableAgentStateUnknownContent.from_unknown_content(content) +class DurableAgentStateRawContent(DurableAgentStateContent): + """Opaque future shared-schema content, preserved without reinterpreting its fields.""" + + def __init__(self, raw: dict[str, Any]) -> None: + self.raw = deepcopy(raw) + + def to_dict(self) -> dict[str, Any]: + return deepcopy(self.raw) + + def to_persisted_dict(self) -> dict[str, Any]: + """Preserve even null fields belonging to an unknown content kind.""" + return _json_snapshot(self.raw) + + def to_core_content(self) -> Content: + """Do not interpret an unknown writer's extension conventions.""" + return self.to_ai_content() + + def to_ai_content(self) -> Content: + return Content(type="unknown", additional_properties={"content": deepcopy(self.raw)}) # type: ignore[arg-type] + + # Core state classes @@ -326,40 +547,157 @@ class DurableAgentStateData: Attributes: conversation_history: Ordered list of conversation entries (requests and responses) + session: Serialized ``AgentSession`` from the previous turn - the context provider state + bag plus any service-issued conversation id. Core treats session state as durable + across turns, so it is persisted here rather than discarded with the per-operation + session. + ingested_positions: Legacy per-producer maxima, retained for read compatibility. + Migration requires delivery evidence because a maximum does not identify skipped positions. + ingested_messages: Actual source identities and content fingerprints, independent of transcript pruning. + response_mailbox: Original serializable results with their delivery expiry. + completed_correlations: Completion evidence retained after mailbox expiry. + truncation: What retention has removed, if anything. A log line is only visible to whoever + was watching at the time, so the fact that this conversation is no longer complete is + recorded in the state itself. Absent until the first eviction, so its absence is a + positive statement that nothing has been dropped. extension_data: Optional dictionary for custom metadata (not part of core schema) """ conversation_history: list[DurableAgentStateEntry] + session: dict[str, Any] | None + ingested_positions: dict[str, int] | None + truncation: dict[str, Any] | None extension_data: dict[str, Any] | None + response_mailbox: dict[str, dict[str, Any]] + completed_correlations: dict[str, dict[str, Any]] + ingested_messages: dict[str, list[str] | None] + unknown_fields: dict[str, Any] def __init__( self, conversation_history: list[DurableAgentStateEntry] | None = None, extension_data: dict[str, Any] | None = None, + session: dict[str, Any] | None = None, + ingested_positions: dict[str, int] | None = None, + truncation: dict[str, Any] | None = None, + response_mailbox: dict[str, dict[str, Any]] | None = None, + completed_correlations: dict[str, dict[str, Any]] | None = None, + ingested_messages: dict[str, list[str] | None] | None = None, ) -> None: """Initialize the data container. Args: conversation_history: Initial conversation history (defaults to empty list) extension_data: Optional custom metadata + session: Optional serialized ``AgentSession`` from the previous turn + ingested_positions: Legacy scalar ingestion state, not exact delivery evidence. + truncation: Record of what retention has removed, absent until something is + response_mailbox: Original response snapshots with independent delivery expiry. + completed_correlations: Completion evidence retained after result expiry. + ingested_messages: Exact message fingerprints or legacy identity-only markers. """ self.conversation_history = conversation_history or [] self.extension_data = extension_data + self.session = session + self.ingested_positions = ingested_positions + self.truncation = truncation + self.response_mailbox = response_mailbox or {} + self.completed_correlations = completed_correlations or {} + self.ingested_messages = ingested_messages or {} + self.unknown_fields = {} def to_dict(self) -> dict[str, Any]: + _validate_delivery_layout(self.unknown_fields) + _validate_completion_outcomes(self.completed_correlations) result: dict[str, Any] = { + **deepcopy(self.unknown_fields), DurableStateFields.CONVERSATION_HISTORY: [entry.to_dict() for entry in self.conversation_history], } if self.extension_data is not None: result[DurableStateFields.EXTENSION_DATA] = self.extension_data - return result + if self.session is not None: + result[DurableStateFields.SESSION] = self.session + if self.ingested_positions: + result[DurableStateFields.INGESTED_POSITIONS] = self.ingested_positions + if self.truncation: + result[DurableStateFields.TRUNCATION] = self.truncation + if self.response_mailbox: + result[DurableStateFields.RESPONSE_MAILBOX] = deepcopy(self.response_mailbox) + if self.completed_correlations: + result[DurableStateFields.COMPLETED_CORRELATIONS] = deepcopy(self.completed_correlations) + if self.ingested_messages: + result[DurableStateFields.INGESTED_MESSAGES] = deepcopy(self.ingested_messages) + return _json_snapshot(result) @classmethod def from_dict(cls, data_dict: dict[str, Any]) -> DurableAgentStateData: - return cls( + _validate_delivery_layout(data_dict) + for name in ( + DurableStateFields.RESPONSE_MAILBOX, + DurableStateFields.COMPLETED_CORRELATIONS, + DurableStateFields.INGESTED_MESSAGES, + ): + if name in data_dict and not isinstance(data_dict[name], dict): + raise ValueError(f"{name} must be an object.") + result = cls( conversation_history=_parse_history_entries(data_dict), extension_data=data_dict.get(DurableStateFields.EXTENSION_DATA), + session=data_dict.get(DurableStateFields.SESSION), + ingested_positions=data_dict.get(DurableStateFields.INGESTED_POSITIONS), + truncation=data_dict.get(DurableStateFields.TRUNCATION), + response_mailbox=deepcopy(data_dict.get(DurableStateFields.RESPONSE_MAILBOX, {})), + completed_correlations=deepcopy(data_dict.get(DurableStateFields.COMPLETED_CORRELATIONS, {})), + ingested_messages=deepcopy(data_dict.get(DurableStateFields.INGESTED_MESSAGES, {})), ) + known = { + DurableStateFields.CONVERSATION_HISTORY, + DurableStateFields.EXTENSION_DATA, + DurableStateFields.SESSION, + DurableStateFields.INGESTED_POSITIONS, + DurableStateFields.TRUNCATION, + DurableStateFields.RESPONSE_MAILBOX, + DurableStateFields.COMPLETED_CORRELATIONS, + DurableStateFields.INGESTED_MESSAGES, + } + result.unknown_fields = {key: deepcopy(value) for key, value in data_dict.items() if key not in known} + for name, records in ( + (DurableStateFields.RESPONSE_MAILBOX, result.response_mailbox), + (DurableStateFields.COMPLETED_CORRELATIONS, result.completed_correlations), + ): + if any(not isinstance(value, dict) for value in records.values()): + raise ValueError(f"{name} must contain objects keyed by correlation ID.") + for correlation_id, record in records.items(): + if not isinstance(correlation_id, str) or not correlation_id: + raise ValueError(f"{name} requires non-empty correlation IDs.") + timestamps = ( + (DurableStateFields.CREATED_AT, DurableStateFields.EXPIRES_AT) + if name == DurableStateFields.RESPONSE_MAILBOX + else (DurableStateFields.COMPLETED_AT,) + ) + for field in timestamps: + try: + _parse_delivery_timestamp(record.get(field)) + except ValueError as exc: + raise ValueError(f"{name}.{field} must be an RFC 3339 timestamp with an offset.") from exc + if name == DurableStateFields.RESPONSE_MAILBOX: + response = record.get(DurableStateFields.RESPONSE) + if not isinstance(response, dict): + raise ValueError("responseMailbox.response must be an inline agent response.") + response = cast(dict[str, Any], response) + if response.get("type") != "agent_response" or not isinstance(response.get("messages"), list): + raise ValueError("responseMailbox.response must be an inline agent response.") + for message in response["messages"]: + _validate_core_message(message) + load_agent_response(response) + elif "legacy" in record and not isinstance(record["legacy"], bool): + raise ValueError("completedCorrelations.legacy must be a boolean.") + _validate_completion_outcomes(result.completed_correlations) + if not isinstance(result.ingested_messages, dict) or any( + values is not None and (not isinstance(values, list) or any(not isinstance(v, str) for v in values)) + for values in result.ingested_messages.values() + ): + raise ValueError("ingestedMessages must contain fingerprint lists or legacy identity markers.") + return result class DurableAgentState: @@ -391,8 +729,9 @@ class DurableAgentState: schema_version: Schema version string (defaults to SCHEMA_VERSION) """ - # Durable Agent Schema version - SCHEMA_VERSION: str = "1.1.0" + # New layout requires compatible workers and response consumers. A version number + # does not make legacy .NET workers or older Python writers safe to share this state. + SCHEMA_VERSION: str = "2.0.0" data: DurableAgentStateData schema_version: str = SCHEMA_VERSION @@ -405,16 +744,17 @@ def __init__(self, schema_version: str = SCHEMA_VERSION): """ self.data = DurableAgentStateData() self.schema_version = schema_version + self.unknown_fields: dict[str, Any] = {} def to_dict(self) -> dict[str, Any]: - - return { + return _json_snapshot({ + **deepcopy(self.unknown_fields), DurableStateFields.SCHEMA_VERSION: self.schema_version, DurableStateFields.DATA: self.data.to_dict(), - } + }) def to_json(self) -> str: - return json.dumps(self.to_dict()) + return json.dumps(self.to_dict(), allow_nan=False) @classmethod def from_dict(cls, state: dict[str, Any]) -> DurableAgentState: @@ -423,13 +763,29 @@ def from_dict(cls, state: dict[str, Any]) -> DurableAgentState: Args: state: Dictionary containing schemaVersion and data (full state structure) """ + if not isinstance(state, dict): + raise ValueError("The durable agent state must be a JSON object.") + state = _json_snapshot(state) schema_version = state.get(DurableStateFields.SCHEMA_VERSION) if schema_version is None: - logger.warning("Resetting state as it is incompatible with the current schema, all history will be lost") - return cls() - - instance = cls(schema_version=state.get(DurableStateFields.SCHEMA_VERSION, DurableAgentState.SCHEMA_VERSION)) - instance.data = DurableAgentStateData.from_dict(state.get(DurableStateFields.DATA, {})) + raise ValueError("The durable agent state is missing schemaVersion; refusing to discard existing state.") + if not isinstance(schema_version, str) or not re.fullmatch(r"[12]\.[0-9]+\.[0-9]+", schema_version): + raise ValueError(f"Unsupported durable agent state schemaVersion: {schema_version!r}.") + raw_data = state.get(DurableStateFields.DATA) + if not isinstance(raw_data, dict): + raise ValueError("The durable agent state data must be an object.") + + instance = cls(schema_version=schema_version) + instance.data = DurableAgentStateData.from_dict(cast(dict[str, Any], raw_data)) + if schema_version.startswith("2.") and ( + instance.data.response_mailbox.keys() - instance.data.completed_correlations.keys() + ): + raise ValueError("Every responseMailbox entry requires a matching completedCorrelations receipt.") + instance.unknown_fields = { + key: deepcopy(value) + for key, value in state.items() + if key not in (DurableStateFields.SCHEMA_VERSION, DurableStateFields.DATA) + } return instance @@ -440,7 +796,9 @@ def from_json(cls, json_str: str) -> DurableAgentState: except json.JSONDecodeError as e: raise ValueError("The durable agent state is not valid JSON.") from e - return cls.from_dict(obj) + if not isinstance(obj, dict): + raise ValueError("The durable agent state must be a JSON object.") + return cls.from_dict(cast(dict[str, Any], obj)) @property def message_count(self) -> int: @@ -448,31 +806,151 @@ def message_count(self) -> int: return len(self.data.conversation_history) def try_get_agent_response(self, correlation_id: str) -> AgentResponse | None: - """Try to get an agent response by correlation ID. - - This method searches the conversation history for a response entry matching the given - correlation ID and returns a dictionary suitable for HTTP API responses. + """Read a retained result or explicit completed status using the persisted layout. - Note: The returned dictionary includes computed properties (message_count) that are - NOT part of the persisted state schema. These are derived values included for backward - compatibility with the HTTP API response format and should not be considered part of - the durable state structure. + Version 2 never falls back to transcript responses, even after mailbox expiry. + Version 1 retains its legacy lookup until an operation migrates the state. Args: - correlation_id: The correlation ID to search for + correlation_id: Request correlation ID whose response or completion status to retrieve. Returns: - Response data dict with 'content', 'message_count', and 'correlationId' if found, - None otherwise + Retained response, expired-response status, or None when no matching result exists. """ - # Search through conversation history for a response with this correlationId + _validate_delivery_layout(self.data.unknown_fields) + _validate_completion_outcomes(self.data.completed_correlations) + if self.schema_version.startswith("2."): + mailbox = self.data.response_mailbox.get(correlation_id) + if mailbox is not None: + expiry = _parse_delivery_timestamp(mailbox[DurableStateFields.EXPIRES_AT]) + if datetime.now(timezone.utc) < expiry: + return load_agent_response(mailbox[DurableStateFields.RESPONSE]) + if correlation_id in self.data.completed_correlations or mailbox is not None: + return AgentResponse( + messages=[ + Message( + "system", + [ + Content.from_error( + message="This request completed, but its response delivery window has expired.", + error_code="response_expired", + ) + ], + ) + ], + additional_properties={ + "durable_status": "already_completed", + "correlation_id": correlation_id, + "durable_outcome": self._completion_outcome(correlation_id) or "unknown", + }, + ) + return None for entry in self.data.conversation_history: if entry.correlation_id == correlation_id and isinstance(entry, DurableAgentStateResponse): - # Found the entry, extract response data return DurableAgentStateResponse.to_run_response(entry) return None + def _completion_outcome(self, correlation_id: str) -> str | None: + """Read a receipt or its independent result, never a possibly altered transcript.""" + receipt = self.data.completed_correlations.get(correlation_id, {}) + if DurableStateFields.OUTCOME in receipt: + return receipt[DurableStateFields.OUTCOME] + mailbox = self.data.response_mailbox.get(correlation_id) + if mailbox is None: + return None + return invocation_outcome( + load_agent_response(mailbox[DurableStateFields.RESPONSE]), legacy=receipt.get("legacy", False) + ) + + def _backfill_completion_outcomes(self, *, require_known: bool = False) -> None: + """Enrich old receipts from retained evidence without changing time or availability.""" + _validate_completion_outcomes(self.data.completed_correlations) + for correlation_id, receipt in self.data.completed_correlations.items(): + outcome = self._completion_outcome(correlation_id) + if outcome is not None: + receipt.setdefault(DurableStateFields.OUTCOME, outcome) + elif require_known: + raise ValueError("A known completion outcome requires authoritative retained result evidence.") + + def record_response( + self, + correlation_id: str, + response: AgentResponse, + *, + delivery_window_seconds: int, + now: datetime | None = None, + legacy: bool = False, + ) -> None: + """Stage an independent JSON snapshot and completion receipt, without persisting them. + + Args: + correlation_id: Request correlation ID used to key the snapshot and completion receipt. + response: Agent response to snapshot for delivery. + delivery_window_seconds: Seconds after the recording timestamp when the snapshot expires. + now: Offset-aware recording timestamp, defaulting to the current UTC time. + legacy: Whether this is a possibly altered legacy transcript projection. + A retained failure proves failure, but missing error content cannot prove success. + """ + if correlation_id in self.data.completed_correlations: + return + timestamp = now or datetime.now(timezone.utc) + _parse_delivery_timestamp(timestamp.isoformat()) + payload = _json_snapshot(serialize_agent_response(response)) + outcome = invocation_outcome(load_agent_response(payload), legacy=legacy) + if outcome is None and not legacy: + raise ValueError("A new completion requires a known invocation outcome, not an acknowledgement.") + self.data.response_mailbox[correlation_id] = { + DurableStateFields.RESPONSE: payload, + DurableStateFields.CREATED_AT: timestamp.isoformat(), + DurableStateFields.EXPIRES_AT: (timestamp + timedelta(seconds=delivery_window_seconds)).isoformat(), + } + self.data.completed_correlations[correlation_id] = { + DurableStateFields.COMPLETED_AT: timestamp.isoformat(), + **({DurableStateFields.OUTCOME: outcome} if outcome is not None else {}), + **({"legacy": True} if legacy else {}), + } + + def expire_responses(self, *, now: datetime | None = None) -> None: + """Expire payloads, preserving the original completion time and known outcome. + + Args: + now: Offset-aware expiry-check timestamp, defaulting to the current UTC time. + """ + _validate_completion_outcomes(self.data.completed_correlations) + timestamp = now or datetime.now(timezone.utc) + for correlation_id, mailbox in list(self.data.response_mailbox.items()): + expiry = _parse_delivery_timestamp(mailbox[DurableStateFields.EXPIRES_AT]) + if timestamp >= expiry: + # Older receipts may lack the outcome. Preserve what their independent + # result proves before deleting it, never infer from the transcript. + outcome = self._completion_outcome(correlation_id) + receipt = self.data.completed_correlations.get(correlation_id) + if receipt is not None and outcome is not None: + receipt.setdefault(DurableStateFields.OUTCOME, outcome) + del self.data.response_mailbox[correlation_id] + + def prepare_for_write(self, *, delivery_window_seconds: int) -> None: + """Admit only the revised writer layout, without silently upgrading legacy state. + + Args: + delivery_window_seconds: Retained for source compatibility; migration now + requires an explicit destination operation, including its grace policy. + """ + _validate_delivery_layout(self.data.unknown_fields) + _validate_completion_outcomes(self.data.completed_correlations) + if self.schema_version == self.SCHEMA_VERSION: + return + if re.fullmatch(r"1\.[0-9]+\.[0-9]+", self.schema_version) is None: + raise ValueError( + f"Unsupported durable agent state schemaVersion for writing: {self.schema_version!r}. " + f"Only {self.SCHEMA_VERSION} is writable." + ) + raise ValueError( + "Legacy state is read-only in this runtime. Keep it on its original deployment or use explicit " + "migration into a separate isolated-v2 entity. Legacy ingestedPositions require recorded delivery evidence." + ) + class DurableAgentStateEntry: """Base class for conversation history entries (requests and responses). @@ -486,8 +964,10 @@ class DurableAgentStateEntry: with their originating requests. Common Attributes: - json_type: Discriminator for entry type ("request" or "response") - correlationId: Unique identifier linking requests and responses + json_type: Discriminator for entry type ("request", "response", "errorResponse" or + "compaction") + correlationId: Unique identifier linking requests and responses. Absent on compaction + entries, which answer no request. created_at: Timestamp when the entry was created messages: List of messages in this entry extensionData: Optional additional metadata (not serialized per schema) @@ -500,7 +980,7 @@ class DurableAgentStateEntry: usage: Token usage statistics - only for response entries """ - json_type: DurableAgentStateEntryJsonType + json_type: DurableAgentStateEntryJsonType | str correlation_id: str | None created_at: datetime messages: list[DurableAgentStateMessage] @@ -508,7 +988,7 @@ class DurableAgentStateEntry: def __init__( self, - json_type: DurableAgentStateEntryJsonType, + json_type: DurableAgentStateEntryJsonType | str, correlation_id: str | None, created_at: datetime, messages: list[DurableAgentStateMessage], @@ -519,27 +999,55 @@ def __init__( self.created_at = created_at self.messages = messages self.extension_data = extension_data + self.unknown_fields: dict[str, Any] = {} def to_dict(self) -> dict[str, Any]: - return { + result: dict[str, Any] = { + **deepcopy(self.unknown_fields), DurableStateFields.TYPE_DISCRIMINATOR: self.json_type, - DurableStateFields.CORRELATION_ID: self.correlation_id, DurableStateFields.CREATED_AT: self.created_at.isoformat(), DurableStateFields.MESSAGES: [m.to_dict() for m in self.messages], } + if self.correlation_id is not None: + # Omitted rather than written as null. A compaction entry answers no request and so has + # no correlation, and "absent" says that where an explicit null only says the field + # exists and is empty. It also keeps the persisted shape a string wherever it appears, + # which is what the schema and the .NET reader both expect. + result[DurableStateFields.CORRELATION_ID] = self.correlation_id + if self.extension_data is not None: + result[DurableStateFields.EXTENSION_DATA] = deepcopy(self.extension_data) + return _json_snapshot(result) @classmethod def from_dict(cls, data: dict[str, Any]) -> DurableAgentStateEntry: created_at = _parse_created_at(data.get(DurableStateFields.CREATED_AT)) messages = _parse_messages(data) - return cls( + entry = cls( json_type=DurableAgentStateEntryJsonType(data.get(DurableStateFields.TYPE_DISCRIMINATOR)), correlation_id=data.get(DurableStateFields.CORRELATION_ID), created_at=created_at, messages=messages, extension_data=data.get(DurableStateFields.EXTENSION_DATA), ) + entry.unknown_fields = _entry_unknown_fields(entry, data) + return entry + + +class DurableAgentStateUnknownEntry(DurableAgentStateEntry): + """Opaque future entry preserved for round-trip, never converted into model context.""" + + def __init__(self, raw: dict[str, Any]) -> None: + self.raw = deepcopy(raw) + super().__init__( + json_type=str(raw.get(DurableStateFields.TYPE_DISCRIMINATOR, "unknown")), + correlation_id=raw.get(DurableStateFields.CORRELATION_ID), + created_at=datetime.min.replace(tzinfo=timezone.utc), + messages=[], + ) + + def to_dict(self) -> dict[str, Any]: + return deepcopy(self.raw) class DurableAgentStateRequest(DurableAgentStateEntry): @@ -591,7 +1099,7 @@ def to_dict(self) -> dict[str, Any]: if self.response_type is not None: data[DurableStateFields.RESPONSE_TYPE] = self.response_type if self.response_schema is not None: - data[DurableStateFields.RESPONSE_SCHEMA] = self.response_schema + data[DurableStateFields.RESPONSE_SCHEMA] = deepcopy(self.response_schema) return data @classmethod @@ -599,7 +1107,7 @@ def from_dict(cls, data: dict[str, Any]) -> DurableAgentStateRequest: created_at = _parse_created_at(data.get(DurableStateFields.CREATED_AT)) messages = _parse_messages(data) - return cls( + entry = cls( correlation_id=data.get(DurableStateFields.CORRELATION_ID), created_at=created_at, messages=messages, @@ -608,13 +1116,21 @@ def from_dict(cls, data: dict[str, Any]) -> DurableAgentStateRequest: response_schema=data.get(DurableStateFields.RESPONSE_SCHEMA), orchestration_id=data.get(DurableStateFields.ORCHESTRATION_ID), ) + entry.unknown_fields = _entry_unknown_fields(entry, data) + return entry @staticmethod def from_run_request(request: RunRequest) -> DurableAgentStateRequest: + # A workflow may deliver the upstream conversation instead of a single message. + if request.context_messages is not None: + messages = [DurableAgentStateMessage.from_core_dict(raw) for raw in request.context_messages] + else: + messages = [DurableAgentStateMessage.from_run_request(request)] + # Determine response_type based on response_format return DurableAgentStateRequest( correlation_id=request.correlation_id, - messages=[DurableAgentStateMessage.from_run_request(request)], + messages=messages, created_at=_parse_created_at(request.created_at), response_type=request.request_response_format, response_schema=serialize_response_format(request.response_format), @@ -631,15 +1147,15 @@ class DurableAgentStateResponse(DurableAgentStateEntry): Attributes: usage: Token usage statistics for this response (input, output, and total tokens) - is_error: Flag indicating if this response represents an error (not persisted in schema) correlation_id: Unique identifier linking this response to its request created_at: Timestamp when the response was created messages: List of assistant messages in this response - json_type: Always "response" for this class + json_type: "response", or "errorResponse" for the failed-turn subclass """ + JSON_TYPE: ClassVar[DurableAgentStateEntryJsonType] = DurableAgentStateEntryJsonType.RESPONSE + usage: DurableAgentStateUsage | None = None - is_error: bool = False def __init__( self, @@ -648,17 +1164,15 @@ def __init__( messages: list[DurableAgentStateMessage], extension_data: dict[str, Any] | None = None, usage: DurableAgentStateUsage | None = None, - is_error: bool = False, ) -> None: super().__init__( - json_type=DurableAgentStateEntryJsonType.RESPONSE, + json_type=type(self).JSON_TYPE, correlation_id=correlation_id, created_at=created_at, messages=messages, extension_data=extension_data, ) self.usage = usage - self.is_error = is_error def to_dict(self) -> dict[str, Any]: data = super().to_dict() @@ -673,21 +1187,29 @@ def from_dict(cls, data: dict[str, Any]) -> DurableAgentStateResponse: usage_dict = data.get(DurableStateFields.USAGE) usage: DurableAgentStateUsage | None = None - if usage_dict and isinstance(usage_dict, dict): + if isinstance(usage_dict, dict): usage = DurableAgentStateUsage.from_dict(cast(dict[str, Any], usage_dict)) + elif usage_dict is not None: + raise ValueError("Response usage must be an object.") - return cls( + entry = cls( correlation_id=data.get(DurableStateFields.CORRELATION_ID), created_at=created_at, messages=messages, extension_data=data.get(DurableStateFields.EXTENSION_DATA), usage=usage, ) + entry.unknown_fields = _entry_unknown_fields(entry, data) + return entry - @staticmethod - def from_run_response(correlation_id: str, response: AgentResponse) -> DurableAgentStateResponse: - """Creates a DurableAgentStateResponse from an AgentResponse.""" - return DurableAgentStateResponse( + @classmethod + def from_run_response(cls, correlation_id: str, response: AgentResponse) -> DurableAgentStateResponse: + """Creates a response entry of this class from an AgentResponse. + + A classmethod rather than a staticmethod so the error subclass produces an error entry + without the caller having to set anything afterwards. + """ + return cls( correlation_id=correlation_id, created_at=_parse_created_at(response.created_at), messages=[DurableAgentStateMessage.from_chat_message(m) for m in response.messages], @@ -707,7 +1229,65 @@ def to_run_response( created_at=response_entry.created_at.isoformat(), messages=messages, usage_details=usage_details, + additional_properties=( + {"durable_status": "error"} if isinstance(response_entry, DurableAgentStateErrorResponse) else None + ), + ) + + +class DurableAgentStateErrorResponse(DurableAgentStateResponse): + """A turn that failed, recorded so the waiting caller can be told why. + + Deliberately a response, because a caller polling its correlation id still needs an answer and + an error is the answer. Deliberately not replayable, because the reason a turn failed is for + the caller, not for the model, and feeding it back would present an exception as something the + assistant said. + + That second part used to be a boolean on the response, which was never serialized. The failure + survived a reload looking like an ordinary reply. Being a distinct type means the distinction + cannot be lost in transit. + + Not to be confused with ``DurableAgentStateErrorContent``, which is error content inside a + single message. This is the entry recording that a whole turn failed. + """ + + JSON_TYPE: ClassVar[DurableAgentStateEntryJsonType] = DurableAgentStateEntryJsonType.ERROR_RESPONSE + + +class DurableAgentStateCompaction(DurableAgentStateEntry): + """A message compaction produced, such as a summary standing in for turns it replaced. + + The exact opposite of an error entry. It belongs to the model's transcript and takes its place + in conversation order, but it answers no request, so it is not a response and can never be + returned to a caller polling for one. Previously these were inserted into whichever entry they + followed, which meant a poll could hand back a summary alongside the real answer. + """ + + def __init__( + self, + created_at: datetime, + messages: list[DurableAgentStateMessage], + correlation_id: str | None = None, + extension_data: dict[str, Any] | None = None, + ) -> None: + super().__init__( + json_type=DurableAgentStateEntryJsonType.COMPACTION, + correlation_id=correlation_id, + created_at=created_at, + messages=messages, + extension_data=extension_data, + ) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> DurableAgentStateCompaction: + entry = cls( + created_at=_parse_created_at(data.get(DurableStateFields.CREATED_AT)), + messages=_parse_messages(data), + correlation_id=data.get(DurableStateFields.CORRELATION_ID), + extension_data=data.get(DurableStateFields.EXTENSION_DATA), ) + entry.unknown_fields = _entry_unknown_fields(entry, data) + return entry class DurableAgentStateMessage: @@ -722,14 +1302,22 @@ class DurableAgentStateMessage: contents: List of content items (text, function calls, errors, etc.) author_name: Optional name of the message author (typically set for assistant messages) created_at: Optional timestamp when the message was created - extension_data: Optional additional metadata (not serialized per schema) + message_id: Optional stable identifier for the message. Persisted so context-management + state (for example compaction summaries that reference the messages they replace) + can be reconciled across entity operations. + extension_data: Optional additional metadata. Carries a message's + ``additional_properties``, including compaction annotations, so that context + management state survives across entity operations. """ role: str contents: list[DurableAgentStateContent] author_name: str | None = None created_at: datetime | None = None + message_id: str | None = None extension_data: dict[str, Any] | None = None + ingestion_identity: str | None = None + ingestion_occurrence: str | None = None def __init__( self, @@ -738,45 +1326,56 @@ def __init__( author_name: str | None = None, created_at: datetime | None = None, extension_data: dict[str, Any] | None = None, + message_id: str | None = None, ) -> None: self.role = role self.contents = contents self.author_name = author_name self.created_at = created_at + self.message_id = message_id self.extension_data = extension_data + self.unknown_fields: dict[str, Any] = {} def to_dict(self) -> dict[str, Any]: result: dict[str, Any] = { + **deepcopy(self.unknown_fields), DurableStateFields.ROLE: self.role, - DurableStateFields.CONTENTS: [ - { - DurableStateFields.TYPE_DISCRIMINATOR: c.to_dict().get( - DurableStateFields.TYPE_INTERNAL, ContentTypes.TEXT - ), - **{k: v for k, v in c.to_dict().items() if k != DurableStateFields.TYPE_INTERNAL}, - } - for c in self.contents - ], + DurableStateFields.CONTENTS: [c.to_persisted_dict() for c in self.contents], } # Only include optional fields if they have values if self.created_at is not None: result[DurableStateFields.CREATED_AT] = self.created_at.isoformat() if self.author_name is not None: result[DurableStateFields.AUTHOR_NAME] = self.author_name - return result + if self.message_id is not None: + result[DurableStateFields.MESSAGE_ID] = self.message_id + if self.extension_data is not None: + result[DurableStateFields.EXTENSION_DATA] = self.extension_data + return _json_snapshot(result) @classmethod def from_dict(cls, data: dict[str, Any]) -> DurableAgentStateMessage: data_created_at = data.get(DurableStateFields.CREATED_AT) created_at = _parse_created_at(data_created_at) if data_created_at else None - return cls( + message = cls( role=data.get(DurableStateFields.ROLE, ""), contents=_parse_contents(data), author_name=data.get(DurableStateFields.AUTHOR_NAME), created_at=created_at, + message_id=data.get(DurableStateFields.MESSAGE_ID), extension_data=data.get(DurableStateFields.EXTENSION_DATA), ) + known = { + DurableStateFields.ROLE, + DurableStateFields.CONTENTS, + DurableStateFields.AUTHOR_NAME, + DurableStateFields.CREATED_AT, + DurableStateFields.MESSAGE_ID, + DurableStateFields.EXTENSION_DATA, + } + message.unknown_fields = {key: deepcopy(value) for key, value in data.items() if key not in known} + return message @property def text(self) -> str: @@ -802,6 +1401,35 @@ def from_run_request(request: RunRequest) -> DurableAgentStateMessage: created_at=_parse_created_at(request.created_at) if request.created_at else None, ) + @staticmethod + def from_core_dict(data: dict[str, Any]) -> DurableAgentStateMessage: + """Keep unknown core fields before consumer filtering can discard them. + + Args: + data: Serialized core message containing content envelopes and optional metadata. + + Returns: + Durable message preserving unknown message and content fields. + """ + raw = _json_snapshot(data) + _validate_core_message(raw) + message = load_agent_response({"messages": [raw]}).messages[0] + stored = DurableAgentStateMessage.from_chat_message(message) + for content, original in zip(stored.contents, raw.get("contents", []), strict=True): + if not isinstance(original, dict): + raise ValueError("Core contents must contain content objects.") + original = cast(dict[str, Any], original) + if isinstance(content, DurableAgentStateUnknownContent): + content.content = original + else: + mapped = content.core_projection() + content.extensionData = { + "coreContent": {key: value for key, value in original.items() if key not in mapped} + } + known = {"type", "role", "contents", "author_name", "message_id", "additional_properties"} + stored.unknown_fields = {key: value for key, value in raw.items() if key not in known} + return stored + @staticmethod def from_chat_message(chat_message: Message) -> DurableAgentStateMessage: """Converts an Agent Framework chat message to a durable state message. @@ -816,12 +1444,19 @@ def from_chat_message(chat_message: Message) -> DurableAgentStateMessage: DurableAgentStateContent.from_ai_content(c) for c in chat_message.contents ] - return DurableAgentStateMessage( + stored = DurableAgentStateMessage( role=chat_message.role if hasattr(chat_message.role, "value") else str(chat_message.role), contents=contents_list, author_name=chat_message.author_name, - extension_data=dict(chat_message.additional_properties) if chat_message.additional_properties else None, + message_id=getattr(chat_message, "message_id", None), + extension_data=deepcopy(chat_message.additional_properties) if chat_message.additional_properties else None, ) + stored.ingestion_identity = message_identity(chat_message) + known = {"type", "role", "contents", "author_name", "message_id", "additional_properties"} + stored.unknown_fields = { + key: value for key, value in _json_snapshot(chat_message.to_dict()).items() if key not in known + } + return stored def to_chat_message(self) -> Any: """Converts this DurableAgentStateMessage back to an agent framework Message. @@ -830,7 +1465,7 @@ def to_chat_message(self) -> Any: Message object with role, contents, and metadata converted back to agent framework types """ # Convert DurableAgentStateContent objects back to agent_framework content objects - ai_contents = [c.to_ai_content() for c in self.contents] + ai_contents = [c.to_core_content() for c in self.contents] # Build kwargs for Message kwargs: dict[str, Any] = { @@ -841,8 +1476,16 @@ def to_chat_message(self) -> Any: if self.author_name is not None: kwargs["author_name"] = self.author_name + if self.message_id is not None: + kwargs["message_id"] = self.message_id + if self.extension_data is not None: - kwargs["additional_properties"] = self.extension_data + # Copied, not shared. Callers treat the result as detached and mutate it: retention + # pops compaction annotations off the copies it measures. Handing out the stored dict + # would make that erase those annotations from durable state. Core does copy this + # during validation today, but that is its internal business, and quietly depending on + # it would mean a change there costs us the user's compaction work. + kwargs["additional_properties"] = deepcopy(self.extension_data) return Message(**kwargs) @@ -935,16 +1578,16 @@ class DurableAgentStateFunctionCallContent(DurableAgentStateContent): Attributes: call_id: Unique identifier for this function call (used to match with results) name: Name of the function/tool to execute - arguments: Dictionary of argument names to values for the function call + arguments: Original argument string or mapping, without lossy reparsing """ call_id: str name: str - arguments: dict[str, Any] + arguments: dict[str, Any] | str | None type: str = ContentTypes.FUNCTION_CALL - def __init__(self, call_id: str, name: str, arguments: dict[str, Any]) -> None: + def __init__(self, call_id: str, name: str, arguments: dict[str, Any] | str | None) -> None: self.call_id = call_id self.name = name self.arguments = arguments @@ -963,22 +1606,13 @@ def from_function_call_content(content: Content) -> DurableAgentStateFunctionCal raise ValueError("call_id is required for function call content") if content.name is None: raise ValueError("name is required for function call content") - # Ensure arguments is a dict; parse string if needed - arguments: dict[str, Any] = {} - if content.arguments: - if isinstance(content.arguments, dict): - arguments = content.arguments - elif isinstance(content.arguments, str): - # Parse JSON string to dict - try: - arguments = json.loads(content.arguments) - except json.JSONDecodeError: - arguments = {} - - return DurableAgentStateFunctionCallContent(call_id=content.call_id, name=content.name, arguments=arguments) + return DurableAgentStateFunctionCallContent( + call_id=content.call_id, name=content.name, arguments=_json_snapshot(content.to_dict().get("arguments")) + ) def to_ai_content(self) -> Content: - return Content.from_function_call(call_id=self.call_id, name=self.name, arguments=json.dumps(self.arguments)) + arguments = json.dumps(self.arguments) if isinstance(self.arguments, dict) else self.arguments + return Content.from_function_call(call_id=self.call_id, name=self.name, arguments=arguments) class DurableAgentStateFunctionResultContent(DurableAgentStateContent): @@ -998,6 +1632,8 @@ class DurableAgentStateFunctionResultContent(DurableAgentStateContent): type: str = ContentTypes.FUNCTION_RESULT + _NULLABLE_FIELDS: ClassVar[frozenset[str]] = frozenset({DurableStateFields.RESULT}) + def __init__(self, call_id: str, result: Any | None = None) -> None: self.call_id = call_id self.result = result @@ -1013,7 +1649,9 @@ def to_dict(self) -> dict[str, Any]: def from_function_result_content(content: Content) -> DurableAgentStateFunctionResultContent: if content.call_id is None: raise ValueError("call_id is required for function result content") - return DurableAgentStateFunctionResultContent(call_id=content.call_id, result=content.result) + return DurableAgentStateFunctionResultContent( + call_id=content.call_id, result=_json_snapshot(content.to_dict().get("result")) + ) def to_ai_content(self) -> Content: return Content.from_function_result(call_id=self.call_id, result=self.result) @@ -1103,6 +1741,12 @@ def __init__(self, text: str | None) -> None: def to_dict(self) -> dict[str, Any]: return {DurableStateFields.TYPE_DISCRIMINATOR: self.type, DurableStateFields.TEXT: self.text} + def to_persisted_dict(self) -> dict[str, Any]: + """Require the schema's text string rather than emit an invalid content item.""" + if not isinstance(self.text, str): + raise ValueError("Text content requires a text string for persistence.") + return super().to_persisted_dict() + @staticmethod def from_text_content(content: Content) -> DurableAgentStateTextContent: return DurableAgentStateTextContent(text=content.text) @@ -1149,11 +1793,11 @@ class DurableAgentStateUriContent(DurableAgentStateContent): """ uri: str - media_type: str + media_type: str | None type: str = ContentTypes.URI - def __init__(self, uri: str, media_type: str) -> None: + def __init__(self, uri: str, media_type: str | None = None) -> None: self.uri = uri self.media_type = media_type @@ -1168,8 +1812,6 @@ def to_dict(self) -> dict[str, Any]: def from_uri_content(content: Content) -> DurableAgentStateUriContent: if content.uri is None: raise ValueError("uri is required for uri content") - if content.media_type is None: - raise ValueError("media_type is required for uri content") return DurableAgentStateUriContent(uri=content.uri, media_type=content.media_type) def to_ai_content(self) -> Content: @@ -1214,25 +1856,35 @@ def __init__( self.output_token_count = output_token_count self.total_token_count = total_token_count self.extensionData = extensionData + self.unknown_fields: dict[str, Any] = {} def to_dict(self) -> dict[str, Any]: - result: dict[str, Any] = { + counts: dict[str, Any] = { DurableStateFields.INPUT_TOKEN_COUNT: self.input_token_count, DurableStateFields.OUTPUT_TOKEN_COUNT: self.output_token_count, DurableStateFields.TOTAL_TOKEN_COUNT: self.total_token_count, } + result = {**self.unknown_fields, **{key: value for key, value in counts.items() if value is not None}} if self.extensionData is not None: result[DurableStateFields.EXTENSION_DATA] = self.extensionData - return result + return _json_snapshot(result) @classmethod def from_dict(cls, data: dict[str, Any]) -> DurableAgentStateUsage: - return cls( + usage = cls( input_token_count=data.get(DurableStateFields.INPUT_TOKEN_COUNT), output_token_count=data.get(DurableStateFields.OUTPUT_TOKEN_COUNT), total_token_count=data.get(DurableStateFields.TOTAL_TOKEN_COUNT), extensionData=data.get(DurableStateFields.EXTENSION_DATA), ) + known = { + DurableStateFields.INPUT_TOKEN_COUNT, + DurableStateFields.OUTPUT_TOKEN_COUNT, + DurableStateFields.TOTAL_TOKEN_COUNT, + DurableStateFields.EXTENSION_DATA, + } + usage.unknown_fields = {key: deepcopy(value) for key, value in data.items() if key not in known} + return usage @staticmethod def from_usage(usage: UsageDetails | MutableMapping[str, Any] | None) -> DurableAgentStateUsage | None: @@ -1253,13 +1905,20 @@ def from_usage(usage: UsageDetails | MutableMapping[str, Any] | None) -> Durable def to_usage_details(self) -> UsageDetails: # Convert back to AI SDK UsageDetails - result = UsageDetails( - input_token_count=self.input_token_count, - output_token_count=self.output_token_count, - total_token_count=self.total_token_count, + result = cast( + UsageDetails, + { + key: value + for key, value in ( + (self._INPUT_TOKEN_COUNT, self.input_token_count), + (self._OUTPUT_TOKEN_COUNT, self.output_token_count), + (self._TOTAL_TOKEN_COUNT, self.total_token_count), + ) + if value is not None + }, ) if self.extensionData: - result.update(self.extensionData) # type: ignore[typeddict-item] + result.update(deepcopy(self.extensionData)) # type: ignore[typeddict-item] return result @@ -1310,6 +1969,8 @@ class DurableAgentStateUnknownContent(DurableAgentStateContent): type: str = ContentTypes.UNKNOWN + _NULLABLE_FIELDS: ClassVar[frozenset[str]] = frozenset({DurableStateFields.CONTENT}) + def __init__(self, content: Any) -> None: self.content = content @@ -1322,13 +1983,16 @@ def from_unknown_content(content: Any) -> DurableAgentStateUnknownContent: return DurableAgentStateUnknownContent(content=content.to_dict()) return DurableAgentStateUnknownContent(content=content) + def to_core_content(self) -> Content: + """Leave unknown content extension conventions opaque, as for future raw kinds.""" + return self.to_ai_content() + def to_ai_content(self) -> Content: - if not self.content: - raise Exception("The content is missing and cannot be converted to valid AI content.") content_value: Any = self.content if isinstance(content_value, dict) and "type" in content_value: - try: - return Content.from_dict(cast(dict[str, Any], content_value)) - except (ValueError, TypeError): - pass - return Content(type=self.type, additional_properties={"content": self.content}) # type: ignore + return ( + load_agent_response({"messages": [{"role": "assistant", "contents": [content_value]}]}) + .messages[0] + .contents[0] + ) + return Content(type=self.type, additional_properties={"content": deepcopy(self.content)}) # type: ignore diff --git a/python/packages/durabletask/agent_framework_durabletask/_entities.py b/python/packages/durabletask/agent_framework_durabletask/_entities.py index 63f7098..75d7366 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_entities.py +++ b/python/packages/durabletask/agent_framework_durabletask/_entities.py @@ -4,34 +4,162 @@ from __future__ import annotations +import asyncio import inspect +import json import logging import warnings +from collections.abc import Mapping, Sequence +from copy import copy, deepcopy from datetime import datetime, timezone from typing import Any, cast from agent_framework import ( + Agent, AgentResponse, AgentResponseUpdate, + AgentSession, Content, Message, ResponseStream, SupportsAgentRun, + register_state_type, ) from durabletask.entities import DurableEntity from ._callbacks import AgentCallbackContext, AgentResponseCallbackProtocol +from ._configuration import validate_response_delivery_window from ._durable_agent_state import ( DurableAgentState, DurableAgentStateEntry, + DurableAgentStateErrorResponse, DurableAgentStateMessage, DurableAgentStateRequest, DurableAgentStateResponse, + DurableAgentStateUnknownEntry, ) +from ._history_provider import ( + DurableHistoryBinding, + DurableHistoryProvider, + bind_durable_history, + ensure_durable_history, + prepare_history_owner, + service_stores_history, + unbind_durable_history, +) +from ._invocation_safety import DurableToolGuard, InvocationProgress +from ._message_identity import message_identity from ._models import RunRequest +from ._response_utils import is_terminal_agent_response, load_agent_response +from ._retention import ( + DEFAULT_MAX_STATE_BYTES, + DEFAULT_RETENTION, + DELIVERY_WINDOW_SECONDS, + HIGH_WATERMARK, + LOW_WATERMARK, + RetentionMode, + StateBudget, + enforce_budget, + prunes_excluded, + resolve_state_budget, + validate_retention, +) +from ._retention_telemetry import record_write, retention_operation +from ._state_migration import migrate_legacy_state, state_snapshot_digest logger = logging.getLogger("agent_framework.durabletask") +# Key produced by core's ``AgentSession.to_dict()``. +_SESSION_ID_KEY = "session_id" + +try: + # Root of core's serializable state types. Not part of core's public surface, so a move must + # not break the entity: without it, restored provider state simply stays as plain dicts, + # which is core's own behavior. + from agent_framework._serialization import SerializationMixin + + _SerializableStateRoot: type | None = SerializationMixin +except ImportError: # pragma: no cover - depends on the installed core version + _SerializableStateRoot = None + +_registered_state_types: set[type] = set() + +# Provider error code for a conversation id the service will not accept as a parent turn. +_MISSING_PREVIOUS_RESPONSE_CODE = "previous_response_not_found" + +_REJECTED_ID_RETRIES = 3 +"""How many times to re-send a request whose conversation id the service would not accept. + +Few, because a retry only helps when the id is late rather than gone, and the two are +indistinguishable from the error alone. Enough to cover the gap that was measured, which was +under a second on the chaining path. +""" + +_REJECTED_ID_BACKOFF_SECONDS = 0.5 +"""Multiplied by the attempt number, so the waits are 0.5s, 1s, 1.5s.""" + + +def _is_missing_previous_response(exc: BaseException, *, prior_error: BaseException | None = None) -> bool: + """Return whether the service refused the conversation id from the previous turn. + + A service that keeps the conversation can hand back the id of a finished response before that + response is durably readable, so the next turn is refused even though the id is genuine and + was captured correctly. Bounded identical-request retries may recover visibility delays; + genuinely expired IDs still fail. No transcript recovery is attempted. + + Match only the structured error code, including wrapped causes, so unrelated + request failures are not retried as conversation visibility failures. + """ + seen: set[int] = set() + current: BaseException | None = exc + while current is not None and id(current) not in seen: + if current is prior_error: + return False + seen.add(id(current)) + code = getattr(current, "code", None) + if code is not None: + return code == _MISSING_PREVIOUS_RESPONSE_CODE + body = getattr(current, "body", None) + if isinstance(body, Mapping): + details = cast("Mapping[str, Any]", body) + if "code" in details: + return details["code"] == _MISSING_PREVIOUS_RESPONSE_CODE + current = current.__cause__ or current.__context__ + return False + + +def _register_loaded_state_types() -> None: + """Let core restore session state values as their own classes after a cold start. + + Core deserializes session state through a type registry that it seeds with exactly one entry + (``Message``). Anything else must be registered explicitly, and the registry is process-local. + A durable entity routinely restores state in a process that never serialized it, so without + this a provider's state comes back as a plain dict rather than its own class. + + Only classes already imported in this process are registered - nothing is imported from + persisted data - so this cannot load code the application has not already loaded itself. That + is enough in practice, because whoever put a value in the state bag had to import its class to + construct it. + """ + if _SerializableStateRoot is None: + return + + seen: set[type] = set() + pending: list[type] = [_SerializableStateRoot] + while pending: + for subclass in pending.pop().__subclasses__(): + if subclass in seen: + continue + seen.add(subclass) + pending.append(subclass) + if subclass in _registered_state_types: + continue + _registered_state_types.add(subclass) + try: + register_state_type(subclass) + except Exception: + logger.debug("Could not register session state type %s", subclass, exc_info=True) + class AgentEntityStateProviderMixin: """Mixin implementing durable agent state caching + (de)serialization + persistence. @@ -63,10 +191,33 @@ def _get_session_id_from_entity(self) -> str: return cast(str, legacy_hook()) raise NotImplementedError + def _get_entity_name_from_entity(self) -> str: + """Return the entity name, when the host exposes one. + + Optional, so state providers written before this hook existed keep working. They fall + back to a core session id built from the key alone. + """ + return "" + @property def session_id(self) -> str: return self._get_session_id_from_entity() + @property + def core_session_id(self) -> str: + """Identity handed to core's ``create_session``, unique to this entity. + + ``session_id`` is only the entity key, which is not unique on its own. Every agent node + in one workflow run shares a key (the orchestration instance id) and is told apart by + entity name, so an external history provider keyed on the key alone would mix the + histories of different nodes. The name is included here to keep them separate. + + Uses the same ``@name@key`` form as :class:`AgentSessionId`, so the result parses back. + """ + name = self._get_entity_name_from_entity() + key = self.session_id + return f"@{name}@{key}" if name else key + @property def thread_id(self) -> str: """Deprecated alias for :attr:`session_id`.""" @@ -90,10 +241,25 @@ def state(self, value: DurableAgentState) -> None: self.persist_state() def persist_state(self) -> None: - """Persist the current state to the underlying storage provider.""" + """Pass state to the host, which may stage rather than confirm a durable write.""" if self._state_cache is None: self._state_cache = DurableAgentState() - self._set_state_dict(self._state_cache.to_dict()) + state = self._state_cache + try: + payload = state.to_dict() + except BaseException: + record_write(state, stage="serialization", outcome="failed") + raise + try: + self._set_state_dict(payload) + except BaseException: + record_write(state, stage="set_state", outcome="failed") + raise + record_write(state, stage="set_state", outcome="returned") + + def replace_cached_state(self, state: DurableAgentState) -> None: + """Stage or restore an operation snapshot without writing to the backend.""" + self._state_cache = state def reset(self) -> None: """Clear conversation history by resetting state to a fresh DurableAgentState.""" @@ -117,10 +283,24 @@ def __init__( callback: AgentResponseCallbackProtocol | None = None, *, state_provider: AgentEntityStateProviderMixin, + retention: RetentionMode = DEFAULT_RETENTION, + max_state_bytes: StateBudget = DEFAULT_MAX_STATE_BYTES, + high_watermark: float = HIGH_WATERMARK, + low_watermark: float = LOW_WATERMARK, + response_delivery_window_seconds: int = DELIVERY_WINDOW_SECONDS, ) -> None: - self.agent = agent + validate_retention(retention, high_watermark, low_watermark) + validate_response_delivery_window(response_delivery_window_seconds) + # Back the agent's conversation history with durable entity state so an agent that + # already works in core runs durably without any configuration change. + self.agent = ensure_durable_history(agent, prune_excluded=prunes_excluded(retention)) self.callback = callback self._state_provider = state_provider + self._retention = retention + self._max_state_bytes = resolve_state_budget(max_state_bytes) + self._high_watermark = high_watermark + self._low_watermark = low_watermark + self._response_delivery_window_seconds = response_delivery_window_seconds logger.debug("[AgentEntity] Initialized with agent type: %s", type(agent).__name__) @@ -135,14 +315,129 @@ def state(self, value: DurableAgentState) -> None: def persist_state(self) -> None: self._state_provider.persist_state() + def expire_responses(self) -> int: + """Remove expired delivery payloads without model execution or deleting receipts. + + Hosts expose this maintenance operation for an application-owned schedule. + An idle entity has no timer of its own; availability expires independently. + + Returns: + The number of payloads removed by this operation. + """ + original = self.state + original.prepare_for_write(delivery_window_seconds=self._response_delivery_window_seconds) + staged = deepcopy(original) + before = len(staged.data.response_mailbox) + staged.expire_responses() + removed = before - len(staged.data.response_mailbox) + if not removed: + return 0 + self._state_provider.replace_cached_state(staged) + try: + self._validate_control_budget() + self.persist_state() + except BaseException: + self._state_provider.replace_cached_state(original) + raise + return removed + + def migrate(self, request: dict[str, Any]) -> dict[str, str]: + """Import a quiesced legacy snapshot into an empty, separately addressed entity. + + This privileged backend operation is not exposed through the generated HTTP + or MCP routes. The deployment owner must authorize the source export, journal + and ownership transfer. No runtime can inspect or fence a legacy deployment. + Retries with the exact same request return the recorded migration, even after + subsequent runs, without rewriting state or refreshing response grace. + + Args: + request: Source snapshot/digest, sourceSessionId, destinationSessionId, + migrationId, ownershipTransferId and optional deliveryEvidence and + requireKnownOutcomes, which rejects imports without outcome evidence. + + Returns: + The committed migration ID and destination session identity. + """ + required = { + "source", + "sourceDigest", + "sourceSessionId", + "destinationSessionId", + "migrationId", + "ownershipTransferId", + } + if ( + not isinstance(request, dict) + or not required <= request.keys() + or request.keys() - required - {"deliveryEvidence", "requireKnownOutcomes"} + ): + raise ValueError("Migration requires a complete explicit source and destination request.") + for name in required - {"source"}: + if not isinstance(request[name], str) or not request[name].strip(): + raise ValueError(f"Migration {name} must be a nonblank string.") + destination = self._state_provider.core_session_id + if request["destinationSessionId"] != destination: + raise ValueError("Migration destinationSessionId does not match this entity.") + if request["sourceSessionId"] == destination: + raise ValueError("Migration requires a separately addressed destination, never an in-place rewrite.") + if not isinstance(request["source"], dict): + raise ValueError("Migration source must be an exported state object.") + digest = state_snapshot_digest(request) + original = self.state + existing = original.data.unknown_fields.get("migration") + if isinstance(existing, dict) and cast("dict[str, Any]", existing).get("requestDigest") == digest: + return {"status": "migrated", "migrationId": request["migrationId"], "sessionId": destination} + if original.to_dict() != DurableAgentState().to_dict(): + raise ValueError( + "Migration destination must be empty; an existing or different migration cannot be replaced." + ) + staged = migrate_legacy_state( + cast("dict[str, Any]", request["source"]), + source_digest=request["sourceDigest"], + source_session_id=request["sourceSessionId"], + migration_id=request["migrationId"], + ownership_transfer_id=request["ownershipTransferId"], + delivery_window_seconds=self._response_delivery_window_seconds, + delivery_evidence=request.get("deliveryEvidence"), + require_known_outcomes=request.get("requireKnownOutcomes", False), + ) + staged.data.unknown_fields["migration"].update({"requestDigest": digest, "destinationSessionId": destination}) + self._state_provider.replace_cached_state(staged) + try: + self._validate_control_budget() + self.persist_state() + except BaseException: + self._state_provider.replace_cached_state(original) + raise + return {"status": "migrated", "migrationId": request["migrationId"], "sessionId": destination} + + def _validate_control_budget(self) -> None: + """Reject an oversized maintenance commit, without pruning any protected state.""" + if self._max_state_bytes is not None: + size = len(json.dumps(self.state.to_dict(), allow_nan=False)) + if size > self._max_state_bytes: + raise ValueError("Retained delivery/control state cannot fit within max_state_bytes.") + def reset(self) -> None: - self._state_provider.reset() + """Clear local history/session context without erasing execution receipts.""" + if self._has_context_pipeline() and self._find_durable_history_provider() is None: + raise NotImplementedError("Reset of external history requires a provider-owned clear operation.") + original = self.state + self._state_provider.replace_cached_state(deepcopy(original)) + try: + self.state.prepare_for_write(delivery_window_seconds=self._response_delivery_window_seconds) + self.state.data.conversation_history.clear() + self.state.data.session = None + self.state.expire_responses() + self._validate_control_budget() + self.persist_state() + except BaseException: + self._state_provider.replace_cached_state(original) + raise def _is_error_response(self, entry: DurableAgentStateEntry) -> bool: - """Check if a conversation history entry is an error response.""" - if isinstance(entry, DurableAgentStateResponse): - return entry.is_error - return False + """Check if a conversation history entry records a failed turn.""" + return isinstance(entry, (DurableAgentStateErrorResponse, DurableAgentStateUnknownEntry)) async def run( self, @@ -156,6 +451,29 @@ async def run( else: run_request = request + # A read-compatible legacy layout is not permission to run a new writer. + self.state.prepare_for_write(delivery_window_seconds=self._response_delivery_window_seconds) + already_answered = self.state.try_get_agent_response(run_request.correlation_id) + if already_answered is not None: + self.expire_responses() + return already_answered + original = self.state + self._state_provider.replace_cached_state(deepcopy(original)) + with retention_operation(self.state): + try: + self.state.expire_responses() + response = await self._execute_request(run_request) + await self._enforce_retention() + self.persist_state() + return response + except BaseException: + # A failed commit must not leave a warm worker with staged completion or + # ingestion receipts. External effects are outside this local rollback. + self._state_provider.replace_cached_state(original) + raise + + async def _execute_request(self, run_request: RunRequest) -> AgentResponse: + """Stage a turn without committing until every local slice and budget is valid.""" message = run_request.message session_id = self._state_provider.session_id correlation_id = run_request.correlation_id @@ -163,55 +481,499 @@ async def run( raise ValueError("Entity State Provider must provide a session_id") options: dict[str, Any] = dict(run_request.options) options.setdefault("response_format", run_request.response_format) - if not run_request.enable_tool_calls: - options.setdefault("tools", None) logger.debug("[AgentEntity.run] Received SessionId %s Message: %s", session_id, run_request) + durable_history = self._find_durable_history_provider() + uses_context_pipeline = self._has_context_pipeline() + # A property of the run rather than of the registration, since ``store`` is an ordinary + # run option. The provider stays attached either way so core never injects one of its own. + service_owns_history = service_stores_history(self.agent, options) + prior_receipts = deepcopy(self.state.data.ingested_messages) state_request = DurableAgentStateRequest.from_run_request(run_request) - self.state.data.conversation_history.append(state_request) - - try: - chat_messages: list[Message] = [ - replayable_message - for entry in self.state.data.conversation_history - if not self._is_error_response(entry) - for m in entry.messages - if (replayable_message := self._to_replayable_message(m)) is not None - ] - - run_kwargs: dict[str, Any] = {"messages": chat_messages, "options": options} - - agent_run_response: AgentResponse = await self._invoke_agent( - run_kwargs=run_kwargs, - correlation_id=correlation_id, - session_id=session_id, - request_message=message, + if run_request.context_messages is not None: + state_request.messages = self._drop_already_stored( + state_request.messages, occurrence_ids=run_request.context_message_ids + ) + if not uses_context_pipeline: + self.state.data.conversation_history.append(state_request) + + binding_token = ( + bind_durable_history( + DurableHistoryBinding( + state_provider=self._state_provider, + correlation_id=correlation_id, + # The provider stays attached either way so core never injects one of its own, + # but it must not load history on a turn the service is already carrying. + service_owns_history=service_owns_history, + ) ) + if durable_history is not None + else None + ) - state_response = DurableAgentStateResponse.from_run_response(correlation_id, agent_run_response) - self.state.data.conversation_history.append(state_response) - self.persist_state() + # Bound before the try so the failure path can always reach it. ``_create_session`` can + # raise, and referencing an unbound name while handling that would replace the agent's + # error with a NameError. + session: Any = None + inactive_service_id: Any = None + succeeded = False + original_agent = self.agent + progress = InvocationProgress() - return agent_run_response + try: + self.agent = prepare_history_owner(self.agent, service_owns_history) + if not run_request.enable_tool_calls: + invocation_agent = copy(self.agent) + # Core merges default, context-provider, MCP and additional tools. A + # model option alone cannot disable the local invocation loop. + defaults = getattr(invocation_agent, "default_options", None) + if isinstance(defaults, Mapping): + invocation_agent.default_options = { # type: ignore[attr-defined] + **cast("Mapping[str, Any]", defaults), + "tools": [], + "tool_choice": "none", + } + client = getattr(invocation_agent, "client", None) + invocation_configuration = getattr(client, "function_invocation_configuration", None) + if client is not None and isinstance(invocation_configuration, Mapping): + invocation_client = copy(client) + invocation_client.function_invocation_configuration = { + **cast("Mapping[str, Any]", invocation_configuration), + "enabled": False, + } + invocation_agent.client = invocation_client # type: ignore[attr-defined] + if isinstance(getattr(invocation_agent, "mcp_tools", None), list): + invocation_agent.mcp_tools = [] # type: ignore[attr-defined] + self.agent = invocation_agent + options["tools"] = [] + options["tool_choice"] = "none" + if uses_context_pipeline: + # The agent's own context providers supply prior turns - durable-backed history, + # an external store (Cosmos/Redis/file), or the model service itself. Only the + # newly received request messages are passed as run input, so history lives in + # its selected store, subject to the service-owned branch's inactive-primary gate. + session = self._create_session() + if not service_owns_history: + inactive_service_id = getattr(session, "service_session_id", None) + session.service_session_id = None + # A conversation ID supplied through defaults/options must not + # override the client-owned branch either. Copy, never mutate + # the agent the application may be using elsewhere. + defaults = getattr(self.agent, "default_options", None) + if isinstance(defaults, Mapping) and "conversation_id" in defaults: + invocation_agent = copy(self.agent) + invocation_agent.default_options = { # type: ignore[attr-defined] + key: value + for key, value in cast("Mapping[str, Any]", defaults).items() + if key != "conversation_id" + } + self.agent = invocation_agent + options.pop("conversation_id", None) + chat_messages: list[Message] = [] + # Core's operation-local copies retain private attributes, while its + # serializers exclude them. This receipt follows the actual appended + # input, even when two equal inputs carry different transport IDs. + for stored in state_request.messages: + current = self._to_current_message(stored, run_request) + if current is None: + continue + if stored.ingestion_occurrence and stored.ingestion_identity: + current._durable_ingestion_receipt = ( # type: ignore[attr-defined] + stored.ingestion_occurrence, + stored.ingestion_identity, + ) + chat_messages.append(current) + run_kwargs: dict[str, Any] = { + "messages": chat_messages, + "session": session, + "options": options, + } + else: + # Fallback for agents without the core context pipeline (for example a fully + # custom agent): the entity replays the persisted conversation on every turn. + session = None + chat_messages = self._replay_all_messages() + run_kwargs = {"messages": chat_messages, "options": options} + + if isinstance(self.agent, Agent): + run_kwargs["client_kwargs"] = { + "middleware": [DurableToolGuard(progress, enabled=run_request.enable_tool_calls)] + } + original_service_id = getattr(session, "service_session_id", None) + try: + agent_run_response: AgentResponse = await self._invoke_agent( + run_kwargs=run_kwargs, + correlation_id=correlation_id, + session_id=session_id, + request_message=message, + progress=progress, + ) + except Exception as exc: + if ( + session is None + or not service_owns_history + or not _is_missing_previous_response(exc) + or progress.stream_started + or progress.function_started + or getattr(session, "service_session_id", None) != original_service_id + ): + raise + retried = await self._retry_rejected_conversation_id( + run_kwargs=run_kwargs, + correlation_id=correlation_id, + session_id=session_id, + request_message=message, + cause=exc, + progress=progress, + original_service_id=original_service_id, + ) + if retried is None: + raise + agent_run_response = retried + + # Resolve structured output inside the runtime-error boundary. A parsing + # error is a committed error result, not an invisible post-run failure. + succeeded = not is_terminal_agent_response(agent_run_response) + if ( + succeeded + and not agent_run_response.user_input_requests + and agent_run_response.additional_properties.get("durable_status") != "accepted" + ): + _ = agent_run_response.value except Exception as exc: + succeeded = False logger.exception("[AgentEntity.run] Agent execution failed.") + # The entity absorbs failures rather than faulting, so the session survives and the + # caller can take the next turn. That is only reasonable if the caller can tell what + # happened: error content alone leaves ``response.text`` empty, which reads as the + # agent having nothing to say. The text carries the same message the error content + # already holds, so callers inspecting contents see no change. + detail = f"{type(exc).__name__}: {exc}" error_message = Message( - role="assistant", contents=[Content.from_error(message=str(exc), error_code=type(exc).__name__)] + role="assistant", + contents=[ + Content.from_error(message=str(exc), error_code=type(exc).__name__), + Content.from_text(detail), + ], ) - error_response = AgentResponse( + agent_run_response = AgentResponse( messages=[error_message], created_at=datetime.now(tz=timezone.utc).isoformat(), + additional_properties={"durable_status": "error", "correlation_id": correlation_id}, ) - error_state_response = DurableAgentStateResponse.from_run_response(correlation_id, error_response) - error_state_response.is_error = True - self.state.data.conversation_history.append(error_state_response) - self.persist_state() + finally: + try: + if session is not None and durable_history is not None and not service_owns_history: + if not succeeded: + durable_history.finalize_failed_run(session.state.get(durable_history.source_id, {})) + durable_history.flush(session.state.get(durable_history.source_id, {})) + finally: + if session is not None and not service_owns_history: + session.service_session_id = inactive_service_id + if binding_token is not None: + unbind_durable_history(binding_token) + self.agent = original_agent + + if not succeeded and uses_context_pipeline: + # A failed pre-invocation/provider load did not deliver these messages. + # Retain receipts only for inputs actually staged by durable history; + # no portable external provider API proves an interrupted append. + staged_inputs = { + (stored.ingestion_occurrence, stored.ingestion_identity) + for entry in self.state.data.conversation_history + if isinstance(entry, DurableAgentStateRequest) and entry.correlation_id == correlation_id + for stored in entry.messages + } + self.state.data.ingested_messages = prior_receipts + for stored in state_request.messages: + identity = stored.ingestion_occurrence or stored.message_id + if identity and (identity, stored.ingestion_identity) in staged_inputs: + fingerprints = self.state.data.ingested_messages.get(identity, []) + if fingerprints is not None and stored.ingestion_identity: + if stored.ingestion_identity not in fingerprints: + fingerprints.append(stored.ingestion_identity) + self.state.data.ingested_messages[identity] = fingerprints + self.state.record_response( + correlation_id, + agent_run_response, + delivery_window_seconds=self._response_delivery_window_seconds, + ) + if not uses_context_pipeline and succeeded: + self.state.data.conversation_history.append( + DurableAgentStateResponse.from_run_response(correlation_id, agent_run_response) + ) + self._capture_session(session) + return agent_run_response + + async def _retry_rejected_conversation_id( + self, + *, + run_kwargs: dict[str, Any], + correlation_id: str, + session_id: str, + request_message: Any, + cause: BaseException, + progress: InvocationProgress, + original_service_id: Any, + ) -> AgentResponse | None: + """Re-send an identical request whose conversation id the service refused. + + The refusal we are recovering from is a read-after-write gap rather than a lost + conversation. Measured against Azure OpenAI, a streamed response reports its id in the + completion event before that response is retrievable, so the very next turn can be + rejected for naming an id that is perfectly valid and simply not readable yet. Waiting + briefly and asking again is the cheapest thing that works, and it leaves the conversation + continuing from the same point rather than restarting it from a resent transcript. + + A retry cannot rescue an id that has genuinely expired, and the error is identical either + way, so the attempts are few and short. Exhausting them fails the turn without + reconstructing a transcript or starting a different service conversation. + + Args: + run_kwargs: The unchanged arguments of the request that was refused. + correlation_id: Correlation id of the in-flight request. + session_id: Session the request belongs to. + request_message: The originating message, for logging. + cause: The refusal that triggered this, so a give-up is reported with its reason. + progress: Run-local observations that prohibit restarting after stream or tool progress. + original_service_id: Session continuation before the first attempt; retries must not advance it. + + Returns: + The response, or None when every attempt was refused the same way. + """ + for attempt in range(1, _REJECTED_ID_RETRIES + 1): + await asyncio.sleep(_REJECTED_ID_BACKOFF_SECONDS * attempt) + try: + response: AgentResponse = await self._invoke_agent( + run_kwargs=run_kwargs, + correlation_id=correlation_id, + session_id=session_id, + request_message=request_message, + progress=progress, + ) + except Exception as retry_exc: + if ( + not _is_missing_previous_response(retry_exc, prior_error=cause) + or progress.stream_started + or progress.function_started + or getattr(run_kwargs.get("session"), "service_session_id", None) != original_service_id + ): + raise + logger.debug( + "[AgentEntity.run] Conversation id still not accepted for session %s (attempt %d of %d).", + session_id, + attempt, + _REJECTED_ID_RETRIES, + ) + continue + logger.info( + "[AgentEntity.run] Conversation id for session %s was accepted on attempt %d, " + "so the turn continued without resending the transcript.", + session_id, + attempt, + ) + return response + + logger.debug( + "[AgentEntity.run] Conversation id for session %s was refused on every attempt. %s", + session_id, + cause, + ) + return None + + async def _enforce_retention(self) -> None: + """Apply optional whole-state pressure budgeting independently of eager pruning.""" + if self._max_state_bytes is None: + return + await enforce_budget( + self.state, + max_state_bytes=self._max_state_bytes, + high_watermark=self._high_watermark, + low_watermark=self._low_watermark, + ) + + def _has_context_pipeline(self) -> bool: + """Whether the agent exposes core's context-provider pipeline. + + When it does, the providers own conversation context and the entity delivers only the + new messages. Agents without it fall back to replaying persisted history. + """ + return isinstance(getattr(self.agent, "context_providers", None), (list, tuple)) - return error_response + def _capture_session(self, session: Any) -> None: + """Persist the session so provider state survives to the next turn. + + The entity creates a fresh session per operation, so anything the context providers keep + in the session state bag - tool approval rules and queued approval requests, todo lists, + memory extraction state - would otherwise be discarded at the end of every turn. Core + documents that state as durable for the life of the session, so agents that rely on it + must behave the same way here. The serialized session also carries the service-issued + conversation id, so service-backed agents continue the same thread. + + Omit only the durable provider's working message buffer and position index, which are + rebuilt from ``conversation_history`` each turn. Keep its other JSON-compatible state. + Removing those transient fields before serialization avoids encoding a second transcript. + + Core can return live objects from session serialization. Validate the payload before + staging it so an unusable session fails the operation without replacing committed state. + """ + if session is None: + return + to_dict = getattr(session, "to_dict", None) + if not callable(to_dict): + return + + durable_history = self._find_durable_history_provider() + session_state = getattr(session, "state", None) + transient: Any = None + has_transient = False + if durable_history is not None and isinstance(session_state, dict): + bag = cast("dict[str, Any]", session_state) + if durable_history.source_id in bag: + transient = bag.pop(durable_history.source_id) + has_transient = True + if isinstance(transient, dict): + persistent = { + key: value + for key, value in cast("dict[str, Any]", transient).items() + if key not in ("messages", "_positions") + } + if persistent: + bag[durable_history.source_id] = persistent + try: + payload = cast("dict[str, Any]", to_dict()) + finally: + if has_transient: + cast("dict[str, Any]", session_state)[durable_history.source_id] = transient # type: ignore[union-attr] + + try: + json.dumps(payload, allow_nan=False) + except (TypeError, ValueError) as exc: + raise ValueError("Agent session state is not JSON-compatible; the operation cannot commit.") from exc + previous = self.state.data.session + if isinstance(previous, dict): + opaque = { + key: deepcopy(value) + for key, value in previous.items() + if key not in {"type", "session_id", "service_session_id", "state"} and key not in payload + } + payload = {**opaque, **payload} + self.state.data.session = payload + + def _drop_already_stored( + self, messages: list[DurableAgentStateMessage], *, occurrence_ids: list[str] | None = None + ) -> list[DurableAgentStateMessage]: + """Remember actual identities, including skipped positions and content revisions. + + Receipts outlive transcript eviction. Anonymous direct inputs are not content- + deduplicated; the workflow sender supplies scoped IDs for anonymous projections. + Entirely repeated projections stay empty instead of re-ingesting their last item. + """ + receipts = self.state.data.ingested_messages + kept: list[DurableAgentStateMessage] = [] + for index, message in enumerate(messages): + identity = occurrence_ids[index] if occurrence_ids is not None else message.message_id + if identity: + fingerprint = message.ingestion_identity or message_identity(message.to_chat_message()) + known = receipts.get(identity, []) + if known is None or fingerprint in known: + continue + known.append(fingerprint) + receipts[identity] = known + message.ingestion_occurrence = identity + kept.append(message) + return kept + + def _find_durable_history_provider(self) -> DurableHistoryProvider | None: + """Return the agent's :class:`DurableHistoryProvider`, if it is configured with one.""" + providers = getattr(self.agent, "context_providers", None) + if not isinstance(providers, (list, tuple)): + return None + for provider in cast("Sequence[Any]", providers): + if isinstance(provider, DurableHistoryProvider): + return provider + return None + + def _create_session(self) -> Any: + """Create the session for this operation and restore what the last turn left on it. + + Conversation history lives in the agent's context providers (durable entity state, an + external store, or the model service), so a fresh session per operation is enough - but it + must carry the entity's **stable** session id. External history providers (Cosmos, Redis, + file) key their storage on ``session.session_id``, and with a freshly generated id they would + read and write a different key every turn and never see prior history. + + The id is qualified with the entity name (see ``core_session_id``) because the key alone + collides across the agent nodes of one workflow run. + """ + create_session = getattr(self.agent, "create_session", None) + if not callable(create_session): + raise TypeError( + f"Agent {type(self.agent).__name__} exposes context providers but does not support create_session()." + ) + migration = self.state.data.unknown_fields.get("migration") + logical_session_id = self._state_provider.core_session_id + if isinstance(migration, dict): + source_session_id = cast("dict[str, Any]", migration).get("sourceSessionId") + if not isinstance(source_session_id, str) or not source_session_id.strip(): + raise ValueError("Migration sourceSessionId must preserve the original logical session identity.") + logical_session_id = source_session_id + session: Any = create_session(session_id=logical_session_id) + self._restore_session(session) + return session + + def _restore_session(self, session: Any) -> None: + """Apply the previous turn's session state onto a freshly created session. + + The agent's own ``create_session`` is used so its session type is preserved. Only the + state bag and the service conversation id are carried over. + """ + stored = self.state.data.session + if not stored or _SESSION_ID_KEY not in stored: + return + + # Done here rather than at import: by now the agent and its providers are built, so the + # classes their state uses are loaded and can be resolved. + _register_loaded_state_types() + + restored = AgentSession.from_dict(dict(stored)) + session.state.update(restored.state) + if getattr(session, "service_session_id", None) is None: + session.service_session_id = restored.service_session_id + + def _replay_all_messages(self) -> list[Message]: + """Build run input from the whole persisted transcript. + + Used only for agents without the core context pipeline. Service conversation + errors do not trigger local transcript reconstruction. + + Failed turns are skipped so an error reply is never presented back to the model as + something it said. + """ + return [ + replayable_message + for entry in self.state.data.conversation_history + if not self._is_error_response(entry) + for m in entry.messages + if (replayable_message := self._to_replayable_message(m)) is not None + ] + + @staticmethod + def _to_current_message(message: DurableAgentStateMessage, request: RunRequest) -> Message | None: + """Preserve core input content metadata rather than round-tripping through legacy types.""" + if request.context_messages is not None and message.ingestion_identity: + for raw in request.context_messages: + original = load_agent_response({"messages": [raw]}).messages[0] + if ( + original.message_id == message.message_id + and message_identity(original) == message.ingestion_identity + ): + return original + return AgentEntity._to_replayable_message(message) @staticmethod def _to_replayable_message(message: DurableAgentStateMessage) -> Message | None: @@ -225,6 +987,7 @@ def _to_replayable_message(message: DurableAgentStateMessage) -> Message | None: role=chat_message.role, contents=replayable_contents, author_name=chat_message.author_name, + message_id=chat_message.message_id, additional_properties=chat_message.additional_properties, ) @@ -234,6 +997,7 @@ async def _invoke_agent( correlation_id: str, session_id: str, request_message: str, + progress: InvocationProgress | None = None, ) -> AgentResponse: """Execute the agent, preferring streaming when available.""" callback_context: AgentCallbackContext | None = None @@ -246,28 +1010,32 @@ async def _invoke_agent( run_callable = self.agent.run - # Try streaming first with run(stream=True) + # Only negotiate an unsupported streaming signature before consuming a stream. + # Errors raised while consuming it must never restart model/tool execution. try: stream_candidate = run_callable(stream=True, **run_kwargs) if inspect.isawaitable(stream_candidate): stream_candidate = await stream_candidate - - return await self._consume_stream( - stream=stream_candidate, - callback_context=callback_context, - ) except TypeError as type_error: - if "__aiter__" not in str(type_error) and "stream" not in str(type_error): + detail = str(type_error) + if not ( + "stream is not supported" in detail + or "streaming not supported" in detail + or "unexpected keyword argument 'stream'" in detail + or 'unexpected keyword argument "stream"' in detail + ): raise logger.debug( - "run(stream=True) returned a non-async result; falling back to run(): %s", + "Agent does not support streaming; invoking non-streaming run(): %s", type_error, ) - except Exception as stream_error: - logger.warning( - "run(stream=True) failed; falling back to run(): %s", - stream_error, - exc_info=True, + else: + if isinstance(stream_candidate, AgentResponse): + direct_response = cast(AgentResponse, stream_candidate) + await self._notify_final_response(direct_response, callback_context) + return direct_response + return await self._consume_stream( + stream=stream_candidate, callback_context=callback_context, progress=progress ) agent_run_response = run_callable(**run_kwargs) if inspect.isawaitable(agent_run_response): @@ -284,12 +1052,12 @@ async def _consume_stream( self, stream: ResponseStream[AgentResponseUpdate, AgentResponse], callback_context: AgentCallbackContext | None = None, + progress: InvocationProgress | None = None, ) -> AgentResponse: """Consume streaming responses and build the final AgentResponse.""" - updates: list[AgentResponseUpdate] = [] - async for update in stream: - updates.append(update) + if progress is not None: + progress.stream_started = True await self._notify_stream_update(update, callback_context) response = await stream.get_final_response() @@ -307,7 +1075,7 @@ async def _notify_stream_update( return try: - callback_result = self.callback.on_streaming_response_update(update, context) + callback_result = self.callback.on_streaming_response_update(deepcopy(update), context) if inspect.isawaitable(callback_result): await callback_result except Exception as exc: @@ -327,7 +1095,14 @@ async def _notify_final_response( return try: - callback_result = self.callback.on_agent_response(response, context) + snapshot = deepcopy(response) + # Core deliberately shares opaque SDK representations during deepcopy. + # Detach them when possible, otherwise omit only that opaque field. + try: + snapshot.raw_representation = deepcopy(response.raw_representation) + except Exception: + snapshot.raw_representation = None + callback_result = self.callback.on_agent_response(snapshot, context) if inspect.isawaitable(callback_result): await callback_result except Exception as exc: @@ -372,3 +1147,6 @@ def _set_state_dict(self, state: dict[str, Any]) -> None: def _get_session_id_from_entity(self) -> str: return self.entity_context.entity_id.key + + def _get_entity_name_from_entity(self) -> str: + return self.entity_context.entity_id.entity diff --git a/python/packages/durabletask/agent_framework_durabletask/_executors.py b/python/packages/durabletask/agent_framework_durabletask/_executors.py index eea17ef..912d55f 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_executors.py +++ b/python/packages/durabletask/agent_framework_durabletask/_executors.py @@ -87,12 +87,11 @@ def on_child_completed(self, task: Task[Any]) -> None: try: response = load_agent_response(raw_result) - if self._response_format is not None: - ensure_response_format( - self._response_format, - self._correlation_id, - response, - ) + ensure_response_format( + self._response_format, + self._correlation_id, + response, + ) # Set the typed AgentResponse as this task's result self.complete(response) @@ -155,11 +154,21 @@ def generate_unique_id(self) -> str: """Generate a new Unique ID.""" return uuid.uuid4().hex + def _orchestration_id(self) -> str | None: + """Return the orchestration instance that issued this request. + + Overridden by executors that run inside an orchestration. Client-side executors + have no orchestration, so the default is ``None``. + """ + return None + def get_run_request( self, message: str, *, options: dict[str, Any] | None = None, + context_messages: list[dict[str, Any]] | None = None, + context_message_ids: list[str] | None = None, ) -> RunRequest: """Create a RunRequest from message and options.""" correlation_id = self.generate_unique_id() @@ -179,6 +188,9 @@ def get_run_request( wait_for_response=wait_for_response, correlation_id=correlation_id, options=opts, + context_messages=context_messages, + context_message_ids=context_message_ids, + orchestration_id=self._orchestration_id(), ) def _create_acceptance_response(self, correlation_id: str) -> AgentResponse: @@ -203,6 +215,7 @@ def _create_acceptance_response(self, correlation_id: str) -> AgentResponse: return AgentResponse( messages=[acceptance_message], created_at=datetime.now(timezone.utc).isoformat(), + additional_properties={"durable_status": "accepted", "correlation_id": correlation_id}, ) @@ -357,12 +370,11 @@ def _handle_agent_response( if agent_response is not None: try: # Validate response format if specified - if response_format is not None: - ensure_response_format( - response_format, - correlation_id, - agent_response, - ) + ensure_response_format( + response_format, + correlation_id, + agent_response, + ) return agent_response @@ -449,23 +461,8 @@ def generate_unique_id(self) -> str: """Create a new UUID that is safe for replay within an orchestration or operation.""" return self._context.new_uuid() - def get_run_request( - self, - message: str, - *, - options: dict[str, Any] | None = None, - ) -> RunRequest: - """Get the current run request from the orchestration context. - - Returns: - RunRequest: The current run request - """ - request = super().get_run_request( - message, - options=options, - ) - request.orchestration_id = self._context.instance_id - return request + def _orchestration_id(self) -> str | None: + return self._context.instance_id def run_durable_agent( self, diff --git a/python/packages/durabletask/agent_framework_durabletask/_history_provider.py b/python/packages/durabletask/agent_framework_durabletask/_history_provider.py new file mode 100644 index 0000000..2fc830d --- /dev/null +++ b/python/packages/durabletask/agent_framework_durabletask/_history_provider.py @@ -0,0 +1,1043 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""A core ``HistoryProvider`` backed by durable entity state. + +Core history hooks stage transcript appends, while compaction annotations and summaries +are reconciled from a transient working buffer. The entity commits the transcript together +with its independent delivery and control state at the operation boundary. + +See ADR-0032 (durable thread compaction). +""" + +from __future__ import annotations + +import copy +import logging +from collections.abc import Iterator, Mapping, Sequence +from contextvars import ContextVar, Token +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Any, cast + +from agent_framework import ( + GROUP_ANNOTATION_KEY, + GROUP_ID_KEY, + SUMMARIZED_BY_SUMMARY_ID_KEY, + SUMMARY_OF_MESSAGE_IDS_KEY, + AgentResponse, + HistoryProvider, + InMemoryHistoryProvider, + Message, + SessionContext, + SupportsAgentRun, + annotate_message_groups, +) + +from ._durable_agent_state import ( + DurableAgentStateCompaction, + DurableAgentStateEntry, + DurableAgentStateEntryJsonType, + DurableAgentStateErrorResponse, + DurableAgentStateFunctionCallContent, + DurableAgentStateFunctionResultContent, + DurableAgentStateMessage, + DurableAgentStateRequest, + DurableAgentStateResponse, + DurableAgentStateUnknownEntry, + DurableAgentStateUsage, +) +from ._response_utils import is_terminal_agent_response +from ._retention_telemetry import eager_state_size, record_retention + +if TYPE_CHECKING: + from ._entities import AgentEntityStateProviderMixin + +logger = logging.getLogger("agent_framework.durabletask") + +WORKING_BUFFER_KEY = "messages" +POSITIONS_KEY = "_positions" +EXCLUDED_KEY = "_excluded" + + +@dataclass +class DurableHistoryBinding: + """Per-operation binding between a durable entity and the history provider.""" + + state_provider: AgentEntityStateProviderMixin + """The entity state provider whose conversation history backs the agent.""" + + correlation_id: str | None = None + """Owner of every request and response appended during this operation.""" + + service_owns_history: bool = False + """Whether the model service is holding the conversation for *this* run. + + The provider stays available when no external primary was selected, so that core never injects + a separate in-memory transcript outside retention's reach. A + service-backed run continues the conversation by id rather than by resending it, so loading + history here as well would hand the model the whole transcript on top of the copy the service + already has. Whoever owns a given run is only known once its options are resolved, which is + why this rides on the binding rather than on the provider. + """ + + append_ordinal: int = 0 + """Operation-local append counter used to give anonymous messages stable stored identities.""" + + pending_inputs: list[Message] = field(default_factory=lambda: list[Message](), repr=False) + """Detached inputs of the latest per-service call, never serialized into session state.""" + + +_current_binding: ContextVar[DurableHistoryBinding | None] = ContextVar( + "durable_history_binding", + default=None, +) + + +def bind_durable_history(binding: DurableHistoryBinding) -> Token[DurableHistoryBinding | None]: + """Bind the durable entity state for the current operation. + + Returns a token that must be passed to :func:`unbind_durable_history`. + """ + return _current_binding.set(binding) + + +def unbind_durable_history(token: Token[DurableHistoryBinding | None]) -> None: + """Release a binding created by :func:`bind_durable_history`.""" + _current_binding.reset(token) + + +def current_durable_history_binding() -> DurableHistoryBinding | None: + """Return the binding for the current durable operation, if any.""" + return _current_binding.get() + + +class DurableHistoryProvider(HistoryProvider): + """Core history hooks backed by the entity's staged transcript. + + Loading restores persisted messages, IDs and compaction annotations. After-run hooks append + the configured inputs, context and outputs, once per core hook. Reconciliation writes working + buffer annotations and summaries back to the transcript without committing to the backend. + The entity owns response delivery and the final flush after all core after-run providers. + + Attributes: + skip_excluded: When True, messages marked ``_excluded`` by compaction are omitted + from the context loaded for the model. The messages remain in durable storage. + prune_excluded: When True, excluded messages are physically removed from durable + storage on flush, preserving system messages and the newest/current exchange. + This is **lossy** and opt-in, independently of any configured pressure budget. + """ + + DEFAULT_SOURCE_ID = "durable_history" + + def __init__( + self, + source_id: str | None = None, + *, + store_inputs: bool = True, + store_outputs: bool = True, + store_context_messages: bool = False, + store_context_from: set[str] | None = None, + skip_excluded: bool = True, + prune_excluded: bool | None = None, + ) -> None: + """Initialize the durable history provider. + + Args: + source_id: Unique identifier for this provider instance. + store_inputs: Store each hook's input messages. + store_outputs: Store each hook's response messages. + store_context_messages: Store context contributed by other providers. + store_context_from: Restrict stored context to these source identifiers, when set. + skip_excluded: Omit compaction-excluded messages from loaded context. + prune_excluded: Physically delete excluded messages from durable storage on flush. + Lossy, so it is off unless asked for. Leaving it unset defers to the entity's + ``retention`` mode, which resolves it when the provider is prepared for a run. + Passing it explicitly pins the behaviour and retention will not override it, which + is what lets a caller who wires this provider by hand opt in or out independently + of the mode. Unset and unresolved, as when this provider is not the one the entity + prepared, it does not prune. + """ + super().__init__( + source_id=source_id or self.DEFAULT_SOURCE_ID, + load_messages=True, + store_inputs=store_inputs, + store_outputs=store_outputs, + store_context_messages=store_context_messages, + store_context_from=set(store_context_from) if store_context_from is not None else None, + ) + self.skip_excluded = skip_excluded + self.prune_excluded = prune_excluded + + def _binding(self) -> DurableHistoryBinding | None: + binding = current_durable_history_binding() + if binding is None: + logger.warning( + "[DurableHistoryProvider] No durable binding is active, so the provider yields no history. " + "This provider only works inside a durable agent entity operation." + ) + return binding + + def _replayable_entries(self, binding: DurableHistoryBinding) -> Iterator[tuple[DurableAgentStateEntry, int]]: + """Yield (entry, message_index) pairs that participate in model context.""" + # A tool loop must see messages saved by earlier calls in this same operation. + # The entity does not pre-append pipeline inputs, so there is no current request to hide. + yield from replayable_entries(binding.state_provider.state.data.conversation_history) + + @staticmethod + def _synthetic_message_id(entry: DurableAgentStateEntry, index: int) -> str: + """Build a deterministic ID candidate for a legacy message. + + The id comes from persisted fields, so a cold start or a retried flush regenerates the + same value. An id derived from object identity would not, and a recycled address could + collide with an id an earlier run already persisted. + + Args: + entry: History entry holding the message. + index: Position of the message within that entry. + + Returns: + An ID candidate, disambiguated against the current history by ``_positions``. + """ + # A request and its response share a correlation id, so the entry type is what tells the + # two sides of an exchange apart. + scope = entry.correlation_id or entry.created_at.isoformat() + kind = entry.json_type.value if isinstance(entry.json_type, DurableAgentStateEntryJsonType) else entry.json_type + return f"durable_{kind}_{scope}_{index}" + + @staticmethod + def _to_message(stored: DurableAgentStateMessage) -> Message | None: + """Convert a persisted message into one that is safe to replay to a chat client.""" + chat_message: Message = copy.deepcopy(stored).to_chat_message() + replayable = [content for content in chat_message.contents if content.type != "reasoning"] + if not replayable: + return None + return Message( + role=chat_message.role, + contents=replayable, + author_name=chat_message.author_name, + message_id=stored.message_id, + additional_properties=chat_message.additional_properties, + ) + + @staticmethod + def _unique_message_id(candidate: str, reserved: set[str]) -> str: + """Disambiguate generated identities, including collisions with caller-supplied IDs.""" + message_id = candidate + revision = 0 + while message_id in reserved: + revision += 1 + message_id = f"{candidate}_{revision}" + reserved.add(message_id) + return message_id + + def _positions(self, binding: DurableHistoryBinding) -> dict[str, tuple[DurableAgentStateEntry, int]]: + """Index current storage, repairing anonymous or duplicate identities in legacy entries.""" + history = binding.state_provider.state.data.conversation_history + reserved = {message.message_id for entry in history for message in entry.messages if message.message_id} + positions: dict[str, tuple[DurableAgentStateEntry, int]] = {} + for entry, index in self._replayable_entries(binding): + stored = entry.messages[index] + if not stored.message_id or stored.message_id in positions: + stored.message_id = self._unique_message_id(self._synthetic_message_id(entry, index), reserved) + positions[stored.message_id] = (entry, index) + return positions + + async def get_messages( + self, session_id: str | None, *, state: dict[str, Any] | None = None, **kwargs: Any + ) -> list[Message]: + """Load conversation history from durable entity state.""" + binding = self._binding() + if binding is None: + return [] + + if binding.service_owns_history: + # The service is holding this conversation and core will continue it by id. Returning + # history as well would send the model everything twice. The provider is still + # attached, which is what keeps core from injecting one whose state nothing bounds. + return [] + + id_map = self._positions(binding) + loaded: list[Message] = [] + for entry, index in self._replayable_entries(binding): + stored = entry.messages[index] + message = self._to_message(stored) + if message is None: + continue + loaded.append(message) + + if state is not None: + # Expose the loaded messages as the working buffer so CompactionProvider's + # after_strategy can annotate them (core reads session.state[source_id]["messages"]). + state[WORKING_BUFFER_KEY] = loaded + state[POSITIONS_KEY] = id_map + + if self.skip_excluded: + return [m for m in loaded if not m.additional_properties.get(EXCLUDED_KEY)] + return list(loaded) + + async def save_messages( + self, + session_id: str | None, + messages: Sequence[Message], + *, + state: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + """Stage a generic message batch as a request entry, without committing entity state.""" + binding = self._binding() + if binding is None or binding.service_owns_history: + return + if state is not None: + self.flush(state) + self._append_messages(binding, messages, state=state) + + def _append_messages( + self, + binding: DurableHistoryBinding, + messages: Sequence[Message], + *, + state: dict[str, Any] | None, + response: AgentResponse | None = None, + ) -> None: + """Append one hook batch, exposing only nonterminal batches to core compaction. + + Terminal outputs are excluded from this provider's local model history, not from + opaque external transcripts. Inputs remain separate accepted receipts, and earlier + per-call appends are not rewritten when a later response is terminal. + """ + if not messages or binding.service_owns_history: + return + if state is not None and WORKING_BUFFER_KEY not in state: + # Direct save_messages callers need the same complete compaction buffer as callers + # that loaded through before_run. Do not replace an already annotated buffer. + state[POSITIONS_KEY] = self._positions(binding) + state[WORKING_BUFFER_KEY] = [ + message + for entry, index in self._replayable_entries(binding) + if (message := self._to_message(entry.messages[index])) is not None + ] + created_at = datetime.now(tz=timezone.utc) + response_type: type[DurableAgentStateResponse] = ( + DurableAgentStateErrorResponse + if response is not None and is_terminal_agent_response(response) + else DurableAgentStateResponse + ) + kind = DurableAgentStateEntryJsonType.REQUEST if response is None else response_type.JSON_TYPE + stored_messages, working_messages = self._copy_append_messages(binding, messages, kind, created_at) + entry: DurableAgentStateEntry + if response is None: + entry = DurableAgentStateRequest(binding.correlation_id, created_at, stored_messages) + else: + entry = response_type( + binding.correlation_id, + created_at, + stored_messages, + usage=copy.deepcopy(DurableAgentStateUsage.from_usage(response.usage_details)), + ) + binding.state_provider.state.data.conversation_history.append(entry) + if state is not None: + buffer = cast("list[Message]", state.setdefault(WORKING_BUFFER_KEY, [])) + if not isinstance(entry, DurableAgentStateErrorResponse): + # Otherwise flush would mistake these non-replayable outputs for new summaries. + buffer.extend(working_messages) + state[POSITIONS_KEY] = self._positions(binding) + + def _copy_append_messages( + self, + binding: DurableHistoryBinding, + messages: Sequence[Message], + kind: DurableAgentStateEntryJsonType, + created_at: datetime, + ) -> tuple[list[DurableAgentStateMessage], list[Message]]: + """Allocate stored identities without changing input messages or caller responses.""" + history = binding.state_provider.state.data.conversation_history + used = {message.message_id for entry in history for message in entry.messages if message.message_id} + reserved = used | {message.message_id for message in messages if message.message_id} + scope = binding.correlation_id or created_at.isoformat() + ordinal = binding.append_ordinal + binding.append_ordinal += 1 + stored_messages: list[DurableAgentStateMessage] = [] + working_messages: list[Message] = [] + for index, message in enumerate(messages): + # Conversion can retain nested tool payloads, so neither stored content nor the + # compaction working copy may share those objects with the caller or each other. + stored = DurableAgentStateMessage.from_chat_message(copy.deepcopy(message)) + receipt = getattr(message, "_durable_ingestion_receipt", None) + if ( + kind == DurableAgentStateEntryJsonType.REQUEST + and isinstance(receipt, tuple) + and len(cast("tuple[Any, ...]", receipt)) == 2 + ): + occurrence, fingerprint = cast("tuple[str, str]", receipt) + stored.ingestion_occurrence = occurrence + stored.ingestion_identity = fingerprint + if not stored.message_id or stored.message_id in used: + prefix = "durable_revision" if stored.message_id else "durable" + candidate = f"{prefix}_{kind.value}_{scope}_{ordinal}_{index}" + stored.message_id = self._unique_message_id(candidate, reserved) + used.add(stored.message_id) + working = copy.deepcopy(message) + working.message_id = stored.message_id + stored_messages.append(stored) + working_messages.append(working) + return stored_messages, working_messages + + async def before_run( + self, + *, + agent: Any, + session: Any, + context: Any, + state: dict[str, Any], + ) -> None: + """Load durable history into context, unless the service owns the conversation.""" + binding = current_durable_history_binding() + if binding is not None: + binding.pending_inputs.clear() + if binding.service_owns_history: + return + if self.store_inputs and getattr(agent, "require_per_service_call_history_persistence", False): + # Capture only. Core still decides whether the after-run persistence hook runs. + binding.pending_inputs = copy.deepcopy(context.input_messages) + await super().before_run(agent=agent, session=session, context=context, state=state) + + def _get_context_messages_to_store(self, context: SessionContext) -> list[Message]: + # Our own contribution is already persisted. Core's in-memory save deduplicates it, + # but durable appends allocate new identities, so exclude it even from explicit masks. + if not self.store_context_messages: + return [] + return context.get_messages(sources=self.store_context_from, exclude_sources={self.source_id}) + + async def after_run( + self, + *, + agent: Any, + session: Any, + context: Any, + state: dict[str, Any], + ) -> None: + """Reconcile compaction, then append exactly the messages selected for this core hook.""" + binding = self._binding() + if binding is None: + return + if binding.service_owns_history: + binding.pending_inputs.clear() + return + self.flush(state) + request_messages = self._get_context_messages_to_store(context) + if self.store_inputs: + request_messages.extend(context.input_messages) + self._append_messages(binding, request_messages, state=state) + if self.store_outputs and context.response and context.response.messages: + self._append_messages(binding, context.response.messages, state=state, response=context.response) + binding.pending_inputs.clear() + + def finalize_failed_run(self, state: dict[str, Any]) -> None: + """Stage actual tool results left unsaved when a later service call fails. + + Call from the entity before its final flush, with the operation binding still active. + Only result-only tool messages answering unresolved calls already stored under this + correlation are eligible. Fresh requests, results for earlier correlations and calls whose + persistence core deferred or disabled do not authorize an append. No backend write occurs. + + Args: + state: The provider-scoped session state holding the working buffer. + """ + binding = current_durable_history_binding() + if binding is None: + return + pending_inputs = binding.pending_inputs + binding.pending_inputs = [] + if ( + not pending_inputs + or not self.store_inputs + or binding.service_owns_history + or binding.correlation_id is None + ): + return + + pending_calls: set[str] = set() + for entry, index in self._replayable_entries(binding): + if entry.correlation_id != binding.correlation_id: + continue + for content in entry.messages[index].contents: + if isinstance(content, DurableAgentStateFunctionCallContent): + pending_calls.add(content.call_id) + elif isinstance(content, DurableAgentStateFunctionResultContent): + pending_calls.discard(content.call_id) + + messages: list[Message] = [] + for message in pending_inputs: + if message.role != "tool": + continue + result_ids = { + content.call_id for content in message.contents if content.type == "function_result" and content.call_id + } + if result_ids and len(result_ids) == len(message.contents) and result_ids <= pending_calls: + # Keep the entire original message so its ingestion hash still identifies the + # caller's input, even if append allocates a different stored message ID. + messages.append(message) + pending_calls.difference_update(result_ids) + self._append_messages(binding, messages, state=state) + + def flush(self, state: dict[str, Any]) -> None: + """Apply compaction results to durable entity state. + + Reconciliation is by ``message_id`` rather than position, so strategies that + *insert* messages (for example ``ToolResultCompactionStrategy``, which replaces a + tool-call group with a summary) are handled as well as ones that only annotate. + Previously loaded messages removed from the list become exclusions, not implicit + permission to physically delete them. Opt-in pruning still applies its atomic-group floor. + + These edits affect only cached state. Repeated flushes reconcile annotations without + repeating appends or strategies. The entity performs the final flush while this operation's + binding is still active, after all core after-run providers and before its single commit. + + Args: + state: The provider-scoped session state holding the working buffer. + """ + binding = current_durable_history_binding() + if binding is None or binding.service_owns_history: + return + + raw_buffer = state.get(WORKING_BUFFER_KEY) + raw_positions = state.get(POSITIONS_KEY) + if not isinstance(raw_buffer, list): + return + buffer = cast("list[Message]", raw_buffer) + previous_ids: set[str] = set() + if isinstance(raw_positions, dict): + previous_ids.update(cast("dict[str, Any]", raw_positions)) + # Positions can refer to entries replaced by pressure eviction, or indices invalidated + # by an earlier flush. Only the previous keys are useful for recognizing removed messages. + stored_by_id = self._positions(binding) + + # Messages that compaction added (summaries) are inserted right after the last + # known message so ordering in durable state matches the compacted conversation. + last_known: tuple[DurableAgentStateEntry, int] | None = None + + for message in buffer: + position = stored_by_id.get(message.message_id) if message.message_id else None + summary_ids = self._summary_original_ids(message) + summary_revision = False + if position is not None and summary_ids is not None: + owner, index = position + original = self._to_message(owner.messages[index]) + # Compare the replayed storage shape on both sides. Metadata discarded by + # content conversion must not make every later flush look like a new summary. + working = self._to_message(DurableAgentStateMessage.from_chat_message(copy.deepcopy(message))) + original_payload = original.to_dict() if original is not None else {} + working_payload = working.to_dict() if working is not None else {} + original_payload.pop("additional_properties", None) + working_payload.pop("additional_properties", None) + # Core's summary_{len(messages)} can recur after pruning. Different source + # messages identify a new summary even when the generated body is identical. + original_ids = self._summary_original_ids(original) if original is not None else None + summary_revision = original_payload != working_payload or original_ids != summary_ids + + if position is None or summary_revision: + if not summary_revision and message.message_id in previous_ids: + # This was persisted before, not a newly generated summary. Never resurrect + # a message removed since the working buffer was assembled. + continue + original_id = message.message_id + position = self._insert_new_message(binding, message, after=last_known) + if original_id and original_id != message.message_id and summary_ids is not None: + group = message.additional_properties.get(GROUP_ANNOTATION_KEY) + if isinstance(group, dict): + group[GROUP_ID_KEY] = f"group_{message.message_id}" + # Only the sources named by this summary now point to its new ID. Older + # sources may still point to the old summary, including when that summary + # is itself a source here. Its forward ID must therefore remain unchanged. + for source in buffer: + if source is message or source.message_id not in summary_ids: + continue + source_group = source.additional_properties.get(GROUP_ANNOTATION_KEY) + if ( + isinstance(source_group, dict) + and cast("dict[str, Any]", source_group).get(SUMMARIZED_BY_SUMMARY_ID_KEY) == original_id + ): + source_group[SUMMARIZED_BY_SUMMARY_ID_KEY] = message.message_id + if source.additional_properties.get(SUMMARIZED_BY_SUMMARY_ID_KEY) == original_id: + source.additional_properties[SUMMARIZED_BY_SUMMARY_ID_KEY] = message.message_id + stored_by_id = self._positions(binding) + last_known = position + + # Link repair can touch sources that precede an inserted summary. Persist annotations + # only after every insertion, using the final positions after any entry splits. + for message in buffer: + position = stored_by_id.get(message.message_id) if message.message_id else None + if position is None: + continue + entry, index = position + stored = entry.messages[index] + stored.extension_data = ( + copy.deepcopy(message.additional_properties) if message.additional_properties else None + ) + + # A strategy may shrink the list without setting _excluded. Compare final identities + # after summary revision IDs have been allocated, so replacing a summary does not leave + # its old revision active. Preserve stored metadata and backlinks on absent messages. + remaining_ids = {message.message_id for message in buffer if message.message_id} + for message_id in previous_ids - remaining_ids: + position = stored_by_id.get(message_id) + if position is not None: + entry, index = position + stored = entry.messages[index] + if self._to_message(stored) is None: + # Empty/non-replayable payloads were never exposed to the strategy. + continue + stored.extension_data = {**(stored.extension_data or {}), EXCLUDED_KEY: True} + + if self.prune_excluded: + # Resolve owners after all insertions. Splitting a multi-message entry may have + # moved a previously annotated message into the tail entry. + self._prune( + binding, + [ + (entry, entry.messages[index]) + for entry, index in stored_by_id.values() + if (entry.messages[index].extension_data or {}).get(EXCLUDED_KEY) + ], + ) + stored_by_id = self._positions(binding) + buffer[:] = [message for message in buffer if message.message_id in stored_by_id] + state[POSITIONS_KEY] = stored_by_id + + @staticmethod + def _summary_original_ids(message: Message) -> list[str] | None: + """Recognize Core's summary links, also accepting top-level custom-strategy links.""" + group = message.additional_properties.get(GROUP_ANNOTATION_KEY) + original_ids = ( + cast("Mapping[str, Any]", group).get(SUMMARY_OF_MESSAGE_IDS_KEY) if isinstance(group, Mapping) else None + ) + if original_ids is None: + original_ids = message.additional_properties.get(SUMMARY_OF_MESSAGE_IDS_KEY) + return cast("list[str]", original_ids) if isinstance(original_ids, list) else None + + def _insert_new_message( + self, + binding: DurableHistoryBinding, + message: Message, + *, + after: tuple[DurableAgentStateEntry, int] | None, + ) -> tuple[DurableAgentStateEntry, int]: + """Persist a message compaction produced, such as a summary, as an entry of its own. + + It takes its place in conversation order, but as a compaction entry rather than inside + whichever request or response it happened to follow. Folding it into a response made it + part of that response, so a caller polling that correlation was handed back a summary the + agent never produced. + + Having its own entry also means nothing downstream has to be told to skip it. It is not a + response, so the lookup that serves waiting callers cannot match it. + """ + history = binding.state_provider.state.data.conversation_history + created_at = datetime.now(tz=timezone.utc) + stored, _ = self._copy_append_messages( + binding, + [message], + DurableAgentStateEntryJsonType.COMPACTION, + created_at, + ) + message.message_id = stored[0].message_id + entry = DurableAgentStateCompaction( + created_at=created_at, + messages=stored, + ) + + if after is not None: + owner, message_index = after + position = next(index for index, candidate in enumerate(history) if candidate is owner) + 1 + if message_index + 1 < len(owner.messages): + # [a, b] + a summary after a must become [a], [summary], [b], not + # [a, b], [summary]. Keep message identities while detaching envelope metadata. + tail = copy.copy(owner) + tail.messages = owner.messages[message_index + 1 :] + tail.extension_data = copy.deepcopy(owner.extension_data) + tail.unknown_fields = copy.deepcopy(owner.unknown_fields) + if isinstance(tail, DurableAgentStateRequest): + tail.response_schema = copy.deepcopy(tail.response_schema) + if isinstance(tail, DurableAgentStateResponse): + tail.usage = copy.deepcopy(tail.usage) + owner.messages = owner.messages[: message_index + 1] + history[position:position] = [entry, tail] + else: + history.insert(position, entry) + return entry, 0 + + history.insert(0, entry) + return entry, 0 + + @staticmethod + def _prune( + binding: DurableHistoryBinding, + pruned: list[tuple[DurableAgentStateEntry, DurableAgentStateMessage]], + ) -> None: + """Remove eligible exclusions and record only actual removals. + + Removal is by identity rather than index, since insertions earlier in this flush may have + moved messages within their entry. System messages and the current exchange are a floor. + """ + if not pruned: + return + + from ._retention import ( + _detached_message, # pyright: ignore[reportPrivateUsage] + _link_atomic_groups, # pyright: ignore[reportPrivateUsage] + _newest_exchange, # pyright: ignore[reportPrivateUsage] + _saved_group_id, # pyright: ignore[reportPrivateUsage] + record_truncation, + ) + + state = binding.state_provider.state + history = state.data.conversation_history + protected = {id(entry) for entry in _newest_exchange(history)} + protected.update( + id(entry) + for entry in history + if binding.correlation_id is not None and entry.correlation_id == binding.correlation_id + ) + # Protect the whole atomic group when a system/current message holds any member, + # including non-contiguous tool results and persisted links beyond Core's grouping. + originals = [(entry, entry.messages[index]) for entry, index in replayable_entries(history)] + messages = [_detached_message(stored) for _, stored in originals] + annotate_message_groups(messages, force_reannotate=True) + groups = _link_atomic_groups(messages, [_saved_group_id(stored) for _, stored in originals]) + protected_groups = { + group + for (entry, stored), group in zip(originals, groups) + if id(entry) in protected or stored.role == "system" + } + # Eager pruning may delete only exclusions, not the included partners of a tool or + # reasoning group. Defer the whole group until all its members are excluded. + excluded_messages = {id(stored) for _, stored in pruned} + protected_groups.update( + group for (_, stored), group in zip(originals, groups) if id(stored) not in excluded_messages + ) + protected_messages = {id(stored) for (_, stored), group in zip(originals, groups) if group in protected_groups} + eligible = [ + (entry, stored) + for entry, stored in pruned + if id(entry) not in protected and stored.role != "system" and id(stored) not in protected_messages + ] + before = sum(len(entry.messages) for entry in history) + before_entries = len(history) + before_bytes = eager_state_size(state) if eligible else None + prune_messages(history, eligible) + removed = before - sum(len(entry.messages) for entry in history) + if removed: + record_truncation(state, removed) + record_retention( + state, + mechanism="eager", + outcome="staged" if removed else "protected", + before_bytes=before_bytes, + after_bytes=eager_state_size(state) if before_bytes is not None else None, + removed_messages=removed, + removed_entries=before_entries - len(history), + ) + + +def replayable_entries( + history: list[DurableAgentStateEntry], + *, + correlation_id: str | None = None, +) -> Iterator[tuple[DurableAgentStateEntry, int]]: + """Yield (entry, message_index) pairs that participate in model context. + + Storage eviction has separate eligibility rules. A failed turn is not model + context, but its expired transcript payload can still consume evictable storage. + + Args: + history: The entity's conversation history. + correlation_id: Optional legacy exclusion. The core pipeline includes current-correlation appends. + + Yields: + Each replayable message as its owning entry and its index within that entry. + """ + for entry in history: + if ( + isinstance(entry, (DurableAgentStateErrorResponse, DurableAgentStateUnknownEntry)) + or entry.json_type not in tuple(DurableAgentStateEntryJsonType) + or entry.json_type == DurableAgentStateEntryJsonType.ERROR_RESPONSE + ): + # Runtime-error entries and opaque future entries are not model messages. + continue + if correlation_id is not None and entry.correlation_id == correlation_id: + continue + for index in range(len(entry.messages)): + yield entry, index + + +def prune_messages( + history: list[DurableAgentStateEntry], + pruned: list[tuple[DurableAgentStateEntry, DurableAgentStateMessage]], +) -> None: + """Remove messages, dropping only changed, bare, known transcript envelopes. + + Removal is by identity rather than index, since an insertion elsewhere in the same pass may + have moved messages within their entry. + + Args: + history: The entity's conversation history, modified in place. + pruned: The messages to remove, each with the entry that owns it. + """ + from ._retention import _can_drop_entry # pyright: ignore[reportPrivateUsage] + + live_entries = {id(entry) for entry in history} + changed: set[int] = set() + for entry, stored in pruned: + if ( + id(entry) not in live_entries + or isinstance(entry, DurableAgentStateUnknownEntry) + or entry.json_type not in tuple(DurableAgentStateEntryJsonType) + ): + continue + for index, candidate in enumerate(entry.messages): + if candidate is stored: + del entry.messages[index] + changed.add(id(entry)) + break + + remaining = [entry for entry in history if entry.messages or id(entry) not in changed or not _can_drop_entry(entry)] + if len(remaining) != len(history): + history[:] = remaining + + +def service_stores_history(agent: Any, options: Mapping[str, Any] | None = None) -> bool: + """Return whether the service keeps conversation history for this run. + + Mirrors core's precedence, most specific first: the option passed on the run itself, then an + explicit ``store`` in the agent's default options, and only when both are unset does the + client's ``STORES_BY_DEFAULT`` apply. Clients that store by default (such as the Responses + API) can therefore be put back in client-side mode either permanently or for a single run, and + in that case durable history is what makes the conversation survive. + + Resolved per run rather than once at registration because ``store`` is an ordinary run option. + An agent registered against a storing client can still be asked to keep one turn client-side, + and whoever answers that turn's history has to be decided at that point. + + Args: + agent: The agent being run. + options: The effective options for this run, when there is a run in progress. + + Returns: + True when the model service is holding this conversation. + """ + if options is not None: + run_store = options.get("store") + if run_store is not None: + return bool(run_store) + default_options = getattr(agent, "default_options", None) + if isinstance(default_options, Mapping): + explicit_store = cast("Mapping[str, Any]", default_options).get("store") + if explicit_store is not None: + return bool(explicit_store) + client = getattr(agent, "client", None) + return bool(getattr(client, "STORES_BY_DEFAULT", False)) + + +def validate_history_providers(agent: SupportsAgentRun) -> None: + """Reject competing primaries and shared state namespaces, allowing distinct store-only sinks.""" + providers = getattr(agent, "context_providers", None) + if not isinstance(providers, (list, tuple)): + return + primaries = [p for p in cast("Sequence[Any]", providers) if isinstance(p, HistoryProvider) and p.load_messages] + if len(primaries) > 1: + raise ValueError("A durable agent supports only one load-enabled primary history provider.") + sources: set[str] = set() + for provider in cast("Sequence[Any]", providers): + source_id = provider.source_id + if source_id in sources: + raise ValueError( + f"Context providers must have unique source_id values; {source_id!r} is duplicated. " + "Assign distinct source_id values to history, audit and other context providers." + ) + sources.add(source_id) + if not primaries and InMemoryHistoryProvider.DEFAULT_SOURCE_ID in sources: + raise ValueError( + "Cannot inject durable history: 'in_memory' is already used by a context provider or store-only sink. " + "Set that provider's source_id to a unique value such as 'audit', or explicitly configure a " + "DurableHistoryProvider with a distinct source_id and matching compaction history_source_id." + ) + + +class _ServiceOwnedHistoryProvider(HistoryProvider): + """Occupy the primary slot without loading or saving the inactive external branch.""" + + def __init__(self, provider: HistoryProvider) -> None: + """Borrow the original provider without modifying its configuration or lifecycle.""" + super().__init__( + source_id=provider.source_id, + load_messages=provider.load_messages, + store_inputs=provider.store_inputs, + store_outputs=provider.store_outputs, + store_context_messages=provider.store_context_messages, + store_context_from=provider.store_context_from, + ) + self.__wrapped__ = provider + # Core 1.13 predates this optional hook-cadence hint. + if hasattr(provider, "after_run_once_per_turn"): + self.after_run_once_per_turn = provider.after_run_once_per_turn + + def __getattr__(self, name: str) -> Any: + # Expose the original configuration/resources without copying or taking their ownership. + return getattr(self.__wrapped__, name) + + async def get_messages( + self, session_id: str | None, *, state: dict[str, Any] | None = None, **kwargs: Any + ) -> list[Message]: + """Do not load the external transcript into a service-owned invocation.""" + return [] + + async def save_messages( + self, + session_id: str | None, + messages: Sequence[Message], + *, + state: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + """Do not append a service-owned turn to the external primary.""" + + async def before_run(self, *, agent: Any, session: Any, context: Any, state: dict[str, Any]) -> None: + """Suppress custom loading hooks as well as the base implementation.""" + # Do not call custom hooks either: an ordinary primary need not know about ownership. + + async def after_run(self, *, agent: Any, session: Any, context: Any, state: dict[str, Any]) -> None: + """Suppress custom persistence hooks for the inactive primary.""" + + +def prepare_history_owner(agent: SupportsAgentRun, service_owns_history: bool) -> SupportsAgentRun: + """Return a per-run view that silences only a service-owned run's external primary. + + Call after registration and ownership resolution. Client-owned runs keep the original + provider, including custom hooks and context attribution. Store-only sinks are never wrapped. + The view borrows the agent's resources: preparation neither enters nor closes them. Keep the + registered agent for lifecycle/reset decisions; wrappers also expose ``__wrapped__``. + """ + providers = getattr(agent, "context_providers", None) + if not isinstance(providers, (list, tuple)): + return agent + updated: list[Any] = [] + changed = False + for provider in cast("Sequence[Any]", providers): + original = provider.__wrapped__ if isinstance(provider, _ServiceOwnedHistoryProvider) else provider + replacement = original + if ( + service_owns_history + and isinstance(original, HistoryProvider) + and original.load_messages + and not isinstance(original, DurableHistoryProvider) + and type(original) is not InMemoryHistoryProvider + ): + replacement = ( + provider + if isinstance(provider, _ServiceOwnedHistoryProvider) + else _ServiceOwnedHistoryProvider(original) + ) + changed |= replacement is not provider + updated.append(replacement) + if not changed: + return agent + return _copy_with_history_providers(agent, updated) + + +def _copy_with_history_providers(agent: SupportsAgentRun, providers: list[Any]) -> SupportsAgentRun: + try: + clone = copy.copy(agent) + clone.context_providers = providers # type: ignore[attr-defined] + except Exception as exc: + raise ValueError( + f"Could not attach durable history to agent {getattr(agent, 'name', type(agent).__name__)}. " + "Configure a supported history provider explicitly." + ) from exc + return clone + + +def ensure_durable_history(agent: SupportsAgentRun, *, prune_excluded: bool = False) -> SupportsAgentRun: + """Back an agent's conversation history with durable entity state. + + Lets a user register an agent that already works in core and get durable behavior with no + configuration change. The agent is never mutated: when a substitution is needed a shallow + copy is returned with its own provider list. + + With no load-enabled primary, append a :class:`DurableHistoryProvider` after existing + providers using core's default history source. This matches core's automatic injection order, + so default compaction resolves that source and its reverse-order after hook sees this turn. + Explicit registration order is preserved. Only exact built-in :class:`InMemoryHistoryProvider` + instances are replaced, preserving source, storage flags, exclusion policy and once-per-turn + hook setting. A hand-configured durable provider keeps explicit ``prune_excluded`` values; + otherwise a shallow copy inherits the registration retention policy. + + Other primaries, including in-memory subclasses, keep their original hooks and state without + an additional durable provider. Their session state may contain a transcript. Such state is + part of the non-evictable floor, not managed by durable transcript retention. + + Service ownership is resolved per run. Without a custom primary, durable history remains + available for client-owned runs and silent for service-owned runs, preventing core from + injecting a separate unmanaged history slice. Agents without the core context pipeline are + left alone and the entity falls back to replaying its own persisted history. + + Args: + agent: The agent being registered with the durable runtime. + + Keyword Args: + prune_excluded: When True, the injected provider physically deletes messages that + compaction excluded, bounding durable storage. This is a **lossy retention policy** + and is off by default. + + Returns: + The agent to run, either unchanged or a shallow copy with durable-backed history. + """ + validate_history_providers(agent) + providers = getattr(agent, "context_providers", None) + if not isinstance(providers, (list, tuple)): + return agent + + provider_list = list(cast("Sequence[Any]", providers)) + existing = next( + (p for p in provider_list if isinstance(p, HistoryProvider) and p.load_messages), + None, + ) + + if existing is None: + # Match core's source_id and append order. After hooks run in reverse, so automatic + # history must save this turn before an earlier compaction provider reads its buffer. + updated = [ + *provider_list, + DurableHistoryProvider( + source_id=InMemoryHistoryProvider.DEFAULT_SOURCE_ID, + prune_excluded=prune_excluded, + ), + ] + elif isinstance(existing, DurableHistoryProvider): + # Already durable. If the caller pinned ``prune_excluded`` themselves that decision + # stands, but an unset one means they never expressed a preference, and leaving it unset + # would make the entity's retention mode silently do nothing. + if existing.prune_excluded is not None: + return agent + replacement = copy.copy(existing) + replacement.prune_excluded = prune_excluded + if existing.store_context_from is not None: + replacement.store_context_from = set(existing.store_context_from) + updated = [replacement if p is existing else p for p in provider_list] + elif type(existing) is InMemoryHistoryProvider: + replacement = DurableHistoryProvider( + source_id=existing.source_id, + store_inputs=existing.store_inputs, + store_outputs=existing.store_outputs, + store_context_messages=existing.store_context_messages, + store_context_from=existing.store_context_from, + skip_excluded=existing.skip_excluded, + prune_excluded=prune_excluded, + ) + if hasattr(existing, "after_run_once_per_turn"): + replacement.after_run_once_per_turn = existing.after_run_once_per_turn + updated = [replacement if p is existing else p for p in provider_list] + else: + # A deliberate storage choice (external or custom), so do not override it. + return agent + + return _copy_with_history_providers(agent, updated) diff --git a/python/packages/durabletask/agent_framework_durabletask/_invocation_safety.py b/python/packages/durabletask/agent_framework_durabletask/_invocation_safety.py new file mode 100644 index 0000000..2e2c42b --- /dev/null +++ b/python/packages/durabletask/agent_framework_durabletask/_invocation_safety.py @@ -0,0 +1,36 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Run-local safeguards at core's function invocation boundary.""" + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass + +from agent_framework import FunctionInvocationContext, FunctionMiddleware + + +@dataclass +class InvocationProgress: + """Track observable progress that makes restarting a whole agent run unsafe.""" + + stream_started: bool = False + function_started: bool = False + + +class DurableToolGuard(FunctionMiddleware): + """Prevent callable execution even when a wrapper delegates to an inner core loop. + + This uses core's public per-run middleware contract. Arbitrary custom agents or + clients that execute tools outside that contract remain responsible for their own + side effects; no portable wrapper can sandbox their implementation. + """ + + def __init__(self, progress: InvocationProgress, *, enabled: bool) -> None: + self.progress = progress + self.enabled = enabled + + async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None: + if not self.enabled: + context.result = "Tool execution is disabled for this invocation." + return + self.progress.function_started = True + await call_next() diff --git a/python/packages/durabletask/agent_framework_durabletask/_message_identity.py b/python/packages/durabletask/agent_framework_durabletask/_message_identity.py new file mode 100644 index 0000000..840fed3 --- /dev/null +++ b/python/packages/durabletask/agent_framework_durabletask/_message_identity.py @@ -0,0 +1,22 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Content-sensitive identities shared by workflow transport and ingestion.""" + +import hashlib +import json + +from agent_framework import Message + + +def message_identity(message: Message) -> str: + """Hash a message's complete wire representation, including its supplied ID. + + Dictionary ordering is immaterial; content ordering, role, author and additional + properties are meaningful. Core's ``to_dict`` already excludes raw SDK objects + and absent optional fields. Do not use this alone to identify anonymous requests: + the workflow sender assigns those a deterministic, source-scoped ID first. + """ + canonical = json.dumps( + message.to_dict(), sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False + ) + return hashlib.sha256(canonical.encode("utf-8")).hexdigest() diff --git a/python/packages/durabletask/agent_framework_durabletask/_models.py b/python/packages/durabletask/agent_framework_durabletask/_models.py index e8eabca..fbdfbe7 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_models.py +++ b/python/packages/durabletask/agent_framework_durabletask/_models.py @@ -109,6 +109,11 @@ class RunRequest: created_at: Optional timestamp when the request was created orchestration_id: Optional ID of the orchestration that initiated this request options: Optional options dictionary forwarded to the agent + context_messages: Optional upstream conversation (serialized ``Message`` dicts) that should + be delivered to the agent as the request's messages. Workflows use this to give a + downstream agent the conversation produced by upstream nodes, matching the in-process + ``AgentExecutor`` context behavior. When set, it replaces ``message`` as the + request payload; ``message`` still carries the latest text for logging. """ message: str @@ -121,6 +126,8 @@ class RunRequest: created_at: datetime | None = None orchestration_id: str | None = None options: dict[str, Any] = field(default_factory=lambda: {}) + context_messages: list[dict[str, Any]] | None = None + context_message_ids: list[str] | None = None def __init__( self, @@ -134,7 +141,11 @@ def __init__( created_at: datetime | None = None, orchestration_id: str | None = None, options: dict[str, Any] | None = None, + context_messages: list[dict[str, Any]] | None = None, + context_message_ids: list[str] | None = None, ) -> None: + if not isinstance(correlation_id, str) or not correlation_id.strip(): + raise ValueError("correlationId must be a non-empty string.") self.message = message self.correlation_id = correlation_id self.role = self.coerce_role(role) @@ -145,6 +156,19 @@ def __init__( self.created_at = created_at if created_at is not None else datetime.now(tz=timezone.utc) self.orchestration_id = orchestration_id self.options = options if options is not None else {} + if context_messages is not None and ( + not isinstance(context_messages, list) or any(not isinstance(message, dict) for message in context_messages) + ): + raise ValueError("contextMessages must be a list of message objects.") + self.context_messages = context_messages + if context_message_ids is not None and ( + context_messages is None + or not isinstance(context_message_ids, list) + or len(context_message_ids) != len(context_messages) + or any(not isinstance(identity, str) or not identity for identity in context_message_ids) + ): + raise ValueError("contextMessageIds must contain one non-empty occurrence ID per context message.") + self.context_message_ids = context_message_ids @staticmethod def coerce_role(value: str | None) -> str: @@ -158,7 +182,7 @@ def coerce_role(value: str | None) -> str: def to_dict(self) -> dict[str, Any]: """Convert to dictionary for JSON serialization.""" - result = { + result: dict[str, Any] = { "message": self.message, "enable_tool_calls": self.enable_tool_calls, "wait_for_response": self.wait_for_response, @@ -173,6 +197,10 @@ def to_dict(self) -> dict[str, Any]: result["created_at"] = self.created_at.isoformat() if self.orchestration_id: result["orchestrationId"] = self.orchestration_id + if self.context_messages is not None: + result["contextMessages"] = self.context_messages + if self.context_message_ids is not None: + result["contextMessageIds"] = self.context_message_ids return result @classmethod @@ -183,7 +211,9 @@ def from_json(cls, data: str) -> RunRequest: except json.JSONDecodeError as e: raise ValueError("The durable agent state is not valid JSON.") from e - return cls.from_dict(dict_data) + if not isinstance(dict_data, dict): + raise ValueError("RunRequest must be a JSON object.") + return cls.from_dict(cast("dict[str, Any]", dict_data)) @classmethod def from_dict(cls, data: dict[str, Any]) -> RunRequest: @@ -200,6 +230,13 @@ def from_dict(cls, data: dict[str, Any]) -> RunRequest: raise ValueError("correlationId is required in RunRequest data") options = data.get("options") + raw_context = data.get("contextMessages") + if raw_context is not None and ( + not isinstance(raw_context, list) + or any(not isinstance(message, dict) for message in cast("list[Any]", raw_context)) + ): + raise ValueError("contextMessages must be a list of message objects.") + context_messages = cast("list[dict[str, Any]] | None", raw_context) return cls( message=data.get("message", ""), @@ -212,6 +249,8 @@ def from_dict(cls, data: dict[str, Any]) -> RunRequest: created_at=created_at, orchestration_id=data.get("orchestrationId"), options=cast(dict[str, Any], options) if isinstance(options, dict) else {}, + context_messages=context_messages, + context_message_ids=data.get("contextMessageIds"), ) diff --git a/python/packages/durabletask/agent_framework_durabletask/_response_utils.py b/python/packages/durabletask/agent_framework_durabletask/_response_utils.py index 2d0ee84..1f34d2a 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_response_utils.py +++ b/python/packages/durabletask/agent_framework_durabletask/_response_utils.py @@ -2,14 +2,178 @@ """Shared utilities for handling AgentResponse parsing and validation.""" +import json import logging -from typing import Any +from collections.abc import Mapping, Sequence +from copy import copy, deepcopy +from functools import lru_cache +from inspect import Parameter, signature +from typing import Any, Literal, cast -from agent_framework import AgentResponse -from pydantic import BaseModel +from agent_framework import AgentResponse, Content, Message +from pydantic import BaseModel, ValidationError logger = logging.getLogger("agent_framework.durabletask") +# Optional reader marker; the serializer does not add it to response payloads. +_DELIVERY_VERSION_KEY = "_durable_response_version" +_DELIVERY_VERSION = 1 +_VALUE_BY_NAME_KEY = "_durable_value_by_name" + + +@lru_cache(maxsize=3) +def _constructor_fields(cls: type[AgentResponse[Any]] | type[Message] | type[Content]) -> tuple[str, ...]: + """Cache explicit public parameters, never names supplied by a stored type.""" + return tuple( + name + for name, parameter in signature(cls).parameters.items() + if not name.startswith("_") and parameter.kind in (Parameter.POSITIONAL_OR_KEYWORD, Parameter.KEYWORD_ONLY) + ) + + +def _constructor_kwargs( + data: Mapping[str, Any], cls: type[AgentResponse[Any]] | type[Message] | type[Content] +) -> dict[str, Any]: + return {name: data[name] for name in _constructor_fields(cls) if name in data} + + +def _load_content(data: Any) -> Any: + """Decode only Content envelope edges, not arbitrary dictionaries with a type key.""" + if not isinstance(data, Mapping): + return data + fields = _constructor_kwargs(cast(Mapping[str, Any], data), Content) + if not isinstance(fields.get("type"), str) or not fields["type"]: + raise ValueError("Content mapping requires 'type' to be a non-empty string") + if isinstance(fields.get("function_call"), Mapping): + fields["function_call"] = _load_content(fields["function_call"]) + for name in ("items", "inputs"): + if isinstance(fields.get(name), list): + fields[name] = [_load_content(item) for item in fields[name]] + # Unlike code/shell outputs, image-generation outputs are arbitrary application data. + if fields["type"] in ("code_interpreter_tool_result", "shell_tool_result") and isinstance( + fields.get("outputs"), list + ): + fields["outputs"] = [_load_content(item) for item in fields["outputs"]] + # arguments, result, output, annotations and additional_properties stay opaque. + return Content(**fields) + + +def _load_message(data: Any) -> Message: + if isinstance(data, Message): + return data + if not isinstance(data, Mapping): + raise TypeError("Agent response messages must be Message instances or mappings") + fields = _constructor_kwargs(cast(Mapping[str, Any], data), Message) + if fields.get("contents") is not None: + fields["contents"] = [_load_content(content) for content in fields["contents"]] + return Message(**fields) + + +def _serialize_model_value(value: BaseModel) -> tuple[Any, bool]: + """Prefer alias JSON; use field-name JSON when serialization aliases are not inputs.""" + payload = value.model_dump(mode="json", by_alias=True, round_trip=True) + field_payload = value.model_dump(mode="json", by_alias=False, round_trip=True) + try: + restored = type(value).model_validate_json(json.dumps(payload)) + except ValidationError: + pass + else: + if restored.model_dump(mode="json", by_alias=False, round_trip=True) == field_payload: + return payload, False + # A serialization alias may be ignored in favor of a default without raising an error. + # Record the input mode, not a Python model name, for the caller's declared format. + restored = type(value).model_validate_json(json.dumps(field_payload), by_alias=False, by_name=True) + if restored.model_dump(mode="json", by_alias=False, round_trip=True) != field_payload: + raise ValueError("Structured response value cannot round-trip through its declared model") + return field_payload, True + + +def is_terminal_agent_response(response: AgentResponse[Any]) -> bool: + """Identify durable failures/completions, retaining the legacy non-tool error fallback. + + Args: + response: Agent response whose durable status and non-tool error contents to inspect. + + Returns: + Whether the response reports a durable failure, completion, or non-tool error. + """ + return response.additional_properties.get("durable_status") in ("error", "already_completed") or any( + content.type == "error" + for message in response.messages + if message.role != "tool" + for content in message.contents + ) + + +def invocation_outcome(response: AgentResponse[Any], *, legacy: bool = False) -> Literal["succeeded", "failed"] | None: + """Classify invocation evidence, not delivery availability or an approval's pending action. + + Legacy transcript projections can have lost their error contents. Their absence + does not prove success. An independent original mailbox does not have that loss. + Accepted or already-unavailable replies likewise cannot establish a new outcome. + """ + status = response.additional_properties.get("durable_status") + if status == "accepted": + return None + if status == "already_completed" or any( + content.type == "error" and content.error_code == "response_expired" + for message in response.messages + if message.role != "tool" + for content in message.contents + ): + outcome = response.additional_properties.get("durable_outcome") + return outcome if outcome in ("succeeded", "failed") else None + if is_terminal_agent_response(response): + return "failed" + return None if legacy else "succeeded" + + +def serialize_agent_response(response: AgentResponse) -> dict[str, Any]: + """Snapshot a response as inline base-response JSON for durable delivery. + + Public base fields are authoritative even for subclasses. Serializable extra + fields may remain in the raw snapshot, but are not constructor arguments when + delivering it. No response-format class or provider raw representation is stored. + The containing entity schema versions delivery; a response version is not added. + + Args: + response: Agent response whose public fields and structured value to snapshot. + + Returns: + Detached response payload with canonical base-response fields. + """ + base = AgentResponse(**{ + name: getattr(response, name) + for name in _constructor_fields(AgentResponse) + if name not in ("value", "response_format", "raw_representation") and hasattr(response, name) + }) + # Use the base serializer, not an override that may omit or replace public response fields. + payload = AgentResponse.to_dict(response) + payload.update(base.to_dict()) + payload.pop("response_format", None) + payload.pop("raw_representation", None) + payload.pop("value", None) + payload.pop(_VALUE_BY_NAME_KEY, None) + payload["type"] = "agent_response" + + # Core's lazy value getter changes its cache. Parse a copy so recording is observational. + source = copy(response) + value = source._value # pyright: ignore[reportPrivateUsage] + if ( + not is_terminal_agent_response(source) + and not source.user_input_requests + and source.additional_properties.get("durable_status") != "accepted" + ): + value = source.value + if value is not None or source._value_parsed: # pyright: ignore[reportPrivateUsage] + by_name = getattr(source, _VALUE_BY_NAME_KEY, False) + if isinstance(value, BaseModel): + value, by_name = _serialize_model_value(value) + payload["value"] = value + if by_name: + payload[_VALUE_BY_NAME_KEY] = True + return deepcopy(payload) + def load_agent_response(agent_response: AgentResponse | dict[str, Any] | None) -> AgentResponse: """Convert raw payloads into AgentResponse instance. @@ -21,8 +185,9 @@ def load_agent_response(agent_response: AgentResponse | dict[str, Any] | None) - AgentResponse: The converted response object Raises: - ValueError: If agent_response is None - TypeError: If agent_response is an unsupported type + ValueError: If agent_response is None, its optional delivery version is unsupported, + or a response or content envelope is malformed. + TypeError: If the input type or required constructor fields are invalid. """ if agent_response is None: raise ValueError("agent_response cannot be None") @@ -32,8 +197,35 @@ def load_agent_response(agent_response: AgentResponse | dict[str, Any] | None) - if isinstance(agent_response, AgentResponse): return agent_response if isinstance(agent_response, dict): - logger.debug("[load_agent_response] Converting dict payload using AgentResponse.from_dict") - return AgentResponse.from_dict(agent_response) + logger.debug("[load_agent_response] Constructing a base response from delivery fields") + if _DELIVERY_VERSION_KEY in agent_response: + version = agent_response[_DELIVERY_VERSION_KEY] + if type(version) is not int or version != _DELIVERY_VERSION: + raise ValueError("Unsupported durable response version") + response_type = agent_response.get("type") + if "type" in agent_response and (not isinstance(response_type, str) or not response_type): + raise ValueError("Agent response type must be a non-empty string") + # Internal callers supply messages without a type. Custom response types are + # projected onto the base class, never imported, but must still be response-like. + if response_type != "agent_response" and agent_response.get("messages") is None: + raise ValueError("Agent response mapping requires a response type or messages") + # Filtering is consumer-only. Neither construction nor subsequent consumer mutations + # may remove or change unknown fields in the raw mailbox payload. + data = deepcopy(agent_response) + fields = _constructor_kwargs(data, AgentResponse) + fields.pop("response_format", None) + messages = fields.get("messages") + if messages is not None and not isinstance(messages, Message): + if not isinstance(messages, Sequence) or isinstance(messages, (str, bytes, bytearray)): + raise TypeError("Agent response messages must be a sequence of messages") + fields["messages"] = [_load_message(message) for message in cast("Sequence[Any]", messages)] + response = AgentResponse(**fields) + if "value" in data: + # Core sets this to False for None, losing the distinction between absent and null. + response._value_parsed = True # pyright: ignore[reportPrivateUsage] + if data.get(_VALUE_BY_NAME_KEY) is True: + setattr(response, _VALUE_BY_NAME_KEY, True) + return response raise TypeError(f"Unsupported type for agent_response: {type(agent_response)}") @@ -41,12 +233,14 @@ def load_agent_response(agent_response: AgentResponse | dict[str, Any] | None) - def ensure_response_format( response_format: type[BaseModel] | None, correlation_id: str, - response: AgentResponse, + response: AgentResponse[Any], ) -> None: """Ensure the AgentResponse value is parsed into the expected response_format. This function modifies the response in-place by parsing its value attribute - into the specified Pydantic model format. + into the specified Pydantic model format. Terminal responses and accepted + acknowledgements are left unchanged. A retained value, including null, + takes precedence over parsing message text again. Args: response_format: Optional Pydantic model class to parse the response value into @@ -57,9 +251,32 @@ def ensure_response_format( ValueError: If response_format is specified but response.value cannot be parsed """ if response_format is not None: + if ( + is_terminal_agent_response(response) + or response.user_input_requests + or response.additional_properties.get("durable_status") == "accepted" + ): + return + + # Only reuse a retained value; an unparsed response must use the requested format. + value = response._value # pyright: ignore[reportPrivateUsage] + value_present = value is not None or response._value_parsed # pyright: ignore[reportPrivateUsage] # Set the response format on the response so .value knows how to parse response._response_format = response_format # pyright: ignore[reportPrivateUsage] - response._value_parsed = False # pyright: ignore[reportPrivateUsage] # Reset to allow re-parsing with new format + if value_present: + if not isinstance(value, response_format): + # Retained values crossed a JSON boundary, just like structured message text. + by_name = getattr(response, _VALUE_BY_NAME_KEY, False) + if isinstance(value, BaseModel): + value, by_name = _serialize_model_value(value) + if by_name: + value = response_format.model_validate_json(json.dumps(value), by_alias=False, by_name=True) + else: + value = response_format.model_validate_json(json.dumps(value)) + response._value = value # pyright: ignore[reportPrivateUsage] + response._value_parsed = True # pyright: ignore[reportPrivateUsage] + else: + response._value_parsed = False # pyright: ignore[reportPrivateUsage] # Access response.value to trigger parsing (may raise ValidationError) # Validate that parsing succeeded diff --git a/python/packages/durabletask/agent_framework_durabletask/_retention.py b/python/packages/durabletask/agent_framework_durabletask/_retention.py new file mode 100644 index 0000000..4e22e31 --- /dev/null +++ b/python/packages/durabletask/agent_framework_durabletask/_retention.py @@ -0,0 +1,559 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Independent eager-pruning policy and opt-in whole-entity pressure eviction.""" + +from __future__ import annotations + +import json +import logging +import math +from collections.abc import Mapping +from copy import deepcopy +from datetime import datetime, timedelta, timezone +from typing import Any, Literal, TypeAlias, cast + +from agent_framework import ( + EXCLUDED_KEY, + GROUP_ANNOTATION_KEY, + GROUP_ID_KEY, + GROUP_INDEX_KEY, + CharacterEstimatorTokenizer, + Message, + TokenBudgetComposedStrategy, + annotate_message_groups, + included_token_count, +) + +from ._constants import DurableStateFields +from ._durable_agent_state import ( + DurableAgentState, + DurableAgentStateEntry, + DurableAgentStateEntryJsonType, + DurableAgentStateMessage, +) +from ._retention_telemetry import record_retention + +__all__ = [ + "DEFAULT_MAX_STATE_BYTES", + "DEFAULT_RETENTION", + "DELIVERY_WINDOW_SECONDS", + "DTS_MAX_STATE_BYTES", + "HIGH_WATERMARK", + "LOW_WATERMARK", + "RetentionMode", + "StateBudget", + "StateCapacityError", + "enforce_budget", + "prunes_excluded", + "resolve_state_budget", + "validate_retention", +] + +logger = logging.getLogger("agent_framework.durabletask") + +RetentionMode: TypeAlias = Literal["keep_all", "follow_compaction"] +"""Whether to eagerly prune compaction exclusions, independently of a pressure budget.""" + +StateBudget: TypeAlias = int | Literal["backend_limit"] | None +"""An explicit byte budget, a host-resolved limit, or disabled pressure eviction.""" + +DEFAULT_RETENTION: RetentionMode = "keep_all" +DEFAULT_MAX_STATE_BYTES: StateBudget = None +DTS_MAX_STATE_BYTES = 1_048_576 +HIGH_WATERMARK = 0.85 +LOW_WATERMARK = 0.70 +DELIVERY_WINDOW_SECONDS = 60 +"""Legacy response protection when independent completion bookkeeping is absent.""" + +_SYSTEM_ROLE = "system" +_MAX_PASSES = 3 +_Origin: TypeAlias = tuple[int, int] + +_EXCHANGE_KINDS = { + DurableAgentStateEntryJsonType.REQUEST, + DurableAgentStateEntryJsonType.RESPONSE, + DurableAgentStateEntryJsonType.ERROR_RESPONSE, +} +_TRANSCRIPT_KINDS = _EXCHANGE_KINDS | {DurableAgentStateEntryJsonType.COMPACTION} +_BARE_ENTRY_FIELDS = { + DurableStateFields.TYPE_DISCRIMINATOR, + DurableStateFields.CORRELATION_ID, + DurableStateFields.CREATED_AT, + DurableStateFields.MESSAGES, +} + + +class StateCapacityError(ValueError): + """The protected state or an unreachable retention target prevents a safe commit.""" + + def __init__(self, *, size_bytes: int, max_state_bytes: int, floor_bytes: int, target_bytes: int) -> None: + """Describe the measured state, configured budget, protected floor and target.""" + self.size_bytes = size_bytes + self.max_state_bytes = max_state_bytes + self.floor_bytes = floor_bytes + self.target_bytes = target_bytes + super().__init__( + f"Durable state capacity cannot meet the {target_bytes}-byte retention target: " + f"serialized size is {size_bytes} bytes, budget is {max_state_bytes} bytes, " + f"and the protected floor is {floor_bytes} bytes. No transcript changes were applied." + ) + + +def _positive_budget(value: object, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, not a boolean or another value type.") + return value + + +def resolve_state_budget(value: StateBudget, *, backend_limit: int | None = None) -> int | None: + """Resolve a pressure budget without enabling eager pruning or assuming a backend. + + Raises: + ValueError: The value is invalid, or ``backend_limit`` is requested but unresolved. + """ + if backend_limit is not None: + _positive_budget(backend_limit, "backend_limit") + if value is None: + return None + if isinstance(value, str) and value == "backend_limit": + if backend_limit is None: + raise ValueError("max_state_bytes='backend_limit' requires a known backend_limit from the host.") + return backend_limit + return _positive_budget(value, "max_state_bytes") + + +def validate_retention( + retention: RetentionMode, + high_watermark: float = HIGH_WATERMARK, + low_watermark: float = LOW_WATERMARK, +) -> None: + """Validate the eager-pruning mode and finite, ordered numeric watermarks. + + Raises: + ValueError: The mode or watermarks do not satisfy the retention contract. + """ + if not isinstance(retention, str) or retention not in ("keep_all", "follow_compaction"): + raise ValueError("retention must be 'keep_all' or 'follow_compaction'.") + for name, value in (("high_watermark", high_watermark), ("low_watermark", low_watermark)): + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not 0 < value <= 1 + or not math.isfinite(value) + ): + raise ValueError(f"{name} must be a finite number in (0, 1], not a boolean.") + if low_watermark >= high_watermark: + raise ValueError("watermarks must satisfy 0 < low_watermark < high_watermark <= 1.") + + +def prunes_excluded(retention: RetentionMode) -> bool: + """Whether compaction exclusions should be deleted as they are made.""" + validate_retention(retention) + return retention == "follow_compaction" + + +async def enforce_budget( + state: DurableAgentState, + *, + max_state_bytes: int, + high_watermark: float = HIGH_WATERMARK, + low_watermark: float = LOW_WATERMARK, +) -> int: + """Evict eligible oldest atomic groups using detached, byte-checked plans. + + Args: + state: Modified only after a plan fits, including its truncation record. + + Keyword Args: + max_state_bytes: An already resolved positive budget. Callers skip this function for None. + high_watermark: The fraction at which pressure eviction starts. + low_watermark: The desired retained fraction, raised to the protected floor if necessary. + + Returns: + The number of transcript messages removed. + + Raises: + ValueError: A budget or watermark is invalid. + StateCapacityError: No safe target is reachable. The input state remains unchanged. + """ + _positive_budget(max_state_bytes, "max_state_bytes") + validate_retention(DEFAULT_RETENTION, high_watermark, low_watermark) + high = int(max_state_bytes * high_watermark) + size = _serialized_size(state) + if size < high: + record_retention( + state, + mechanism="pressure", + outcome="below_threshold", + before_bytes=size, + after_bytes=size, + budget_bytes=max_state_bytes, + ) + return 0 + + baseline = deepcopy(state) + now = datetime.now(tz=timezone.utc) + messages, origins = _candidates(baseline, now=now) + floor_state = _stage_eviction(baseline, set(origins)) + floor_without_record = _serialized_size(floor_state) + if origins: + record_truncation(floor_state, len(origins), now=now) + floor = _serialized_size(floor_state) + target = max(int(max_state_bytes * low_watermark), floor) + if floor >= high: + record_retention( + state, + mechanism="pressure", + outcome="protected_floor", + before_bytes=size, + after_bytes=size, + budget_bytes=max_state_bytes, + ) + raise StateCapacityError( + size_bytes=size, max_state_bytes=max_state_bytes, floor_bytes=floor, target_bytes=high - 1 + ) + + groups: dict[str, list[int]] = {} + for index, message in enumerate(messages): + groups.setdefault(_group_id(message), []).append(index) + ordered_groups = list(groups.values()) + group_tokens = [included_token_count([messages[index] for index in group]) for group in ordered_groups] + group_sizes = _prefix_sizes( + baseline, + origins, + ordered_groups, + size=size, + record_cost=floor - floor_without_record, + ) + stored_origins = [ + (baseline.data.conversation_history[entry], baseline.data.conversation_history[entry].messages[message]) + for entry, message in origins + ] + evictable_bytes = sum(_message_size(stored) for _, stored in stored_origins) + planning_target = target + + for _ in range(_MAX_PASSES): + cutoff = next( + (index + 1 for index, projected_size in enumerate(group_sizes) if projected_size <= planning_target), + len(ordered_groups), + ) + retained_tokens = sum(group_tokens[cutoff:]) + estimate = _token_budget( + stored_origins, + serialized_size=size, + evictable_bytes=evictable_bytes, + target_bytes=planning_target, + floor_bytes=floor, + evictable_tokens=sum(group_tokens), + ) + # Align the estimate to a byte-measured group boundary. A global bytes/token ratio + # alone can over-delete mixed Unicode, tool payloads and small prose messages. + token_budget = min(max(estimate, retained_tokens), retained_tokens + group_tokens[cutoff - 1] - 1) + planned = deepcopy(messages) + # Core 1.16 retains its last non-system group even above budget. A detached, empty + # user anchor occupies that slot, so the last eligible OLD group is not pinned. + anchor = Message("user", [], message_id="retention_anchor") + annotate_message_groups([anchor], tokenizer=CharacterEstimatorTokenizer()) + planned.append(anchor) + strategy = TokenBudgetComposedStrategy( + token_budget=token_budget + included_token_count([anchor]), + tokenizer=CharacterEstimatorTokenizer(), + strategies=[], + ) + await strategy(planned) + removed = { + origin + for origin, message in zip(origins, planned) + if message.additional_properties.get(EXCLUDED_KEY, False) + } + staged = _stage_eviction(baseline, removed) + if removed: + record_truncation(staged, len(removed), now=now) + measured = _serialized_size(staged) + if measured <= target and measured < high: + state.data.conversation_history[:] = staged.data.conversation_history + state.data.truncation = staged.data.truncation + record_retention( + state, + mechanism="pressure", + outcome="staged", + before_bytes=size, + after_bytes=measured, + budget_bytes=max_state_bytes, + removed_messages=len(removed), + removed_entries=len(baseline.data.conversation_history) - len(staged.data.conversation_history), + ) + logger.warning( + "[Retention] Evicted %d oldest transcript message(s), leaving %d serialized bytes " + "against a %d-byte budget. Set max_state_bytes=None to disable pressure eviction.", + len(removed), + measured, + max_state_bytes, + ) + return len(removed) + # Correct the observed planning error, not an arbitrary fraction of the target. + # Subtracting only the excess over target can select the same group boundary again. + planning_error = max(measured - group_sizes[cutoff - 1], 1) + planning_target = max(floor, target - planning_error) + + record_retention( + state, + mechanism="pressure", + outcome="unreachable_target", + before_bytes=size, + after_bytes=size, + budget_bytes=max_state_bytes, + ) + raise StateCapacityError(size_bytes=size, max_state_bytes=max_state_bytes, floor_bytes=floor, target_bytes=target) + + +def record_truncation(state: DurableAgentState, removed: int, *, now: datetime | None = None) -> None: + """Accumulate bounded eviction evidence without discarding unknown metadata.""" + timestamp = (now or datetime.now(tz=timezone.utc)).isoformat() + existing = state.data.truncation or {} + state.data.truncation = { + **existing, + DurableStateFields.EVICTED_MESSAGE_COUNT: int(existing.get(DurableStateFields.EVICTED_MESSAGE_COUNT, 0)) + + removed, + DurableStateFields.FIRST_EVICTED_AT: existing.get(DurableStateFields.FIRST_EVICTED_AT, timestamp), + DurableStateFields.LAST_EVICTED_AT: timestamp, + } + + +def _serialized_size(state: DurableAgentState) -> int: + """Measure default JSON serialization, including ASCII escapes but excluding transport framing.""" + return len(json.dumps(state.to_dict())) + + +def _detached_message(stored: DurableAgentStateMessage) -> Message: + message: Message = deepcopy(stored).to_chat_message() + message.additional_properties.pop(EXCLUDED_KEY, None) + # Recount with this tokenizer rather than trusting another strategy's cached token count. + message.additional_properties.pop(GROUP_ANNOTATION_KEY, None) + return message + + +def _group_id(message: Message) -> str: + return cast("str", message.additional_properties[GROUP_ANNOTATION_KEY][GROUP_ID_KEY]) + + +def _saved_group_id(stored: DurableAgentStateMessage) -> str | None: + annotation = (stored.extension_data or {}).get(GROUP_ANNOTATION_KEY) + if isinstance(annotation, Mapping): + group_id = cast("Mapping[str, object]", annotation).get(GROUP_ID_KEY) + if isinstance(group_id, str): + return group_id + return None + + +def _link_atomic_groups(messages: list[Message], saved_ids: list[str | None]) -> list[int]: + """Unite core-inferred groups with persisted atomic links, including non-contiguous spans.""" + parents = list(range(len(messages))) + + def root(index: int) -> int: + while parents[index] != index: + parents[index] = parents[parents[index]] + index = parents[index] + return index + + for group_ids in (saved_ids, [_group_id(message) for message in messages]): + first: dict[str, int] = {} + for index, group_id in enumerate(group_ids): + if group_id is not None: + left, right = root(first.setdefault(group_id, index)), root(index) + parents[max(left, right)] = min(left, right) + + roots = [root(index) for index in range(len(messages))] + for message, group in zip(messages, roots): + annotation = cast("dict[str, Any]", message.additional_properties[GROUP_ANNOTATION_KEY]) + annotation[GROUP_ID_KEY] = f"retention_group_{group}" + annotation[GROUP_INDEX_KEY] = group + return roots + + +def _candidates(state: DurableAgentState, *, now: datetime) -> tuple[list[Message], list[_Origin]]: + history = state.data.conversation_history + completed = cast("Mapping[str, object] | None", getattr(state.data, "completed_correlations", None)) + protected = {id(entry) for entry in _protected_entries(history, completed_correlations=completed, now=now)} + messages: list[Message] = [] + origins: list[_Origin | None] = [] + saved_ids: list[str | None] = [] + held: set[int] = set() + reserved = {stored.message_id for entry in history for stored in entry.messages if stored.message_id} + seen: set[str] = set() + + for entry_index, entry in enumerate(history): + known = entry.json_type in _TRANSCRIPT_KINDS + # Unknown entries are opaque barriers, not model-conversion inputs or deletion candidates. + for message_index, stored in enumerate(entry.messages if known else (entry.messages or [None])): + index = len(messages) + eligible = known and stored is not None and bool(stored.contents) + message = _detached_message(stored) if known and stored is not None else Message(_SYSTEM_ROLE, []) + if not eligible or id(entry) in protected or message.role == _SYSTEM_ROLE: + held.add(index) + message_id = message.message_id + if not message_id or message_id in seen: + suffix = index + message_id = f"retention_message_{suffix}" + while message_id in reserved: + suffix += 1 + message_id = f"retention_message_{suffix}" + message.message_id = message_id + reserved.add(message_id) + seen.add(message_id) + messages.append(message) + origins.append((entry_index, message_index) if eligible else None) + saved_ids.append(_saved_group_id(stored) if stored is not None else None) + + annotate_message_groups(messages, force_reannotate=True, tokenizer=CharacterEstimatorTokenizer()) + roots = _link_atomic_groups(messages, saved_ids) + protected_groups = {roots[index] for index in held} + candidates: list[Message] = [] + candidate_origins: list[_Origin] = [] + for index, origin in enumerate(origins): + if origin is not None and roots[index] not in protected_groups: + candidates.append(messages[index]) + candidate_origins.append(origin) + return candidates, candidate_origins + + +def _can_drop_entry(entry: DurableAgentStateEntry) -> bool: + # Only a bare transcript envelope may disappear with its final message. Keep usage, + # response schemas, orchestration metadata and unknown fields in the protected floor. + return ( + entry.json_type in _TRANSCRIPT_KINDS + and not entry.extension_data + and entry.to_dict().keys() <= _BARE_ENTRY_FIELDS + ) + + +def _stage_eviction(state: DurableAgentState, removed: set[_Origin]) -> DurableAgentState: + staged = deepcopy(state) + history: list[DurableAgentStateEntry] = [] + for entry_index, entry in enumerate(staged.data.conversation_history): + remaining = [message for index, message in enumerate(entry.messages) if (entry_index, index) not in removed] + changed = len(remaining) != len(entry.messages) + entry.messages = remaining + # Do not incidentally remove an already-empty or unknown entry. + if remaining or not changed or not _can_drop_entry(entry): + history.append(entry) + staged.data.conversation_history = history + return staged + + +def _prefix_sizes( + state: DurableAgentState, + origins: list[_Origin], + groups: list[list[int]], + *, + size: int, + record_cost: int, +) -> list[int]: + """Compute default-JSON byte costs at core group boundaries without repeated whole-state copies.""" + history = state.data.conversation_history + remaining = [len(entry.messages) for entry in history] + entry_sizes = [len(json.dumps(entry.to_dict())) for entry in history] + droppable = [_can_drop_entry(entry) for entry in history] + entry_count = len(history) + previous_count = int((state.data.truncation or {}).get(DurableStateFields.EVICTED_MESSAGE_COUNT, 0)) + final_count_digits = len(str(previous_count + len(origins))) + removed = 0 + sizes: list[int] = [] + for group in groups: + for index in group: + entry_index, message_index = origins[index] + if remaining[entry_index] == 1 and droppable[entry_index]: + saved = entry_sizes[entry_index] + (2 if entry_count > 1 else 0) + entry_count -= 1 + else: + stored = history[entry_index].messages[message_index] + saved = _message_size(stored) + (2 if remaining[entry_index] > 1 else 0) + entry_sizes[entry_index] -= saved + remaining[entry_index] -= 1 + size -= saved + removed += 1 + # The timestamp and unknown truncation fields are fixed across plans. Only the + # decimal width of the aggregate count varies with the chosen prefix. + count_correction = len(str(previous_count + removed)) - final_count_digits + sizes.append(size + record_cost + count_correction) + return sizes + + +def _newest_exchange(history: list[DurableAgentStateEntry]) -> list[DurableAgentStateEntry]: + """Return the entries belonging to the most recent exchange. + + Grouped by correlation id, so a request and the response it produced are protected together. + + Compaction entries answer no request and carry no correlation, so they are skipped when + deciding which exchange is newest. Taking the last entry blindly would let a summary appended + at the end stand in for the turn that actually just happened, leaving that turn unprotected. + """ + for entry in reversed(history): + if entry.json_type in _EXCHANGE_KINDS and entry.correlation_id is not None: + newest = entry.correlation_id + return [candidate for candidate in history if candidate.correlation_id == newest] + return [history[-1]] if history else [] + + +def _as_utc(value: datetime) -> datetime: + """Persisted timestamps can come back without a timezone, so read those as UTC.""" + return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc) + + +def _protected_entries( + history: list[DurableAgentStateEntry], + *, + completed_correlations: Mapping[str, object] | None = None, + now: datetime | None = None, +) -> list[DurableAgentStateEntry]: + """Protect the newest exchange and recent responses lacking independent completion records.""" + protected = list(_newest_exchange(history)) + completed = completed_correlations or {} + cutoff = (now or datetime.now(tz=timezone.utc)) - timedelta(seconds=DELIVERY_WINDOW_SECONDS) + responses = [ + entry + for entry in history + if entry.json_type in (DurableAgentStateEntryJsonType.RESPONSE, DurableAgentStateEntryJsonType.ERROR_RESPONSE) + and entry.correlation_id not in completed + and _as_utc(entry.created_at) > cutoff + ] + undelivered = {entry.correlation_id for entry in responses if entry.correlation_id is not None} + response_ids = {id(entry) for entry in responses} + protected_ids = {id(entry) for entry in protected} + protected.extend( + entry + for entry in history + if (id(entry) in response_ids or entry.correlation_id in undelivered) and id(entry) not in protected_ids + ) + return protected + + +def _message_size(stored: DurableAgentStateMessage) -> int: + """The persisted message payload size, including non-text contents and metadata.""" + return len(json.dumps(stored.to_dict())) + + +def _token_budget( + origins: list[tuple[DurableAgentStateEntry, DurableAgentStateMessage]], + *, + serialized_size: int, + evictable_bytes: int, + target_bytes: int, + floor_bytes: int | None = None, + evictable_tokens: int | None = None, +) -> int: + """Estimate tokens from persisted candidate bytes and core's actual token annotations. + + The optional measurements let the engine reuse its detached grouping pass and exact floor. + The four original arguments remain usable by callers that only need a conservative estimate. + """ + if evictable_bytes <= 0: + return 1 + if floor_bytes is None: + floor_bytes = max(serialized_size - evictable_bytes, 0) + allowed_bytes = max(target_bytes - floor_bytes, 0) + if evictable_tokens is None: + messages = [_detached_message(stored) for _, stored in origins] + annotate_message_groups(messages, tokenizer=CharacterEstimatorTokenizer()) + evictable_tokens = included_token_count(messages) + return max(allowed_bytes * evictable_tokens // evictable_bytes, 1) diff --git a/python/packages/durabletask/agent_framework_durabletask/_retention_telemetry.py b/python/packages/durabletask/agent_framework_durabletask/_retention_telemetry.py new file mode 100644 index 0000000..24a50e7 --- /dev/null +++ b/python/packages/durabletask/agent_framework_durabletask/_retention_telemetry.py @@ -0,0 +1,197 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Bounded observations of local retention, never proof of a durable commit. + +Deletion measurements describe staged state at the retention boundary. Operation and +write measurements describe the later host call, which may itself only stage a write. +Only the OpenTelemetry API is required. No provider or exporter is configured here. +""" + +from __future__ import annotations + +import json +from collections.abc import Generator +from contextlib import contextmanager, suppress +from contextvars import ContextVar +from dataclasses import dataclass +from functools import lru_cache +from typing import TYPE_CHECKING, Literal + +from opentelemetry.metrics import NoOpMeter, get_meter + +if TYPE_CHECKING: + from ._durable_agent_state import DurableAgentState + +_Mechanism = Literal["eager", "pressure"] +_Outcome = Literal["below_threshold", "staged", "protected_floor", "unreachable_target", "protected"] + + +class _Instruments: + def __init__(self) -> None: + meter = get_meter("agent_framework.durabletask") + self.noop = isinstance(meter, NoOpMeter) + self.evaluations = meter.create_counter( + "durable.retention.evaluations", unit="{evaluation}", description="Local retention evaluations." + ) + self.budget = meter.create_histogram( + "durable.retention.budget", unit="By", description="Requested resolved whole-entity pressure budget." + ) + self.size = meter.create_histogram( + "durable.retention.state.size", unit="By", description="Serialized entity JSON at a retention boundary." + ) + self.messages = meter.create_counter( + "durable.retention.removed_messages", unit="{message}", description="Messages removed from staged state." + ) + self.entries = meter.create_counter( + "durable.retention.removed_entries", unit="{entry}", description="Entries removed from staged state." + ) + self.reclaimed = meter.create_counter( + "durable.retention.reclaimed_bytes", unit="By", description="Nonnegative byte reduction in staged state." + ) + self.capacity_failures = meter.create_counter( + "durable.retention.capacity_failures", unit="{failure}", description="Unreachable pressure targets." + ) + self.writes = meter.create_counter( + "durable.retention.write_attempts", + unit="{attempt}", + description="State serialization or host set_state outcomes, not durable commit confirmation.", + ) + self.operations = meter.create_counter( + "durable.retention.operations", + unit="{operation}", + description="Run operations with retention observations and their host write status.", + ) + + +@lru_cache(maxsize=1) +def _instruments() -> _Instruments: + # Cache the API's proxy too: it can bind to an SDK installed after import. + return _Instruments() + + +@dataclass +class _Operation: + state: DurableAgentState + active: bool = True + observed: bool = False + removed_messages: int = 0 + removed_entries: int = 0 + commit_status: Literal["not_attempted", "unknown"] = "not_attempted" + + +_operation: ContextVar[_Operation | None] = ContextVar("durable_retention_operation", default=None) + + +def _current(state: DurableAgentState) -> _Operation | None: + operation = _operation.get() + if operation is not None and operation.active and operation.state is state: + return operation + return None + + +@contextmanager +def retention_operation(state: DurableAgentState) -> Generator[None]: + """Isolate run observations through rollback and the host write attempt. + + The identity check prevents attributing another state's retention to this run. + Closing the object also invalidates contexts inherited by unfinished child tasks. + """ + operation = _Operation(state) + token = _operation.set(operation) + failed = False + try: + yield + except BaseException: + failed = True + raise + finally: + operation.active = False + _operation.reset(token) + if operation.observed: + # Optional telemetry must not replace the operation's result or error. + with suppress(Exception): + _instruments().operations.add( + 1, + { + "outcome": "failed" if failed else "returned", + "commit_status": operation.commit_status, + "deletion_staged": operation.removed_messages > 0, + }, + ) + + +def eager_state_size(state: DurableAgentState) -> int | None: + """Measure only an eligible eager-prune boundary, skipping an explicit no-op meter. + + The API has no portable enabled check for a proxy or an SDK without readers. + Those meters still measure eligible eager deletions, but never ordinary flushes. + """ + try: + if _instruments().noop: + return None + return len(json.dumps(state.to_dict())) + except Exception: + return None + + +def record_retention( + state: DurableAgentState, + *, + mechanism: _Mechanism, + outcome: _Outcome, + before_bytes: int | None = None, + after_bytes: int | None = None, + budget_bytes: int | None = None, + removed_messages: int = 0, + removed_entries: int = 0, +) -> None: + """Record actual staged changes, never the exclusions on a detached trial plan.""" + operation = _current(state) + if operation is not None: + operation.observed = True + operation.removed_messages += removed_messages + operation.removed_entries += removed_entries + attributes = {"mechanism": mechanism, "outcome": outcome, "commit_status": "not_attempted"} + with suppress(Exception): + instruments = _instruments() + instruments.evaluations.add(1, attributes) + if budget_bytes is not None: + instruments.budget.record(budget_bytes, attributes) + if before_bytes is not None: + instruments.size.record(before_bytes, {**attributes, "phase": "before"}) + if after_bytes is not None: + instruments.size.record(after_bytes, {**attributes, "phase": "after"}) + if removed_messages: + instruments.messages.add(removed_messages, attributes) + if before_bytes is not None and after_bytes is not None: + instruments.reclaimed.add(max(0, before_bytes - after_bytes), attributes) + if removed_entries: + instruments.entries.add(removed_entries, attributes) + if outcome in ("protected_floor", "unreachable_target"): + instruments.capacity_failures.add(1, attributes) + + +def record_write( + state: DurableAgentState, + *, + stage: Literal["serialization", "set_state"], + outcome: Literal["returned", "failed"], +) -> None: + """Observe a host write boundary without interpreting its return as persistence.""" + operation = _current(state) + if operation is None: + return + if stage == "set_state": + # Even a failed host call may have staged work before it raised. + operation.commit_status = "unknown" + if operation.observed: + with suppress(Exception): + _instruments().writes.add( + 1, + { + "stage": stage, + "outcome": outcome, + "commit_status": operation.commit_status, + "deletion_staged": operation.removed_messages > 0, + }, + ) diff --git a/python/packages/durabletask/agent_framework_durabletask/_shim.py b/python/packages/durabletask/agent_framework_durabletask/_shim.py index ed8a752..0c365a0 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_shim.py +++ b/python/packages/durabletask/agent_framework_durabletask/_shim.py @@ -18,13 +18,47 @@ from ._executors import DurableAgentExecutor from ._feature_usage import FeatureIndex -from ._models import DurableAgentSession +from ._models import AgentSessionId, DurableAgentSession # TypeVar for the task type returned by executors # Covariant because TaskT only appears in return positions (output) TaskT = TypeVar("TaskT", covariant=True) +def build_agent_task( + executor: DurableAgentExecutor[Any], + executor_id: str, + message: str, + orchestration_instance_id: str, + context_messages: list[dict[str, Any]] | None = None, + context_message_ids: list[str] | None = None, +) -> Any: + """Create the yieldable task that runs a workflow's agent node. + + Shared by every host adapter: the only host-specific part of dispatching an agent is + which :class:`DurableAgentExecutor` drives it, so the surrounding session/agent wiring + lives here rather than being repeated per host. + + Args: + executor: The host's executor, which knows how to reach the agent entity. + executor_id: The workflow-scoped agent identity to dispatch to. + message: The text message for this turn. + orchestration_instance_id: Used as the entity session key, keeping conversation + state isolated per workflow run. + context_messages: Optional upstream conversation delivered as prior context. + context_message_ids: Durable occurrence IDs, separate from application message IDs. + + Returns: + A yieldable task whose result is an ``AgentResponse``. + """ + session_id = AgentSessionId(name=executor_id, key=orchestration_instance_id) + session = DurableAgentSession(durable_session_id=session_id) + agent = DurableAIAgent(executor, executor_id) + return agent.run( + message, session=session, context_messages=context_messages, context_message_ids=context_message_ids + ) + + class DurableAgentProvider(ABC, Generic[TaskT]): """Abstract provider for constructing durable agent proxies. @@ -94,6 +128,8 @@ def run( # type: ignore[override] stream: Literal[False] = False, session: AgentSession | None = None, options: dict[str, Any] | None = None, + context_messages: list[dict[str, Any]] | None = None, + context_message_ids: list[str] | None = None, ) -> TaskT: """Execute the agent via the injected provider. @@ -105,6 +141,10 @@ def run( # type: ignore[override] options: Optional options dictionary. Supported keys include ``response_format``, ``enable_tool_calls``, and ``wait_for_response``. Additional keys are forwarded to the agent execution. + context_messages: Optional upstream conversation (serialized ``Message`` dicts) + delivered to the agent as prior context. Workflows use this to give a + downstream agent the conversation produced by upstream nodes. + context_message_ids: Durable occurrence identities paired with context messages. Note: This method overrides SupportsAgentRun.run() with a different return type: @@ -122,11 +162,23 @@ def run( # type: ignore[override] """ if stream is not False: raise ValueError("DurableAIAgent does not support streaming mode (stream must be False)") - message_str = self._normalize_messages(messages) + # Explicit context is the invocation payload, including an empty delta or + # tool-only messages. The separate workflow string is just a logging preview. + message_str = ( + messages + if context_messages is not None and isinstance(messages, str) + else self._normalize_messages(messages) + ) + # Only forward context messages when a workflow supplied them, so executors that do + # not implement the parameter keep working unchanged. + extra: dict[str, Any] = {"context_messages": context_messages} if context_messages is not None else {} + if context_message_ids is not None: + extra["context_message_ids"] = context_message_ids run_request = self._executor.get_run_request( message=message_str, options=options, + **extra, ) mark_feature_used(FeatureIndex.DURABLETASK) diff --git a/python/packages/durabletask/agent_framework_durabletask/_state_migration.py b/python/packages/durabletask/agent_framework_durabletask/_state_migration.py new file mode 100644 index 0000000..e164faa --- /dev/null +++ b/python/packages/durabletask/agent_framework_durabletask/_state_migration.py @@ -0,0 +1,350 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Explicit, detached legacy-to-v2 state migration, with no storage or provider access.""" + +from __future__ import annotations + +import hashlib +import json +import re +from collections.abc import Iterator +from datetime import datetime, timedelta, timezone +from typing import Any, cast + +from ._constants import DurableStateFields +from ._durable_agent_state import ( + DurableAgentState, + DurableAgentStateRequest, + DurableAgentStateResponse, + _validate_json, # pyright: ignore[reportPrivateUsage] +) +from ._message_identity import message_identity +from ._response_utils import load_agent_response +from ._retention import StateCapacityError +from ._workflows.naming import parse_workflow_message_id + +__all__ = ["migrate_legacy_state", "state_snapshot_digest"] + +_SHA256 = re.compile(r"[0-9a-f]{64}") +_EVIDENCE_FIELDS = {"sourceDigest", "evidenceId", "complete", "messages"} +_JOURNAL_REQUIRED = ( + "Legacy ingestedPositions require recorded delivery evidence: a complete authoritative accepted-message " + "journal from the quiesced legacy deployment, including evicted messages. If that journal is unavailable, " + "keep the old session on the old engine rather than guessing delivery receipts." +) + + +def _canonical_json(value: Any) -> str: + try: + _validate_json(value) + return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False) + except (TypeError, ValueError, RecursionError) as exc: + raise ValueError("Migration inputs must be strict JSON with string keys and finite numbers.") from exc + + +def state_snapshot_digest(source: dict[str, Any]) -> str: + """Return the SHA-256 of the complete strict-JSON source snapshot encoded as UTF-8. + + Object keys are sorted, separators are compact, Unicode is not ASCII-escaped, + and non-finite numbers, non-string keys and non-JSON values are rejected. + Array order and all unknown fields participate in the digest. + + Args: + source: The unmodified exported legacy state, not a parsed or upgraded state. + + Returns: + A lowercase hexadecimal SHA-256 digest. + """ + if not isinstance(source, dict): + raise ValueError("source must be a JSON object.") + return hashlib.sha256(_canonical_json(source).encode("utf-8")).hexdigest() + + +def _nonblank(value: Any, name: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{name} must be a nonblank string.") + return value + + +def _positive_int(value: Any, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, not a boolean.") + return value + + +def _workflow_position(identity: str) -> tuple[str, int] | None: + parsed = parse_workflow_message_id(identity) + if identity.startswith("wf_") and (parsed is None or identity != identity.strip() or not parsed[0].strip()): + raise ValueError("Malformed workflow message ID in recorded delivery evidence or legacy receipts.") + return parsed + + +def _legacy_positions(data: dict[str, Any]) -> dict[str, int]: + raw = data.get("ingestedPositions", {}) + if not isinstance(raw, dict): + raise ValueError("Legacy ingestedPositions must be an object of nonnegative integer producer positions.") + positions: dict[str, int] = {} + for producer, position in cast(dict[str, Any], raw).items(): + _nonblank(producer, "ingestedPositions producer") + if isinstance(position, bool) or not isinstance(position, int) or position < 0: + raise ValueError("Every ingestedPositions position must be a nonnegative integer, not a boolean.") + positions[producer] = position + return positions + + +def _validate_receipts(receipts: dict[str, list[str] | None]) -> None: + for identity, fingerprints in receipts.items(): + _nonblank(identity, "ingestedMessages ID") + workflow = _workflow_position(identity) + if fingerprints is None: + if workflow is not None: + raise ValueError("Workflow identity-only markers are not exact recorded delivery evidence.") + continue + if not fingerprints or any(_SHA256.fullmatch(value) is None for value in fingerprints): + raise ValueError("ingestedMessages requires nonempty lists of lowercase SHA-256 fingerprints.") + if len(set(fingerprints)) != len(fingerprints): + raise ValueError("ingestedMessages must not contain duplicate fingerprints.") + + +def _journal_receipts( + evidence: dict[str, Any], *, source_digest: str, positions: dict[str, int] +) -> tuple[str, dict[str, list[str]]]: + if not isinstance(evidence, dict) or evidence.keys() != _EVIDENCE_FIELDS: + raise ValueError("Recorded delivery evidence requires exactly sourceDigest, evidenceId, complete and messages.") + # Detach before constructing any core object. The loader must never touch the caller's journal. + journal: dict[str, Any] = json.loads(_canonical_json(evidence)) + if journal["sourceDigest"] != source_digest: + raise ValueError("Recorded delivery evidence sourceDigest does not match the source snapshot.") + evidence_id = _nonblank(journal["evidenceId"], "Recorded delivery evidence evidenceId") + if journal["complete"] is not True: + raise ValueError("Recorded delivery evidence requires the explicit complete=True operator assertion.") + messages = journal["messages"] + if not isinstance(messages, list): + raise ValueError("Recorded delivery evidence messages must be a list of complete canonical message objects.") + + receipts: dict[str, list[str]] = {} + maxima: dict[str, int] = {} + for raw in cast(list[Any], messages): + if not isinstance(raw, dict): + raise ValueError("Recorded delivery evidence messages must contain canonical message objects.") + raw = cast(dict[str, Any], raw) + identity = _nonblank(raw.get("message_id"), "Recorded delivery evidence message_id") + workflow = _workflow_position(identity) + _nonblank(raw.get("role"), "Recorded delivery evidence message role") + if not isinstance(raw.get("contents"), list): + raise ValueError("Recorded delivery evidence message contents must be a canonical array.") + try: + message = load_agent_response({"messages": [raw]}).messages[0] + # Do not hash a projection which silently lost unknown fields or changed their types. + if _canonical_json(message.to_dict()) != _canonical_json(raw): + raise ValueError("The message does not round-trip as a complete canonical input.") + fingerprint = message_identity(message) + except (TypeError, ValueError, AttributeError) as exc: + raise ValueError("Recorded delivery evidence requires lossless complete canonical message inputs.") from exc + revisions = receipts.setdefault(identity, []) + if fingerprint in revisions: + raise ValueError("Recorded delivery evidence contains a duplicate message ID/fingerprint pair.") + revisions.append(fingerprint) + if workflow is not None: + producer, position = workflow + maxima[producer] = max(maxima.get(producer, position), position) + + # This is only a consistency check. Sparse positions are valid; a maximum is never + # proof of a complete prefix, nor proof that the operator's journal is complete. + if maxima != positions: + raise ValueError( + "Recorded delivery evidence workflow producers and maximum positions must match ingestedPositions." + ) + return evidence_id, receipts + + +def _retained_custom_request_ids(state: DurableAgentState) -> Iterator[str]: + """Yield legacy lookup identities, never fingerprints of possibly pruned content.""" + for entry in state.data.conversation_history: + if isinstance(entry, DurableAgentStateRequest): + for message in entry.messages: + if message.message_id is not None: + if isinstance(message.message_id, str) and not message.message_id.strip(): + continue + identity = _nonblank(message.message_id, "Legacy request message ID") + if _workflow_position(identity) is None: + yield identity + + +def _apply_journal(state: DurableAgentState, journal: dict[str, list[str]]) -> None: + receipts = state.data.ingested_messages + for identity, existing in receipts.items(): + recorded = journal.get(identity) + if recorded is None or (existing is not None and not set(existing).issubset(recorded)): + raise ValueError("Recorded delivery evidence is inconsistent with existing ingestedMessages receipts.") + for identity in _retained_custom_request_ids(state): + if identity not in journal: + raise ValueError("Recorded delivery evidence must include every retained legacy custom request message ID.") + for identity, recorded in journal.items(): + existing = receipts.get(identity) + if existing is None: + receipts[identity] = list(recorded) + else: + existing.extend(fingerprint for fingerprint in recorded if fingerprint not in existing) + + +def _preserve_session(state: DurableAgentState, source_session_id: str) -> None: + session = state.data.session + if session is None: + state.data.session = {"session_id": source_session_id, "state": {}} + return + if not isinstance(session, dict): + raise ValueError("Legacy session must be an AgentSession-like object or null.") + existing_id = session.get("session_id") + if existing_id is not None and not isinstance(existing_id, str): + raise ValueError("Legacy session.session_id must be a string or null.") + if isinstance(existing_id, str) and existing_id.strip() and existing_id != source_session_id: + raise ValueError("Legacy session.session_id must match source_session_id to preserve external-store identity.") + session["session_id"] = source_session_id + session.setdefault("state", {}) + + +def migrate_legacy_state( + source: dict[str, Any], + *, + source_digest: str, + source_session_id: str, + migration_id: str, + ownership_transfer_id: str, + delivery_window_seconds: int, + max_state_bytes: int | None = None, + delivery_evidence: dict[str, Any] | None = None, + require_known_outcomes: bool = False, + now: datetime | None = None, +) -> DurableAgentState: + """Stage a detached legacy migration for an explicit entity migrate operation. + + The parent must enforce an EMPTY, separate destination on an isolated v2 hub, + quiesce the legacy owner, authorize ownership transfer, and atomically commit + once with idempotency keyed by the migration request. This function performs + no model, tool or provider calls, backend writes, or provider-transcript import. + It does not authorize the supplied IDs or prove source ownership. + + A nonempty scalar ingestedPositions map requires privileged operator-provided + recorded delivery evidence, not cryptographically proven history. The operator + must obtain the COMPLETE authoritative accepted-message journal from a quiesced + legacy deployment, including retained and evicted inputs and every accepted + revision of a message ID. Migration cannot independently verify the evidence's + authority or completeness. If that journal is unavailable, keep the old session + on the old engine rather than guessing. Equal maxima check consistency only; + no contiguous positions or prefixes are required or inferred. + + Only recorded responses receive completion/mailbox backfill. Surviving legacy + responses may be partial, not immutable originals. Existing delivery records + are preserved, never reopened. A fresh grace timestamp is captured once per + call, not once per migration ID: the parent owns retry idempotency and must not + repeatedly migrate the same source to refresh grace. Fixed now gives fixed + backfill timestamps. Existing state parsing owns legacy transcript conversion. + + Args: + source: Raw exported version-1 state. The caller's object remains untouched. + source_digest: Lowercase SHA-256 returned by state_snapshot_digest(source). + source_session_id: Original logical session identity, including its existing namespace. + migration_id: Nonblank parent-managed idempotency identifier. + ownership_transfer_id: Nonblank parent-authorized ownership transfer identifier. + delivery_window_seconds: Positive bounded grace period for legacy response backfill. + max_state_bytes: Optional positive resolved budget, measured with default ASCII JSON. + All migrated data and metadata are protected; oversize states fail without pruning. + delivery_evidence: Exactly sourceDigest, nonblank evidenceId, complete=True and + messages, a list of complete canonical Message.to_dict() inputs. Unsupported + or lossy canonical inputs and duplicate ID/fingerprint pairs are rejected. + require_known_outcomes: Reject imports with unknown invocation outcomes instead + of using legacy-compatible receipts. Neither mode discards completion evidence. + now: Offset-aware timestamp for this staging call, defaulting to UTC now. + + Returns: + A detached version-2 DurableAgentState ready for parent validation and commit. + + Raises: + ValueError: Invalid input, version, digest, evidence, identity, or timestamp. + StateCapacityError: The complete staged result exceeds max_state_bytes. + """ + _nonblank(source_session_id, "source_session_id") + _nonblank(migration_id, "migration_id") + _nonblank(ownership_transfer_id, "ownership_transfer_id") + _positive_int(delivery_window_seconds, "delivery_window_seconds") + if not isinstance(require_known_outcomes, bool): + raise ValueError("require_known_outcomes must be a boolean.") + if max_state_bytes is not None: + _positive_int(max_state_bytes, "max_state_bytes") + if not isinstance(source_digest, str) or _SHA256.fullmatch(source_digest) is None: + raise ValueError("source_digest must be a lowercase SHA-256 snapshot digest.") + if state_snapshot_digest(source) != source_digest: + raise ValueError("source_digest does not match the canonical source snapshot.") + snapshot: dict[str, Any] = json.loads(_canonical_json(source)) + version = snapshot.get("schemaVersion") + if not isinstance(version, str) or re.fullmatch(r"1\.[0-9]+\.[0-9]+", version) is None: + raise ValueError("Explicit migration accepts only legacy version-1 state, never a v2 source.") + raw_data = snapshot.get("data") + if not isinstance(raw_data, dict): + raise ValueError("Legacy state data must be an object.") + positions = _legacy_positions(cast(dict[str, Any], raw_data)) + if positions and delivery_evidence is None: + raise ValueError(_JOURNAL_REQUIRED) + timestamp = now if now is not None else datetime.now(timezone.utc) + if not isinstance(timestamp, datetime) or timestamp.utcoffset() is None: + raise ValueError("now must be an offset-aware datetime.") + timestamp = timestamp.astimezone(timezone.utc) + try: + _ = timestamp + timedelta(seconds=delivery_window_seconds) + except OverflowError as exc: + raise ValueError("delivery_window_seconds exceeds the representable bounded grace period.") from exc + + state = DurableAgentState.from_dict(snapshot) + if "migration" in state.data.unknown_fields: + raise ValueError("Legacy state already contains reserved migration metadata; refusing to overwrite it.") + _validate_receipts(state.data.ingested_messages) + _preserve_session(state, source_session_id) + evidence_id: str | None = None + if delivery_evidence is not None: + evidence_id, journal = _journal_receipts(delivery_evidence, source_digest=source_digest, positions=positions) + _apply_journal(state, journal) + else: + for identity in _retained_custom_request_ids(state): + state.data.ingested_messages.setdefault(identity, None) + + # A mailbox is itself a recorded response. Do not replace it from a transcript + # or refresh its expiry, even if its matching completion receipt was absent. + for correlation_id, mailbox in state.data.response_mailbox.items(): + # An original mailbox establishes its completion time, unlike a legacy + # transcript's created_at. Backfill before marking new receipts as legacy. + state.data.completed_correlations.setdefault( + correlation_id, {DurableStateFields.COMPLETED_AT: mailbox[DurableStateFields.CREATED_AT]} + ) + state._backfill_completion_outcomes(require_known=False) # pyright: ignore[reportPrivateUsage] + for correlation_id in state.data.response_mailbox: + if correlation_id not in cast(dict[str, Any], raw_data).get(DurableStateFields.COMPLETED_CORRELATIONS, {}): + state.data.completed_correlations[correlation_id]["legacy"] = True + for entry in state.data.conversation_history: + if isinstance(entry, DurableAgentStateResponse) and entry.correlation_id is not None: + correlation_id = _nonblank(entry.correlation_id, "Legacy response correlation ID") + if correlation_id not in state.data.completed_correlations: + state.record_response( + correlation_id, + entry.to_run_response(entry), + delivery_window_seconds=delivery_window_seconds, + now=timestamp, + legacy=True, + ) + + state._backfill_completion_outcomes(require_known=require_known_outcomes) # pyright: ignore[reportPrivateUsage] + state.schema_version = DurableAgentState.SCHEMA_VERSION + state.data.unknown_fields["migration"] = { + "id": migration_id, + "sourceDigest": source_digest, + "sourceSessionId": source_session_id, + "ownershipTransferId": ownership_transfer_id, + "createdAt": timestamp.isoformat(), + **({"evidenceId": evidence_id} if evidence_id is not None else {}), + } + size = len(json.dumps(state.to_dict(), allow_nan=False)) + if max_state_bytes is not None and size > max_state_bytes: + raise StateCapacityError( + size_bytes=size, max_state_bytes=max_state_bytes, floor_bytes=size, target_bytes=max_state_bytes + ) + return state diff --git a/python/packages/durabletask/agent_framework_durabletask/_worker.py b/python/packages/durabletask/agent_framework_durabletask/_worker.py index 77fae09..a791891 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_worker.py +++ b/python/packages/durabletask/agent_framework_durabletask/_worker.py @@ -14,13 +14,37 @@ from agent_framework import SupportsAgentRun, Workflow from agent_framework._telemetry import mark_feature_used +from durabletask.azuremanaged.worker import DurableTaskSchedulerWorker from durabletask.task import ActivityContext, OrchestrationContext from durabletask.worker import TaskHubGrpcWorker from ._async_bridge import run_agent_coroutine from ._callbacks import AgentResponseCallbackProtocol +from ._configuration import ( + INHERIT, + AgentRegistrationSettings, + RegistrationIdentity, + StateBudgetOverride, + resolve_state_budget_override, + validate_agent_configuration, + validate_response_delivery_window, + validate_runtime_deployment, +) from ._entities import AgentEntity, DurableTaskEntityStateProvider from ._feature_usage import FeatureIndex +from ._response_utils import serialize_agent_response +from ._retention import ( + DEFAULT_MAX_STATE_BYTES, + DEFAULT_RETENTION, + DELIVERY_WINDOW_SECONDS, + DTS_MAX_STATE_BYTES, + HIGH_WATERMARK, + LOW_WATERMARK, + RetentionMode, + StateBudget, + resolve_state_budget, + validate_retention, +) from ._workflows.activity import execute_workflow_activity from ._workflows.dt_context import DurableTaskWorkflowContext from ._workflows.naming import ( @@ -31,6 +55,7 @@ workflow_scoped_executor_id, ) from ._workflows.orchestrator import run_workflow_orchestrator +from ._workflows.protocol import unwrap_workflow_input from ._workflows.registration import collect_hosted_workflows, plan_workflow_registration logger = logging.getLogger("agent_framework.durabletask") @@ -53,6 +78,11 @@ class DurableAIAgentWorker: surfaces are split into :class:`DurableAIAgentClient` and ``DurableWorkflowClient``, because a caller invokes one or the other.) + Set ``deployment_mode="isolated_v2"`` or ``DURABLE_AGENTS_DEPLOYMENT_MODE=isolated_v2`` + to acknowledge an isolated schema 2 task hub/deployment with upgraded clients. + Old workflow histories must remain on the old engine. This acknowledgement is + not runtime proof of isolation and cannot detect peer workers. + Example: ```python from durabletask.worker import TaskHubGrpcWorker @@ -63,8 +93,8 @@ class DurableAIAgentWorker: # Create the underlying worker worker = TaskHubGrpcWorker(host_address="localhost:4001") - # Wrap it with the agent worker - agent_worker = DurableAIAgentWorker(worker) + # Acknowledge that this is an isolated schema 2 deployment + agent_worker = DurableAIAgentWorker(worker, deployment_mode="isolated_v2") # Register agents (or call configure_workflow(workflow) to host a workflow) client = OpenAIChatCompletionClient() @@ -80,15 +110,44 @@ def __init__( self, worker: TaskHubGrpcWorker, callback: AgentResponseCallbackProtocol | None = None, + *, + deployment_mode: str | None = None, + retention: RetentionMode = DEFAULT_RETENTION, + max_state_bytes: StateBudget = DEFAULT_MAX_STATE_BYTES, + high_watermark: float = HIGH_WATERMARK, + low_watermark: float = LOW_WATERMARK, + response_delivery_window_seconds: int = DELIVERY_WINDOW_SECONDS, ): """Initialize the worker wrapper. Args: worker: The durabletask worker instance to wrap callback: Optional callback for agent response notifications + deployment_mode: Exactly ``isolated_v2`` to acknowledge an isolated schema 2 + deployment with upgraded clients. None reads ``DURABLE_AGENTS_DEPLOYMENT_MODE``. + Old workflow histories stay on the old engine. This is not runtime proof of isolation. + retention: Eager pruning policy. ``keep_all`` does not prune compaction exclusions; + ``follow_compaction`` does. Pressure eviction is controlled separately by the budget. + max_state_bytes: Optional serialized-state budget. None disables pressure eviction; + ``backend_limit`` requires a DurableTaskSchedulerWorker. An explicit positive + integer works with any backend. + high_watermark: Budget fraction at which pressure eviction starts. + low_watermark: Target budget fraction after pressure eviction. + response_delivery_window_seconds: Positive integer response delivery window in seconds. """ + validate_runtime_deployment(deployment_mode) + validate_retention(retention, high_watermark, low_watermark) + self._backend_limit = DTS_MAX_STATE_BYTES if isinstance(worker, DurableTaskSchedulerWorker) else None + resolved_max_state_bytes = resolve_state_budget(max_state_bytes, backend_limit=self._backend_limit) + validate_response_delivery_window(response_delivery_window_seconds) + self._worker = worker self._callback = callback + self._retention: RetentionMode = retention + self._max_state_bytes = resolved_max_state_bytes + self._high_watermark = high_watermark + self._low_watermark = low_watermark + self._response_delivery_window_seconds = response_delivery_window_seconds self._registered_agents: dict[str, SupportsAgentRun] = {} self._workflows: dict[str, Workflow] = {} # Every workflow whose orchestration has been registered (top-level plus nested @@ -96,6 +155,8 @@ def __init__( # sub-workflow shared across the tree is registered once while two different # workflows whose names collide (including case-only differences) are rejected. self._registered_orchestrations: dict[str, Workflow] = {} + self._registration_identities: dict[tuple[str, str], RegistrationIdentity] = {} + self._registration_failed = False logger.debug("[DurableAIAgentWorker] Initialized with worker type: %s", type(worker).__name__) def add_agent( @@ -104,6 +165,11 @@ def add_agent( callback: AgentResponseCallbackProtocol | None = None, *, entity_id: str | None = None, + retention: RetentionMode | None = None, + max_state_bytes: StateBudgetOverride = INHERIT, + high_watermark: float | None = None, + low_watermark: float | None = None, + response_delivery_window_seconds: int | None = None, ) -> None: """Register an agent with the worker. @@ -117,33 +183,73 @@ def add_agent( entity_id: Optional identity to register the entity under instead of ``agent.name``. Workflow hosting passes the executor's ``id`` so the entity matches the identity the orchestrator dispatches to. + retention: Per-agent retention override. When None, the worker-level setting is used. + max_state_bytes: Per-agent budget. INHERIT uses the worker default; None disables it. + high_watermark: Pressure trigger override, or None to inherit the worker default. + low_watermark: Pressure target override, or None to inherit the worker default. + response_delivery_window_seconds: Delivery window override, or None to inherit. Raises: - ValueError: If the agent doesn't have a name or is already registered + ValueError: If the name, retention settings, or history-provider composition is invalid, + or the agent is already registered. """ + self._ensure_registration_usable() registration_name = entity_id or agent.name - if not registration_name: + if not isinstance(registration_name, str) or not registration_name: raise ValueError("Agent must have a name to be registered") if registration_name in self._registered_agents: raise ValueError(f"Agent '{registration_name}' is already registered") + effective_retention = self._retention if retention is None else retention + effective_budget = resolve_state_budget_override( + max_state_bytes, self._max_state_bytes, backend_limit=self._backend_limit + ) + effective_high = self._high_watermark if high_watermark is None else high_watermark + effective_low = self._low_watermark if low_watermark is None else low_watermark + effective_window = ( + self._response_delivery_window_seconds + if response_delivery_window_seconds is None + else response_delivery_window_seconds + ) + validate_retention(effective_retention, effective_high, effective_low) + validate_response_delivery_window(effective_window) + validate_agent_configuration(agent, retention=effective_retention) + effective_callback = self._callback if callback is None else callback + settings = AgentRegistrationSettings( + effective_retention, effective_budget, effective_high, effective_low, effective_window, effective_callback + ) + identities = dict(self._registration_identities) + RegistrationIdentity(agent, agent, "entity", settings, f"agent '{registration_name}'").reserve( + identities, f"dafx-{registration_name}", namespace="entity-name" + ) + logger.info( "[DurableAIAgentWorker] Registering agent: %s as entity: dafx-%s", registration_name, registration_name ) - # Store the agent reference - self._registered_agents[registration_name] = agent - - # Use agent-specific callback if provided, otherwise use worker-level callback - effective_callback = callback or self._callback - # Create a configured entity class using the factory - entity_class = self.__create_agent_entity(agent, effective_callback, entity_id=registration_name) + entity_class = self.__create_agent_entity( + agent, + effective_callback, + entity_id=registration_name, + retention=effective_retention, + max_state_bytes=effective_budget, + high_watermark=effective_high, + low_watermark=effective_low, + response_delivery_window_seconds=effective_window, + ) # Register the entity class with the worker # The worker.add_entity method takes a class - entity_registered: str = self._worker.add_entity(entity_class) + try: + entity_registered: str = self._worker.add_entity(entity_class) + except Exception: + # A backend can fail after mutating its registry, with no public rollback API. + self._registration_failed = True + raise + self._registered_agents[registration_name] = agent + self._registration_identities = identities logger.debug( "[DurableAIAgentWorker] Successfully registered entity class %s for agent: %s", @@ -151,6 +257,13 @@ def add_agent( registration_name, ) + def _ensure_registration_usable(self) -> None: + if self._registration_failed: + raise RuntimeError( + "Backend registration failed; this host may be partially registered. " + "Create a new host with a new underlying worker before registering or starting." + ) + def start(self) -> None: """Start the worker to begin processing tasks. @@ -158,6 +271,7 @@ def start(self) -> None: This method delegates to the underlying worker's start method. The worker will block until stopped. """ + self._ensure_registration_usable() logger.info("[DurableAIAgentWorker] Starting worker with %d registered agents", len(self._registered_agents)) mark_feature_used(FeatureIndex.DURABLETASK) self._worker.start() @@ -198,6 +312,12 @@ def configure_workflow( self, workflow: Workflow, callback: AgentResponseCallbackProtocol | None = None, + *, + retention: RetentionMode | None = None, + max_state_bytes: StateBudgetOverride = INHERIT, + high_watermark: float | None = None, + low_watermark: float | None = None, + response_delivery_window_seconds: int | None = None, ) -> None: """Register a :class:`Workflow` for automatic orchestration. @@ -208,8 +328,8 @@ def configure_workflow( Multiple workflows can be hosted on one worker: call this method once per workflow. Each workflow is keyed by its :attr:`Workflow.name`, and its durable primitives are scoped by that name (orchestration - ``dafx-{name}``; activities/entities ``dafx-{name}-{executorId}``) so two - co-hosted workflows that reuse an executor id do not collide. + ``dafx-{name}``; activities/entities ``dafx-{name}-{executorId}``). Ambiguous + derived names are rejected rather than renamed, preserving deployment compatibility. Sub-workflows nest: if the workflow contains :class:`~agent_framework.WorkflowExecutor` nodes, each inner workflow's @@ -223,58 +343,110 @@ def configure_workflow( across restarts and would break durable resume). Every nested sub-workflow must likewise be named. callback: Optional callback for agent response notifications. + retention: Retention for this workflow's agent nodes. When None, the worker-level + setting is used. Worth setting separately, since a workflow node's entity lives + for one orchestration while a standalone agent's can live indefinitely. + max_state_bytes: Budget for newly registered agent nodes in this workflow and its nested + workflows. INHERIT uses the worker default; None disables pressure eviction. + high_watermark: Pressure trigger override, or None to inherit the worker default. + low_watermark: Pressure target override, or None to inherit the worker default. + response_delivery_window_seconds: Delivery window override, or None to inherit. Raises: ValueError: If the workflow (or a nested sub-workflow) name is missing, - invalid, or auto-generated, or if the top-level workflow name is - already registered on this worker. + invalid, or auto-generated, a derived name has a different owner, + a shared workflow has different settings, or history preparation fails. """ + self._ensure_registration_usable() workflow_name = workflow.name validate_workflow_name(workflow_name) - if any(name.casefold() == workflow_name.casefold() for name in self._workflows): - raise ValueError( - f"Workflow '{workflow_name}' is already registered on this worker " - "(workflow names are compared case-insensitively)." - ) - # Validate the whole composition (top-level plus every nested sub-workflow) - # up front, so an invalid/auto-generated nested name (or an executor id that - # would break durable naming / nested-HITL addressing) fails before any - # registration side effects leave the worker partially configured. + effective_retention = self._retention if retention is None else retention + effective_budget = resolve_state_budget_override( + max_state_bytes, self._max_state_bytes, backend_limit=self._backend_limit + ) + effective_high = self._high_watermark if high_watermark is None else high_watermark + effective_low = self._low_watermark if low_watermark is None else low_watermark + effective_window = ( + self._response_delivery_window_seconds + if response_delivery_window_seconds is None + else response_delivery_window_seconds + ) + validate_retention(effective_retention, effective_high, effective_low) + validate_response_delivery_window(effective_window) + settings = AgentRegistrationSettings( + effective_retention, + effective_budget, + effective_high, + effective_low, + effective_window, + self._callback if callback is None else callback, + ) + + # Reserve the actual derived identities for the entire composition before any SDK calls. hosted_workflows = list(collect_hosted_workflows(workflow)) + identities = dict(self._registration_identities) for hosted in hosted_workflows: validate_workflow_name(hosted.name) for executor_id in hosted.executors: validate_executor_id(executor_id) - - # Check every cross-call collision *before* mutating any state, so a clash - # between a nested sub-workflow and an already-registered orchestration cannot - # leave the worker partially configured (e.g. the top-level name added to - # ``_workflows`` while a later child fails). Registration below is then a pure - # commit step. - for hosted in hosted_workflows: - existing = self._registered_orchestrations.get(hosted.name.casefold()) - if existing is not None and existing is not hosted: - raise ValueError( - f"A different workflow named '{hosted.name}' collides with already-registered " - f"'{existing.name}' on this worker. A workflow name maps to a single durable " - f"orchestration ('dafx-{hosted.name}'), compared case-insensitively; rename one " - "of them." + label = f"workflow '{hosted.name}'" + RegistrationIdentity(hosted, hosted, "orchestration", settings, label).reserve( + identities, workflow_orchestrator_name(hosted.name), namespace="orchestrator-name" + ) + plan = plan_workflow_registration(hosted) + for agent_executor in plan.agent_executors: + validate_executor_id(agent_executor.id) + validate_agent_configuration(agent_executor.agent, retention=effective_retention) + RegistrationIdentity( + hosted, agent_executor.agent, "entity", settings, f"{label} executor '{agent_executor.id}'" + ).reserve( + identities, + f"dafx-{workflow_scoped_executor_id(hosted.name, agent_executor.id)}", + namespace="entity-name", + ) + for executor in plan.activity_executors: + validate_executor_id(executor.id) + RegistrationIdentity( + hosted, executor, "activity", settings, f"{label} executor '{executor.id}'" + ).reserve( + identities, workflow_executor_activity_name(hosted.name, executor.id), namespace="activity-name" ) + previous_agents = dict(self._registered_agents) + previous_identities = self._registration_identities + try: + for hosted in hosted_workflows: + if hosted.name.casefold() in self._registered_orchestrations: + continue + self._register_single_workflow( + hosted, + callback, + effective_retention, + max_state_bytes=effective_budget, + high_watermark=effective_high, + low_watermark=effective_low, + response_delivery_window_seconds=effective_window, + ) + except Exception: + self._registration_failed = True + self._registered_agents = previous_agents + self._registration_identities = previous_identities + raise + self._registration_identities = identities + self._registered_orchestrations.update({hosted.name.casefold(): hosted for hosted in hosted_workflows}) self._workflows[workflow_name] = workflow - # Commit: register the top-level workflow and every nested sub-workflow (deduped - # by name), so the parent can drive sub-workflows as durable child orchestrations. - for hosted in hosted_workflows: - if hosted.name.casefold() in self._registered_orchestrations: - continue - self._register_single_workflow(hosted, callback) - def _register_single_workflow( self, workflow: Workflow, callback: AgentResponseCallbackProtocol | None, + retention: RetentionMode | None = None, + *, + max_state_bytes: StateBudgetOverride = INHERIT, + high_watermark: float | None = None, + low_watermark: float | None = None, + response_delivery_window_seconds: int | None = None, ) -> None: """Register one workflow's durable primitives (no recursion into sub-workflows). @@ -283,18 +455,24 @@ def _register_single_workflow( via ``plan_workflow_registration``. """ validate_workflow_name(workflow.name) - self._registered_orchestrations[workflow.name.casefold()] = workflow plan = plan_workflow_registration(workflow) - # Register agent executors as durable entities, scoped by workflow name so - # two workflows that reuse an executor id register distinct entities. The + # Register agent executors under the names validated by composition preflight. The # entity is keyed by the scoped identity (the same identity the orchestrator # dispatches to); the entity *key* at run time is the orchestration instance # id, which keeps conversation state isolated per run. for agent_executor in plan.agent_executors: scoped_id = workflow_scoped_executor_id(workflow.name, agent_executor.id) - if scoped_id not in self._registered_agents: - self.add_agent(agent_executor.agent, callback=callback, entity_id=scoped_id) + self.add_agent( + agent_executor.agent, + callback=callback, + entity_id=scoped_id, + retention=retention, + max_state_bytes=max_state_bytes, + high_watermark=high_watermark, + low_watermark=low_watermark, + response_delivery_window_seconds=response_delivery_window_seconds, + ) # Register non-agent executors as durable activities, scoped by workflow name. # WorkflowExecutor nodes are intentionally not registered as activities: their @@ -338,9 +516,8 @@ def _register_workflow_orchestrator(self, workflow: Workflow) -> None: orchestrator_name = workflow_orchestrator_name(workflow.name) def workflow_orchestrator(context: OrchestrationContext, input_data: Any) -> Any: - # Pass the deserialized client input straight to the shared engine, which - # reconstructs the start executor's declared type (see _coerce_initial_input). - initial_message = input_data + # Never replay the changed engine against a legacy recorded start. + initial_message = unwrap_workflow_input(input_data) shared_state: dict[str, Any] = {} dt_ctx = DurableTaskWorkflowContext(context) @@ -359,6 +536,11 @@ def __create_agent_entity( callback: AgentResponseCallbackProtocol | None = None, *, entity_id: str | None = None, + retention: RetentionMode = DEFAULT_RETENTION, + max_state_bytes: int | None = None, + high_watermark: float = HIGH_WATERMARK, + low_watermark: float = LOW_WATERMARK, + response_delivery_window_seconds: int = DELIVERY_WINDOW_SECONDS, ) -> type[DurableTaskEntityStateProvider]: """Factory function to create a DurableEntity class configured with an agent. @@ -371,6 +553,11 @@ def __create_agent_entity( entity_id: Optional identity to register the entity under instead of ``agent.name`` (used by workflow hosting to key entities by executor id). + retention: How much of the conversation durable state may discard. + max_state_bytes: Resolved pressure budget, or None to disable pressure eviction. + high_watermark: Budget fraction at which pressure eviction starts. + low_watermark: Target budget fraction after pressure eviction. + response_delivery_window_seconds: Response delivery window in seconds. Returns: A new DurableEntity subclass configured for this agent @@ -388,6 +575,11 @@ def __init__(self) -> None: agent=agent, callback=callback, state_provider=self, + retention=retention, + max_state_bytes=max_state_bytes, + high_watermark=high_watermark, + low_watermark=low_watermark, + response_delivery_window_seconds=response_delivery_window_seconds, ) logger.debug( "[ConfiguredAgentEntity] Initialized entity for agent: %s (entity name: %s)", @@ -409,13 +601,21 @@ def run(self, request: Any) -> Any: # shared agent clients/credentials stay bound to a live loop across # successive entity invocations (avoids cross-loop hangs). response = run_agent_coroutine(self._agent_entity.run(request)) - return response.to_dict() + return serialize_agent_response(response) def reset(self) -> None: - """Reset the agent's conversation history.""" + """Delegate reset to the configured AgentEntity.""" logger.debug("[ConfiguredAgentEntity.reset] Resetting agent: %s", agent_name) self._agent_entity.reset() + def expire_responses(self) -> int: + """Remove expired payloads when signaled by application-owned maintenance.""" + return self._agent_entity.expire_responses() + + def migrate(self, request: dict[str, Any]) -> dict[str, str]: + """Import an authorized legacy export into an empty destination entity.""" + return self._agent_entity.migrate(request) + # Set the entity name to match the prefixed agent name # This is used by durabletask to register the entity ConfiguredAgentEntity.__name__ = entity_name diff --git a/python/packages/durabletask/agent_framework_durabletask/_workflows/client.py b/python/packages/durabletask/agent_framework_durabletask/_workflows/client.py index b68265f..98965f1 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_workflows/client.py +++ b/python/packages/durabletask/agent_framework_durabletask/_workflows/client.py @@ -24,6 +24,7 @@ split_subworkflow_request_id, workflow_orchestrator_name, ) +from .protocol import wrap_workflow_input from .serialization import ( deserialize_workflow_event, deserialize_workflow_output, @@ -125,7 +126,7 @@ def start_workflow( # internal child dispatch (post trust boundary) may carry those reserved # keys, so stripping them here keeps untrusted input off the orchestrator's # trusted-deserialization path even if start_workflow is exposed remotely. - input=strip_subworkflow_markers(input), + input=wrap_workflow_input(strip_subworkflow_markers(input)), instance_id=instance_id, ) logger.debug("[DurableWorkflowClient] Started workflow instance: %s", new_instance_id) diff --git a/python/packages/durabletask/agent_framework_durabletask/_workflows/context.py b/python/packages/durabletask/agent_framework_durabletask/_workflows/context.py index d757d00..96ec2d8 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_workflows/context.py +++ b/python/packages/durabletask/agent_framework_durabletask/_workflows/context.py @@ -73,13 +73,23 @@ def current_utc_datetime(self) -> datetime: """The current replay-safe UTC datetime.""" ... - def prepare_agent_task(self, executor_id: str, message: str, orchestration_instance_id: str) -> Any: + def prepare_agent_task( + self, + executor_id: str, + message: str, + orchestration_instance_id: str, + context_messages: list[dict[str, Any]] | None = None, + context_message_ids: list[str] | None = None, + ) -> Any: """Create a yieldable task that runs an agent executor. Args: executor_id: Agent name / executor ID. message: The text message to send to the agent. orchestration_instance_id: Instance ID used as the entity session key. + context_messages: Optional upstream conversation (serialized ``Message`` dicts) + delivered to the agent as prior context. + context_message_ids: Occurrence identities, without changing application-visible IDs. Returns: A yieldable task whose result is an ``AgentResponse``. diff --git a/python/packages/durabletask/agent_framework_durabletask/_workflows/dt_context.py b/python/packages/durabletask/agent_framework_durabletask/_workflows/dt_context.py index 7388a0a..35ba696 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_workflows/dt_context.py +++ b/python/packages/durabletask/agent_framework_durabletask/_workflows/dt_context.py @@ -20,8 +20,7 @@ ) from .._executors import OrchestrationAgentExecutor -from .._models import AgentSessionId, DurableAgentSession -from .._shim import DurableAIAgent +from .._shim import build_agent_task from .context import WorkflowOrchestrationContext logger = logging.getLogger(__name__) @@ -57,11 +56,22 @@ def current_utc_datetime(self) -> datetime: # -- Agent / Activity dispatch -------------------------------------------- - def prepare_agent_task(self, executor_id: str, message: str, orchestration_instance_id: str) -> Any: - session_id = AgentSessionId(name=executor_id, key=orchestration_instance_id) - session = DurableAgentSession(durable_session_id=session_id) - agent = DurableAIAgent(self._executor, executor_id) - return agent.run(message, session=session) + def prepare_agent_task( + self, + executor_id: str, + message: str, + orchestration_instance_id: str, + context_messages: list[dict[str, Any]] | None = None, + context_message_ids: list[str] | None = None, + ) -> Any: + return build_agent_task( + self._executor, + executor_id, + message, + orchestration_instance_id, + context_messages, + context_message_ids, + ) def prepare_activity_task(self, activity_name: str, input_json: str) -> Any: return cast(Any, self._context.call_activity(activity_name, input=input_json)) diff --git a/python/packages/durabletask/agent_framework_durabletask/_workflows/naming.py b/python/packages/durabletask/agent_framework_durabletask/_workflows/naming.py index b1b9072..7780857 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_workflows/naming.py +++ b/python/packages/durabletask/agent_framework_durabletask/_workflows/naming.py @@ -31,12 +31,15 @@ "DURABLE_NAME_PREFIX", "MAX_EXECUTOR_ID_LENGTH", "SUBWORKFLOW_REQUEST_SEPARATOR", + "WORKFLOW_INPUT_EXECUTOR_ID", "is_auto_generated_workflow_name", + "parse_workflow_message_id", "qualify_subworkflow_request_id", "split_subworkflow_request_id", "validate_executor_id", "validate_workflow_name", "workflow_executor_activity_name", + "workflow_message_id", "workflow_name_from_orchestrator", "workflow_orchestrator_name", "workflow_scoped_executor_id", @@ -47,6 +50,50 @@ # ``AgentSessionId.ENTITY_NAME_PREFIX``. DURABLE_NAME_PREFIX = "dafx-" +# Identifies the workflow's own input in the conversation chained between agent nodes. It has no +# producing executor, so it carries a reserved id in that position. +WORKFLOW_INPUT_EXECUTOR_ID = "input" + +_WORKFLOW_MESSAGE_ID_PREFIX = "wf_" +_WORKFLOW_MESSAGE_ID_RE = re.compile(rf"^{_WORKFLOW_MESSAGE_ID_PREFIX}(?P.+)_(?P\d+)$") + + +def workflow_message_id(executor_id: str, position: int) -> str: + """Build the id for a message the workflow itself puts in the chained conversation. + + Core leaves ``message_id`` unset, so without this an agent node cannot tell context it has + already recorded from genuinely new input. The position is the message's index in the chained + conversation, which is fixed once the message joins it and is reproduced identically when the + orchestrator replays. + + Args: + executor_id: The node that produced the message, or ``WORKFLOW_INPUT_EXECUTOR_ID``. + position: The message's index in the chained conversation. + + Returns: + An id unique within one workflow run. + """ + return f"{_WORKFLOW_MESSAGE_ID_PREFIX}{executor_id}_{position}" + + +def parse_workflow_message_id(message_id: str | None) -> tuple[str, int] | None: + """Recover the producing executor and conversation position from a message id. + + Args: + message_id: The id to parse, if the message has one. + + Returns: + The executor id and position, or None when the id was not produced by + :func:`workflow_message_id`. + """ + if not message_id: + return None + match = _WORKFLOW_MESSAGE_ID_RE.match(message_id) + if match is None: + return None + return match.group("executor"), int(match.group("position")) + + # Separator used to qualify a nested sub-workflow's pending HITL request when it is # bubbled up to the top-level instance (one top-level addressing surface). A qualified id # is a path of ``{executorId}~{ordinal}`` hops ending in the leaf's bare request id, diff --git a/python/packages/durabletask/agent_framework_durabletask/_workflows/orchestrator.py b/python/packages/durabletask/agent_framework_durabletask/_workflows/orchestrator.py index 3116ab9..054fc85 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_workflows/orchestrator.py +++ b/python/packages/durabletask/agent_framework_durabletask/_workflows/orchestrator.py @@ -19,12 +19,14 @@ from __future__ import annotations +import hashlib import inspect import json import logging -from collections import defaultdict -from collections.abc import Generator -from dataclasses import dataclass +from collections import Counter, defaultdict +from collections.abc import Generator, Mapping +from copy import copy +from dataclasses import dataclass, field, replace from enum import Enum from typing import Any, cast @@ -33,10 +35,12 @@ AgentExecutorRequest, AgentExecutorResponse, AgentResponse, + Content, Executor, Message, Workflow, WorkflowConvergenceException, + WorkflowEvent, WorkflowExecutor, ) from agent_framework._workflows._edge import ( @@ -47,15 +51,22 @@ SingleEdgeGroup, SwitchCaseEdgeGroup, ) +from agent_framework._workflows._message_utils import normalize_messages_input from agent_framework._workflows._state import State +from pydantic import BaseModel +from .._message_identity import message_identity +from .._response_utils import ensure_response_format, load_agent_response from .context import WorkflowOrchestrationContext from .naming import ( + WORKFLOW_INPUT_EXECUTOR_ID, qualify_subworkflow_request_id, workflow_executor_activity_name, + workflow_message_id, workflow_orchestrator_name, workflow_scoped_executor_id, ) +from .protocol import wrap_workflow_input from .runner_context import ( HOST_METADATA_INSTANCE_ID, HOST_METADATA_REQUEST_PATH_PREFIX, @@ -69,6 +80,8 @@ reconstruct_to_type, resolve_type, serialize_value, + serialize_workflow_agent_response, + serialize_workflow_event, strip_pickle_markers, ) @@ -83,6 +96,9 @@ SOURCE_ORCHESTRATOR = "__orchestrator__" SOURCE_HITL_RESPONSE = "__hitl_response__" +# Private checkpoint provenance on dispatch copies, never application message IDs. +_FORWARDING_PROVENANCE = "_durable_workflow_forwarding" + # A WorkflowExecutor node runs its inner workflow as a durable child orchestration. # The parent wraps the node's input in SUBWORKFLOW_INPUT_KEY (defined alongside the # trust-boundary sanitizer in serialization.py) so the child orchestrator can tell a @@ -121,6 +137,11 @@ class TaskMetadata: # parent records these in its custom status before awaiting the child so the read # side can reach nested pending HITL requests while the parent is suspended. child_instance_id: str | None = None + selected_context: list[Message] | None = None + selected_context_ids: list[str] | None = None + invocation_ordinal: int = 0 + response_format: type[BaseModel] | None = None + skip_dispatch: bool = False @dataclass @@ -131,6 +152,8 @@ class ExecutorResult: output_message: AgentExecutorResponse | None activity_result: dict[str, Any] | None task_type: TaskType + source_message: Any = None + child_instance_id: str | None = None @dataclass @@ -142,6 +165,259 @@ class PendingHITLRequest: request_data: Any request_type: str | None response_type: str | None + task_type: TaskType = TaskType.ACTIVITY + + +@dataclass +class _WorkflowDeliveryLedger: + """Logical conversations and occurrence receipts rebuilt by deterministic replay. + + Application IDs are opaque. Object addresses only look up live envelopes and + aliases in this episode; retained references prevent address reuse. Wire IDs + contain only deterministic structural addresses, never memory addresses. + """ + + instance_id: str = "" + sent: dict[str, set[tuple[str, str]]] = field(default_factory=lambda: dict[str, set[tuple[str, str]]]()) + handoffs: dict[str, int] = field(default_factory=lambda: dict[str, int]()) + completions: int = 0 + envelopes: dict[int, tuple[AgentExecutorResponse, list[str], list[str]]] = field( + default_factory=lambda: dict[int, tuple[AgentExecutorResponse, list[str], list[str]]]() + ) + aliases: dict[int, tuple[Message, set[str]]] = field(default_factory=lambda: dict[int, tuple[Message, set[str]]]()) + cached: dict[str, tuple[list[Message], list[str]]] = field( + default_factory=lambda: dict[str, tuple[list[Message], list[str]]]() + ) + pending_agent_requests: dict[str, dict[str, Content]] = field( + default_factory=lambda: dict[str, dict[str, Content]]() + ) + pending_agent_responses: dict[str, list[Content]] = field(default_factory=lambda: dict[str, list[Content]]()) + + def occurrence(self, *address: Any) -> str: + """Identify an occurrence without changing its application message.""" + framed = json.dumps([self.instance_id, *address], ensure_ascii=False) + return "wf:occurrence:" + hashlib.sha256(framed.encode("utf-8")).hexdigest() + + def fork(self) -> _WorkflowDeliveryLedger: + """Stage new associations until projection and task preparation succeed.""" + return replace( + self, + sent=dict(self.sent), + handoffs=dict(self.handoffs), + envelopes=dict(self.envelopes), + aliases=dict(self.aliases), + cached=dict(self.cached), + pending_agent_requests={key: dict(value) for key, value in self.pending_agent_requests.items()}, + pending_agent_responses={key: list(value) for key, value in self.pending_agent_responses.items()}, + ) + + def remember(self, response: AgentExecutorResponse, ids: list[str], latest_ids: list[str]) -> None: + """Associate a logical envelope with parallel full/latest occurrence lists.""" + self.envelopes[id(response)] = (response, ids, latest_ids) + for message, occurrence in zip(response.full_conversation, ids, strict=True): + previous = self.aliases.get(id(message)) + # Two distinct witnesses are enough to make this alias ambiguous. + # Do not retain every later occurrence of a reused application object. + if previous is None: + self.aliases[id(message)] = (message, {occurrence}) + elif len(previous[1]) < 2 and occurrence not in previous[1]: + self.aliases[id(message)] = (message, previous[1] | {occurrence}) + + def identify( + self, response: AgentExecutorResponse, source: Any = None, *, scope: str | None = None + ) -> tuple[list[str], list[str]]: + """Register a new producer envelope once, before fan-out or projection. + + New response outputs are new events, even with equal application IDs or + contents. Forwarded history may reuse aliases or positions from the + explicitly associated activity input. Child outputs use their child scope. + """ + known = self.envelopes.get(id(response)) + if known is not None: + return known[1], known[2] + ordinal = self.completions + self.completions += 1 + latest = list(response.agent_response.messages) if response.agent_response else [] + latest_ids = [self.occurrence(scope, response.executor_id, ordinal, "output", i) for i in range(len(latest))] + # Locate each output once, preferring the appended turn when an object is + # also present earlier in the history. Equal text is not an output marker. + output_positions = _match_occurrences( + list(reversed(latest)), + list(reversed(response.full_conversation)), + [str(i) for i in reversed(range(len(response.full_conversation)))], + ) + # Core appends the latest turn. Its live suffix is positional evidence even + # when an application reuses the same Message object in earlier positions. + if ( + latest + and len(latest) <= len(response.full_conversation) + and all(a is b for a, b in zip(latest, response.full_conversation[-len(latest) :], strict=True)) + ): + output_positions = [ + str(i) + for i in reversed(range(len(response.full_conversation) - len(latest), len(response.full_conversation))) + ] + source_messages: list[Message] = [] + source_ids: list[str] = [] + provenance = cast( + tuple[str, list[Message]] | None, getattr(response.agent_response, _FORWARDING_PROVENANCE, None) + ) + for prior in _upstream_responses(source) or []: + prior_ids, prior_latest_ids = self.identify(prior) + source_messages.extend(prior.full_conversation) + source_ids.extend(prior_ids) + prior_latest = list(prior.agent_response.messages) if prior.agent_response else [] + # Equality is not producer identity. Only a dispatch witness that + # survives the activity checkpoint round trip can identify forwarding. + if ( + isinstance(provenance, tuple) + and len(provenance) == 2 + and provenance[0] == self.forwarding_key(prior, prior_ids, prior_latest_ids) + and response.executor_id == prior.executor_id + and isinstance(provenance[1], list) + and len(latest) == len(prior_latest) == len(provenance[1]) + and all(a is b for a, b in zip(latest, provenance[1], strict=True)) + and _same_message_values(latest, prior_latest) + ): + latest_ids = list(prior_latest_ids) + if _same_message_values(response.full_conversation, prior.full_conversation): + full_matches = _match_occurrences(response.full_conversation, prior.full_conversation, prior_ids) + if all(occurrence is not None for occurrence in full_matches): + full_ids = cast(list[str], full_matches) + self.remember(response, full_ids, latest_ids) + return full_ids, latest_ids + output_matches = { + int(position): occurrence + for position, occurrence in zip(output_positions, reversed(latest_ids), strict=True) + if position is not None + } + history_positions = [i for i in range(len(response.full_conversation)) if i not in output_matches] + history_matches = _match_occurrences( + [response.full_conversation[i] for i in history_positions], source_messages, source_ids + ) + forwarded = dict(zip(history_positions, history_matches, strict=True)) + alias_counts = Counter(id(message) for message in response.full_conversation) + ids: list[str] = [] + for index, message in enumerate(response.full_conversation): + occurrence = output_matches.get(index) + if occurrence is None: + occurrence = forwarded.get(index) + alias = self.aliases.get(id(message)) + if ( + occurrence is None + and scope is None + and alias_counts[id(message)] == 1 + and alias is not None + and len(alias[1]) == 1 + ): + occurrence = next(iter(alias[1])) + ids.append(occurrence or self.occurrence(scope, response.executor_id, ordinal, "context", index)) + self.remember(response, ids, latest_ids) + return ids, latest_ids + + def forwarding_key(self, prior: AgentExecutorResponse, ids: list[str], latest_ids: list[str]) -> str: + """Retain an inherited witness while forwarding through a child workflow.""" + provenance = cast(tuple[str, list[Message]] | None, getattr(prior.agent_response, _FORWARDING_PROVENANCE, None)) + latest = list(prior.agent_response.messages) if prior.agent_response else [] + if ( + isinstance(provenance, tuple) + and len(provenance) == 2 + and isinstance(provenance[0], str) + and isinstance(provenance[1], list) + and len(latest) == len(provenance[1]) + and all(a is b for a, b in zip(latest, provenance[1], strict=True)) + ): + return provenance[0] + return self.occurrence("forward", ids, latest_ids) + + def forwarding_input(self, message: Any) -> Any: + """Copy upstream envelopes with replay-stable, checkpoint-only witnesses.""" + upstream = _upstream_responses(message) + if upstream is None: + return message + forwarded: list[AgentExecutorResponse] = [] + for prior in upstream: + ids, latest_ids = self.identify(prior) + response = copy(prior.agent_response) + # Pickle preserves these references alongside response.messages. + # A new AgentResponse or replacement message has no such witness. + setattr( + response, + _FORWARDING_PROVENANCE, + (self.forwarding_key(prior, ids, latest_ids), list(response.messages)), + ) + forwarded.append(replace(prior, agent_response=response)) + return forwarded[0] if isinstance(message, AgentExecutorResponse) else forwarded + + +def _same_message_values(left: list[Message], right: list[Message]) -> bool: + """Compare JSON message values without requiring excluded data to serialize.""" + try: + return len(left) == len(right) and all( + message_identity(a) == message_identity(b) for a, b in zip(left, right, strict=True) + ) + except (TypeError, ValueError): + return False + + +def _match_occurrences( + selected: list[Message], originals: list[Message], ids: list[str], *, allow_positional: bool = True +) -> list[str | None]: + """Match aliases and unambiguous detached copies within one source list. + + A unique application ID also identifies a redacted version of that occurrence. + Ambiguous detached selections are new handoff occurrences, not global ID guesses. + Only detached matching needs fingerprints; excluded non-JSON data stays local. + """ + aliases: dict[int, list[int]] = defaultdict(list) + application_ids: dict[str, list[int]] = defaultdict(list) + for index, original in enumerate(originals): + aliases[id(original)].append(index) + if original.message_id is not None: + application_ids[original.message_id].append(index) + # Whole-list positions are evidence for both aliases and detached copies, but + # a reused alias in the wrong position must not masquerade as an equal copy. + if allow_positional and selected and len(selected) == len(originals): + copied_aliases: dict[int, int] = {} + try: + if all( + (a is b or (id(a) not in aliases and message_identity(a) == message_identity(b))) + and copied_aliases.setdefault(id(a), id(b)) == id(b) + for a, b in zip(selected, originals, strict=True) + ): + return list(ids) + except (TypeError, ValueError): + pass + fingerprints: dict[str, list[int]] | None = None + used: set[int] = set() + matches: list[str | None] = [] + for message in selected: + candidates = aliases.get(id(message), []) + if not candidates and message.message_id is not None: + candidates = application_ids.get(message.message_id, []) + if len(candidates) != 1: + candidates = [] + if not candidates and originals: + try: + fingerprint = message_identity(message) + except (TypeError, ValueError): + fingerprint = None + if fingerprint is not None: + if fingerprints is None: + fingerprints = defaultdict(list) + for i, original in enumerate(originals): + try: + fingerprints[message_identity(original)].append(i) + except (TypeError, ValueError): + continue + candidates = fingerprints.get(fingerprint, []) + if len(candidates) != 1: + candidates = [] + position = candidates[0] if len(candidates) == 1 and candidates[0] not in used else None + matches.append(ids[position] if position is not None else None) + if position is not None: + used.add(position) + return matches # ============================================================================ @@ -218,8 +494,14 @@ def build_agent_executor_response( response_text: str | None, structured_response: dict[str, Any] | None, previous_message: Any, + *, + position: int | None = None, ) -> AgentExecutorResponse: - """Build an AgentExecutorResponse from entity response data.""" + """Build a legacy text response, leaving upstream application messages untouched. + + Production agent completions retain the actual AgentResponse instead. This + compatibility helper assigns IDs only to the messages it creates itself. + """ final_text: str = response_text or "" if structured_response: final_text = json.dumps(structured_response) @@ -228,10 +510,29 @@ def build_agent_executor_response( agent_response = AgentResponse(messages=[assistant_message]) full_conversation: list[Message] = [] - if isinstance(previous_message, AgentExecutorResponse) and previous_message.full_conversation: - full_conversation.extend(previous_message.full_conversation) + upstream = _upstream_responses(previous_message) + if upstream is not None: + for prior in upstream: + full_conversation.extend(prior.full_conversation) elif isinstance(previous_message, str): - full_conversation.append(Message(role="user", contents=[previous_message])) + full_conversation.append( + Message( + role="user", + contents=[previous_message], + message_id=workflow_message_id(WORKFLOW_INPUT_EXECUTOR_ID, 0), + ) + ) + else: + full_conversation.extend( + normalize_messages_input( + previous_message.messages if isinstance(previous_message, AgentExecutorRequest) else previous_message + ) + ) + # Keep the assigned identity when the conversation is forwarded. Conversation length + # alone is insufficient when a producer receives another short, independent input. + assistant_message.message_id = workflow_message_id( + executor_id, len(full_conversation) if position is None else position + ) full_conversation.append(assistant_message) return AgentExecutorResponse( @@ -246,11 +547,103 @@ def build_agent_executor_response( # ============================================================================ +def _upstream_responses(message: Any) -> list[AgentExecutorResponse] | None: + """Recognize a chained response or a fan-in batch of chained responses.""" + if isinstance(message, AgentExecutorResponse): + return [message] + if isinstance(message, list): + items = cast(list[Any], message) + if all(isinstance(item, AgentExecutorResponse) for item in items): + return cast(list[AgentExecutorResponse], items) + return None + + +def _select_context_messages(executor: AgentExecutor, message: AgentExecutorResponse) -> list[Message]: + """Apply core's projection before assigning any transport-only identities.""" + mode = getattr(executor, "_context_mode", "full") + if mode == "last_agent": + return list(message.agent_response.messages) if message.agent_response else [] + if mode == "custom": + context_filter = getattr(executor, "_context_filter", None) + if context_filter is None: + raise ValueError("context_filter must be provided for 'custom' context_mode.") + return list(context_filter(list(message.full_conversation))) + return list(message.full_conversation) + + +def _build_context_messages( # pyright: ignore[reportUnusedFunction] + executor: AgentExecutor, message: Any +) -> list[dict[str, Any]] | None: + """Project the upstream conversation into messages for a downstream agent. + + Mirrors the in-process :class:`AgentExecutor` context behavior so a workflow behaves the + same way durably: ``full`` forwards the whole upstream conversation, ``last_agent`` only the + previous agent's messages, and ``custom`` applies the executor's ``context_filter``. + + Returns ``None`` when there is no upstream response (for example the first node, + which receives raw input instead). An empty projection is ``[]``, never a fallback + to unfiltered input. Fan-in responses are projected in their aggregation order. + This helper is stateless: delta selection belongs to agent task preparation. + + The mode and filter are read off private attributes because core takes them as constructor + arguments and exposes no public accessor for either. Reading them is therefore the only way + to match in-process behavior. The coupling is deliberate rather than accidental, and it is + covered: the projection tests build a real ``AgentExecutor`` for each mode, so if core ever + renames these the fallback to ``full`` changes the projection and those tests fail. + """ + upstream = _upstream_responses(message) + if upstream is None: + return None + return [m.to_dict() for prior in upstream for m in _select_context_messages(executor, prior)] + + +def _identify_context_messages( + prior: AgentExecutorResponse, + selected: list[Message], + target: str, + handoff: int, + response_ordinal: int, + ledger: _WorkflowDeliveryLedger, + *, + latest_only: bool = False, +) -> list[str]: + """Associate selected copies with source occurrences, never rewrite their IDs.""" + ids, latest_ids = ledger.identify(prior) + if latest_only: + return list(latest_ids) + matches = _match_occurrences(selected, prior.full_conversation, ids) + latest = list(prior.agent_response.messages) if prior.agent_response else [] + unmatched = [index for index, occurrence in enumerate(matches) if occurrence is None] + if unmatched: + # Do not resolve an ambiguous full-history alias by searching only the + # latest turn. Keep every source candidate in the fallback's evidence. + combined_messages = list(prior.full_conversation) + combined_ids = list(ids) + full_ids = set(ids) + for message, occurrence in zip(latest, latest_ids, strict=True): + if occurrence not in full_ids: + combined_messages.append(message) + combined_ids.append(occurrence) + latest_matches = _match_occurrences(selected, combined_messages, combined_ids, allow_positional=False) + for index in unmatched: + matches[index] = latest_matches[index] + return [ + occurrence or ledger.occurrence("selection", target, handoff, response_ordinal, index) + for index, occurrence in enumerate(matches) + ] + + +_AGENT_TASK_MESSAGE_PREVIEW_LIMIT = 1024 + + def _prepare_agent_task( ctx: WorkflowOrchestrationContext, + executor: AgentExecutor, executor_id: str, message: Any, workflow_name: str, + delivery_ledger: _WorkflowDeliveryLedger | None = None, + metadata: TaskMetadata | None = None, ) -> Any: """Prepare an agent task for execution via the context adapter. @@ -259,10 +652,101 @@ def _prepare_agent_task( executor id dispatch to distinct entities (the entity layer prefixes this with ``dafx-``). The session *key* stays the orchestration instance id, so conversation state remains isolated per run. + + Project first, then send only identities not yet dispatched to this target. The + caller shares a replay-local ledger across all dispatch paths, never on an executor + retained between workflow runs. A standalone helper call gets a fresh ledger. """ - message_content = _extract_message_content(message) - scoped_id = workflow_scoped_executor_id(workflow_name, executor_id) - return ctx.prepare_agent_task(scoped_id, message_content, ctx.instance_id) + if delivery_ledger is None: + delivery_ledger = _WorkflowDeliveryLedger(instance_id=ctx.instance_id) + staged = delivery_ledger.fork() + staged.instance_id = ctx.instance_id + if metadata is not None: + options = getattr(executor.agent, "default_options", None) + response_format = ( + cast(Mapping[str, Any], options).get("response_format") if isinstance(options, Mapping) else None + ) + if isinstance(response_format, type) and issubclass(response_format, BaseModel): + metadata.response_format = response_format + if metadata.source_executor_id.startswith(SOURCE_HITL_RESPONSE): + message = _prepare_agent_hitl_message(executor_id, message, staged) + if message is None: + metadata.skip_dispatch = True + delivery_ledger.__dict__.update(staged.__dict__) + return None + upstream = _upstream_responses(message) + handoff = staged.handoffs.get(executor_id, 0) + cached_messages, cached_ids = staged.cached.get(executor_id, ([], [])) + selected_context = list(cached_messages) + selected_ids = list(cached_ids) + if upstream is None: + inputs = normalize_messages_input(message.messages if isinstance(message, AgentExecutorRequest) else message) + selected_context.extend(inputs) + selected_ids.extend(staged.occurrence("input", executor_id, handoff, i) for i in range(len(inputs))) + else: + for response_ordinal, prior in enumerate(upstream): + selected = _select_context_messages(executor, prior) + selected_context.extend(selected) + selected_ids.extend( + _identify_context_messages( + prior, + selected, + executor_id, + handoff, + response_ordinal, + staged, + latest_only=getattr(executor, "_context_mode", "full") == "last_agent", + ) + ) + + # Cache-only input is replay-local control state, not an entity/model task. + cache_only = isinstance(message, AgentExecutorRequest) and not message.should_respond + context_messages: list[dict[str, Any]] | None = [] + context_message_ids: list[str] | None = [] + pending_keys: set[tuple[str, str]] = set() + message_content = "" + sent = staged.sent.get(executor_id, set()) + for selected, occurrence in zip(selected_context, selected_ids, strict=True): + key = (occurrence, message_identity(selected)) + if key in sent or key in pending_keys: + continue + context_messages.append(selected.to_dict()) + context_message_ids.append(occurrence) + pending_keys.add(key) + message_content = selected.text[:_AGENT_TASK_MESSAGE_PREVIEW_LIMIT] + + # Preserve the legacy nonempty-string adapter contract. Its occurrence still + # accompanies the logical outgoing conversation, never a shared wf_input_0. + if isinstance(message, str) and message and not cached_messages: + context_messages = None + context_message_ids = None + message_content = message + pending_keys.clear() + + task = None + if not cache_only: + scoped_id = workflow_scoped_executor_id(workflow_name, executor_id) + if context_message_ids is None: + task = ctx.prepare_agent_task(scoped_id, message_content, ctx.instance_id, context_messages) + else: + task = ctx.prepare_agent_task( + scoped_id, message_content, ctx.instance_id, context_messages, context_message_ids=context_message_ids + ) + # Preparation/serialization can fail before a task is scheduled. Do not record + # those messages or consume a synthetic identity until the adapter accepts it. + if cache_only: + staged.cached[executor_id] = (selected_context, selected_ids) + else: + staged.cached.pop(executor_id, None) + if pending_keys: + staged.sent[executor_id] = sent | pending_keys + staged.handoffs[executor_id] = handoff + 1 + delivery_ledger.__dict__.update(staged.__dict__) + if metadata is not None: + metadata.selected_context = selected_context + metadata.selected_context_ids = selected_ids + metadata.invocation_ordinal = handoff + return task def _prepare_activity_task( @@ -273,6 +757,7 @@ def _prepare_activity_task( shared_state_snapshot: dict[str, Any] | None, workflow_name: str, address: dict[str, str], + delivery_ledger: _WorkflowDeliveryLedger | None = None, ) -> Any: """Prepare an activity task for execution via the context adapter. @@ -280,9 +765,10 @@ def _prepare_activity_task( ``dafx-{workflow_name}-{executor_id}`` so two co-hosted workflows that reuse an executor id register and dispatch to distinct activity functions. """ + staged = delivery_ledger.fork() if delivery_ledger is not None else None activity_input = { "executor_id": executor_id, - "message": serialize_value(message), + "message": serialize_value(staged.forwarding_input(message) if staged else message), "shared_state_snapshot": shared_state_snapshot, "source_executor_ids": [source_executor_id], # host_context addresses the *root* (HTTP-routable) orchestration so an executor @@ -299,7 +785,10 @@ def _prepare_activity_task( } activity_input_json = json.dumps(activity_input) activity_name = workflow_executor_activity_name(workflow_name, executor_id) - return ctx.prepare_activity_task(activity_name, activity_input_json) + task = ctx.prepare_activity_task(activity_name, activity_input_json) + if delivery_ledger is not None and staged is not None: + delivery_ledger.__dict__.update(staged.__dict__) + return task def _prepare_subworkflow_task( @@ -308,6 +797,7 @@ def _prepare_subworkflow_task( message: Any, child_instance_id: str, child_address: dict[str, str], + delivery_ledger: _WorkflowDeliveryLedger | None = None, ) -> Any: """Prepare a child-orchestration task that runs a ``WorkflowExecutor``'s inner workflow. @@ -319,12 +809,18 @@ def _prepare_subworkflow_task( executor inside the child can build a respond URL that targets the top-level instance with a qualified request id. """ + staged = delivery_ledger.fork() if delivery_ledger is not None else None inner_orchestration_name = workflow_orchestrator_name(executor.workflow.name) child_input = { - SUBWORKFLOW_INPUT_KEY: serialize_value(message), + SUBWORKFLOW_INPUT_KEY: serialize_value(staged.forwarding_input(message) if staged else message), SUBWORKFLOW_ADDRESS_KEY: child_address, } - return ctx.call_sub_orchestrator(inner_orchestration_name, child_input, instance_id=child_instance_id) + task = ctx.call_sub_orchestrator( + inner_orchestration_name, wrap_workflow_input(child_input), instance_id=child_instance_id + ) + if delivery_ledger is not None and staged is not None: + delivery_ledger.__dict__.update(staged.__dict__) + return task # ============================================================================ @@ -332,30 +828,111 @@ def _prepare_subworkflow_task( # ============================================================================ +def _raise_for_agent_failure(agent_response: AgentResponse | dict[str, Any], executor_id: str) -> None: + """Reject terminal durable results before reducing them to downstream text. + + Entities should mark runtime failures with response-level ``durable_status=error``. + Direct non-tool error content is the legacy fallback, only within AgentResponse + envelopes. Tool results (including nested errors) and application dicts are data. + Unmarked direct non-tool errors cannot distinguish application errors from legacy + entity failures, so that fallback treats them as terminal. + """ + if isinstance(agent_response, AgentResponse): + properties: dict[str, Any] = agent_response.additional_properties + error_codes = [ + content.error_code + for message in agent_response.messages + if message.role != "tool" + for content in message.contents + if content.type == "error" + ] + elif isinstance(agent_response, dict) and agent_response.get("type") == "agent_response": + properties = cast(dict[str, Any], agent_response.get("additional_properties") or {}) + messages = cast(list[dict[str, Any]], agent_response.get("messages") or []) + # Inspect the wire envelope directly, without deserializing unknown fields. + error_codes = [ + content.get("error_code") + for message in messages + if isinstance(message, dict) and message.get("role") != "tool" + for content in cast(list[dict[str, Any]], message.get("contents") or []) + if isinstance(content, dict) and content.get("type") == "error" + ] + else: + return + + status = properties.get("durable_status") + # Do not include response text, error details or the request in the exception. + if status == "already_completed" or "response_expired" in error_codes: + raise RuntimeError(f"Agent executor {executor_id!r} returned an expired durable response.") + if status == "error" or error_codes: + raise RuntimeError(f"Agent executor {executor_id!r} returned a terminal runtime error.") + + def _process_agent_response( - agent_response: AgentResponse, + agent_response: AgentResponse | dict[str, Any], executor_id: str, message: Any, + delivery_ledger: _WorkflowDeliveryLedger, + metadata: TaskMetadata | None = None, ) -> ExecutorResult: - """Process an agent response into an ExecutorResult.""" - response_text = agent_response.text if agent_response else None - structured_response: dict[str, Any] | None = None - - if agent_response and agent_response.value is not None: - model_dump = getattr(agent_response.value, "model_dump", None) + """Emit core's selected cache plus the unaltered agent response messages.""" + _raise_for_agent_failure(agent_response, executor_id) + if isinstance(agent_response, dict) and agent_response.get("type") == "agent_response": + agent_response = load_agent_response(agent_response) + if isinstance(agent_response, dict): + # Lightweight text/value payloads are data, not durable response envelopes. + value = agent_response.get("value") + model_dump = getattr(value, "model_dump", None) if callable(model_dump): - dumped = model_dump() - if isinstance(dumped, dict): - structured_response = dumped # type: ignore[assignment] - elif isinstance(agent_response.value, dict): - structured_response = agent_response.value + value = model_dump() + text = json.dumps(value) if isinstance(value, dict) else agent_response.get("text") or "" + agent_response = AgentResponse(messages=[Message("assistant", [text])]) + + # Core does not yield or send a partial response while approval is pending. + # These dictionaries belong to this replay, never the registered executor. + requests = agent_response.user_input_requests + if requests: + pending = dict(delivery_ledger.pending_agent_requests.get(executor_id, {})) + events: list[dict[str, Any]] = [] + for request in requests: + request_id = request.id + if not isinstance(request_id, str) or not request_id: + raise ValueError(f"Agent executor {executor_id!r} returned a user input request without an id.") + if request_id in pending: + raise ValueError(f"Agent executor {executor_id!r} returned a duplicate user input request id.") + pending[request_id] = request + event = serialize_workflow_event( + WorkflowEvent.request_info( + request_id=request_id, source_executor_id=executor_id, request_data=request, response_type=Content + ) + ) + event["request_type"] = f"{Content.__module__}:{Content.__name__}" + events.append(event) + delivery_ledger.pending_agent_requests[executor_id] = pending + return ExecutorResult( + executor_id=executor_id, + output_message=None, + activity_result={"pending_request_info_events": events, "events": events}, + task_type=TaskType.AGENT, + ) - output_message = build_agent_executor_response( - executor_id=executor_id, - response_text=response_text, - structured_response=structured_response, - previous_message=message, + if metadata is None or metadata.selected_context is None or metadata.selected_context_ids is None: + raise ValueError("Agent completion requires its prepared logical context.") + if metadata.response_format is not None: + # The entity wire carries values, not model classes. Only the locally + # registered agent's declared format can restore the structured value. + agent_response = copy(agent_response) + ensure_response_format(metadata.response_format, f"{executor_id}:{metadata.invocation_ordinal}", agent_response) + latest_ids = [ + delivery_ledger.occurrence("agent", executor_id, metadata.invocation_ordinal, index) + for index in range(len(agent_response.messages)) + ] + output_message = AgentExecutorResponse( + executor_id, + agent_response, + full_conversation=[*metadata.selected_context, *agent_response.messages], ) + delivery_ledger.remember(output_message, [*metadata.selected_context_ids, *latest_ids], latest_ids) return ExecutorResult( executor_id=executor_id, @@ -419,23 +996,35 @@ def _unpack_subworkflow_result(child_result: Any) -> tuple[list[Any], list[dict[ return [child_result], [] +def _classify_workflow_output(workflow: Workflow, executor_id: str) -> str | None: + """Use core's yield designation for both agent and direct child outputs.""" + # A truthy mock return is not an explicit designation. + if workflow.is_terminal_executor(executor_id) is True: + return "output" + if workflow.is_intermediate_executor(executor_id) is True: + return "intermediate" + return None + + def _process_subworkflow_result( child_result: Any, executor: WorkflowExecutor, workflow_outputs: list[Any], + workflow: Workflow | None = None, ) -> ExecutorResult: """Process a child orchestration's result into an ``ExecutorResult``. The child orchestration returns a result envelope (see :data:`SUBWORKFLOW_RESULT_KEY`) carrying the inner workflow's outputs (a list of - values already encoded by the inner activity via ``serialize_value``) plus its - accumulated event timeline. Mirroring the in-process + already encoded activity values or generated agent response envelopes) plus + its accumulated event timeline. Mirroring the in-process :class:`~agent_framework.WorkflowExecutor`: * ``allow_direct_output`` is ``False`` (default): each inner output becomes a message routed through the ``WorkflowExecutor`` node's outgoing edges. - * ``allow_direct_output`` is ``True``: each inner output becomes one of the - parent workflow's own outputs. + * ``allow_direct_output`` is ``True``: each inner output follows the parent + workflow's yield designation for this node (output, intermediate, or hidden). + Omitting ``workflow`` retains the helper's legacy direct-output behavior. The inner workflow's *intermediate* events are bubbled into the parent's event stream **re-tagged with this node's id** (``executor.id``), matching the @@ -449,10 +1038,15 @@ def _process_subworkflow_result( outputs, child_events = _unpack_subworkflow_result(child_result) sent_messages: list[dict[str, Any]] = [] + output_events: list[dict[str, Any]] = [] if executor.allow_direct_output: - # Inner outputs are already serialized (serialize_value); workflow_outputs - # holds serialized values, so they are directly compatible. - workflow_outputs.extend(outputs) + event_type = _classify_workflow_output(workflow, executor.id) if workflow is not None else "output" + # Inner outputs are already encoded. Reuse them without decoding/re-pickling + # a portable agent response or a checkpoint value. + if event_type == "output": + workflow_outputs.extend(outputs) + if workflow is not None and event_type is not None: + output_events = [{"type": event_type, "executor_id": executor.id, "data": output} for output in outputs] else: # Route each inner output as a message from the node; _route_result_messages # deserializes each "message" value before routing through edge groups. @@ -468,7 +1062,7 @@ def _process_subworkflow_result( return ExecutorResult( executor_id=executor.id, output_message=None, - activity_result={"sent_messages": sent_messages, "outputs": [], "events": bubbled_events}, + activity_result={"sent_messages": sent_messages, "outputs": [], "events": [*output_events, *bubbled_events]}, task_type=TaskType.SUBWORKFLOW, ) @@ -483,6 +1077,7 @@ def _route_result_messages( workflow: Workflow, next_pending_messages: dict[str, list[tuple[Any, str]]], fan_in_pending: dict[str, dict[str, list[tuple[Any, str]]]], + delivery_ledger: _WorkflowDeliveryLedger | None = None, ) -> None: """Route messages from an executor result to their targets.""" executor_id = result.executor_id @@ -503,6 +1098,9 @@ def _route_result_messages( for msg_to_route, explicit_target in messages_to_route: logger.debug("Routing output from %s", executor_id) + if delivery_ledger is not None: + for response in _upstream_responses(msg_to_route) or []: + delivery_ledger.identify(response, result.source_message, scope=result.child_instance_id) if explicit_target: if explicit_target not in next_pending_messages: @@ -568,17 +1166,21 @@ def _collect_hitl_requests( result: ExecutorResult, pending_hitl_requests: dict[str, PendingHITLRequest], ) -> None: - """Collect pending HITL requests from an activity result.""" + """Collect pending HITL requests from executor results without losing agent requests.""" if result.activity_result and result.activity_result.get("pending_request_info_events"): for req_data in result.activity_result["pending_request_info_events"]: request_id = req_data.get("request_id") if request_id: + existing = pending_hitl_requests.get(request_id) + if existing is not None and TaskType.AGENT in (existing.task_type, result.task_type): + raise ValueError("Agent user input request id collides with an outstanding workflow request.") pending_hitl_requests[request_id] = PendingHITLRequest( request_id=request_id, source_executor_id=req_data.get("source_executor_id", result.executor_id), request_data=req_data.get("data"), request_type=req_data.get("request_type"), response_type=req_data.get("response_type"), + task_type=result.task_type, ) logger.debug( "Collected HITL request %s from executor %s", @@ -614,29 +1216,6 @@ def _route_hitl_response( ) -# ============================================================================ -# Message Content Extraction -# ============================================================================ - - -def _extract_message_content(message: Any) -> str: - """Extract text content from various message types.""" - message_content = "" - if isinstance(message, AgentExecutorResponse) and message.agent_response: - if message.agent_response.text: - message_content = message.agent_response.text - elif message.agent_response.messages: - message_content = message.agent_response.messages[-1].text or "" - elif isinstance(message, AgentExecutorRequest) and message.messages: - message_content = message.messages[-1].text or "" - elif isinstance(message, dict): - key_names = list(message.keys()) # type: ignore[union-attr] - logger.warning("Unexpected dict message in _extract_message_content. Keys: %s", key_names) # type: ignore - elif isinstance(message, str): - message_content = message - return message_content - - def _select_primary_input_type(executor: Executor) -> type | None: """Return the executor's primary concrete declared input type, if any. @@ -717,7 +1296,8 @@ def _coerce_initial_input(workflow: Workflow, raw_value: Any) -> Any: reconstruction to mirror in-process delivery, where the start executor receives its declared type: - * Agent start executors only consume text, so non-text input is stringified. + * Agent start executors preserve core's typed inputs and message lists. Other + JSON payloads retain the legacy stringification fallback. * Other executors get their primary declared input type reconstructed (``dict`` -> Pydantic/dataclass, ``str`` -> ``str``, ...) via :func:`reconstruct_to_type`; union/unannotated types pass through unchanged. @@ -738,8 +1318,12 @@ def _coerce_initial_input(workflow: Workflow, raw_value: Any) -> Any: return raw_value if isinstance(start_executor, AgentExecutor): - if isinstance(raw_value, str): + if raw_value is None or isinstance(raw_value, (str, Message, AgentExecutorRequest, AgentExecutorResponse)): return raw_value + if isinstance(raw_value, list): + items = cast(list[Any], raw_value) + if all(isinstance(item, (str, Message)) for item in items): + return items if isinstance(raw_value, (dict, list)): return json.dumps(raw_value) return str(raw_value) @@ -758,6 +1342,49 @@ def _coerce_initial_input(workflow: Workflow, raw_value: Any) -> Any: # ============================================================================ +def _load_agent_hitl_content(request_id: str, original_request: Content, raw_response: Any) -> Content: + """Rebuild a reply using the fixed local Content type, never a supplied type name.""" + sanitized = strip_pickle_markers(raw_response) + response = Content.from_text(sanitized) if isinstance(sanitized, str) else reconstruct_to_type(sanitized, Content) + if not isinstance(response, Content): + raise TypeError("Agent user input responses must be Content objects or Content mappings.") + if response.type == "function_approval_response" and response.id != request_id: + raise ValueError("Agent approval response does not match the pending request id.") + if response.type == "function_result": + call = original_request.function_call + call_id = call.call_id if isinstance(call, Content) else original_request.call_id + if call_id is not None and response.call_id != call_id: + raise ValueError("Agent function result does not match the pending call id.") + return response + + +def _prepare_agent_hitl_message(executor_id: str, message: Any, ledger: _WorkflowDeliveryLedger) -> Message | None: + """Accumulate replies like core's response handler before scheduling one agent turn.""" + if not isinstance(message, dict): + raise TypeError("Agent HITL message must be a response envelope.") + envelope = cast("dict[str, Any]", message) + request_id = envelope.get("request_id") + pending = ledger.pending_agent_requests.get(executor_id, {}) + if not isinstance(request_id, str) or request_id not in pending: + # Duplicate or unknown responses must not resume the agent or erase replies. + logger.warning("Ignoring unknown or already-handled agent response for executor %s", executor_id) + return None + response = _load_agent_hitl_content(request_id, pending[request_id], envelope.get("response")) + responses = ledger.pending_agent_responses.setdefault(executor_id, []) + responses.append(response) + del pending[request_id] + if pending: + return None + role = "tool" if all(reply.type == "function_result" for reply in responses) else "user" + combined = Message(role=role, contents=list(responses)) + ledger.pending_agent_requests.pop(executor_id, None) + ledger.pending_agent_responses.pop(executor_id, None) + # Core replaces its cache on resumption. Durable service/session state stays + # in the same entity; only this new combined reply is dispatched as a delta. + ledger.cached.pop(executor_id, None) + return combined + + async def execute_hitl_response_handler( executor: Any, hitl_message: dict[str, Any], @@ -784,19 +1411,17 @@ async def execute_hitl_response_handler( handler = executor._find_response_handler(original_request, response) if handler is None: - logger.warning( - "No response handler found for HITL response in executor %s. Request type: %s, Response type: %s", - executor.id, - type(original_request).__name__, - type(response).__name__, + raise ValueError( + f"No response handler found for HITL response in executor {executor.id!r}. " + f"Request type: {type(original_request).__name__}, Response type: {type(response).__name__}" ) - return ctx = WorkflowContext( executor=executor, source_executor_ids=[SOURCE_HITL_RESPONSE], runner_context=runner_context, state=shared_state, + request_id=hitl_message.get("request_id"), ) logger.debug( @@ -850,6 +1475,7 @@ def _prepare_all_tasks( shared_state: dict[str, Any] | None, subworkflow_counter: list[int], address: dict[str, str], + delivery_ledger: _WorkflowDeliveryLedger | None = None, ) -> tuple[list[Any], list[TaskMetadata], list[tuple[str, Any, str]]]: """Prepare all pending tasks for parallel execution. @@ -872,7 +1498,11 @@ def _prepare_all_tasks( (``{root_instance_id, root_workflow_name, request_path_prefix}``). Surfaced to activity executors via ``host_context`` and extended by one ``{executor}~{ordinal}~`` hop for each dispatched sub-workflow child. + delivery_ledger: Replay-local agent delivery receipts shared with sequential + dispatch and later supersteps. Standalone calls default to a fresh ledger. """ + if delivery_ledger is None: + delivery_ledger = _WorkflowDeliveryLedger(instance_id=ctx.instance_id) all_tasks: list[Any] = [] task_metadata_list: list[TaskMetadata] = [] remaining_agent_messages: list[tuple[str, Any, str]] = [] @@ -913,7 +1543,9 @@ def _prepare_all_tasks( + qualify_subworkflow_request_id(executor_id, ordinal, ""), } logger.debug("Preparing sub-workflow task: %s -> %s", executor_id, child_instance_id) - task = _prepare_subworkflow_task(ctx, executor, message, child_instance_id, child_address) + task = _prepare_subworkflow_task( + ctx, executor, message, child_instance_id, child_address, delivery_ledger + ) all_tasks.append(task) task_metadata_list.append( TaskMetadata( @@ -928,7 +1560,7 @@ def _prepare_all_tasks( for message, source_executor_id in messages_with_sources: logger.debug("Preparing activity task: %s", executor_id) task = _prepare_activity_task( - ctx, executor_id, message, source_executor_id, shared_state, workflow.name, address + ctx, executor_id, message, source_executor_id, shared_state, workflow.name, address, delivery_ledger ) all_tasks.append(task) task_metadata_list.append( @@ -941,22 +1573,24 @@ def _prepare_all_tasks( ) for executor_id, messages_list in agent_messages_by_executor.items(): - first_msg = messages_list[0] - remaining = messages_list[1:] - - logger.debug("Preparing agent task: %s", executor_id) - task = _prepare_agent_task(ctx, first_msg[0], first_msg[1], workflow.name) - all_tasks.append(task) - task_metadata_list.append( - TaskMetadata( - executor_id=first_msg[0], - message=first_msg[1], - source_executor_id=first_msg[2], - task_type=TaskType.AGENT, + for index, (_, message, source_executor_id) in enumerate(messages_list): + metadata = TaskMetadata(executor_id, message, source_executor_id, TaskType.AGENT) + logger.debug("Preparing agent task: %s", executor_id) + task = _prepare_agent_task( + ctx, + cast(AgentExecutor, workflow.executors[executor_id]), + executor_id, + message, + workflow.name, + delivery_ledger, + metadata, ) - ) - - remaining_agent_messages.extend(remaining) + if metadata.skip_dispatch or (isinstance(message, AgentExecutorRequest) and not message.should_respond): + continue + all_tasks.append(task) + task_metadata_list.append(metadata) + remaining_agent_messages.extend(messages_list[index + 1 :]) + break return all_tasks, task_metadata_list, remaining_agent_messages @@ -1013,7 +1647,7 @@ def run_workflow_orchestrator( Returns: For a top-level run, the list of workflow outputs collected from executor - activities. For a sub-workflow run (``initial_message`` carries + activities and designated agents. For a sub-workflow run (``initial_message`` carries :data:`SUBWORKFLOW_INPUT_KEY`), a :data:`SUBWORKFLOW_RESULT_KEY` envelope ``{"outputs": [...], "events": [...]}`` so the parent can bubble nested progress. @@ -1042,11 +1676,16 @@ def run_workflow_orchestrator( # persists across supersteps so repeated sub-workflow invocations never collide. subworkflow_counter: list[int] = [0] + # Rebuilt by executing this generator on replay, not checkpointed separately or + # attached to the shared Workflow/AgentExecutor objects. Survives cycles and HITL + # waits within this invocation and is shared by parallel and sequential dispatch. + delivery_ledger = _WorkflowDeliveryLedger(instance_id=ctx.instance_id) + # Accumulate workflow events and publish them to the orchestration custom status # after each superstep so an external client can stream progress by polling. # Non-agent executors are run inside a durable activity that captures their events - # with data payloads (replayed via append_activity_events); agents contribute only - # synthesized invoked/completed lifecycle events. Events are per executor / per + # with data payloads (replayed via append_activity_events); agents contribute + # lifecycle, request-info and designated output events. Events are per executor / per # yielded output, not token-level, and accumulate for the run. # # Only hosts that stream this timeline (ctx.supports_event_streaming) accumulate @@ -1073,6 +1712,19 @@ def append_activity_events(activity_result: dict[str, Any] | None) -> None: enriched["iteration"] = iteration live_events.append(enriched) + def record_agent_result(result: ExecutorResult) -> None: + append_activity_events(result.activity_result) + if result.output_message is not None: + event_type = _classify_workflow_output(workflow, result.executor_id) + if event_type == "output" or (event_type == "intermediate" and ctx.supports_event_streaming): + encoded = serialize_workflow_agent_response(result.output_message.agent_response) + if event_type == "output": + workflow_outputs.append(encoded) + append_activity_events({ + "events": [{"type": event_type, "executor_id": result.executor_id, "data": encoded}] + }) + emit_event("executor_completed", result.executor_id) + def publish_live_status( state: str, pending_requests: dict[str, Any] | None = None, @@ -1105,13 +1757,28 @@ def publish_live_status( pending_hitl_requests: dict[str, PendingHITLRequest] = {} + def publish_pending_status() -> None: + publish_live_status( + "waiting_for_human_input", + pending_requests={ + req_id: { + "request_id": req.request_id, + "source_executor_id": req.source_executor_id, + "data": req.request_data, + "request_type": req.request_type, + "response_type": req.response_type, + } + for req_id, req in pending_hitl_requests.items() + }, + ) + while pending_messages and iteration < workflow.max_iterations: logger.debug("Orchestrator iteration %d", iteration) next_pending_messages: dict[str, list[tuple[Any, str]]] = {} # Phase 1: Prepare all tasks all_tasks, task_metadata_list, remaining_agent_messages = _prepare_all_tasks( - ctx, workflow, pending_messages, shared_state, subworkflow_counter, workflow_address + ctx, workflow, pending_messages, shared_state, subworkflow_counter, workflow_address, delivery_ledger ) # Agents and sub-workflows bypass the per-executor activity, so synthesize their @@ -1120,8 +1787,6 @@ def publish_live_status( for task_meta in task_metadata_list: if task_meta.task_type in (TaskType.AGENT, TaskType.SUBWORKFLOW): emit_event("executor_invoked", task_meta.executor_id) - for invoked_executor_id, _invoked_message, _invoked_source in remaining_agent_messages: - emit_event("executor_invoked", invoked_executor_id) # Phase 2: Execute all tasks in parallel all_results: list[ExecutorResult] = [] @@ -1141,31 +1806,46 @@ def publish_live_status( for idx, raw_result in enumerate(raw_results): metadata = task_metadata_list[idx] if metadata.task_type == TaskType.AGENT: - result = _process_agent_response(raw_result, metadata.executor_id, metadata.message) - emit_event("executor_completed", metadata.executor_id) + result = _process_agent_response( + raw_result, metadata.executor_id, metadata.message, delivery_ledger, metadata + ) + record_agent_result(result) elif metadata.task_type == TaskType.SUBWORKFLOW: subworkflow_executor = cast(WorkflowExecutor, workflow.executors[metadata.executor_id]) - result = _process_subworkflow_result(raw_result, subworkflow_executor, workflow_outputs) - # Bubble the child's (re-tagged) intermediate events into this - # parent's timeline before the node's completed event, preserving - # chronological order: node invoked -> child progress -> completed. + result = _process_subworkflow_result(raw_result, subworkflow_executor, workflow_outputs, workflow) + # Publish classified direct outputs and re-tagged child progress + # before the node's completed event, as in core WorkflowExecutor. append_activity_events(result.activity_result) emit_event("executor_completed", metadata.executor_id) else: result = _process_activity_result(raw_result, metadata.executor_id, shared_state, workflow_outputs) append_activity_events(result.activity_result) + result.source_message = metadata.message + result.child_instance_id = metadata.child_instance_id all_results.append(result) # Phase 3: Process sequential agent messages - for executor_id, message, _source_executor_id in remaining_agent_messages: + for executor_id, message, source_executor_id in remaining_agent_messages: logger.debug("Processing sequential message for agent: %s", executor_id) - task = _prepare_agent_task(ctx, executor_id, message, workflow.name) - agent_response: AgentResponse = yield task + metadata = TaskMetadata(executor_id, message, source_executor_id, TaskType.AGENT) + task = _prepare_agent_task( + ctx, + cast(AgentExecutor, workflow.executors[executor_id]), + executor_id, + message, + workflow.name, + delivery_ledger, + metadata, + ) + if metadata.skip_dispatch or (isinstance(message, AgentExecutorRequest) and not message.should_respond): + continue + emit_event("executor_invoked", executor_id) + agent_response: AgentResponse | dict[str, Any] = yield task logger.debug("Agent %s sequential response completed", executor_id) - result = _process_agent_response(agent_response, executor_id, message) + result = _process_agent_response(agent_response, executor_id, message, delivery_ledger, metadata) all_results.append(result) - emit_event("executor_completed", executor_id) + record_agent_result(result) # Phase 4: Collect HITL requests for result in all_results: @@ -1173,7 +1853,7 @@ def publish_live_status( # Phase 5: Route results for result in all_results: - _route_result_messages(result, workflow, next_pending_messages, fan_in_pending) + _route_result_messages(result, workflow, next_pending_messages, fan_in_pending, delivery_ledger) # Phase 6: Check fan-in readiness _check_fan_in_ready(workflow, fan_in_pending, next_pending_messages) @@ -1190,19 +1870,7 @@ def publish_live_status( if not pending_messages and pending_hitl_requests: logger.debug("Workflow paused for HITL - %d pending requests", len(pending_hitl_requests)) - publish_live_status( - "waiting_for_human_input", - pending_requests={ - req_id: { - "request_id": req.request_id, - "source_executor_id": req.source_executor_id, - "data": req.request_data, - "request_type": req.request_type, - "response_type": req.response_type, - } - for req_id, req in pending_hitl_requests.items() - }, - ) + publish_pending_status() for request_id, hitl_request in list(pending_hitl_requests.items()): # Wait indefinitely for the human response, matching MAF core's @@ -1241,12 +1909,26 @@ def publish_live_status( ) continue + if isinstance(workflow.executors[hitl_request.source_executor_id], AgentExecutor): + original_request = delivery_ledger.pending_agent_requests[hitl_request.source_executor_id][ + request_id + ] + try: + sanitized_response = _load_agent_hitl_content( + request_id, original_request, sanitized_response + ) + except (TypeError, ValueError): + logger.warning("Rejected malformed agent HITL response for request %s", request_id) + continue + del pending_hitl_requests[request_id] _route_hitl_response( hitl_request, sanitized_response, pending_messages, ) + if pending_hitl_requests: + publish_pending_status() break publish_live_status("running") diff --git a/python/packages/durabletask/agent_framework_durabletask/_workflows/protocol.py b/python/packages/durabletask/agent_framework_durabletask/_workflows/protocol.py new file mode 100644 index 0000000..be69138 --- /dev/null +++ b/python/packages/durabletask/agent_framework_durabletask/_workflows/protocol.py @@ -0,0 +1,48 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Explicit start-envelope versioning for the incompatible workflow execution engine.""" + +from typing import Any, cast + +WORKFLOW_ENGINE_VERSION = 2 +_VERSION_KEY = "_durable_workflow_version" + + +def wrap_workflow_input(value: Any) -> dict[str, Any]: + """Mark a newly scheduled workflow input without changing its application payload. + + This envelope is not an authorization boundary. It distinguishes recorded starts + from the prior engine, which must remain on their original deployment. + + Args: + value: The application input, or trusted internal child input. + + Returns: + The versioned scheduling envelope. + """ + return {_VERSION_KEY: WORKFLOW_ENGINE_VERSION, "input": value} + + +def unwrap_workflow_input(envelope: Any) -> Any: + """Reject old starts before a hosted orchestrator executes any revised actions. + + Rewrapping recorded history does not migrate it. Deploy the old engine to finish + old instances, and schedule only new instances with this engine's clients. + + Args: + envelope: The durable orchestration's recorded start input. + + Returns: + The original application payload for a supported new start. + """ + data = cast("dict[str, Any]", envelope) if isinstance(envelope, dict) else {} + if ( + data.keys() != {_VERSION_KEY, "input"} + or type(data[_VERSION_KEY]) is not int + or data[_VERSION_KEY] != WORKFLOW_ENGINE_VERSION + ): + raise ValueError( + "This workflow start belongs to an unsupported execution protocol. Keep old workflow histories on their " + "original deployment; use the v2 client/start route for new instances in an isolated-v2 deployment." + ) + return data["input"] diff --git a/python/packages/durabletask/agent_framework_durabletask/_workflows/serialization.py b/python/packages/durabletask/agent_framework_durabletask/_workflows/serialization.py index cabd2c4..8f38c7a 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_workflows/serialization.py +++ b/python/packages/durabletask/agent_framework_durabletask/_workflows/serialization.py @@ -21,7 +21,9 @@ :mod:`agent_framework._workflows._checkpoint_encoding` for the full security model. Contents: -- ``serialize_value`` / ``deserialize_value``: internal codec aliases for encode/decode. +- ``serialize_value`` / ``deserialize_value``: internal checkpoint encoding/decoding. +- ``serialize_workflow_agent_response``: portable JSON for generated agent yields, + recognized by ``deserialize_value`` without loading worker response-format types. - ``reconstruct_to_type``: rebuilds HITL response data (which arrives without type markers) to a known type. - ``resolve_type``: resolves 'module:class' type keys to Python types. @@ -37,7 +39,7 @@ from dataclasses import is_dataclass from typing import Any, cast -from agent_framework import WorkflowEvent +from agent_framework import AgentResponse, Content, Message, WorkflowEvent from agent_framework._workflows._checkpoint_encoding import ( _PICKLE_MARKER, # pyright: ignore[reportPrivateUsage] _TYPE_MARKER, # pyright: ignore[reportPrivateUsage] @@ -47,8 +49,13 @@ from agent_framework._workflows._events import WorkflowEventType from pydantic import BaseModel +from .._response_utils import load_agent_response, serialize_agent_response + logger = logging.getLogger(__name__) +_WORKFLOW_AGENT_RESPONSE_KEY = "_durable_agent_response" +_WORKFLOW_AGENT_RESPONSE_VERSION = 1 + def resolve_type(type_key: str) -> type | None: """Resolve a 'module:class' type key to its Python type. @@ -171,6 +178,14 @@ def strip_subworkflow_markers(data: Any) -> Any: # ============================================================================ +def serialize_workflow_agent_response(response: AgentResponse) -> dict[str, Any]: + """Encode a generated agent yield as base-response JSON, without worker types.""" + return { + _WORKFLOW_AGENT_RESPONSE_KEY: _WORKFLOW_AGENT_RESPONSE_VERSION, + "response": serialize_agent_response(response), + } + + def serialize_value(value: Any) -> Any: """Encode a value for JSON-compatible cross-activity communication (internal). @@ -188,13 +203,12 @@ def serialize_value(value: Any) -> Any: def deserialize_value(value: Any) -> Any: - """Decode a value previously encoded with :func:`serialize_value` (internal). + """Decode checkpoint values and known generated-agent response envelopes. - Framework-internal codec. Delegates to core checkpoint decoding which - unpickles base64-encoded values and verifies type integrity. Not part of the - public API: callers only ever hand it values that the framework produced - itself or that have already passed the :func:`strip_pickle_markers` trust - boundary, so untrusted markers can never reach ``pickle.loads()`` here. + Generated agent yields contain base-response JSON, not persisted Python type + names. Ordinary checkpoint envelopes still delegate to core decoding. Callers + must supply framework-produced data or values that have already passed the + :func:`strip_pickle_markers` trust boundary. Args: value: The serialized data (dict with pickle markers, list, or primitive) @@ -202,16 +216,36 @@ def deserialize_value(value: Any) -> Any: Returns: Reconstructed typed object if type metadata found, otherwise original value. """ + if isinstance(value, dict): + data = cast(dict[str, Any], value) + if _WORKFLOW_AGENT_RESPONSE_KEY in data: + version = data[_WORKFLOW_AGENT_RESPONSE_KEY] + if ( + type(version) is not int + or version != _WORKFLOW_AGENT_RESPONSE_VERSION + or set(data) != {_WORKFLOW_AGENT_RESPONSE_KEY, "response"} + or not isinstance(data["response"], dict) + ): + raise ValueError("Invalid or unsupported workflow agent response envelope") + # The response loader follows only known envelope fields. In particular, + # value/additional_properties remain application JSON, not codec input. + return load_agent_response(cast("dict[str, Any]", data["response"])) + if _PICKLE_MARKER in data and _TYPE_MARKER in data: + # Do not walk the restored object: the core codec also pickles ordinary + # application dictionaries that contain reserved checkpoint keys. + return decode_checkpoint_value(data) + return {key: deserialize_value(item) for key, item in data.items()} + if isinstance(value, list): + return [deserialize_value(item) for item in cast(list[Any], value)] return decode_checkpoint_value(value) def deserialize_workflow_output(output: Any) -> Any: - """Reconstruct the workflow outputs produced by the shared activity. + """Reconstruct activity and generated agent outputs from the shared engine. - Each value an executor yields is encoded with :func:`serialize_value` before - it reaches the orchestrator, so typed objects (dataclasses, Pydantic models, - ``AgentResponse``, ...) are stored as checkpoint-marker dicts. This reverses - that encoding so callers receive the original objects. + Activity yields retain their checkpoint encoding. Generated agent yields use + a known response envelope and restore as base ``AgentResponse`` objects with + JSON structured values, without requiring the worker's response-format class. This is the single decode path shared by every host (the in-process :class:`DurableWorkflowClient` and the Azure Functions status endpoint) so @@ -226,8 +260,8 @@ def deserialize_workflow_output(output: Any) -> Any: of yielded outputs or a single value). Returns: - The output with every checkpoint-encoded value reconstructed; primitives - and plain JSON structures pass through unchanged. + The output with checkpoint values and known response envelopes reconstructed; + primitives and other plain JSON structures pass through unchanged. """ return deserialize_value(output) @@ -323,8 +357,9 @@ def reconstruct_to_type(value: Any, target_type: type) -> Any: Tries strategies in order: 1. Return as-is if already the correct type 2. deserialize_value (for data with any type markers) - 3. Pydantic model_validate (for Pydantic models) - 4. Dataclass constructor (for dataclasses) + 3. Safe base Content/Message construction (for those exact declared types) + 4. Pydantic model_validate (for Pydantic models) + 5. Dataclass constructor (for dataclasses) Args: value: The value to reconstruct (typically a dict from JSON) @@ -332,6 +367,10 @@ def reconstruct_to_type(value: Any, target_type: type) -> Any: Returns: Reconstructed value if possible, otherwise the original value + + Raises: + TypeError: If a declared Content or Message payload has invalid constructor fields. + ValueError: If a declared Content or Message payload has a malformed envelope. """ if value is None: return None @@ -351,6 +390,13 @@ def reconstruct_to_type(value: Any, target_type: type) -> Any: if not isinstance(decoded, dict): return decoded + # The declared type is trusted, but nested payload type names are not. Use + # the fixed envelope loader, leaving arbitrary application data opaque. + if target_type is Message: + return load_agent_response({"messages": [value]}).messages[0] + if target_type is Content: + return load_agent_response({"messages": [{"role": "user", "contents": [value]}]}).messages[0].contents[0] + # Try Pydantic model validation (for unmarked dicts, e.g., external HITL data) if issubclass(target_type, BaseModel): try: diff --git a/python/packages/durabletask/pyproject.toml b/python/packages/durabletask/pyproject.toml index 6e5ca54..8a4a89b 100644 --- a/python/packages/durabletask/pyproject.toml +++ b/python/packages/durabletask/pyproject.toml @@ -25,6 +25,8 @@ dependencies = [ "agent-framework-core>=1.13.0,<2", "durabletask>=1.5.0,<2", "durabletask-azuremanaged>=1.4.0,<2", + "opentelemetry-api>=1.39.0,<2", + "pydantic>=2.11,<3", "python-dateutil>=2.8.0,<3", ] diff --git a/python/packages/durabletask/tests/integration_tests/conftest.py b/python/packages/durabletask/tests/integration_tests/conftest.py index d25fa36..58a7e65 100644 --- a/python/packages/durabletask/tests/integration_tests/conftest.py +++ b/python/packages/durabletask/tests/integration_tests/conftest.py @@ -401,6 +401,8 @@ class TestSingleAgent: env = os.environ.copy() env["ENDPOINT"] = dts_endpoint env["TASKHUB"] = unique_taskhub + # Opt in only for the subprocess using this isolated test hub. + env["DURABLE_AGENTS_DEPLOYMENT_MODE"] = "isolated_v2" # Start worker subprocess try: diff --git a/python/packages/durabletask/tests/integration_tests/live_retention_worker.py b/python/packages/durabletask/tests/integration_tests/live_retention_worker.py new file mode 100644 index 0000000..b4fcf63 --- /dev/null +++ b/python/packages/durabletask/tests/integration_tests/live_retention_worker.py @@ -0,0 +1,185 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Test-only DTS host with real core history and a deterministic, recording model. + +The stdin/stdout protocol carries bounded control records only. Full model inputs +and simulated external effects stay in the parent test's temporary directory. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import logging +import sys +from collections.abc import AsyncIterable, Awaitable, Generator, Mapping, Sequence +from pathlib import Path +from threading import Event, Lock +from typing import Any + +from agent_framework import Agent, BaseChatClient, ChatResponse, ChatResponseUpdate, Content, Message, ResponseStream +from durabletask.azuremanaged.worker import DurableTaskSchedulerWorker +from durabletask.entities import EntityInstanceId +from durabletask.task import OrchestrationContext +from opentelemetry.metrics import set_meter_provider +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import InMemoryMetricReader, Sum + +import agent_framework_durabletask +from agent_framework_durabletask import DurableAIAgentWorker, DurableHistoryProvider + +AGENT_NAME = "live-retention" +MAX_STATE_BYTES = 50_000 +DELIVERY_WINDOW_SECONDS = 3600 +CONTROL_TIMEOUT = 45 +_output_lock = Lock() + + +def emit(event: str, **fields: Any) -> None: + record = json.dumps({"event": event, **fields}, ensure_ascii=True) + if len(record) > 2048: + raise ValueError("Control record exceeds its bound") + with _output_lock: + sys.stdout.write(record + "\n") + sys.stdout.flush() + + +class RecordingModel(BaseChatClient): + """Only model I/O is replaced. Agent streaming and history hooks are real.""" + + def __init__(self, artifacts: Path, blocked_message_id: str) -> None: + super().__init__() + self.artifacts = artifacts + self.blocked_message_id = blocked_message_id + self.release = Event() + self.calls = 0 + + def _capture(self, messages: Sequence[Message], current_id: str) -> None: + self.calls += 1 + ordinal = self.calls + captured = [message.to_dict() for message in messages] + (self.artifacts / f"model-{ordinal}.json").write_text(json.dumps(captured), encoding="utf-8") + # This file is the simulated nontransactional external effect, not entity state. + (self.artifacts / f"effect-{ordinal}.json").write_text( + json.dumps({"message_id": current_id, "ordinal": ordinal}), encoding="utf-8" + ) + emit("model_entered", ordinal=ordinal, message_id=current_id) + + def _inner_get_response( + self, + *, + messages: Sequence[Message], + stream: bool, + options: Mapping[str, Any], + **kwargs: Any, + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + async def response_update() -> ChatResponseUpdate: + await self._validate_options(options) + current_id = next(message.message_id for message in reversed(messages) if message.role == "user") + if not current_id: + raise ValueError("The test requires a current user message ID") + await asyncio.to_thread(self._capture, messages, current_id) + if current_id == self.blocked_message_id: + released = await asyncio.to_thread(self.release.wait, CONTROL_TIMEOUT) + if not released: + # Do not turn a missed test barrier into a committed agent error response. + raise asyncio.CancelledError("Test model barrier timed out") + return ChatResponseUpdate( + role="assistant", + author_name="retention-model", + contents=[Content.from_text(f"answer:{current_id}")], + message_id=f"{current_id}-answer", + response_id=f"response:{current_id}", + finish_reason="stop", + ) + + async def updates() -> AsyncIterable[ChatResponseUpdate]: + yield await response_update() + + async def response() -> ChatResponse: + return ChatResponse.from_updates([await response_update()]) + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) if stream else response() + + +def live_retention_duplicate(context: OrchestrationContext, payload: dict[str, Any]) -> Generator[Any, Any, Any]: + """A same-sender signal then call supplies an acknowledged duplicate barrier.""" + entity = EntityInstanceId(entity=f"dafx-{AGENT_NAME}", key=payload["key"]) + context.signal_entity(entity, "run", payload["request"]) + result = yield context.call_entity(entity, "run", payload["request"]) + return result # noqa: B901 + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--endpoint", required=True) + parser.add_argument("--taskhub", required=True) + parser.add_argument("--artifacts", required=True, type=Path) + parser.add_argument("--block-message-id", default="") + args = parser.parse_args() + expected_package = Path(__file__).resolve().parents[2] / "agent_framework_durabletask" + if Path(agent_framework_durabletask.__file__).resolve().parent != expected_package: + raise RuntimeError("Worker imported durabletask extension from a different checkout") + logging.basicConfig(level=logging.WARNING) + reader = InMemoryMetricReader() + meters = MeterProvider(metric_readers=[reader], shutdown_on_exit=False) + set_meter_provider(meters) + model = RecordingModel(args.artifacts, args.block_message_id) + worker = DurableTaskSchedulerWorker( + host_address=args.endpoint, taskhub=args.taskhub, token_credential=None, secure_channel=False + ) + host = DurableAIAgentWorker( + worker, + deployment_mode="isolated_v2", + retention="keep_all", + max_state_bytes=MAX_STATE_BYTES, + response_delivery_window_seconds=DELIVERY_WINDOW_SECONDS, + ) + host.add_agent( + Agent( + client=model, + name=AGENT_NAME, + id=AGENT_NAME, + default_options={"store": False}, + context_providers=[DurableHistoryProvider()], + ) + ) + worker.add_orchestrator(live_retention_duplicate) + try: + host.start() + # start() launches the SDK background thread. Only a backend receipt proves readiness. + emit("started") + for line in sys.stdin: + command = json.loads(line)["command"] + if command == "release": + model.release.set() + elif command == "metrics": + data = reader.get_metrics_data() + rows: list[dict[str, Any]] = [] + if data is not None: + for resource in data.resource_metrics: + for scope in resource.scope_metrics: + for metric in scope.metrics: + if metric.name == "durable.retention.removed_messages": + assert isinstance(metric.data, Sum) and metric.data.is_monotonic + rows.extend( + {"value": point.value, "attributes": dict(point.attributes or {})} + for point in metric.data.data_points + ) + (args.artifacts / "metrics.json").write_text(json.dumps(rows), encoding="utf-8") + emit("metrics") + elif command == "stop": + break + else: + raise ValueError("Unknown test control command") + finally: + model.release.set() + try: + host.stop() + finally: + meters.shutdown() + + +if __name__ == "__main__": + main() diff --git a/python/packages/durabletask/tests/integration_tests/test_01_dt_single_agent.py b/python/packages/durabletask/tests/integration_tests/test_01_dt_single_agent.py index 0bb5f4b..a546c2f 100644 --- a/python/packages/durabletask/tests/integration_tests/test_01_dt_single_agent.py +++ b/python/packages/durabletask/tests/integration_tests/test_01_dt_single_agent.py @@ -62,21 +62,25 @@ def test_single_interaction(self): assert len(response.text) > 0 def test_conversation_continuity(self): - """Test that conversation context is maintained across turns.""" + """Prior turns must reach the model, not just be recorded. + + The second turn is only answerable from persisted history, so this fails if durable + history is not actually being loaded and delivered to the agent. + """ agent = self.agent_client.get_agent("Joker") session = agent.create_session() - # First turn: Ask for a joke about a specific topic - response1 = agent.run("Tell me a joke about cats.", session=session) + # First turn establishes a fact that exists nowhere else. + response1 = agent.run("My favorite animal is the axolotl. Tell me a joke about it.", session=session) assert response1 is not None assert len(response1.text) > 0 - # Second turn: Ask a follow-up that requires context - response2 = agent.run("Can you make it funnier?", session=session) + # Second turn can only be answered from the conversation history. + response2 = agent.run("What is my favorite animal? Reply with just the animal name.", session=session) assert response2 is not None - assert len(response2.text) > 0 - - # The agent should understand "it" refers to the previous joke + assert "axolotl" in response2.text.lower(), ( + f"Agent lost conversation context across turns. Got: {response2.text!r}" + ) def test_multiple_sessions(self): """Test that different sessions maintain separate contexts.""" diff --git a/python/packages/durabletask/tests/integration_tests/test_06_dt_multi_agent_orchestration_conditionals.py b/python/packages/durabletask/tests/integration_tests/test_06_dt_multi_agent_orchestration_conditionals.py index d50748f..f64a980 100644 --- a/python/packages/durabletask/tests/integration_tests/test_06_dt_multi_agent_orchestration_conditionals.py +++ b/python/packages/durabletask/tests/integration_tests/test_06_dt_multi_agent_orchestration_conditionals.py @@ -63,23 +63,42 @@ def test_agents_registered(self): assert email_agent is not None assert email_agent.name == EMAIL_AGENT_NAME - def test_conditional_branching(self): - """Test that conditional branching works correctly.""" - # Test with obvious spam - spam_payload = { - "email_id": "spam-001", - "email_content": "Buy cheap medications online! No prescription needed! Limited time offer!", - } + def test_conditional_branching(self) -> None: + """Spam takes the spam-handler branch and legitimate mail takes the reply branch. + Asserting only that the orchestration completed would pass even if the condition sent + every email down the same branch, so each case checks the branch-specific output. + """ spam_instance_id = self.dts_client.schedule_new_orchestration( orchestrator="spam_detection_orchestration", - input=spam_payload, + input={ + "email_id": "spam-001", + "email_content": "Buy cheap medications online! No prescription needed! Limited time offer!", + }, ) - - # Both should complete successfully (different branches) - spam_metadata = self.orch_helper.wait_for_orchestration( + spam_metadata, spam_output = self.orch_helper.wait_for_orchestration_with_output( instance_id=spam_instance_id, timeout=300.0, ) assert spam_metadata.runtime_status == OrchestrationStatus.COMPLETED + # The spam handler returns "Email marked as spam: ..."; the other branch returns "Email sent: ...". + assert "marked as spam" in str(spam_output).lower(), f"spam took the wrong branch: {spam_output}" + + legit_instance_id = self.dts_client.schedule_new_orchestration( + orchestrator="spam_detection_orchestration", + input={ + "email_id": "legit-001", + "email_content": ( + "Hi team, please confirm receipt of purchase order PRJ-4417 for the new lab " + "hardware, and let me know the expected delivery date." + ), + }, + ) + legit_metadata, legit_output = self.orch_helper.wait_for_orchestration_with_output( + instance_id=legit_instance_id, + timeout=300.0, + ) + + assert legit_metadata.runtime_status == OrchestrationStatus.COMPLETED + assert "email sent" in str(legit_output).lower(), f"legitimate mail took the wrong branch: {legit_output}" diff --git a/python/packages/durabletask/tests/integration_tests/test_07_dt_single_agent_orchestration_hitl.py b/python/packages/durabletask/tests/integration_tests/test_07_dt_single_agent_orchestration_hitl.py index 8c90d07..49c9c25 100644 --- a/python/packages/durabletask/tests/integration_tests/test_07_dt_single_agent_orchestration_hitl.py +++ b/python/packages/durabletask/tests/integration_tests/test_07_dt_single_agent_orchestration_hitl.py @@ -147,7 +147,13 @@ def test_hitl_orchestration_with_rejection_and_feedback(self): assert metadata.runtime_status == OrchestrationStatus.COMPLETED def test_hitl_orchestration_timeout(self): - """Test HITL orchestration timeout behavior.""" + """With no approval sent, the orchestration fails on its own approval timeout. + + The shared helper raises when an orchestration reaches FAILED, so this waits on the client + directly. Catching and ignoring that exception (as this test used to) also swallowed a + TimeoutError from a hung orchestration, which left no outcome that could fail the test for + the right reason. + """ payload = { "topic": "Cloud computing fundamentals", "max_review_attempts": 1, @@ -160,15 +166,13 @@ def test_hitl_orchestration_timeout(self): input=payload, ) - # Don't send any approval - let it timeout - # The orchestration should fail due to timeout - try: - metadata = self.orch_helper.wait_for_orchestration( - instance_id=instance_id, - timeout=90.0, - ) - # If it completes, it should be failed status due to timeout - assert metadata.runtime_status == OrchestrationStatus.FAILED - except (RuntimeError, TimeoutError): - # Expected - orchestration should timeout and fail - pass + # Don't send any approval - let it hit its own approval timeout. + metadata = self.dts_client.wait_for_orchestration_completion(instance_id=instance_id, timeout=90) + + assert metadata is not None, "orchestration never reached a terminal state" + assert metadata.runtime_status == OrchestrationStatus.FAILED + + # Fail for the right reason: the sample raises TimeoutError("Human approval timed out ..."). + failure = metadata.failure_details + details = getattr(failure, "message", None) or str(failure) + assert "timed out" in details.lower(), f"expected an approval timeout, got: {details}" diff --git a/python/packages/durabletask/tests/integration_tests/test_08_dt_workflow.py b/python/packages/durabletask/tests/integration_tests/test_08_dt_workflow.py index 2aa9a9d..e1b3d03 100644 --- a/python/packages/durabletask/tests/integration_tests/test_08_dt_workflow.py +++ b/python/packages/durabletask/tests/integration_tests/test_08_dt_workflow.py @@ -5,18 +5,17 @@ Exercises the standalone (non-Azure-Functions) workflow path: - ``DurableAIAgentWorker.configure_workflow`` auto-registers the agent entities, non-agent executor activities, and the workflow orchestrator. -- A client starts the workflow by scheduling its ``dafx-{workflow_name}`` orchestration. +- ``DurableWorkflowClient.start_workflow`` schedules the versioned workflow input. - Conditional routing sends spam to a non-agent handler and legitimate email through a second agent and a sender executor. """ import logging -from typing import Any, Protocol import pytest from durabletask.client import OrchestrationStatus -from agent_framework_durabletask import DurableAIAgentClient, workflow_orchestrator_name +from agent_framework_durabletask import DurableWorkflowClient # Must match the workflow name in samples/08_workflow/worker.py WORKFLOW_NAME = "email_triage" @@ -24,13 +23,6 @@ logging.basicConfig(level=logging.WARNING) -class AgentClientFactoryProtocol(Protocol): - """Protocol for the agent client factory fixture.""" - - @classmethod - def create(cls, max_poll_retries: int = 90) -> tuple[Any, DurableAIAgentClient]: ... - - # Module-level markers pytestmark = [ pytest.mark.flaky, @@ -45,15 +37,15 @@ class TestStandaloneWorkflow: """Standalone (non-Azure-Functions) workflow execution on a durabletask worker.""" @pytest.fixture(autouse=True) - def setup(self, agent_client_factory: type[AgentClientFactoryProtocol], orchestration_helper) -> None: - """Provide a DTS client and orchestration helper for each test.""" - self.dts_client, self.agent_client = agent_client_factory.create() + def setup(self, workflow_client: DurableWorkflowClient, orchestration_helper) -> None: + """Provide a workflow client and orchestration helper for each test.""" + self.workflow_client = workflow_client self.orch_helper = orchestration_helper def test_legitimate_email_drafts_response(self) -> None: """A legitimate email routes through the email agent and is 'sent'.""" - instance_id = self.dts_client.schedule_new_orchestration( - orchestrator=workflow_orchestrator_name(WORKFLOW_NAME), + instance_id = self.workflow_client.start_workflow( + workflow_name=WORKFLOW_NAME, input=( "Hi team, just a reminder about our sprint planning meeting tomorrow at 10 AM. " "Please review the agenda in Jira." @@ -69,10 +61,35 @@ def test_legitimate_email_drafts_response(self) -> None: assert output is not None assert "Email sent" in str(output) + def test_downstream_agent_receives_upstream_conversation(self) -> None: + """The email agent can only reference the original email if upstream context reached it. + + The edge into the email agent carries the spam agent's structured verdict, not the email. + A purchase order number is used as the marker because a spam verdict explains *why* a + message is legitimate and would not repeat an arbitrary code, whereas a drafted reply to + the email naturally does. + """ + instance_id = self.workflow_client.start_workflow( + workflow_name=WORKFLOW_NAME, + input=( + "Hi team, please confirm receipt of purchase order PRJ-4417 for the new lab " + "hardware, and let me know the expected delivery date." + ), + ) + + metadata, output = self.orch_helper.wait_for_orchestration_with_output( + instance_id=instance_id, + timeout=180.0, + ) + + assert metadata.runtime_status == OrchestrationStatus.COMPLETED + assert output is not None + assert "PRJ-4417" in str(output), f"drafted reply did not reference the original email: {output}" + def test_spam_email_handled(self) -> None: """A spam email routes to the non-agent spam handler.""" - instance_id = self.dts_client.schedule_new_orchestration( - orchestrator=workflow_orchestrator_name(WORKFLOW_NAME), + instance_id = self.workflow_client.start_workflow( + workflow_name=WORKFLOW_NAME, input="URGENT! You've won $1,000,000! Click here now to claim your prize! Limited time offer!", ) diff --git a/python/packages/durabletask/tests/integration_tests/test_13_dt_conversation_compaction.py b/python/packages/durabletask/tests/integration_tests/test_13_dt_conversation_compaction.py new file mode 100644 index 0000000..037ef9f --- /dev/null +++ b/python/packages/durabletask/tests/integration_tests/test_13_dt_conversation_compaction.py @@ -0,0 +1,247 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Integration tests for compaction of client-owned durable history. + +Covers the sample's ``store=False`` agent with input/output storage enabled and explicit +``retention="keep_all", max_state_bytes=None``: + +- provider-selected history reaches the model on later turns, +- compaction annotations and message identities survive entity state serialization, +- excluded local history is retained without coupling delivery to transcript entries, +- original response payloads live in the correlation-keyed mailbox with completion receipts. + +This is not a bounded-capacity stress test. Live mailbox payloads and completion receipts still +consume state, and no delivery window is shortened to make the sample fit a small budget. +""" + +import json +from datetime import datetime +from pathlib import Path +from typing import Any, Protocol + +import pytest +from durabletask.entities import EntityInstanceId + +from agent_framework_durabletask import ( + DurableAgentState, + DurableAgentStateRequest, + DurableAgentStateResponse, + DurableAIAgentClient, + serialize_agent_response, +) + +# Matches worker.py: only the most recent groups stay in the model's context. +KEEP_LAST_GROUPS = 4 + + +class AgentClientFactoryProtocol(Protocol): + """Protocol for the agent client factory fixture.""" + + @classmethod + def create(cls, max_poll_retries: int = 90) -> tuple[Any, DurableAIAgentClient]: ... + + +pytestmark = [ + pytest.mark.flaky, + pytest.mark.integration, + pytest.mark.sample("13_conversation_compaction"), + pytest.mark.integration_test, + pytest.mark.requires_foundry, + pytest.mark.requires_dts, +] + + +class TestConversationCompaction: + """Local provider history compacts without changing the sample's keep-all policy.""" + + @pytest.fixture(autouse=True) + def setup(self, agent_client_factory: type[AgentClientFactoryProtocol]) -> None: + """Setup test fixtures.""" + self.dts_client, self.agent_client = agent_client_factory.create() + + def _read_state(self, session_id: Any) -> DurableAgentState: + """Load the agent entity's persisted state straight from the scheduler.""" + entity_id = EntityInstanceId(entity=session_id.entity_name, key=session_id.key) + metadata = self.dts_client.get_entity(entity_id) + assert metadata is not None, f"no durable state found for {entity_id}" + + raw = metadata.get_state() + # The scheduler returns the entity payload as serialized JSON. + if isinstance(raw, str): + return DurableAgentState.from_json(raw) + assert isinstance(raw, dict), f"unexpected entity state payload: {type(raw)}" + return DurableAgentState.from_dict(raw) + + def test_agent_registration(self) -> None: + """The compacting agent is registered like any other agent.""" + agent = self.agent_client.get_agent("Historian") + assert agent is not None + assert agent.name == "Historian" + + def test_session_is_persisted_and_scoped(self) -> None: + """The serialized session survives real entity storage with the right shape. + + Unit tests keep the session dict in memory, so they cannot show that the blob survives the + entity's JSON encoding, that it carries the entity's **own** session id, or that the durable + history provider's slice really is kept out of it. + """ + agent = self.agent_client.get_agent("Historian") + session = agent.create_session() + + assert agent.run("Name a color.", session=session) is not None + assert agent.run("Name a fruit.", session=session) is not None + + stored = self._read_state(session.durable_session_id).data.session + assert stored is not None, "the session was not persisted" + + # The entity's own identity rather than a per-operation id. External history providers key + # their storage on this, so a generated id would restart their conversation every turn. + # It carries the entity name as well as the key, because agent nodes in one workflow run + # share a key and would otherwise all resolve to the same conversation. + assert session.durable_session_id is not None + key = session.durable_session_id.key + assert stored["session_id"].endswith(f"@{key}"), ( + f"expected the session id to end with the entity key {key}, got {stored['session_id']}" + ) + # The runtime lowercases entity names, so compare that way. + entity_name = session.durable_session_id.entity_name.lower() + assert entity_name in stored["session_id"].lower(), ( + f"expected the entity name in the session id, got {stored['session_id']}" + ) + + slices = stored["state"] + # The compaction provider's own slice is carried across turns... + assert "compaction" in slices, f"expected provider state to be persisted, got {slices}" + # ...but the durable history provider's is not, since it is derived from + # conversationHistory and would otherwise duplicate the transcript. "in_memory" is the + # source_id the sample's provider keeps after the durable swap. + assert "in_memory" not in slices, f"durable history slice leaked into the session: {slices}" + + def test_persisted_state_matches_the_shared_schema(self) -> None: + """Real scheduler round-tripped state must satisfy the cross-language contract. + + Unit tests validate a synthetic dict. This validates what the entity actually wrote and + the scheduler actually stored, which is where drift between the two would show up. + """ + jsonschema = pytest.importorskip("jsonschema") + schema_path = Path(__file__).resolve().parents[5] / "schemas" / "durable-agent-entity-state.json" + schema = json.loads(schema_path.read_text(encoding="utf-8")) + + agent = self.agent_client.get_agent("Historian") + session = agent.create_session() + assert agent.run("Name a city.", session=session) is not None + # A second turn exercises loading persisted ids and annotations as well as assigning + # identities to new messages in the provider's append hooks. + assert agent.run("Name another.", session=session) is not None + + state = self._read_state(session.durable_session_id) + jsonschema.Draft202012Validator(schema).validate(state.to_dict()) + + # The fields compaction depends on must actually be present, not merely permitted. + stored = [m for entry in state.data.conversation_history for m in entry.messages] + assert any(m.message_id for m in stored), "no message carried an id through real storage" + + def test_recent_context_survives_compaction(self) -> None: + """A fact inside the retained window is still answerable after several turns.""" + agent = self.agent_client.get_agent("Historian") + session = agent.create_session() + + for filler in ("Name a color.", "Name a country.", "Name a fruit."): + assert agent.run(filler, session=session) is not None + + agent.run("My project codename is BLUEHERON.", session=session) + answer = agent.run("What is my project codename? Reply with just the codename.", session=session) + + assert "blueheron" in answer.text.lower(), ( + f"Recent context was lost despite being inside the retained window. Got: {answer.text!r}" + ) + + def test_compaction_annotations_are_persisted(self) -> None: + """Compaction state must survive durable state serialization. + + The strategy still runs each turn. Persisted message metadata and ids let it operate on + the annotated history rather than losing prior exclusions across entity operations. + """ + agent = self.agent_client.get_agent("Historian") + session = agent.create_session() + + # Run enough turns that the sliding window must exclude earlier ones. + for index in range(KEEP_LAST_GROUPS + 3): + assert agent.run(f"Name animal number {index + 1}.", session=session) is not None + + state = self._read_state(session.durable_session_id) + + stored = [message for entry in state.data.conversation_history for message in entry.messages] + assert stored, "expected the conversation to be persisted" + + # Compaction excluded older messages, and that annotation round-tripped through storage. + annotated = [m for m in stored if m.extension_data] + assert annotated, "expected compaction annotations to be persisted in durable state" + + excluded = [m for m in annotated if (m.extension_data or {}).get("_excluded")] + assert excluded, "expected the sliding window to exclude older messages" + + # Provider appends assign identities without changing caller messages. Compaction + # reconciles by those ids, including the newest turn, not by transcript position. + assert all(m.message_id for m in stored), "stored messages must carry stable message ids" + assert len({m.message_id for m in stored}) == len(stored), "stored message ids must be unique" + + def test_local_provider_retains_selected_inputs_and_outputs_with_keep_all(self) -> None: + """This local provider stores both sides; compaction alone does not delete them.""" + agent = self.agent_client.get_agent("Historian") + session = agent.create_session() + + turns = KEEP_LAST_GROUPS + 3 + prompts = [f"Name city number {index + 1}." for index in range(turns)] + replies = [agent.run(prompt, session=session) for prompt in prompts] + assert all(reply.text for reply in replies) + assert all( + content.type != "error" for reply in replies for message in reply.messages for content in message.contents + ) + + state = self._read_state(session.durable_session_id) + + # These counts follow this sample's local provider flags and single-call, tool-free turns. + # They are not a delivery invariant for external or service-managed history. + requests = [entry for entry in state.data.conversation_history if isinstance(entry, DurableAgentStateRequest)] + responses = [entry for entry in state.data.conversation_history if isinstance(entry, DurableAgentStateResponse)] + assert len(requests) == len(responses) == turns + assert [message.text for entry in requests for message in entry.messages] == prompts + assert [[message.text for message in entry.messages] for entry in responses] == [ + [message.text for message in reply.messages] for reply in replies + ] + assert len(state.data.completed_correlations) == turns + assert state.data.truncation is None + + def test_mailbox_delivers_original_response_without_transcript_entries(self) -> None: + """A real stored result remains readable without reconstructing a transcript response.""" + agent = self.agent_client.get_agent("Historian") + session = agent.create_session() + response = agent.run("Name a river.", session=session) + assert response.text + assert all(content.type != "error" for message in response.messages for content in message.contents) + expected = json.loads(json.dumps(serialize_agent_response(response))) + assert expected["created_at"], "the Foundry result timestamp was lost" + + state = self._read_state(session.durable_session_id) + assert len(state.data.completed_correlations) == 1 + correlation_id = next(iter(state.data.completed_correlations)) + assert correlation_id + assert set(state.data.response_mailbox) == {correlation_id} + mailbox = state.data.response_mailbox[correlation_id] + assert mailbox["response"] == expected + # The result's date and message count are not the request's date or the transcript count. + assert mailbox["response"]["created_at"] == expected["created_at"] + assert len(mailbox["response"]["messages"]) == len(response.messages) + assert state.data.completed_correlations[correlation_id]["completedAt"] == mailbox["createdAt"] + assert datetime.fromisoformat(mailbox["expiresAt"]) > datetime.fromisoformat(mailbox["createdAt"]) + + # Mutate only this detached read, not scheduler state. Version 2 lookup must still use + # the mailbox even when no local transcript entry can provide an answer. + assert state.data.conversation_history + state.data.conversation_history.clear() + restored = DurableAgentState.from_json(state.to_json()) + assert restored.message_count == 0 + delivered = restored.try_get_agent_response(correlation_id) + assert delivered is not None + assert json.loads(json.dumps(serialize_agent_response(delivered))) == expected diff --git a/python/packages/durabletask/tests/integration_tests/test_14_dt_external_history_redis.py b/python/packages/durabletask/tests/integration_tests/test_14_dt_external_history_redis.py new file mode 100644 index 0000000..58d0b80 --- /dev/null +++ b/python/packages/durabletask/tests/integration_tests/test_14_dt_external_history_redis.py @@ -0,0 +1,183 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Integration tests for agents whose history lives in an external store. + +A user who deliberately configured their own history provider (Redis here, but Cosmos DB or a +file behaves the same) must get the same behavior under the durable runtime as in core: + +- the provider is not swapped out for durable-backed history, +- it participates in the run and its stored history reaches the model on later turns, +- it is handed the entity's stable session id, so its keys line up across turns, +- fresh durable state has no local transcript mirror or metadata-only exchange envelopes, +- responses and completion evidence are stored separately, keyed by correlation id. + +The stable session id matters because the entity builds a fresh session per operation, and if +that session carried a generated id an externally keyed store would silently start over every turn. +This sample uses blind Redis appends. These tests do not assert exactly-once external writes across +an interrupted operation or portable reset support. +""" + +import json +import os +from datetime import datetime +from typing import Any, Protocol + +import pytest +import redis.asyncio as aioredis + +from agent_framework_durabletask import DurableAgentState, DurableAIAgentClient, serialize_agent_response + + +class AgentClientFactoryProtocol(Protocol): + """Protocol for the agent client factory fixture.""" + + @classmethod + def create(cls, max_poll_retries: int = 90) -> tuple[Any, DurableAIAgentClient]: ... + + +pytestmark = [ + pytest.mark.flaky, + pytest.mark.integration, + pytest.mark.sample("14_external_history_redis"), + pytest.mark.integration_test, + pytest.mark.requires_foundry, + pytest.mark.requires_dts, + pytest.mark.requires_redis, +] + +# Matches redis_history_provider.py in the sample. +KEY_PREFIX = "durable_sample:history" + + +class TestExternalHistoryProvider: + """An external history provider works durably with no durable-specific configuration.""" + + @pytest.fixture(autouse=True) + def setup(self, agent_client_factory: type[AgentClientFactoryProtocol]) -> None: + """Setup test fixtures.""" + self.dts_client, self.agent_client = agent_client_factory.create() + self.redis_url = os.environ.get("REDIS_CONNECTION_STRING", "redis://localhost:6379") + + async def _history_entries(self, session_id: Any) -> list[str]: + """Read the raw history entries the sample's provider wrote for a session. + + The provider keys on the core session id, which qualifies the entity key with the entity + name so that agent nodes sharing a key in a workflow run stay separate. The exact name + casing is the runtime's, so the key is discovered rather than reconstructed. + + Args: + session_id: The durable session id used for the conversation. + + Returns: + The serialized messages stored in Redis, oldest first. + """ + client = aioredis.from_url(self.redis_url, decode_responses=True) + try: + matches: Any = await client.keys(f"{KEY_PREFIX}:*{session_id.key}") # type: ignore[misc] + keys = [k if isinstance(k, str) else k.decode() for k in matches] + assert len(keys) <= 1, f"the conversation was scattered across keys: {keys}" + if not keys: + return [] + # The client is configured with decode_responses, so entries come back as strings. + # Coerce anyway, since redis-py types lrange as bytes or str depending on version. + entries: Any = await client.lrange(keys[0], 0, -1) # type: ignore[misc] + return [entry if isinstance(entry, str) else entry.decode() for entry in entries] + finally: + await client.aclose() + + def test_agent_registration(self) -> None: + """The externally backed agent is registered like any other agent.""" + agent = self.agent_client.get_agent("Archivist") + assert agent is not None + assert agent.name == "Archivist" + + def test_history_from_the_external_store_reaches_the_model(self) -> None: + """Nothing else could supply the earlier turn, so recall proves the provider ran.""" + agent = self.agent_client.get_agent("Archivist") + session = agent.create_session() + + assert agent.run("My library card number is 4417.", session=session) is not None + answer = agent.run("What is my library card number? Reply with just the number.", session=session) + + assert answer is not None + assert "4417" in answer.text + + async def test_provider_is_keyed_by_the_stable_session_id(self) -> None: + """All turns must land under one key; a per-operation id would scatter them.""" + agent = self.agent_client.get_agent("Archivist") + session = agent.create_session() + + assert agent.run("Remember that my favorite number is 12.", session=session) is not None + assert agent.run("Remember that my favorite color is teal.", session=session) is not None + + entries = await self._history_entries(session.durable_session_id) + + # Two turns, each storing its input and the model's reply, all under the entity's own id. + assert len(entries) >= 4, f"expected the whole conversation under one key, found {len(entries)}" + assert any("12" in entry for entry in entries) + assert any("teal" in entry for entry in entries) + + def test_external_history_has_no_local_transcript_but_keeps_correlated_delivery(self) -> None: + """Fresh external history needs no local transcript to deliver completed results.""" + agent = self.agent_client.get_agent("Archivist") + session = agent.create_session() + completed: set[str] = set() + + for prompt in ("Note that the archive opens at nine.", "Name a weekday."): + response = agent.run(prompt, session=session) + assert response.text + assert all(content.type != "error" for message in response.messages for content in message.contents) + expected = json.loads(json.dumps(serialize_agent_response(response))) + assert expected["created_at"], "the Foundry result timestamp was lost" + + state = self._read_state(session.durable_session_id) + assert state.data.conversation_history == [], "external history must not create a local transcript mirror" + assert state.message_count == 0, "transcript count is not a delivery or completion count" + + # Discover the new correlation from completion state, never from transcript entries. + correlations = set(state.data.completed_correlations) + assert completed <= correlations, "earlier completion receipts were lost" + new_correlations = correlations - completed + assert len(new_correlations) == 1 + correlation_id = new_correlations.pop() + assert correlation_id + completed = correlations + + # Check this turn immediately rather than assuming older payloads remain within + # their delivery window after another model call. Receipts outlive those payloads. + assert correlation_id in state.data.response_mailbox + assert set(state.data.response_mailbox) <= completed + mailbox = state.data.response_mailbox[correlation_id] + assert mailbox["response"] == expected + assert mailbox["response"]["created_at"] == expected["created_at"] + assert len(mailbox["response"]["messages"]) == len(response.messages) + assert state.data.completed_correlations[correlation_id]["completedAt"] == mailbox["createdAt"] + assert datetime.fromisoformat(mailbox["expiresAt"]) > datetime.fromisoformat(mailbox["createdAt"]) + + delivered = state.try_get_agent_response(correlation_id) + assert delivered is not None + assert json.loads(json.dumps(serialize_agent_response(delivered))) == expected + + assert len(completed) == 2, "two completed turns must not require two local transcript exchanges" + + def _read_state(self, session_id: Any) -> DurableAgentState: + """Load the agent entity's persisted state straight from the scheduler. + + Args: + session_id: The durable session id used for the conversation. + + Returns: + The deserialized durable agent state. + """ + from durabletask.entities import EntityInstanceId + + entity_id = EntityInstanceId(entity=session_id.entity_name, key=session_id.key) + metadata = self.dts_client.get_entity(entity_id) + assert metadata is not None, f"no durable state found for {entity_id}" + + raw = metadata.get_state() + # The scheduler returns the entity payload as serialized JSON. + if isinstance(raw, str): + return DurableAgentState.from_json(raw) + assert isinstance(raw, dict), f"unexpected entity state payload: {type(raw)}" + return DurableAgentState.from_dict(raw) diff --git a/python/packages/durabletask/tests/integration_tests/test_15_dt_live_retention.py b/python/packages/durabletask/tests/integration_tests/test_15_dt_live_retention.py new file mode 100644 index 0000000..3478455 --- /dev/null +++ b/python/packages/durabletask/tests/integration_tests/test_15_dt_live_retention.py @@ -0,0 +1,525 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Live DTS persistence tests, not live LLM or graceful cancellation tests. + +Requires the installed worktree package, pytest, pytest-timeout, redis and +python-dotenv (for the existing conftest), and opentelemetry-sdk. The only required service is DTS at +ENDPOINT (default http://localhost:8080). No model credentials or sample marker. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import struct +import subprocess +import sys +import time +import uuid +import zlib +from collections import namedtuple +from collections.abc import Callable, Iterator +from contextlib import contextmanager, suppress +from copy import deepcopy +from datetime import datetime +from pathlib import Path +from queue import Empty, Queue +from threading import Event, Thread +from typing import Any + +import grpc +import pytest +from agent_framework import Content, Message +from durabletask.azuremanaged.client import DurableTaskSchedulerClient +from durabletask.client import OrchestrationStatus +from durabletask.entities import EntityInstanceId +from live_retention_worker import AGENT_NAME, DELIVERY_WINDOW_SECONDS, MAX_STATE_BYTES + +import agent_framework_durabletask +from agent_framework_durabletask import DurableAgentState, DurableHistoryProvider + +pytestmark = [pytest.mark.integration, pytest.mark.requires_dts, pytest.mark.timeout(150)] +WAIT_SECONDS = 30 +PACKAGE_ROOT = Path(__file__).resolve().parents[2] +WORKER_SCRIPT = Path(__file__).with_name("live_retention_worker.py") + + +class _CallDetails( + namedtuple("CallDetails", "method timeout metadata credentials wait_for_ready compression"), grpc.ClientCallDetails +): + pass + + +class _RpcDeadline(grpc.UnaryUnaryClientInterceptor): + def intercept_unary_unary(self, continuation: Any, details: Any, request: Any) -> Any: + # SDK get_entity/signal_entity have no timeout parameter. Bound the actual RPC, + # not just the polling loop around it, while preserving the DTS routing metadata. + bounded = _CallDetails( + details.method, + min(details.timeout, 3.0) if details.timeout is not None else 3.0, + details.metadata, + details.credentials, + details.wait_for_ready, + details.compression, + ) + return continuation(bounded, request) + + +@pytest.fixture +def live_taskhub(unique_taskhub: str) -> str: + # The existing fixture is module-scoped. Isolate parameter cases too, including + # pending work left behind when a deliberately killed worker's test fails. + return f"{unique_taskhub}-{uuid.uuid4().hex[:8]}" + + +@pytest.fixture +def live_client(dts_available: bool, dts_endpoint: str, live_taskhub: str) -> Iterator[DurableTaskSchedulerClient]: + assert dts_available + loaded_package = Path(agent_framework_durabletask.__file__).resolve().parent + assert loaded_package == PACKAGE_ROOT / "agent_framework_durabletask", ( + "Run with this exact worktree package installed, not an editable install from another checkout" + ) + client = DurableTaskSchedulerClient( + host_address=dts_endpoint, + taskhub=live_taskhub, + token_credential=None, + secure_channel=False, + interceptors=[_RpcDeadline()], + ) + with client: + yield client + + +class _WorkerProcess: + def __init__(self, endpoint: str, taskhub: str, artifacts: Path, block_message_id: str) -> None: + artifacts.mkdir() + self.artifacts = artifacts + self.events: Queue[dict[str, Any]] = Queue(maxsize=128) + self.exited = Event() + self.log = (artifacts / "worker.log").open("w", encoding="utf-8") + env = { + **os.environ, + "ENDPOINT": endpoint, + "TASKHUB": taskhub, + "DURABLE_AGENTS_DEPLOYMENT_MODE": "isolated_v2", + "PYTHONDONTWRITEBYTECODE": "1", + "PYTHONPATH": str(PACKAGE_ROOT) + os.pathsep + os.environ.get("PYTHONPATH", ""), + } + try: + self.process = subprocess.Popen( + [ + sys.executable, + "-B", + "-u", + str(WORKER_SCRIPT), + "--endpoint", + endpoint, + "--taskhub", + taskhub, + "--artifacts", + str(artifacts), + "--block-message-id", + block_message_id, + ], + cwd=artifacts, + env=env, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=self.log, + text=True, + encoding="utf-8", + shell=False, + ) + except BaseException: + self.log.close() + raise + self.reader = Thread(target=self._read_events, name="live-retention-control", daemon=True) + try: + self.reader.start() + except BaseException: + self.hard_stop() + if self.process.stdin is not None: + self.process.stdin.close() + if self.process.stdout is not None: + self.process.stdout.close() + self.log.close() + raise + + def _read_events(self) -> None: + try: + assert self.process.stdout is not None + while line := self.process.stdout.readline(4097): + if len(line) > 4096: + raise ValueError("Oversized worker control record") + self.events.put_nowait(json.loads(line)) + except Exception: + with suppress(Exception): + self.events.put_nowait({"event": "protocol_error"}) + finally: + self.exited.set() + + def event(self, expected: str) -> dict[str, Any]: + deadline = time.monotonic() + WAIT_SECONDS + while time.monotonic() < deadline: + try: + record = self.events.get(timeout=min(0.1, max(0.001, deadline - time.monotonic()))) + except Empty: + self.check_alive() + continue + assert record.get("event") == expected, f"Expected {expected}, received control event {record.get('event')}" + return record + raise TimeoutError(f"No {expected} control record within {WAIT_SECONDS}s. Inspect temporary worker.log") + + def check_alive(self) -> None: + if self.exited.is_set() or self.process.poll() is not None: + raise RuntimeError("Worker exited or control pipe failed. Inspect temporary worker.log") + + def command(self, command: str) -> None: + self.check_alive() + assert self.process.stdin is not None + self.process.stdin.write(json.dumps({"command": command}) + "\n") + self.process.stdin.flush() + + def hard_stop(self) -> None: + self.process.kill() + self.process.wait(timeout=10) + assert self.process.returncode is not None + + def close(self) -> None: + try: + if self.process.poll() is None: + with suppress(BrokenPipeError, OSError, RuntimeError): + self.command("stop") + try: + self.process.wait(timeout=10) + except subprocess.TimeoutExpired: + self.hard_stop() + finally: + if self.process.poll() is None: + self.hard_stop() + if self.process.stdin is not None: + with suppress(BrokenPipeError, OSError): + self.process.stdin.close() + self.reader.join(timeout=5) + if self.process.stdout is not None: + self.process.stdout.close() + self.log.close() + assert not self.reader.is_alive(), "Worker control thread did not terminate" + + def captured(self, message_id: str) -> list[dict[str, Any]]: + record = self.event("model_entered") + assert record["message_id"] == message_id + return json.loads((self.artifacts / f"model-{record['ordinal']}.json").read_text(encoding="utf-8")) + + def removed_measurement(self) -> int: + self.command("metrics") + self.event("metrics") + rows = json.loads((self.artifacts / "metrics.json").read_text(encoding="utf-8")) + for row in rows: + assert row["attributes"] == { + "mechanism": "pressure", + "outcome": "staged", + "commit_status": "not_attempted", + } + return sum(row["value"] for row in rows) + + +@contextmanager +def _worker(endpoint: str, hub: str, artifacts: Path, block: str = "") -> Iterator[_WorkerProcess]: + worker = _WorkerProcess(endpoint, hub, artifacts, block) + try: + worker.event("started") + yield worker + finally: + worker.close() + + +def _poll(worker: _WorkerProcess, probe: Callable[[], Any], description: str) -> Any: + deadline = time.monotonic() + WAIT_SECONDS + while time.monotonic() < deadline: + worker.check_alive() + if result := probe(): + return result + # A bounded backend poll, never a sleep used as evidence of completion. + worker.exited.wait(min(0.1, max(0, deadline - time.monotonic()))) + raise TimeoutError(f"DTS did not expose {description} within {WAIT_SECONDS}s") + + +def _snapshot(client: DurableTaskSchedulerClient, entity: EntityInstanceId) -> dict[str, Any]: + metadata = client.get_entity(entity) + assert metadata is not None, "Expected an existing backend entity" + raw = metadata.get_state() + state = json.loads(raw) if isinstance(raw, str) else raw + assert isinstance(state, dict), "Backend returned no JSON entity state" + return { + "id": str(metadata.id), + "last_modified": metadata.last_modified.isoformat(), + "backlog_queue_size": metadata.backlog_queue_size, + "state": state, + } + + +def _committed( + client: DurableTaskSchedulerClient, entity: EntityInstanceId, correlation: str, worker: _WorkerProcess +) -> dict[str, Any]: + def probe() -> dict[str, Any] | None: + metadata = client.get_entity(entity) + if metadata is None or not metadata.get_state(): + return None + raw = metadata.get_state() + state = json.loads(raw) if isinstance(raw, str) else raw + if correlation not in state.get("data", {}).get("completedCorrelations", {}): + return None + return _snapshot(client, entity) + + snapshot = _poll(worker, probe, f"completion receipt for {correlation}") + (worker.artifacts / f"committed-{correlation}.json").write_text(json.dumps(snapshot), encoding="utf-8") + receipt = snapshot["state"]["data"]["completedCorrelations"][correlation] + assert receipt["outcome"] == "succeeded", f"The real Agent failed for {correlation}" + return snapshot["state"] + + +def _equal(actual: Any, expected: Any, label: str) -> None: + # Compare entire payloads without leaking large media into pytest assertion output. + if actual != expected: + + def digest(value: Any) -> str: + return hashlib.sha256(json.dumps(value, sort_keys=True).encode()).hexdigest() + + pytest.fail(f"{label}: full JSON mismatch ({digest(actual)} != {digest(expected)})") + + +def _stored(raw: dict[str, Any]) -> list[dict[str, Any]]: + state = DurableAgentState.from_json(json.dumps(raw)) + return [ + message.to_chat_message().to_dict() for entry in state.data.conversation_history for message in entry.messages + ] + + +def _model_history(stored: list[dict[str, Any]]) -> list[dict[str, Any]]: + expected = deepcopy(stored) + for message in expected: + message.setdefault("additional_properties", {})["_attribution"] = { + "source_id": DurableHistoryProvider.DEFAULT_SOURCE_ID, + "source_type": "DurableHistoryProvider", + } + return expected + + +def _input(kind: str, turn: str) -> list[Message]: + def chunk(tag: bytes, value: bytes) -> bytes: + return struct.pack(">I", len(value)) + tag + value + struct.pack(">I", zlib.crc32(tag + value)) + + width, height = 128, 64 + pixels = hashlib.shake_256(b"durable-media-pressure").digest(width * height) + rows = b"".join(b"\x00" + pixels[row * width : (row + 1) * width] for row in range(height)) + png = ( + b"\x89PNG\r\n\x1a\n" + + chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 0, 0, 0, 0)) + + chunk(b"IDAT", zlib.compress(rows)) + + chunk(b"IEND", b"") + ) + if kind == "inline-png": + media = Content.from_data(png, "image/png") + elif kind == "inline-file": + media = Content.from_data((f"{turn}: inline document 界\n" * 256).encode(), "text/plain") + else: + raise ValueError(f"Unexpected media case: {kind}") + properties = {"application": {"type": "text", "values": [turn, "界", 0, False, None]}} + return [ + Message( + "user", + [Content.from_text(f"{turn}: " + "context " * 100), media], + message_id=f"{turn}-input", + author_name="media-user", + additional_properties=deepcopy(properties), + ), + Message( + "assistant", + [Content.from_function_call(f"{turn}-call", "lookup", arguments={"query": turn})], + message_id=f"{turn}-call-message", + author_name="planner", + additional_properties=deepcopy(properties), + ), + Message( + "tool", + [Content.from_function_result(f"{turn}-call", result={"records": [turn, "界", False]})], + message_id=f"{turn}-result-message", + author_name="lookup", + additional_properties=deepcopy(properties), + ), + ] + + +def _request(correlation: str, messages: list[Message]) -> dict[str, Any]: + return { + "message": "projected test input", + "correlationId": correlation, + "contextMessages": [message.to_dict() for message in messages], + } + + +def _answer(message_id: str) -> dict[str, Any]: + return Message( + "assistant", [f"answer:{message_id}"], message_id=f"{message_id}-answer", author_name="retention-model" + ).to_dict() + + +@pytest.mark.parametrize("kind", ["inline-png", "inline-file"]) +def test_live_media_pressure_cold_read_and_exact_model_input( + kind: str, live_client: DurableTaskSchedulerClient, dts_endpoint: str, live_taskhub: str, tmp_path: Path +) -> None: + entity = EntityInstanceId(entity=f"dafx-{AGENT_NAME}", key=uuid.uuid4().hex) + originals: dict[str, dict[str, Any]] = {} + previous: list[dict[str, Any]] = [] + raw: dict[str, Any] = {} + previous_removed = 0 + previous_measured = 0 + with _worker(dts_endpoint, live_taskhub, tmp_path / "warm") as warm: + for index in range(8): + correlation = f"turn-{index}" + inputs = _input(kind, correlation) + current_id = f"{correlation}-input" + live_client.signal_entity(entity, "run", _request(correlation, inputs)) + expected_input = [*_model_history(previous), *[message.to_dict() for message in inputs]] + _equal(warm.captured(current_id), expected_input, "model input") + raw = _committed(live_client, entity, correlation, warm) + expected_turn = [*[message.to_dict() for message in inputs], _answer(current_id)] + originals.update({message["message_id"]: message for message in expected_turn}) + retained = _stored(raw) + retained_ids = {message["message_id"] for message in retained} + assert {message["message_id"] for message in expected_turn} <= retained_ids + _equal( + retained, [value for key, value in originals.items() if key in retained_ids], "retained payload/order" + ) + for turn in range(index + 1): + pair = {f"turn-{turn}-call-message", f"turn-{turn}-result-message"} + assert pair <= retained_ids or pair.isdisjoint(retained_ids), "Pressure split an atomic tool pair" + removed = len(originals) - len(retained) + assert (raw["data"].get("truncation") or {}).get("evictedMessageCount", 0) == removed + measured = warm.removed_measurement() + assert measured - previous_measured == removed - previous_removed + assert len(json.dumps(raw)) < int(MAX_STATE_BYTES * 0.85) + previous, previous_removed, previous_measured = retained, removed, measured + assert previous_removed >= 4, "This must exercise real pressure eviction, not merely media serialization" + assert sum(message["role"] == "user" for message in previous) >= 2, "Retain older media for cold model replay" + assert len(raw["data"]["completedCorrelations"]) == len(raw["data"]["responseMailbox"]) == 8 + + with _worker(dts_endpoint, live_taskhub, tmp_path / "cold") as cold: + snapshot = _snapshot(live_client, entity) + (cold.artifacts / "cold-read.json").write_text(json.dumps(snapshot), encoding="utf-8") + _equal(snapshot["state"], raw, "cold backend read") + current = Message("user", ["next turn"], message_id="cold-input") + # Resending an evicted projected input must not resurrect it after a process restart. + evicted = set(originals) - {message["message_id"] for message in previous} + assert "turn-0-input" in evicted + live_client.signal_entity(entity, "run", _request("cold", [*_input(kind, "turn-0"), current])) + _equal(cold.captured("cold-input"), [*_model_history(previous), current.to_dict()], "cold model input") + final = _committed(live_client, entity, "cold", cold) + final_messages = _stored(final) + assert evicted.isdisjoint(message["message_id"] for message in final_messages) + assert {"cold-input", "cold-input-answer"} <= {message["message_id"] for message in final_messages} + originals.update({"cold-input": current.to_dict(), "cold-input-answer": _answer("cold-input")}) + final_ids = {message["message_id"] for message in final_messages} + _equal( + final_messages, [value for key, value in originals.items() if key in final_ids], "cold persisted payloads" + ) + for turn in range(8): + pair = {f"turn-{turn}-call-message", f"turn-{turn}-result-message"} + assert pair <= final_ids or pair.isdisjoint(final_ids), "Cold pressure split an atomic tool pair" + _equal( + {key: final["data"]["ingestedMessages"][key] for key in raw["data"]["ingestedMessages"]}, + raw["data"]["ingestedMessages"], + "cold replay preserves ingestion receipts", + ) + assert set(final["data"]["ingestedMessages"]) == {*raw["data"]["ingestedMessages"], "cold-input"} + total_removed = len(originals) - len(final_messages) + assert final["data"]["truncation"]["evictedMessageCount"] == total_removed + assert cold.removed_measurement() == total_removed - previous_removed + assert len(json.dumps(final)) < int(MAX_STATE_BYTES * 0.85) + _equal( + final["data"]["completedCorrelations"]["turn-0"], + raw["data"]["completedCorrelations"]["turn-0"], + "evicted turn completion receipt", + ) + for correlation, mailbox in raw["data"]["responseMailbox"].items(): + _equal(final["data"]["responseMailbox"][correlation], mailbox, "retained mailbox through pressure") + assert ( + datetime.fromisoformat(mailbox["expiresAt"]) - datetime.fromisoformat(mailbox["createdAt"]) + ).total_seconds() == DELIVERY_WINDOW_SECONDS + assert set(final["data"]["responseMailbox"]) == {*raw["data"]["responseMailbox"], "cold"} + + +def test_live_hard_stop_before_commit_repeats_effect_but_committed_duplicate_does_not( + live_client: DurableTaskSchedulerClient, dts_endpoint: str, live_taskhub: str, tmp_path: Path +) -> None: + entity = EntityInstanceId(entity=f"dafx-{AGENT_NAME}", key=uuid.uuid4().hex) + target = Message("user", ["simulated external effect"], message_id="target-input") + request = _request("target", [target]) + with _worker(dts_endpoint, live_taskhub, tmp_path / "interrupted", "target-input") as first: + seed = Message("user", ["establish committed baseline"], message_id="seed-input") + live_client.signal_entity(entity, "run", _request("seed", [seed])) + first.captured("seed-input") + baseline = _committed(live_client, entity, "seed", first) + live_client.signal_entity(entity, "run", request) # Accepted is not committed. + captured = first.captured("target-input") + _equal(captured, [*_model_history(_stored(baseline)), target.to_dict()], "interrupted model input") + first.hard_stop() # No model response, history after-hook, or set_state can finish. + snapshot = _snapshot(live_client, entity) + (tmp_path / "after-hard-stop.json").write_text(json.dumps(snapshot), encoding="utf-8") + _equal(snapshot["state"], baseline, "backend state after hard stop") + assert "target" not in snapshot["state"]["data"]["completedCorrelations"] + assert "target" not in snapshot["state"]["data"]["responseMailbox"] + + with _worker(dts_endpoint, live_taskhub, tmp_path / "retry", "target-input") as retry: + # The killed work item may redeliver before this explicit retry. Either must use + # committed state, and the same correlation must execute once in this new process. + live_client.signal_entity(entity, "run", request) + _equal(retry.captured("target-input"), captured, "retried model input") + _equal(_snapshot(live_client, entity)["state"], baseline, "blocked retry is still uncommitted") + retry.command("release") + committed = _committed(live_client, entity, "target", retry) + assert committed["data"]["completedCorrelations"]["target"]["outcome"] == "succeeded" + _equal(_stored(committed), [*_stored(baseline), target.to_dict(), _answer("target-input")], "committed retry") + retry.hard_stop() # Only after authoritative scheduler readback, never a warm-cache acknowledgement. + + with _worker(dts_endpoint, live_taskhub, tmp_path / "duplicate") as duplicate: + _equal(_snapshot(live_client, entity)["state"], committed, "post-commit cold read") + instance = live_client.schedule_new_orchestration( + "live_retention_duplicate", input={"key": entity.key, "request": request} + ) + barrier_completed = False + try: + + def completed_barrier() -> Any: + state = live_client.get_orchestration_state(instance) + if state is not None: + state.raise_if_failed() + if state.runtime_status == OrchestrationStatus.COMPLETED: + return state + return None + + barrier = _poll(duplicate, completed_barrier, "acknowledged duplicate signal/call") + barrier_completed = True + _equal( + json.loads(barrier.serialized_output), + committed["data"]["responseMailbox"]["target"]["response"], + "duplicate response", + ) + snapshot = _snapshot(live_client, entity) + (duplicate.artifacts / "duplicate-read.json").write_text(json.dumps(snapshot), encoding="utf-8") + _equal(snapshot["state"], committed, "duplicate must not rewrite completion or expiry") + assert not list(duplicate.artifacts.glob("model-*.json")), "Committed duplicate executed the model" + assert not list(duplicate.artifacts.glob("effect-*.json")), "Committed duplicate repeated the effect" + finally: + if not barrier_completed: + with suppress(grpc.RpcError): + live_client.terminate_orchestration(instance) + + effects = [json.loads(path.read_text(encoding="utf-8")) for path in tmp_path.glob("*/effect-*.json")] + assert sum(effect["message_id"] == "target-input" for effect in effects) == 2 + # The delivery window is not shortened to make duplicate suppression or pressure fit. + mailbox = committed["data"]["responseMailbox"]["target"] + window = datetime.fromisoformat(mailbox["expiresAt"]) - datetime.fromisoformat(mailbox["createdAt"]) + assert window.total_seconds() == DELIVERY_WINDOW_SECONDS diff --git a/python/packages/durabletask/tests/test_cancellation_boundaries.py b/python/packages/durabletask/tests/test_cancellation_boundaries.py new file mode 100644 index 0000000..f0954bf --- /dev/null +++ b/python/packages/durabletask/tests/test_cancellation_boundaries.py @@ -0,0 +1,390 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Deterministic interruption and JSON commit boundaries, not live worker shutdown tests.""" + +import asyncio +import json +from collections.abc import Awaitable, Sequence +from copy import deepcopy +from datetime import datetime, timezone +from types import SimpleNamespace +from typing import Any +from unittest.mock import Mock + +import pytest +from agent_framework import AgentResponse, AgentSession, ChatResponse, HistoryProvider, Message +from test_durable_history_provider import RecordingChatClient +from test_history_pipeline_revision import NonStreamingAgent +from test_revision_contract import JsonStateProvider + +from agent_framework_durabletask import AgentEntity, DurableAgentState, DurableHistoryProvider, RunRequest +from agent_framework_durabletask import _entities as entities +from agent_framework_durabletask._executors import ClientAgentExecutor +from agent_framework_durabletask._history_provider import current_durable_history_binding + + +class SimulatedWorkerStop(BaseException): + """An explicit process-boundary sentinel, not a claim about SDK shutdown plumbing.""" + + +class PhaseBarrier: + def __init__(self) -> None: + self.phase: str | None = None + self.entered = asyncio.Event() + self.release = asyncio.Event() + self.stop: SimulatedWorkerStop | None = None + self.observed_binding: Any = None + + async def wait(self, phase: str) -> None: + if phase != self.phase: + return + self.observed_binding = current_durable_history_binding() + self.entered.set() + await self.release.wait() + if self.stop is not None: + raise self.stop + + +async def await_boundary(task: asyncio.Task[Any], entered: asyncio.Event) -> None: + """Fail promptly if execution ends before its barrier, without clock-based waits.""" + waiter = asyncio.create_task(entered.wait()) + try: + await asyncio.wait({task, waiter}, return_when=asyncio.FIRST_COMPLETED) + if not entered.is_set(): + await task + pytest.fail("execution finished without reaching the required boundary") + finally: + if not waiter.done(): + waiter.cancel() + await asyncio.gather(waiter, return_exceptions=True) + + +class BarrierClient(RecordingChatClient): + def __init__(self, barrier: PhaseBarrier) -> None: + super().__init__() + self.barrier = barrier + self.effects: list[str] = [] + + def get_response(self, messages: Any, *, stream: bool = False, **kwargs: Any) -> Awaitable[ChatResponse]: + if stream: + raise TypeError("stream is not supported") + self.received_messages.append(deepcopy(list(messages))) + + async def get() -> ChatResponse: + # The externally visible attempt precedes the cancellable model await. + self.effects.append("model-side-effect") + await self.barrier.wait("model") + return ChatResponse(messages=[Message("assistant", ["boundary answer"], message_id="boundary-answer")]) + + return get() + + +def _agent(client: BarrierClient, **kwargs: Any) -> NonStreamingAgent: + chat_client: Any = client + return NonStreamingAgent(client=chat_client, **kwargs) + + +class _BoundaryHistory(DurableHistoryProvider): + def __init__(self, barrier: PhaseBarrier) -> None: + super().__init__(prune_excluded=False) + self.barrier = barrier + self.effects: list[str] = [] + self.sessions: list[AgentSession] = [] + self.agents: list[Any] = [] + + async def before_run(self, *, agent: Any, session: AgentSession, state: dict[str, Any], **kwargs: Any) -> None: + self.agents.append(agent) + self.sessions.append(session) + await super().before_run(agent=agent, session=session, state=state, **kwargs) + state["provider_control"] = {"visits": state.get("provider_control", {}).get("visits", 0) + 1} + self.effects.append("external-before-effect") + await self.barrier.wait("before_run") + + async def after_run(self, **kwargs: Any) -> None: + await super().after_run(**kwargs) + self.effects.append("external-after-effect") + await self.barrier.wait("after_run") + + +def projected_request(correlation: str = "interrupted") -> dict[str, Any]: + message = Message( + "user", + ["projected boundary input"], + message_id=f"{correlation}-input", + additional_properties={"json": [1, False]}, + ) + return {"message": message.text, "correlationId": correlation, "contextMessages": [message.to_dict()]} + + +def _initial_state() -> dict[str, Any]: + state = DurableAgentState() + session = AgentSession(session_id="revision-session", service_session_id="saved-service-id") + session.state = {"foreign": {"pending_approval": {"id": "keep", "approved": False}}} + state.data.session = session.to_dict() + state.data.ingested_messages = {"previous-input": ["previous-fingerprint"]} + state.data.ingested_positions = {"source": 4} + state.data.extension_data = {"control": {"keep": [1]}} + state.record_response( + "expired", + AgentResponse(messages=[Message("assistant", ["old delivery payload"])]), + now=datetime(2020, 1, 1, tzinfo=timezone.utc), + delivery_window_seconds=60, + ) + state.record_response( + "previous", + AgentResponse(messages=[Message("assistant", ["retained delivery payload"])]), + delivery_window_seconds=3600, + ) + return json.loads(state.to_json()) + + +@pytest.mark.parametrize("phase", ["before_run", "model", "after_run", "budget"]) +@pytest.mark.parametrize("interruption", ["task-cancel", "base-exception-stop"]) +async def test_interruption_rolls_back_every_local_slice_and_retry_can_repeat_external_effects( + phase: str, interruption: str, monkeypatch: pytest.MonkeyPatch +) -> None: + barrier = PhaseBarrier() + barrier.phase = phase + history = _BoundaryHistory(barrier) + client = BarrierClient(barrier) + agent: Any = _agent( + client=client, + name="boundary", + context_providers=[history], + default_options={"store": False, "conversation_id": "inactive-default-id"}, + ) + provider = JsonStateProvider(_initial_state()) + entity = AgentEntity(agent, state_provider=provider, max_state_bytes=1_000_000) + original_state = entity.state + original_agent = entity.agent + original_defaults = deepcopy(agent.default_options) + before = json.loads(json.dumps(provider.raw)) + request = projected_request() + real_budget = entities.enforce_budget + + async def budget(state: DurableAgentState, **kwargs: Any) -> int: + removed = await real_budget(state, **kwargs) + await barrier.wait("budget") + return removed + + monkeypatch.setattr(entities, "enforce_budget", budget) + task_final_bindings: list[Any] = [] + + async def execute() -> AgentResponse: + try: + return await entity.run(request) + finally: + # Inspect the cancelled task's own context, not merely its parent's ContextVar. + task_final_bindings.append(current_durable_history_binding()) + + task = asyncio.create_task(execute()) + try: + await await_boundary(task, barrier.entered) + assert not task.done() + assert provider.raw == before and provider.writes == 0 + staged = entity.state.to_dict()["data"] + assert "expired" not in staged["responseMailbox"], "TTL cleanup must actually have been staged" + assert "interrupted-input" in staged["ingestedMessages"] + if phase == "budget": + assert "interrupted" in staged["completedCorrelations"] + assert "interrupted" in staged["responseMailbox"] + assert staged["session"] != before["data"]["session"] + assert barrier.observed_binding is None + else: + assert barrier.observed_binding is not None + if phase == "after_run": + assert any(entry.get("correlationId") == "interrupted" for entry in staged["conversationHistory"]) + if interruption == "task-cancel": + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert task.cancelled() + else: + barrier.stop = SimulatedWorkerStop("explicit worker-stop boundary") + barrier.release.set() + with pytest.raises(SimulatedWorkerStop) as error: + await task + assert error.value is barrier.stop + assert not task.cancelled() + finally: + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + assert task_final_bindings == [None] + assert current_durable_history_binding() is None + assert history.agents[-1] is not original_agent, "exercise restoration of a real operation-local Agent clone" + assert history.sessions[-1].service_session_id == "saved-service-id" + assert entity.agent is original_agent and agent.default_options == original_defaults + assert entity.state is original_state and entity.state.to_dict() == before + assert provider.raw == before and provider.writes == 0 + assert entity.state.try_get_agent_response("interrupted") is None + assert "interrupted" not in provider.raw["data"]["completedCorrelations"] + assert "expired" in provider.raw["data"]["responseMailbox"] + assert history.effects.count("external-before-effect") == 1 + assert len(client.effects) == int(phase != "before_run") + + barrier.phase = None + barrier.stop = None + response = await entity.run(request) + assert response.text == "boundary answer" + assert history.effects.count("external-before-effect") == 2 + assert len(client.effects) == 1 + int(phase != "before_run") + assert provider.writes == 1 + assert "expired" not in provider.raw["data"]["responseMailbox"] + assert ( + provider.raw["data"]["completedCorrelations"]["expired"] == before["data"]["completedCorrelations"]["expired"] + ) + assert provider.raw["data"]["responseMailbox"]["previous"] == before["data"]["responseMailbox"]["previous"] + assert provider.raw["data"]["session"]["state"]["foreign"] == before["data"]["session"]["state"]["foreign"] + attempts = (len(client.effects), len(history.effects)) + cold_provider = JsonStateProvider(provider.raw) + cold = AgentEntity(agent, state_provider=cold_provider) + assert (await cold.run(request)).to_dict() == response.to_dict() + assert (len(client.effects), len(history.effects)) == attempts and cold_provider.writes == 0 + + +class _LostAcknowledgementStorage(JsonStateProvider): + def _set_state_dict(self, state: dict[str, Any]) -> None: + super()._set_state_dict(state) + # The JSON write has definitely happened, but the operation cannot know that. + raise OSError("storage acknowledgement lost after write") + + +async def test_unknown_commit_requires_fresh_json_read_and_suppresses_duplicate_execution() -> None: + barrier = PhaseBarrier() + client = BarrierClient(barrier) + agent: Any = _agent(client=client, name="unknown-commit") + provider = _LostAcknowledgementStorage(_initial_state()) + entity = AgentEntity(agent, state_provider=provider) + original = entity.state + before = deepcopy(provider.raw) + request = projected_request("unknown-commit") + + with pytest.raises(OSError, match="acknowledgement lost"): + await entity.run(request) + + assert provider.writes == 1 and len(client.effects) == 1 + assert entity.state is original and entity.state.to_dict() == before + assert entity.state.try_get_agent_response("unknown-commit") is None + assert current_durable_history_binding() is None + raw = json.loads(json.dumps(provider.raw)) + assert raw != before + committed = DurableAgentState.from_json(json.dumps(raw)).try_get_agent_response("unknown-commit") + assert committed is not None and committed.text == "boundary answer" + assert "unknown-commit" in raw["data"]["completedCorrelations"] + assert raw["data"]["responseMailbox"]["unknown-commit"]["response"] == committed.to_dict() + cold_provider = JsonStateProvider(raw) + cold_agent: Any = _agent(client=client, name="unknown-commit") + cold = AgentEntity(cold_agent, state_provider=cold_provider) + duplicate = await cold.run(request) + assert duplicate.to_dict() == committed.to_dict() + assert len(client.effects) == 1 and cold_provider.writes == 0 + assert cold_provider.raw == raw + + +class FailingExternalHistory(HistoryProvider): + def __init__(self, phase: str) -> None: + super().__init__("external-boundary") + self.phase: str | None = phase + self.loads = 0 + self.saved: list[list[Message]] = [] + + async def get_messages(self, session_id: str | None, **kwargs: Any) -> list[Message]: + self.loads += 1 + if self.phase == "load": + raise OSError("external load boundary failed") + return deepcopy([message for batch in self.saved for message in batch]) + + async def save_messages(self, session_id: str | None, messages: Sequence[Message], **kwargs: Any) -> None: + # An external append is not undone by an entity storage failure. + self.saved.append(deepcopy(list(messages))) + if self.phase == "store": + raise OSError("external store boundary failed") + + +class RejectedWriteStorage(JsonStateProvider): + def __init__(self) -> None: + super().__init__(_initial_state()) + self.attempts: list[dict[str, Any]] = [] + self.reject = True + + def _set_state_dict(self, state: dict[str, Any]) -> None: + self.attempts.append(json.loads(json.dumps(state))) + if self.reject: + raise OSError("entity storage write rejected") + super()._set_state_dict(state) + + +def failure_boundary(phase: str) -> tuple[AgentEntity, RejectedWriteStorage, FailingExternalHistory, BarrierClient]: + external = FailingExternalHistory(phase) + client = BarrierClient(PhaseBarrier()) + agent: Any = _agent(client=client, name="failure-boundary", context_providers=[external]) + provider = RejectedWriteStorage() + return AgentEntity(agent, state_provider=provider), provider, external, client + + +def assert_staged_not_committed(provider: RejectedWriteStorage, before: dict[str, Any], phase: str) -> None: + assert len(provider.attempts) == 1 and provider.writes == 0 + assert provider.raw == before + attempted = DurableAgentState.from_json(json.dumps(provider.attempts[0])) + failed = attempted.try_get_agent_response("provider-failed") + assert failed is not None + assert failed.additional_properties["durable_status"] == "error" + assert f"external {phase} boundary failed" in failed.text + assert any(content.error_code == "OSError" for message in failed.messages for content in message.contents) + assert "provider-failed" in attempted.data.completed_correlations + assert DurableAgentState.from_json(json.dumps(provider.raw)).try_get_agent_response("provider-failed") is None + + +class _JsonDTBackend: + """Signal acceptance and state reads only, with no fabricated response or completion.""" + + def __init__(self, provider: JsonStateProvider) -> None: + self.provider = provider + self.signals: list[Any] = [] + self.reads = 0 + + def signal_entity(self, *args: Any) -> None: + self.signals.append(args) + + def get_entity(self, entity_id: Any, *, include_state: bool) -> Any: + assert include_state + self.reads += 1 + return SimpleNamespace(get_state=lambda: json.dumps(self.provider.raw)) + + +@pytest.mark.parametrize("phase", ["load", "store"]) +async def test_external_failure_then_rejected_error_commit_is_invisible_to_real_dt_poller( + phase: str, monkeypatch: pytest.MonkeyPatch +) -> None: + entity, provider, external, client = failure_boundary(phase) + before = deepcopy(provider.raw) + request = projected_request("provider-failed") + with pytest.raises(OSError, match="entity storage write rejected"): + await entity.run(request) + assert_staged_not_committed(provider, before, phase) + assert entity.state.to_dict() == before and current_durable_history_binding() is None + assert external.loads == 1 + assert len(external.saved) == len(client.effects) == int(phase == "store") + backend: Any = _JsonDTBackend(provider) + sleep = Mock() + monkeypatch.setattr("agent_framework_durabletask._executors.time.sleep", sleep) + executor = ClientAgentExecutor(backend, max_poll_retries=3, poll_interval_seconds=0.01) + response = executor.run_durable_agent("failure-boundary", RunRequest.from_dict(request)) + assert [content.error_code for message in response.messages for content in message.contents] == ["response_timeout"] + assert backend.reads == 3 and len(backend.signals) == 1 and sleep.call_count == 3 + assert provider.raw == before and len(provider.attempts) == 1 + assert external.loads == 1, "polling is a reader, not a direct entity execution" + + # Permit a real error commit while the provider outage remains. That changes visibility. + provider.reject = False + failed = await entity.run(request) + assert failed.additional_properties["durable_status"] == "error" and provider.writes == 1 + calls = (external.loads, len(external.saved), len(client.effects)) + delivered = executor.run_durable_agent("failure-boundary", RunRequest.from_dict(request)) + assert delivered.to_dict() == failed.to_dict() and backend.reads == 4 + external.phase = None + cold = AgentEntity(entity.agent, state_provider=JsonStateProvider(provider.raw)) + assert (await cold.run(request)).to_dict() == failed.to_dict() + assert (external.loads, len(external.saved), len(client.effects)) == calls diff --git a/python/packages/durabletask/tests/test_completion_outcomes.py b/python/packages/durabletask/tests/test_completion_outcomes.py new file mode 100644 index 0000000..4e18095 --- /dev/null +++ b/python/packages/durabletask/tests/test_completion_outcomes.py @@ -0,0 +1,309 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Outcome receipts across expiry, old-state handling and trusted migration boundaries.""" + +import json +from copy import deepcopy +from datetime import datetime, timedelta, timezone +from typing import Any + +import pytest +from agent_framework import Agent, AgentResponse, Content, Message +from test_durable_history_provider import RecordingChatClient +from test_revision_contract import JsonStateProvider +from typing_extensions import Self + +from agent_framework_durabletask import AgentEntity, DurableAgentState, migrate_legacy_state, state_snapshot_digest +from agent_framework_durabletask import _durable_agent_state as state_module +from agent_framework_durabletask._response_utils import serialize_agent_response + +NOW = datetime(2026, 9, 11, 12, tzinfo=timezone.utc) +WINDOW = 60 +CORRELATION = "outcome-correlation" + + +class Clock(datetime): + current = NOW + + @classmethod + def now(cls, tz: Any = None) -> Self: + return cls.fromtimestamp(cls.current.timestamp(), tz or timezone.utc) + + +@pytest.fixture(autouse=True) +def clock(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(Clock, "current", NOW) + monkeypatch.setattr(state_module, "datetime", Clock) + + +def _response(kind: str) -> AgentResponse[Any]: + if kind == "error-status": + return AgentResponse(messages=[], additional_properties={"durable_status": "error"}) + if kind == "error-content": + return AgentResponse(messages=[Message("assistant", [Content.from_error(message="provider failed")])]) + if kind == "recovered-tool": + return AgentResponse( + messages=[ + Message("tool", [Content.from_error(message="recoverable lookup failure")]), + Message("assistant", ["recovered answer"]), + ] + ) + if kind == "approval": + return AgentResponse( + messages=[ + Message( + "assistant", + [ + Content.from_function_approval_request( + "approval-1", Content.from_function_call("call-1", "work") + ) + ], + ) + ] + ) + if kind == "empty": + return AgentResponse(messages=[]) + if kind == "structured-false": + return AgentResponse(messages=[], value=False) + return AgentResponse(messages=[Message("assistant", ["original answer"])], response_id="original-response") + + +def _state(kind: str = "success") -> DurableAgentState: + state = DurableAgentState() + state.record_response(CORRELATION, _response(kind), delivery_window_seconds=WINDOW, now=NOW) + return state + + +def _cold(state: DurableAgentState) -> DurableAgentState: + return DurableAgentState.from_json(state.to_json()) + + +@pytest.mark.parametrize( + "kind", ["success", "empty", "structured-false", "recovered-tool", "approval", "error-status", "error-content"] +) +@pytest.mark.parametrize("cleanup", [False, True]) +@pytest.mark.parametrize("cold", [False, True]) +def test_new_receipt_retains_invocation_outcome_without_payload_after_expiry( + kind: str, cleanup: bool, cold: bool +) -> None: + response = _response(kind) + state = _state(kind) + expected = "failed" if kind.startswith("error-") else "succeeded" + receipt = {"completedAt": NOW.isoformat(), "outcome": expected} + assert state.data.completed_correlations[CORRELATION] == receipt + Clock.current = NOW + timedelta(seconds=WINDOW, microseconds=-1) + delivered = state.try_get_agent_response(CORRELATION) + assert delivered is not None + assert serialize_agent_response(delivered) == serialize_agent_response(response) + + Clock.current = NOW + timedelta(seconds=WINDOW) + if cleanup: + state.expire_responses(now=Clock.current) + if cold: + state = _cold(state) + before = state.to_json() + expired = state.try_get_agent_response(CORRELATION) + assert expired is not None + assert expired.additional_properties == { + "durable_status": "already_completed", + "correlation_id": CORRELATION, + "durable_outcome": expected, + } + assert expired.messages[0].contents[0].error_code == "response_expired" + assert expired.response_id is None and expired.value is None and expired.continuation_token is None + assert state.to_json() == before + assert state.data.completed_correlations[CORRELATION] == receipt + assert bool(state.data.response_mailbox) is not cleanup + state.record_response( + CORRELATION, _response("error-status" if expected == "succeeded" else "success"), delivery_window_seconds=999 + ) + assert state.to_json() == before + + +@pytest.mark.parametrize("kind", ["success", "error-content"]) +@pytest.mark.parametrize("legacy", [False, True]) +def test_old_receipt_uses_only_independent_result_evidence_before_cleanup(kind: str, legacy: bool) -> None: + state = _state(kind) + receipt = state.data.completed_correlations[CORRELATION] + receipt.pop("outcome", None) + receipt["future"] = {"keep": [1]} + if legacy: + receipt["legacy"] = True + state = _cold(state) + original_receipt = deepcopy(state.data.completed_correlations[CORRELATION]) + Clock.current = NOW + timedelta(seconds=WINDOW) + expected = "failed" if kind == "error-content" else "unknown" if legacy else "succeeded" + before = state.to_json() + response = state.try_get_agent_response(CORRELATION) + assert response is not None and response.additional_properties["durable_outcome"] == expected + assert state.to_json() == before, "lookup cannot rewrite persisted evidence" + state.expire_responses(now=Clock.current) + state = _cold(state) + assert state.data.response_mailbox == {} + assert state.data.completed_correlations[CORRELATION] == { + **original_receipt, + **({"outcome": expected} if expected != "unknown" else {}), + } + expired = state.try_get_agent_response(CORRELATION) + assert expired is not None and expired.additional_properties["durable_outcome"] == expected + + +@pytest.mark.parametrize("cold", [False, True]) +async def test_old_unknown_receipt_still_suppresses_execution_and_survives_reset(cold: bool) -> None: + state = DurableAgentState() + state.data.completed_correlations[CORRELATION] = {"completedAt": NOW.isoformat(), "future": {"keep": True}} + original = deepcopy(state.to_dict()) + client: Any = RecordingChatClient() + provider = JsonStateProvider(_cold(state).to_dict() if cold else state.to_dict()) + entity = AgentEntity(Agent(client=client), state_provider=provider) + response = await entity.run({"message": "never rerun", "correlationId": CORRELATION}) + assert response.additional_properties["durable_outcome"] == "unknown" + assert response.additional_properties["durable_status"] == "already_completed" + assert client.received_messages == [] and provider.writes == 0 and provider.raw == original + entity.reset() + assert _cold(entity.state).data.completed_correlations == state.data.completed_correlations + assert entity.state.data.response_mailbox == {} + + +@pytest.mark.parametrize("invalid", [None, "", "success", "FAILED", "unknown", 1, False, [], {}]) +def test_invalid_present_outcome_is_rejected_at_read_and_warm_boundaries(invalid: Any) -> None: + state = _state() + state.data.response_mailbox.clear() + state.data.completed_correlations[CORRELATION]["outcome"] = invalid + raw = {"schemaVersion": "2.0.0", "data": {"completedCorrelations": deepcopy(state.data.completed_correlations)}} + with pytest.raises(ValueError, match="outcome"): + DurableAgentState.from_dict(raw) + with pytest.raises(ValueError, match="outcome"): + DurableAgentState.from_json(json.dumps(raw)) + with pytest.raises(ValueError, match="outcome"): + state.to_dict() + with pytest.raises(ValueError, match="outcome"): + state.prepare_for_write(delivery_window_seconds=WINDOW) + with pytest.raises(ValueError, match="outcome"): + state.try_get_agent_response(CORRELATION) + + +def _migrate(source: dict[str, Any], **kwargs: Any) -> DurableAgentState: + return migrate_legacy_state( + source, + source_digest=state_snapshot_digest(source), + source_session_id="original-session", + migration_id="outcome-migration", + ownership_transfer_id="authorized-transfer", + delivery_window_seconds=WINDOW, + now=NOW + timedelta(days=1), + **kwargs, + ) + + +@pytest.mark.parametrize("mailbox", [False, True]) +def test_known_outcome_import_requires_authoritative_evidence_and_preserves_completion_time(mailbox: bool) -> None: + source = _state("error-content").to_dict() + source["schemaVersion"] = "1.1.0" + source["data"]["completedCorrelations"][CORRELATION].pop("outcome", None) + if not mailbox: + source["data"].pop("responseMailbox") + before = deepcopy(source) + if mailbox: + result = _cold(_migrate(source, require_known_outcomes=True)) + assert result.data.completed_correlations[CORRELATION] == {"completedAt": NOW.isoformat(), "outcome": "failed"} + assert result.data.response_mailbox == before["data"]["responseMailbox"] + else: + with pytest.raises(ValueError, match="outcome.*evidence"): + _migrate(source, require_known_outcomes=True) + compatible = _cold(_migrate(source)) + assert compatible.data.completed_correlations == source["data"]["completedCorrelations"] + assert compatible.data.response_mailbox == {} + response = compatible.try_get_agent_response(CORRELATION) + assert response is not None and response.additional_properties["durable_outcome"] == "unknown" + assert source == before + + +def test_existing_mailbox_backfill_uses_original_completion_timestamp_not_migration_time() -> None: + source = _state().to_dict() + source["schemaVersion"] = "1.1.0" + source["data"].pop("completedCorrelations") + result = _cold(_migrate(source, require_known_outcomes=True)) + assert result.data.completed_correlations[CORRELATION] == { + "completedAt": NOW.isoformat(), + "outcome": "succeeded", + "legacy": True, + } + assert result.data.response_mailbox == source["data"]["responseMailbox"] + + +def test_partial_legacy_transcript_is_not_proof_of_success() -> None: + source: dict[str, Any] = { + "schemaVersion": "1.1.0", + "data": { + "conversationHistory": [ + { + "$type": "response", + "correlationId": CORRELATION, + "createdAt": NOW.isoformat(), + "messages": [], + } + ] + }, + } + with pytest.raises(ValueError, match="outcome.*evidence"): + _migrate(source, require_known_outcomes=True) + compatible = _cold(_migrate(source)) + assert "outcome" not in compatible.data.completed_correlations[CORRELATION] + compatible.expire_responses(now=NOW + timedelta(days=2)) + assert compatible.try_get_agent_response(CORRELATION).additional_properties["durable_outcome"] == "unknown" # type: ignore[union-attr] + + +def test_strict_entity_import_rejects_unknown_without_any_write_or_model_call() -> None: + source = DurableAgentState("1.1.0") + source.data.completed_correlations[CORRELATION] = {"completedAt": NOW.isoformat()} + provider = JsonStateProvider() + client: Any = RecordingChatClient() + entity = AgentEntity(Agent(client=client), state_provider=provider) + request = { + "source": source.to_dict(), + "sourceDigest": state_snapshot_digest(source.to_dict()), + "sourceSessionId": "original-session", + "destinationSessionId": provider.core_session_id, + "migrationId": "strict-import", + "ownershipTransferId": "authorized-transfer", + "requireKnownOutcomes": True, + } + with pytest.raises(ValueError, match="outcome.*evidence"): + entity.migrate(request) + assert provider.raw == {} and provider.writes == 0 and client.received_messages == [] + + +@pytest.mark.parametrize("status", ["accepted", "already_completed"]) +def test_new_completion_requires_outcome_not_an_acceptance_or_unavailable_status(status: str) -> None: + state = DurableAgentState() + response = AgentResponse(messages=[], additional_properties={"durable_status": status}) + before = state.to_dict() + with pytest.raises(ValueError, match="completion.*outcome"): + state.record_response(CORRELATION, response, delivery_window_seconds=WINDOW) + assert state.to_dict() == before + + +@pytest.mark.parametrize("formatted", [False, True]) +async def test_inner_acceptance_is_not_committed_as_success(formatted: bool) -> None: + class AcceptanceAgent: + name = "acceptance" + calls = 0 + + async def run(self, *, stream: bool = False, **kwargs: Any) -> AgentResponse: + if stream: + raise TypeError("stream is not supported") + self.calls += 1 + return AgentResponse( + messages=[Message("assistant", ["Request accepted"])] if formatted else [], + response_format={"type": "object"} if formatted else None, + additional_properties={"durable_status": "accepted"}, + ) + + agent: Any = AcceptanceAgent() + provider = JsonStateProvider() + entity = AgentEntity(agent, state_provider=provider) + with pytest.raises(ValueError, match="completion.*outcome"): + await entity.run({"message": "work", "correlationId": CORRELATION}) + assert agent.calls == 1 and provider.raw == {} and provider.writes == 0 + assert entity.state.try_get_agent_response(CORRELATION) is None diff --git a/python/packages/durabletask/tests/test_delivery_consumers_dt.py b/python/packages/durabletask/tests/test_delivery_consumers_dt.py new file mode 100644 index 0000000..317cde2 --- /dev/null +++ b/python/packages/durabletask/tests/test_delivery_consumers_dt.py @@ -0,0 +1,406 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Delivery consumers using real core responses, JSON reloads, and durable tasks.""" + +import json +from copy import deepcopy +from datetime import date, datetime, timezone +from typing import Any, cast +from unittest.mock import Mock + +import pytest +from agent_framework import AgentResponse, Content, ContinuationToken, Message +from durabletask.client import TaskHubGrpcClient +from durabletask.task import CompletableTask +from pydantic import BaseModel, ConfigDict, RootModel + +from agent_framework_durabletask import ( + DurableAgentState, + DurableAgentStateErrorResponse, + DurableAgentStateResponse, + RunRequest, + ensure_response_format, + load_agent_response, + serialize_agent_response, +) +from agent_framework_durabletask._executors import ClientAgentExecutor, DurableAgentTask + +CORRELATION_ID = "consumer-correlation" +HISTORICAL_TIME = datetime(2024, 1, 1, tzinfo=timezone.utc) + + +class Answer(BaseModel): + answer: int + + +def _response(*, value: Any = None, text: str = "Readable answer") -> AgentResponse[Any]: + return AgentResponse( + messages=[ + Message( + "tool", + [Content.from_function_result("call-1", result=[Content.from_text("lookup result")])], + author_name="lookup", + message_id="tool-message", + ), + Message( + "assistant", + [ + Content.from_text( + text, + annotations=[{"type": "citation", "title": "Source", "url": "https://example.test/source"}], + additional_properties={"provider": {"labels": ["content"]}}, + raw_representation=object(), + ) + ], + author_name="writer", + message_id="answer-message", + additional_properties={"provider": {"labels": ["message"]}}, + raw_representation=object(), + ), + ], + response_id="response-1", + agent_id="agent-1", + created_at=HISTORICAL_TIME.isoformat(), + finish_reason="stop", + usage_details={"input_token_count": 3, "output_token_count": 2, "total_token_count": 5}, + continuation_token=cast(ContinuationToken, {"cursor": {"pages": [1, 2]}}), + additional_properties={"provider": {"labels": ["response"]}}, + raw_representation=object(), + value=value, + ) + + +def _mailbox_state(response: AgentResponse[Any], *, expired: bool = False, cleanup: bool = False) -> str: + state = DurableAgentState() + state.data.conversation_history.append(DurableAgentStateResponse.from_run_response(CORRELATION_ID, response)) + state.record_response( + CORRELATION_ID, + response, + delivery_window_seconds=3600, + now=HISTORICAL_TIME if expired else None, + ) + if not expired: + state.data.conversation_history.clear() + if cleanup: + state.expire_responses() + return state.to_json() + + +def _client(state_json: str | None) -> tuple[ClientAgentExecutor, Mock]: + client = Mock(spec=TaskHubGrpcClient) + if state_json is None: + client.get_entity.return_value = None + else: + client.get_entity.return_value.get_state.return_value = state_json + return ClientAgentExecutor(client, max_poll_retries=3, poll_interval_seconds=0.01), client + + +def _task( + payload: dict[str, Any], response_format: type[BaseModel] | None, *, precompleted: bool = False +) -> DurableAgentTask: + child: CompletableTask[Any] = CompletableTask() + if precompleted: + child.complete(payload) + task = DurableAgentTask(child, response_format, CORRELATION_ID) + if not precompleted: + assert not task.is_complete + child.complete(payload) + return task + + +def _assert_expired(response: AgentResponse[Any], outcome: str = "succeeded") -> None: + assert response.additional_properties == { + "durable_status": "already_completed", + "correlation_id": CORRELATION_ID, + "durable_outcome": outcome, + } + content = response.messages[0].contents[0] + assert content.type == "error" + assert content.error_code == "response_expired" + assert content.message == "This request completed, but its response delivery window has expired." + assert response.value is None + + +@pytest.fixture +def sleep(monkeypatch: pytest.MonkeyPatch) -> Mock: + mocked = Mock() + monkeypatch.setattr("agent_framework_durabletask._executors.time.sleep", mocked) + return mocked + + +@pytest.mark.parametrize("value", [None, 0, False, "", [], {}, {"items": [{"answer": 42}]}]) +def test_public_serializer_and_loader_preserve_values_and_core_metadata(value: Any) -> None: + response = _response(value=deepcopy(value)) + expected = response.to_dict() + if value is not None: + expected["value"] = deepcopy(value) + + snapshot = json.loads(json.dumps(serialize_agent_response(response), allow_nan=False)) + + assert snapshot == expected + assert ("value" in snapshot) is (value is not None) + loaded = load_agent_response(snapshot) + assert isinstance(loaded, AgentResponse) + assert loaded.value == value + assert type(loaded.value) is type(value) + assert loaded.to_dict() == response.to_dict() + assert all(isinstance(message, Message) for message in loaded.messages) + assert all(isinstance(content, Content) for message in loaded.messages for content in message.contents) + assert loaded.messages[1].author_name == "writer" + assert loaded.messages[1].message_id == "answer-message" + assert loaded.messages[1].contents[0].annotations == response.messages[1].contents[0].annotations + assert loaded.messages[0].contents[0].items == response.messages[0].contents[0].items + assert "raw_representation" not in snapshot + assert "raw_representation" not in snapshot["messages"][1] + assert "raw_representation" not in snapshot["messages"][1]["contents"][0] + assert load_agent_response(loaded) is loaded + + +def test_public_serializer_uses_json_mode_for_pydantic_values() -> None: + class DatedAnswer(BaseModel): + answer: int + day: date + + response = _response(value=DatedAnswer(answer=42, day=date(2026, 9, 8))) + snapshot = json.loads(json.dumps(serialize_agent_response(response), allow_nan=False)) + + assert snapshot["value"] == {"answer": 42, "day": "2026-09-08"} + assert load_agent_response(snapshot).value == snapshot["value"] + + +@pytest.mark.parametrize("response_format", [Answer, Answer.model_json_schema()]) +def test_public_serializer_captures_a_lazy_structured_value(response_format: Any) -> None: + response = AgentResponse(messages=[Message("assistant", ['{"answer":42}'])], response_format=response_format) + + snapshot = json.loads(json.dumps(serialize_agent_response(response))) + + assert snapshot["value"] == {"answer": 42} + assert load_agent_response(snapshot).value == {"answer": 42} + + +@pytest.mark.parametrize("response_format", [None, Answer]) +def test_client_reads_full_mailbox_response_after_cold_reload_and_transcript_pruning( + response_format: type[BaseModel] | None, sleep: Mock +) -> None: + response = _response(value={"answer": 42}) + state_json = _mailbox_state(response) + assert json.loads(state_json)["data"]["conversationHistory"] == [] + executor, client = _client(state_json) + + result = executor.run_durable_agent( + "consumer", RunRequest(message="question", correlation_id=CORRELATION_ID, response_format=response_format) + ) + + assert result.to_dict() == response.to_dict() + if response_format is None: + assert result.value == {"answer": 42} + else: + assert isinstance(result.value, Answer) + assert result.value.answer == 42 + client.signal_entity.assert_called_once() + entity_id = client.signal_entity.call_args.args[0] + client.get_entity.assert_called_once_with(entity_id, include_state=True) + sleep.assert_called_once_with(0.01) + + +@pytest.mark.parametrize("version", ["1.0.0", "1.1.0"]) +@pytest.mark.parametrize("failed", [False, True]) +def test_client_retains_legacy_lookup_and_does_not_reparse_legacy_errors( + version: str, failed: bool, sleep: Mock +) -> None: + response = _response(text='{"answer":42}') + if failed: + response.messages[1].contents.append(Content.from_error(message="Provider failed", error_code="RuntimeError")) + state = DurableAgentState(schema_version=version) + entry_type = DurableAgentStateErrorResponse if failed else DurableAgentStateResponse + state.data.conversation_history.append(entry_type.from_run_response(CORRELATION_ID, response)) + state_json = state.to_json() + executor, client = _client(state_json) + + result = executor.run_durable_agent( + "consumer", RunRequest(message="question", correlation_id=CORRELATION_ID, response_format=Answer) + ) + + assert result.text == response.text + assert result.messages[1].author_name == "writer" + assert result.messages[1].message_id == "answer-message" + assert result.usage_details == response.usage_details + if failed: + assert result.messages[1].contents[-1].error_code == "RuntimeError" + assert result.value is None + else: + assert isinstance(result.value, Answer) + assert result.value.answer == 42 + client.get_entity.assert_called_once() + sleep.assert_called_once_with(0.01) + assert json.loads(state_json)["schemaVersion"] == version + + +@pytest.mark.parametrize("response_format", [None, Answer]) +@pytest.mark.parametrize("cleanup", [False, True]) +@pytest.mark.parametrize("failed", [False, True]) +def test_expired_client_delivery_is_terminal_on_the_first_read( + response_format: type[BaseModel] | None, cleanup: bool, failed: bool, sleep: Mock +) -> None: + original = _response(value={"answer": 42}) + if failed: + original.additional_properties["durable_status"] = "error" + executor, client = _client(_mailbox_state(original, expired=True, cleanup=cleanup)) + + result = executor.run_durable_agent( + "consumer", RunRequest(message="question", correlation_id=CORRELATION_ID, response_format=response_format) + ) + + _assert_expired(result, "failed" if failed else "succeeded") + client.signal_entity.assert_called_once() + client.get_entity.assert_called_once() + sleep.assert_called_once_with(0.01) + + +@pytest.mark.parametrize("response_format", [None, Answer]) +@pytest.mark.parametrize("precompleted", [False, True]) +def test_task_reconstructs_snapshot_value_and_metadata( + response_format: type[BaseModel] | None, precompleted: bool +) -> None: + response = _response(value={"answer": 42}) + payload = json.loads(_mailbox_state(response))["data"]["responseMailbox"][CORRELATION_ID]["response"] + + task = _task(payload, response_format, precompleted=precompleted) + + assert task.is_complete and not task.is_failed + result = task.get_result() + assert isinstance(result, AgentResponse) + assert result.to_dict() == response.to_dict() + assert result.messages[1].author_name == "writer" + if response_format is None: + assert result.value == {"answer": 42} + else: + assert isinstance(result.value, Answer) + assert result.value.answer == 42 + + +@pytest.mark.parametrize("precompleted", [False, True]) +@pytest.mark.parametrize("cleanup", [False, True]) +def test_task_returns_expired_status_instead_of_failing_schema_validation(precompleted: bool, cleanup: bool) -> None: + state = DurableAgentState.from_json(_mailbox_state(_response(), expired=True, cleanup=cleanup)) + expired = state.try_get_agent_response(CORRELATION_ID) + assert isinstance(expired, AgentResponse) + + task = _task(json.loads(json.dumps(serialize_agent_response(expired))), Answer, precompleted=precompleted) + + assert task.is_complete and not task.is_failed + _assert_expired(task.get_result()) + + +@pytest.mark.parametrize("terminal_kind", ["error", "already_completed"]) +@pytest.mark.parametrize("text", ["not JSON", '{"answer":0}']) +def test_terminal_response_formats_skip_all_messages_and_status_only_results(terminal_kind: str, text: str) -> None: + response = _response(text=text) + if terminal_kind == "error": + response.messages[1].contents.append(Content.from_error(message="Provider failed", error_code="RuntimeError")) + else: + response.additional_properties["durable_status"] = "already_completed" + snapshot = json.loads(json.dumps(serialize_agent_response(response))) + + direct = load_agent_response(deepcopy(snapshot)) + ensure_response_format(Answer, CORRELATION_ID, direct) + executor, _ = _client(None) + polled = executor._handle_agent_response(load_agent_response(deepcopy(snapshot)), Answer, CORRELATION_ID) + task = _task(deepcopy(snapshot), Answer) + + assert task.is_complete and not task.is_failed + for result in (direct, polled, task.get_result()): + assert result.to_dict() == response.to_dict() + assert result.value is None + + +def test_response_format_validates_the_saved_value_not_conflicting_text() -> None: + response = _response(value={"answer": 42}, text='{"answer":0}') + + ensure_response_format(Answer, CORRELATION_ID, response) + + assert isinstance(response.value, Answer) + assert response.value.answer == 42 + assert response.messages[1].text == '{"answer":0}' + + +def test_response_format_does_not_replace_an_invalid_saved_value_with_valid_text() -> None: + response = _response(value={"wrong": 42}, text='{"answer":0}') + + with pytest.raises(ValueError): + ensure_response_format(Answer, CORRELATION_ID, response) + + +def test_response_format_keeps_a_matching_pydantic_value() -> None: + value = Answer(answer=42) + response = _response(value=value) + + ensure_response_format(Answer, CORRELATION_ID, response) + + assert response.value is value + + +def test_response_format_uses_json_validation_for_saved_strict_models() -> None: + class StrictAnswer(BaseModel): + model_config = ConfigDict(strict=True) + day: date + coordinates: tuple[int, int] + + value = StrictAnswer(day=date(2026, 9, 8), coordinates=(1, 2)) + payload = json.loads(json.dumps(serialize_agent_response(_response(value=value)))) + response = load_agent_response(payload) + + ensure_response_format(StrictAnswer, CORRELATION_ID, response) + + assert isinstance(response.value, StrictAnswer) + assert response.value == value + + +@pytest.mark.parametrize("value", [0, False, "", [], {}]) +def test_response_format_preserves_falsey_saved_values(value: Any) -> None: + class SavedValue(RootModel[Any]): + pass + + response = _response(value=deepcopy(value)) + + ensure_response_format(SavedValue, CORRELATION_ID, response) + + assert isinstance(response.value, SavedValue) + assert response.value.root == value + assert type(response.value.root) is type(value) + + +def test_response_format_override_still_controls_unparsed_responses() -> None: + class OtherAnswer(BaseModel): + missing: str + + response = AgentResponse(messages=[Message("assistant", ['{"answer":42}'])], response_format=OtherAnswer) + + ensure_response_format(Answer, CORRELATION_ID, response) + + assert isinstance(response.value, Answer) + assert response.value.answer == 42 + + +def test_successful_invalid_schema_still_fails_validation() -> None: + response = _response(text='{"wrong":42}') + executor, _ = _client(None) + + with pytest.raises(ValueError): + ensure_response_format(Answer, CORRELATION_ID, response) + polled = executor._handle_agent_response(load_agent_response(response.to_dict()), Answer, CORRELATION_ID) + assert polled.messages[0].contents[0].error_code == "response_processing_error" + task = _task(response.to_dict(), Answer) + assert task.is_complete and task.is_failed + + +def test_missing_response_keeps_the_bounded_timeout_behavior(sleep: Mock) -> None: + executor, client = _client(None) + + result = executor.run_durable_agent( + "consumer", RunRequest(message="question", correlation_id=CORRELATION_ID, response_format=Answer) + ) + + assert result.messages[0].contents[0].error_code == "response_timeout" + assert client.get_entity.call_count == executor.max_poll_retries + assert sleep.call_count == executor.max_poll_retries diff --git a/python/packages/durabletask/tests/test_delivery_state.py b/python/packages/durabletask/tests/test_delivery_state.py new file mode 100644 index 0000000..2e871c0 --- /dev/null +++ b/python/packages/durabletask/tests/test_delivery_state.py @@ -0,0 +1,691 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""State-only delivery regressions using core responses and real JSON deserialization.""" + +import json +from copy import deepcopy +from datetime import datetime, timedelta, timezone +from typing import Any, cast + +import pytest +from agent_framework import AgentResponse, Annotation, Content, ContinuationToken, Message +from pydantic import BaseModel + +from agent_framework_durabletask import migrate_legacy_state, state_snapshot_digest +from agent_framework_durabletask._durable_agent_state import ( + DurableAgentState, + DurableAgentStateEntryJsonType, + DurableAgentStateResponse, + DurableAgentStateTextContent, + DurableAgentStateUnknownEntry, +) +from agent_framework_durabletask._history_provider import replayable_entries +from agent_framework_durabletask._message_identity import message_identity + +DELIVERY_WINDOW_SECONDS = 60 +HISTORICAL_TIME = datetime(2024, 1, 1, tzinfo=timezone.utc) +CORRELATION_ID = "correlation-1" +SOURCE_SESSION_ID = "@dafx-delivery@legacy-source" + + +def _migrate_legacy_payload(payload: dict[str, Any]) -> DurableAgentState: + return migrate_legacy_state( + payload, + source_digest=state_snapshot_digest(payload), + source_session_id=SOURCE_SESSION_ID, + migration_id="delivery-migration-1", + ownership_transfer_id="delivery-transfer-1", + delivery_window_seconds=DELIVERY_WINDOW_SECONDS, + ) + + +def _response(*, value: Any = None) -> AgentResponse[Any]: + """Use public core 1.16 constructor arguments, not attributes invented by a mock.""" + annotations: list[Annotation] = [ + { + "type": "citation", + "title": "Source", + "url": "https://example.test/source", + "annotated_regions": [{"type": "text_span", "start_index": 0, "end_index": 6}], + "additional_properties": {"pages": [2, 3]}, + } + ] + return AgentResponse( + messages=[ + Message( + "assistant", + [ + Content.from_text( + "answer", + annotations=annotations, + additional_properties={"nested": {"labels": ["content"]}}, + raw_representation=object(), + ), + Content.from_function_call("call-1", "lookup", arguments={"ids": [1, 2]}, informational_only=True), + Content.from_text_reasoning( + id="reasoning-1", + text="reasoning summary", + protected_data="opaque-protected-payload", + additional_properties={"provider": {"sequence": [1]}}, + ), + ], + author_name="planner", + message_id="message-1", + additional_properties={"nested": {"labels": ["message"]}}, + raw_representation=object(), + ), + Message( + "tool", + [ + Content.from_function_result( + "call-1", + result=[ + Content.from_text("tool result"), + Content.from_data(b"data", "application/octet-stream"), + ], + additional_properties={"provider": {"sequence": [1]}}, + ) + ], + author_name="lookup", + message_id="message-2", + ), + ], + response_id="response-1", + agent_id="agent-1", + created_at=HISTORICAL_TIME.isoformat(), + finish_reason="stop", + usage_details={ + "input_token_count": 12, + "output_token_count": 8, + "total_token_count": 20, + "cache_creation_input_token_count": 2, + "cache_read_input_token_count": 3, + "reasoning_output_token_count": 4, + }, + value=value, + continuation_token=cast(ContinuationToken, {"cursor": {"pages": [1, 2]}}), + additional_properties={"nested": {"labels": ["response"]}}, + raw_representation=object(), + ) + + +def _record(state: DurableAgentState, response: AgentResponse[Any], *, now: datetime | None = None) -> None: + state.record_response(CORRELATION_ID, response, delivery_window_seconds=DELIVERY_WINDOW_SECONDS, now=now) + + +def _legacy_payload(version: str) -> dict[str, Any]: + return { + "schemaVersion": version, + "data": { + "conversationHistory": [ + { + "$type": "request", + "correlationId": CORRELATION_ID, + "createdAt": HISTORICAL_TIME.isoformat(), + "messages": [{"role": "user", "contents": [], "messageId": "legacy-known-id"}], + }, + { + "$type": "response", + "correlationId": CORRELATION_ID, + "createdAt": HISTORICAL_TIME.isoformat(), + "messages": [ + { + "role": "assistant", + "contents": [{"$type": "text", "text": "surviving legacy transcript"}], + "messageId": "legacy-response", + "authorName": "legacy-agent", + } + ], + "usage": {"inputTokenCount": 3, "outputTokenCount": 2, "totalTokenCount": 5}, + }, + ] + }, + } + + +def _assert_expired(response: AgentResponse[Any] | None) -> None: + assert isinstance(response, AgentResponse) + assert response.additional_properties["durable_status"] == "already_completed" + assert response.additional_properties["correlation_id"] == CORRELATION_ID + assert len(response.messages) == 1 + assert response.messages[0].role == "system" + assert len(response.messages[0].contents) == 1 + content = response.messages[0].contents[0] + assert content.type == "error" + assert content.error_code == "response_expired" + assert content.error_details is None + assert response.response_id is None + assert response.agent_id is None + assert response.continuation_token is None + assert response.value is None + + +def test_record_response_snapshots_core_metadata_and_reloads_real_response() -> None: + response = _response() + expected = json.loads(json.dumps(response.to_dict(), allow_nan=False)) + now = datetime.now(timezone.utc) + state = DurableAgentState() + + _record(state, response, now=now) + + payload = json.loads(state.to_json()) + assert payload["schemaVersion"] == "2.0.0" + assert payload["data"]["conversationHistory"] == [] + assert payload["data"]["responseMailbox"][CORRELATION_ID] == { + "response": expected, + "createdAt": now.isoformat(), + "expiresAt": (now + timedelta(seconds=DELIVERY_WINDOW_SECONDS)).isoformat(), + } + assert payload["data"]["completedCorrelations"][CORRELATION_ID] == { + "completedAt": now.isoformat(), + "outcome": "succeeded", + } + assert expected["type"] == "agent_response" + assert expected["response_id"] == "response-1" + assert expected["agent_id"] == "agent-1" + assert expected["created_at"] == HISTORICAL_TIME.isoformat() + assert expected["finish_reason"] == "stop" + assert expected["usage_details"]["cache_read_input_token_count"] == 3 + assert expected["continuation_token"] == {"cursor": {"pages": [1, 2]}} + assert expected["additional_properties"] == {"nested": {"labels": ["response"]}} + assert "raw_representation" not in expected + + # Exercise core's own reader as well as the durable state's reader. + direct = AgentResponse.from_dict(deepcopy(payload["data"]["responseMailbox"][CORRELATION_ID]["response"])) + restored = DurableAgentState.from_json(json.dumps(payload)) + delivered = restored.try_get_agent_response(CORRELATION_ID) + assert isinstance(delivered, AgentResponse) + assert direct.to_dict() == delivered.to_dict() == expected + assert all(isinstance(message, Message) for message in delivered.messages) + assert all(isinstance(content, Content) for message in delivered.messages for content in message.contents) + assert delivered.messages[0].author_name == "planner" + assert delivered.messages[0].message_id == "message-1" + assert delivered.messages[0].contents[0].annotations == response.messages[0].contents[0].annotations + assert delivered.messages[0].contents[1].call_id == "call-1" + assert delivered.messages[0].contents[1].informational_only is True + assert delivered.messages[0].contents[2].id == "reasoning-1" + assert delivered.messages[0].contents[2].protected_data == "opaque-protected-payload" + assert delivered.messages[1].contents[0].items == response.messages[1].contents[0].items + assert restored.try_get_agent_response("unknown-correlation") is None + + +@pytest.mark.parametrize("value", [{"items": [{"answer": 42}]}, {}, [], 0, False, "structured result"]) +def test_record_response_preserves_structured_value_not_just_core_to_dict(value: Any) -> None: + """Core 1.16 keeps value in private state, so to_dict equality alone cannot prove delivery fidelity.""" + response = _response(value=deepcopy(value)) + state = DurableAgentState() + _record(state, response) + + snapshot = json.loads(state.to_json())["data"]["responseMailbox"][CORRELATION_ID]["response"] + assert "value" in snapshot, "record_response lost the public structured result" + assert snapshot["value"] == value + assert type(snapshot["value"]) is type(value) + direct = AgentResponse.from_dict(snapshot) + assert direct.value == value + restored = DurableAgentState.from_json(state.to_json()) + delivered = restored.try_get_agent_response(CORRELATION_ID) + assert isinstance(delivered, AgentResponse) + assert delivered.value == value + + +def test_structured_model_value_is_stored_as_inline_json() -> None: + class Result(BaseModel): + answer: int + citations: list[str] + + value = Result(answer=42, citations=["source-1"]) + expected = value.model_dump(mode="json") + state = DurableAgentState() + _record(state, _response(value=value)) + value.citations.append("caller edit") + + snapshot = json.loads(state.to_json())["data"]["responseMailbox"][CORRELATION_ID]["response"] + assert snapshot["value"] == expected + assert AgentResponse.from_dict(snapshot).value == expected + delivered = DurableAgentState.from_json(state.to_json()).try_get_agent_response(CORRELATION_ID) + assert isinstance(delivered, AgentResponse) + assert delivered.value == expected + + +def test_lazy_structured_value_is_captured_before_caller_text_changes() -> None: + response = AgentResponse( + messages=[Message("assistant", ['{"answer":42}'])], + response_format={"type": "object", "properties": {"answer": {"type": "integer"}}}, + ) + state = DurableAgentState() + # Do not access response.value first: recording must capture the public lazy value itself. + _record(state, response) + response.messages[0].contents[0].text = '{"answer":0}' + + snapshot = json.loads(state.to_json())["data"]["responseMailbox"][CORRELATION_ID]["response"] + assert snapshot["value"] == {"answer": 42} + assert AgentResponse.from_dict(snapshot).value == {"answer": 42} + delivered = DurableAgentState.from_json(state.to_json()).try_get_agent_response(CORRELATION_ID) + assert isinstance(delivered, AgentResponse) + assert delivered.value == {"answer": 42} + + +def test_mutating_caller_response_and_transcript_cannot_change_mailbox() -> None: + response = _response() + expected = json.loads(json.dumps(response.to_dict(), allow_nan=False)) + state = DurableAgentState() + transcript = DurableAgentStateResponse.from_run_response(CORRELATION_ID, response) + state.data.conversation_history.append(transcript) + _record(state, response) + + response.response_id = "changed" + response.agent_id = "changed" + response.created_at = datetime.now(timezone.utc).isoformat() + response.finish_reason = "length" + response.additional_properties["nested"]["labels"].append("changed") + assert response.usage_details is not None + response.usage_details["input_token_count"] = 999 + assert response.continuation_token is not None + cast(dict[str, Any], response.continuation_token)["cursor"]["pages"].append(999) + response.messages[0].author_name = "changed" + response.messages[0].message_id = "changed" + response.messages[0].additional_properties["nested"]["labels"].append("changed") + response.messages[0].contents[0].text = "changed" + response.messages[0].contents[0].additional_properties["nested"]["labels"].append("changed") + assert response.messages[0].contents[0].annotations is not None + response.messages[0].contents[0].annotations[0]["additional_properties"]["pages"].append(999) + cast(dict[str, Any], response.messages[0].contents[1].arguments)["ids"].append(999) + response.messages[0].contents[2].id = "changed" + response.messages[0].contents[2].protected_data = "changed" + assert response.messages[1].contents[0].items is not None + response.messages[1].contents[0].items[0].text = "changed tool result" + response.messages.clear() + transcript.messages[0].contents = [DurableAgentStateTextContent("compacted, not the original answer")] + transcript.messages[0].extension_data = {"_excluded": True} + transcript.messages.clear() + state.data.conversation_history.clear() + + restored = DurableAgentState.from_json(state.to_json()) + assert restored.data.response_mailbox[CORRELATION_ID]["response"] == expected + delivered = restored.try_get_agent_response(CORRELATION_ID) + assert isinstance(delivered, AgentResponse) + assert delivered.to_dict() == expected + + +def test_structured_value_is_detached_from_the_caller() -> None: + value = {"nested": {"items": [1, 2]}} + expected = deepcopy(value) + response = _response(value=value) + state = DurableAgentState() + _record(state, response) + value["nested"]["items"].append(3) + assert response.value is not None + response.value["nested"]["items"].append(4) + + restored = DurableAgentState.from_json(state.to_json()) + delivered = restored.try_get_agent_response(CORRELATION_ID) + assert isinstance(delivered, AgentResponse) + assert delivered.value == expected + + +def test_poll_results_and_serialized_delivery_records_are_detached() -> None: + state = DurableAgentState() + _record(state, _response()) + state.data.ingested_messages = {"message-1": ["a" * 64], "legacy-known-id": None} + state = DurableAgentState.from_json(state.to_json()) + expected = json.loads(state.to_json()) + + delivered = state.try_get_agent_response(CORRELATION_ID) + assert isinstance(delivered, AgentResponse) + delivered.additional_properties["nested"]["labels"].append("caller edit") + delivered.messages[0].contents[0].additional_properties["nested"]["labels"].append("caller edit") + delivered.messages.clear() + exported = state.to_dict() + exported["data"]["responseMailbox"][CORRELATION_ID]["response"]["messages"].clear() + exported["data"]["completedCorrelations"][CORRELATION_ID]["completedAt"] = "changed" + exported["data"]["ingestedMessages"]["message-1"].clear() + + assert state.to_dict() == expected + second = state.try_get_agent_response(CORRELATION_ID) + assert isinstance(second, AgentResponse) + assert second.to_dict() == expected["data"]["responseMailbox"][CORRELATION_ID]["response"] + + +@pytest.mark.parametrize("cleanup", [False, True], ids=["before-cleanup", "after-cleanup"]) +@pytest.mark.parametrize("original_error", [False, True], ids=["success", "error"]) +def test_expiry_returns_completed_status_never_the_surviving_transcript(cleanup: bool, original_error: bool) -> None: + state = DurableAgentState() + response = _response() + if original_error: + response.messages = [ + Message( + "system", + [ + Content.from_error( + message="original provider failure", + error_code="previous_response_not_found", + error_details="original provider details", + ) + ], + ) + ] + state.data.conversation_history.append(DurableAgentStateResponse.from_run_response(CORRELATION_ID, response)) + now = datetime.now(timezone.utc) + _record(state, response, now=now - timedelta(seconds=DELIVERY_WINDOW_SECONDS + 1)) + receipt = deepcopy(state.data.completed_correlations[CORRELATION_ID]) + transcript = deepcopy(state.to_dict()["data"]["conversationHistory"]) + if cleanup: + state.expire_responses(now=now) + + restored = DurableAgentState.from_json(state.to_json()) + assert bool(restored.data.response_mailbox) is not cleanup + before_poll = restored.to_json() + _assert_expired(restored.try_get_agent_response(CORRELATION_ID)) + assert restored.to_json() == before_poll + assert restored.data.completed_correlations[CORRELATION_ID] == receipt + assert restored.to_dict()["data"]["conversationHistory"] == transcript + assert restored.try_get_agent_response("never-completed") is None + + +def test_expiry_boundary_removes_only_due_payloads_not_receipts() -> None: + state = DurableAgentState() + _record(state, _response(), now=HISTORICAL_TIME) + state.record_response( + "later", + _response(), + delivery_window_seconds=DELIVERY_WINDOW_SECONDS, + now=HISTORICAL_TIME + timedelta(seconds=30), + ) + receipts = deepcopy(state.data.completed_correlations) + boundary = HISTORICAL_TIME + timedelta(seconds=DELIVERY_WINDOW_SECONDS) + state.expire_responses(now=boundary - timedelta(microseconds=1)) + assert set(state.data.response_mailbox) == {CORRELATION_ID, "later"} + state.expire_responses(now=boundary) + assert set(state.data.response_mailbox) == {"later"} + assert state.data.completed_correlations == receipts + state.expire_responses(now=boundary + timedelta(seconds=30)) + assert state.data.response_mailbox == {} + assert DurableAgentState.from_json(state.to_json()).data.completed_correlations == receipts + + +@pytest.mark.parametrize("expired", [False, True]) +def test_duplicate_record_does_not_replace_or_reopen_a_completed_response(expired: bool) -> None: + state = DurableAgentState() + now = datetime.now(timezone.utc) + _record(state, _response(), now=now) + if expired: + state.expire_responses(now=now + timedelta(seconds=DELIVERY_WINDOW_SECONDS)) + state = DurableAgentState.from_json(state.to_json()) + before = state.to_json() + replacement = AgentResponse(messages=[Message("assistant", ["must not replace the original"])]) + _record(state, replacement, now=now + timedelta(days=1)) + assert state.to_json() == before + if expired: + _assert_expired(state.try_get_agent_response(CORRELATION_ID)) + + +def test_version_two_does_not_poll_transcript_without_delivery_evidence() -> None: + state = DurableAgentState.from_dict(_legacy_payload("2.0.0")) + assert state.try_get_agent_response(CORRELATION_ID) is None + + +@pytest.mark.parametrize("version", ["1.0.0", "1.1.0"]) +@pytest.mark.parametrize("kind", ["response", "errorResponse"]) +def test_legacy_reader_round_trip_and_polling_do_not_upgrade_state(version: str, kind: str) -> None: + payload = _legacy_payload(version) + payload["data"]["conversationHistory"][1]["$type"] = kind + state = DurableAgentState.from_dict(deepcopy(payload)) + restored = DurableAgentState.from_json(state.to_json()) + assert restored.to_dict() == payload + delivered = restored.try_get_agent_response(CORRELATION_ID) + assert isinstance(delivered, AgentResponse) + assert delivered.text == "surviving legacy transcript" + assert delivered.messages[0].author_name == "legacy-agent" + assert delivered.usage_details == {"input_token_count": 3, "output_token_count": 2, "total_token_count": 5} + assert restored.try_get_agent_response("never-completed") is None + assert restored.to_dict() == payload + + +@pytest.mark.parametrize("version", ["1.0.0", "1.1.0"]) +def test_legacy_conversion_records_a_fresh_grace_window_not_a_historical_original(version: str) -> None: + payload = _legacy_payload(version) + original = deepcopy(payload) + legacy_response = DurableAgentState.from_dict(payload).try_get_agent_response(CORRELATION_ID) + assert isinstance(legacy_response, AgentResponse) + before = datetime.now(timezone.utc) + state = _migrate_legacy_payload(payload) + after = datetime.now(timezone.utc) + + restored = DurableAgentState.from_json(state.to_json()) + assert restored.schema_version == "2.0.0" + mailbox = restored.data.response_mailbox[CORRELATION_ID] + created_at = datetime.fromisoformat(mailbox["createdAt"]) + assert before <= created_at <= after + assert created_at != HISTORICAL_TIME + assert datetime.fromisoformat(mailbox["expiresAt"]) - created_at == timedelta(seconds=DELIVERY_WINDOW_SECONDS) + assert restored.data.completed_correlations[CORRELATION_ID] == {"completedAt": mailbox["createdAt"], "legacy": True} + assert restored.data.ingested_messages == {"legacy-known-id": None} + assert restored.data.unknown_fields["migration"] == { + "id": "delivery-migration-1", + "sourceDigest": state_snapshot_digest(original), + "sourceSessionId": SOURCE_SESSION_ID, + "ownershipTransferId": "delivery-transfer-1", + "createdAt": mailbox["createdAt"], + } + assert restored.data.session == {"session_id": SOURCE_SESSION_ID, "state": {}} + assert payload == original + delivered = restored.try_get_agent_response(CORRELATION_ID) + assert isinstance(delivered, AgentResponse) + assert delivered.to_dict() == legacy_response.to_dict() + assert delivered.created_at == HISTORICAL_TIME.isoformat() + + first_conversion = restored.to_json() + restored.prepare_for_write(delivery_window_seconds=DELIVERY_WINDOW_SECONDS * 2) + assert restored.to_json() == first_conversion + restored.expire_responses(now=datetime.fromisoformat(mailbox["expiresAt"])) + expired = DurableAgentState.from_json(restored.to_json()) + after_expiry = expired.to_json() + expired.prepare_for_write(delivery_window_seconds=DELIVERY_WINDOW_SECONDS * 2) + assert expired.to_json() == after_expiry + _assert_expired(expired.try_get_agent_response(CORRELATION_ID)) + + +@pytest.mark.parametrize("version", ["1.0.0", "1.1.0"]) +@pytest.mark.parametrize("position", [0, 3]) +def test_scalar_legacy_ingestion_cannot_be_migrated_without_evidence(version: str, position: int) -> None: + payload = _legacy_payload(version) + payload["futureRoot"] = {"opaque": [1]} + payload["data"]["ingestedPositions"] = {"source": position} + original = deepcopy(payload) + state = DurableAgentState.from_dict(payload) + before = state.to_json() + for _ in range(2): + with pytest.raises(ValueError, match="ingestedPositions.*recorded delivery evidence"): + state.prepare_for_write(delivery_window_seconds=DELIVERY_WINDOW_SECONDS) + with pytest.raises(ValueError, match="ingestedPositions.*recorded delivery evidence"): + _migrate_legacy_payload(payload) + assert state.to_json() == before + assert payload == original + assert state.data.response_mailbox == {} + assert state.data.completed_correlations == {} + assert state.data.ingested_messages == {} + # Refusing the writer upgrade must not prevent legacy read-only polling. + delivered = DurableAgentState.from_json(state.to_json()).try_get_agent_response(CORRELATION_ID) + assert isinstance(delivered, AgentResponse) + assert delivered.text == "surviving legacy transcript" + + +@pytest.mark.parametrize("version", ["1.0.0", "1.1.0", "2.0.0", "2.7.3"]) +def test_unknown_root_data_and_entry_properties_survive_reload_and_explicit_migration(version: str) -> None: + payload = _legacy_payload(version) + payload["futureRoot"] = {"nested": [1, {"keep": True}]} + payload["data"]["futureData"] = {"nested": [2, {"keep": None}]} + payload["data"]["session"] = { + "session_id": SOURCE_SESSION_ID, + "owner": "custom-provider", + "state": {"external": {"messages": [{"custom": "owned data"}], "cursor": [3, 4]}}, + } + known_entry = deepcopy(payload["data"]["conversationHistory"][1]) + history: list[dict[str, Any]] = [] + for kind in DurableAgentStateEntryJsonType: + entry = deepcopy(known_entry) + entry["$type"] = kind.value + entry["correlationId"] = kind.value + entry["messages"][0]["contents"][0]["text"] = kind.value + entry["futureEntry"] = {"nested": [kind.value, {"keep": False}]} + entry["extensionData"] = {"existing": {"keep": True}} + if kind not in (DurableAgentStateEntryJsonType.RESPONSE, DurableAgentStateEntryJsonType.ERROR_RESPONSE): + entry.pop("usage") + if kind == DurableAgentStateEntryJsonType.COMPACTION: + entry.pop("correlationId") + history.append(entry) + opaque = { + "$type": "future-owner-entry", + "correlationId": "opaque", + "messages": [{"role": "assistant", "contents": [{"$type": "text", "text": "do not replay"}]}], + "futureEntry": {"nested": [None, {"keep": "opaque"}]}, + } + history.insert(1, opaque) + payload["data"]["conversationHistory"] = history + state = DurableAgentState.from_dict(deepcopy(payload)) + state = DurableAgentState.from_json(state.to_json()) + assert state.to_dict() == payload + assert isinstance(state.data.conversation_history[1], DurableAgentStateUnknownEntry) + replayed = [entry.messages[index].text for entry, index in replayable_entries(state.data.conversation_history)] + assert replayed == ["request", "response", "compaction"] + assert state.try_get_agent_response("opaque") is None + + if version.startswith("1."): + state = _migrate_legacy_payload(payload) + elif version == "2.0.0": + state.prepare_for_write(delivery_window_seconds=DELIVERY_WINDOW_SECONDS) + else: + # Future revisions remain readable without permitting a write or downgrading the source. + with pytest.raises(ValueError, match="Only 2.0.0 is writable"): + state.prepare_for_write(delivery_window_seconds=DELIVERY_WINDOW_SECONDS) + assert state.to_dict() == payload + upgraded = DurableAgentState.from_json(state.to_json()).to_dict() + assert upgraded["schemaVersion"] == ("2.0.0" if version.startswith("1.") else version) + assert upgraded["futureRoot"] == payload["futureRoot"] + for key in ("futureData", "session", "conversationHistory"): + assert upgraded["data"][key] == payload["data"][key] + + +def test_ingestion_hash_lists_and_legacy_known_id_markers_survive_json_reload() -> None: + first = Message("user", ["first"], message_id="same-id") + changed = Message("user", ["changed"], message_id="same-id") + hashes = [message_identity(first), message_identity(changed)] + assert hashes[0] != hashes[1] + payload = { + "schemaVersion": "2.0.0", + "data": {"conversationHistory": [], "ingestedMessages": {"same-id": hashes, "legacy-known-id": None}}, + } + state = DurableAgentState.from_dict(payload) + assert DurableAgentState.from_json(state.to_json()).to_dict() == payload + + +@pytest.mark.parametrize("version", [None, False, 2, "", "0.1.0", "3.0.0", "2.0", "2.0.0-preview", "2.0.0\n"]) +def test_unsupported_or_malformed_version_fails_without_resetting_input(version: Any) -> None: + payload = _legacy_payload("1.1.0") + payload["schemaVersion"] = version + original = deepcopy(payload) + with pytest.raises(ValueError, match="schemaVersion"): + DurableAgentState.from_dict(payload) + with pytest.raises(ValueError, match="schemaVersion"): + DurableAgentState.from_json(json.dumps(payload)) + assert payload == original + + +def test_missing_version_fails_without_resetting_existing_history() -> None: + payload = _legacy_payload("1.1.0") + del payload["schemaVersion"] + original = deepcopy(payload) + with pytest.raises(ValueError, match="missing schemaVersion"): + DurableAgentState.from_dict(payload) + with pytest.raises(ValueError, match="missing schemaVersion"): + DurableAgentState.from_json(json.dumps(payload)) + assert payload == original + + +@pytest.mark.parametrize("data", [None, False, 0, "", [], [1]]) +def test_non_object_data_is_not_silently_reset(data: Any) -> None: + payload = {"schemaVersion": "2.0.0", "data": data} + with pytest.raises(ValueError, match="data"): + DurableAgentState.from_dict(payload) + + +@pytest.mark.parametrize("field", ["responseMailbox", "completedCorrelations", "ingestedMessages"]) +@pytest.mark.parametrize("value", [None, False, 0, "", [], "not-an-object", [1]]) +def test_malformed_delivery_containers_fail_on_initial_read_including_falsy_values(field: str, value: Any) -> None: + """An explicitly malformed field must not be normalized to an empty receipt store.""" + payload = {"schemaVersion": "2.0.0", "data": {"conversationHistory": [], field: value}} + original = deepcopy(payload) + with pytest.raises(ValueError): + DurableAgentState.from_dict(payload) + assert payload == original + + +@pytest.mark.parametrize("field", ["responseMailbox", "completedCorrelations"]) +@pytest.mark.parametrize("value", [None, False, 0, "", [], "not-an-entry"]) +def test_delivery_record_must_be_an_object_at_initial_read(field: str, value: Any) -> None: + payload = {"schemaVersion": "2.0.0", "data": {field: {CORRELATION_ID: value}}} + with pytest.raises(ValueError): + DurableAgentState.from_dict(payload) + + +@pytest.mark.parametrize( + ("record_name", "required_field"), + [ + ("responseMailbox", "response"), + ("responseMailbox", "createdAt"), + ("responseMailbox", "expiresAt"), + ("completedCorrelations", "completedAt"), + ], +) +def test_required_delivery_record_fields_are_checked_before_polling(record_name: str, required_field: str) -> None: + state = DurableAgentState() + _record(state, _response()) + payload = json.loads(state.to_json()) + del payload["data"][record_name][CORRELATION_ID][required_field] + original = deepcopy(payload) + with pytest.raises(ValueError): + DurableAgentState.from_dict(payload) + assert payload == original + + +@pytest.mark.parametrize( + ("record_name", "field"), + [("responseMailbox", "createdAt"), ("responseMailbox", "expiresAt"), ("completedCorrelations", "completedAt")], +) +@pytest.mark.parametrize("value", [None, False, 0, [], "", "not-a-timestamp"]) +def test_invalid_delivery_timestamps_fail_at_initial_read(record_name: str, field: str, value: Any) -> None: + state = DurableAgentState() + _record(state, _response()) + payload = json.loads(state.to_json()) + payload["data"][record_name][CORRELATION_ID][field] = value + with pytest.raises(ValueError): + DurableAgentState.from_dict(payload) + + +@pytest.mark.parametrize("response", [None, [], "{}", {}, {"type": "other", "messages": []}]) +def test_invalid_inline_response_fails_at_initial_read(response: Any) -> None: + state = DurableAgentState() + _record(state, _response()) + payload = json.loads(state.to_json()) + payload["data"]["responseMailbox"][CORRELATION_ID]["response"] = response + with pytest.raises(ValueError): + DurableAgentState.from_dict(payload) + + +@pytest.mark.parametrize("legacy", [None, 0, 1, "true", [], {}]) +def test_legacy_receipt_marker_must_be_boolean_at_initial_read(legacy: Any) -> None: + payload = { + "schemaVersion": "2.0.0", + "data": { + "completedCorrelations": {CORRELATION_ID: {"completedAt": HISTORICAL_TIME.isoformat(), "legacy": legacy}} + }, + } + with pytest.raises(ValueError): + DurableAgentState.from_dict(payload) + + +@pytest.mark.parametrize("fingerprints", [False, 0, "a" * 64, {}, [None], [1], ["a" * 64, False]]) +def test_ingestion_record_rejects_anything_but_hash_lists_or_legacy_null(fingerprints: Any) -> None: + payload = {"schemaVersion": "2.0.0", "data": {"ingestedMessages": {"message-id": fingerprints}}} + with pytest.raises(ValueError): + DurableAgentState.from_dict(payload) diff --git a/python/packages/durabletask/tests/test_deployment_gate_review.py b/python/packages/durabletask/tests/test_deployment_gate_review.py new file mode 100644 index 0000000..e003fb1 --- /dev/null +++ b/python/packages/durabletask/tests/test_deployment_gate_review.py @@ -0,0 +1,137 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Deployment acknowledgement validation for the shared configuration and worker.""" + +from typing import Any +from unittest.mock import Mock + +import pytest +from durabletask.worker import TaskHubGrpcWorker + +from agent_framework_durabletask import DurableAIAgentWorker +from agent_framework_durabletask import _configuration as configuration_module +from agent_framework_durabletask import _worker as worker_module +from agent_framework_durabletask._configuration import validate_runtime_deployment + +_ENVIRONMENT_VARIABLE = "DURABLE_AGENTS_DEPLOYMENT_MODE" +_INVALID_MODES = ("", "isolated_v1", "mixed", "ISOLATED_V2", " isolated_v2", "isolated_v2 ", "isolated_v2\n") + + +def test_missing_deployment_mode_explains_the_required_acknowledgement(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv(_ENVIRONMENT_VARIABLE, raising=False) + + with pytest.raises(ValueError) as error: + validate_runtime_deployment() + + message = str(error.value) + assert "Schema 2 requires an isolated task hub/deployment with upgraded clients" in message + assert "Old workflow histories must remain on the old engine" in message + assert "deployment_mode='isolated_v2'" in message + assert _ENVIRONMENT_VARIABLE in message + assert "explicit operator acknowledgement" in message + assert "not runtime proof" in message + assert "cannot detect peer workers" in message + + +def test_none_deployment_mode_reads_the_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(_ENVIRONMENT_VARIABLE, "isolated_v2") + validate_runtime_deployment() + validate_runtime_deployment(deployment_mode=None) + + +@pytest.mark.parametrize("environment_mode", [None, "", "mixed"]) +def test_explicit_valid_mode_overrides_missing_or_invalid_environment( + monkeypatch: pytest.MonkeyPatch, environment_mode: str | None +) -> None: + if environment_mode is None: + monkeypatch.delenv(_ENVIRONMENT_VARIABLE, raising=False) + else: + monkeypatch.setenv(_ENVIRONMENT_VARIABLE, environment_mode) + validate_runtime_deployment(deployment_mode="isolated_v2") + + +def test_explicit_mode_never_reads_the_environment(monkeypatch: pytest.MonkeyPatch) -> None: + getenv = Mock(side_effect=AssertionError("Explicit deployment mode must not read the environment")) + with monkeypatch.context() as scoped: + scoped.setattr(configuration_module.os, "getenv", getenv) + validate_runtime_deployment(deployment_mode="isolated_v2") + with pytest.raises(ValueError, match="isolated_v2"): + validate_runtime_deployment(deployment_mode="") + getenv.assert_not_called() + + +@pytest.mark.parametrize("deployment_mode", _INVALID_MODES) +def test_invalid_environment_mode_is_rejected(monkeypatch: pytest.MonkeyPatch, deployment_mode: str) -> None: + monkeypatch.setenv(_ENVIRONMENT_VARIABLE, deployment_mode) + with pytest.raises(ValueError, match="isolated_v2"): + validate_runtime_deployment() + + +@pytest.mark.parametrize("deployment_mode", [*_INVALID_MODES, False, 2, ["isolated_v2"]]) +def test_invalid_explicit_mode_is_not_overridden_by_valid_environment( + monkeypatch: pytest.MonkeyPatch, deployment_mode: Any +) -> None: + monkeypatch.setenv(_ENVIRONMENT_VARIABLE, "isolated_v2") + with pytest.raises(ValueError, match="isolated_v2"): + validate_runtime_deployment(deployment_mode=deployment_mode) + + +def test_worker_missing_mode_fails_before_configuration_or_registry_changes(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv(_ENVIRONMENT_VARIABLE, raising=False) + native = Mock(spec=TaskHubGrpcWorker) + agent_configuration = Mock(side_effect=AssertionError("Agent configuration ran before deployment validation")) + retention = Mock(side_effect=AssertionError("Retention configuration ran before deployment validation")) + monkeypatch.setattr(worker_module, "validate_agent_configuration", agent_configuration) + monkeypatch.setattr(worker_module, "validate_retention", retention) + host = DurableAIAgentWorker.__new__(DurableAIAgentWorker) + + with pytest.raises(ValueError, match="isolated_v2"): + DurableAIAgentWorker.__init__(host, native) + + assert vars(host) == {} + assert native.mock_calls == [] + agent_configuration.assert_not_called() + retention.assert_not_called() + + +@pytest.mark.parametrize("deployment_mode", _INVALID_MODES) +def test_worker_rejects_explicit_invalid_mode_despite_valid_environment( + monkeypatch: pytest.MonkeyPatch, deployment_mode: str +) -> None: + monkeypatch.setenv(_ENVIRONMENT_VARIABLE, "isolated_v2") + native = Mock(spec=TaskHubGrpcWorker) + with pytest.raises(ValueError, match="isolated_v2"): + DurableAIAgentWorker(native, deployment_mode=deployment_mode) + assert native.mock_calls == [] + + +@pytest.mark.parametrize("deployment_mode", _INVALID_MODES) +def test_worker_rejects_invalid_environment_mode(monkeypatch: pytest.MonkeyPatch, deployment_mode: str) -> None: + monkeypatch.setenv(_ENVIRONMENT_VARIABLE, deployment_mode) + native = Mock(spec=TaskHubGrpcWorker) + with pytest.raises(ValueError, match="isolated_v2"): + DurableAIAgentWorker(native) + assert native.mock_calls == [] + + +@pytest.mark.parametrize("source", ["explicit", "environment"]) +def test_worker_accepts_isolated_mode_and_preserves_entity_names(monkeypatch: pytest.MonkeyPatch, source: str) -> None: + kwargs: dict[str, Any] = {} + if source == "explicit": + monkeypatch.delenv(_ENVIRONMENT_VARIABLE, raising=False) + kwargs["deployment_mode"] = "isolated_v2" + else: + monkeypatch.setenv(_ENVIRONMENT_VARIABLE, "isolated_v2") + native = Mock(spec=TaskHubGrpcWorker) + native.add_entity.return_value = "dafx-assistant" + host = DurableAIAgentWorker(native, **kwargs) + + # The private worker factory relies on the host's completed deployment validation. + monkeypatch.delenv(_ENVIRONMENT_VARIABLE, raising=False) + agent = Mock(context_providers=None) + agent.name = "assistant" + host.add_agent(agent) + + assert host.registered_agent_names == ["assistant"] + native.add_entity.assert_called_once() + assert native.add_entity.call_args.args[0].__name__ == "dafx-assistant" diff --git a/python/packages/durabletask/tests/test_durable_agent_state.py b/python/packages/durabletask/tests/test_durable_agent_state.py index d3a36c9..fc29cd4 100644 --- a/python/packages/durabletask/tests/test_durable_agent_state.py +++ b/python/packages/durabletask/tests/test_durable_agent_state.py @@ -156,7 +156,7 @@ class TestDurableAgentState: def test_schema_version(self) -> None: """Test that schema version is set correctly.""" state = DurableAgentState() - assert state.schema_version == "1.1.0" + assert state.schema_version == "2.0.0" def test_to_dict_serialization(self) -> None: """Test that to_dict produces correct structure.""" @@ -165,13 +165,13 @@ def test_to_dict_serialization(self) -> None: assert "schemaVersion" in data assert "data" in data - assert data["schemaVersion"] == "1.1.0" + assert data["schemaVersion"] == "2.0.0" assert "conversationHistory" in data["data"] def test_from_dict_deserialization(self) -> None: """Test that from_dict restores state correctly.""" original_data = { - "schemaVersion": "1.1.0", + "schemaVersion": "1.2.0", "data": { "conversationHistory": [ { @@ -191,7 +191,7 @@ def test_from_dict_deserialization(self) -> None: state = DurableAgentState.from_dict(original_data) - assert state.schema_version == "1.1.0" + assert state.schema_version == "1.2.0" assert len(state.data.conversation_history) == 1 assert isinstance(state.data.conversation_history[0], DurableAgentStateRequest) @@ -218,15 +218,16 @@ def test_round_trip_serialization(self) -> None: assert len(restored.data.conversation_history) == len(state.data.conversation_history) assert restored.data.conversation_history[0].correlation_id == "test-456" - def test_function_call_round_trip_preserves_string_arguments(self) -> None: - """Function call arguments should remain strings across durable state replay.""" + @pytest.mark.parametrize("arguments", ['{"location":"Chicago"}', '{\n "location": "Chicago"\n}', '{"location":']) + def test_function_call_round_trip_preserves_string_arguments(self, arguments: str) -> None: + """Replay preserves the original argument string, including whitespace or partial JSON.""" original = Message( role="assistant", contents=[ Content.from_function_call( call_id="call-123", name="get_weather", - arguments='{"location":"Chicago"}', + arguments=arguments, ) ], ) @@ -235,7 +236,7 @@ def test_function_call_round_trip_preserves_string_arguments(self) -> None: restored = durable_message.to_chat_message() assert restored.contents[0].type == "function_call" - assert restored.contents[0].arguments == '{"location": "Chicago"}' + assert restored.contents[0].arguments == arguments def test_function_call_content_supports_legacy_mapping_arguments(self) -> None: """Existing persisted mapping arguments should still restore successfully.""" @@ -466,15 +467,23 @@ def test_unknown_content_from_plain_dict_unchanged(self) -> None: assert unknown.content == {"some": "data"} - def test_unknown_content_to_ai_content_fallback_on_invalid_type_dict(self) -> None: - """Test that to_ai_content falls back when dict has 'type' but is not valid Content.""" - invalid = {"type": "bogus_not_a_real_content_type", "extra": "stuff"} - unknown = DurableAgentStateUnknownContent(content=invalid) + def test_unknown_content_to_ai_content_preserves_future_type(self) -> None: + """Core accepts arbitrary content type strings and ignores unknown envelope fields.""" + future = { + "type": "bogus_not_a_real_content_type", + "extra": "stuff", + "additional_properties": {"opaque": [1]}, + } + unknown = DurableAgentStateUnknownContent(content=future) result = unknown.to_ai_content() - assert result.type == "unknown" - assert result.additional_properties == {"content": invalid} + assert result.type == future["type"] + assert result.additional_properties == {"opaque": [1]} + assert not hasattr(result, "extra") + result.additional_properties["opaque"].append(2) + assert unknown.to_dict()["content"] == future + assert future["additional_properties"] == {"opaque": [1]} def test_from_ai_content_unknown_type_produces_serializable_state(self) -> None: """Test that unknown content types in message conversion produce JSON-serializable state.""" diff --git a/python/packages/durabletask/tests/test_durable_entities.py b/python/packages/durabletask/tests/test_durable_entities.py index d8cd5e3..9deab6a 100644 --- a/python/packages/durabletask/tests/test_durable_entities.py +++ b/python/packages/durabletask/tests/test_durable_entities.py @@ -340,7 +340,10 @@ def mock_run(*args, stream=False, **kwargs): # Validate callback arguments stream_calls = callback.stream_mock.await_args_list for expected_update, recorded_call in zip(updates, stream_calls, strict=True): - assert recorded_call.args[0] is expected_update + recorded_update = recorded_call.args[0] + assert recorded_update is not expected_update + assert recorded_update.to_dict() == expected_update.to_dict() + assert recorded_update.contents[0] is not expected_update.contents[0] context = recorded_call.args[1] assert context.agent_name == "StreamingAgent" assert context.correlation_id == "corr-stream-1" @@ -350,6 +353,9 @@ def mock_run(*args, stream=False, **kwargs): final_call = callback.response_mock.await_args assert final_call is not None final_response, final_context = final_call.args + assert final_response is not result + assert final_response.to_dict() == result.to_dict() + assert final_response.messages[0] is not result.messages[0] assert final_context.agent_name == "StreamingAgent" assert final_context.correlation_id == "corr-stream-1" assert final_context.session_id == "session-1" @@ -380,7 +386,12 @@ async def test_run_agent_final_callback_without_streaming(self) -> None: final_call = callback.response_mock.await_args assert final_call is not None - assert final_call.args[0] is agent_response + final_response = final_call.args[0] + assert final_response is not agent_response and final_response is not result + assert final_response.to_dict() == agent_response.to_dict() == result.to_dict() + assert final_response.messages[0] is not agent_response.messages[0] + final_response.messages[0].contents[0].text = "callback mutation" + assert agent_response.text == result.text == "Final response" final_context = final_call.args[1] assert final_context.agent_name == "NonStreamingAgent" assert final_context.correlation_id == "corr-final-1" @@ -641,6 +652,38 @@ async def test_run_agent_preserves_message_on_error(self) -> None: content = result.messages[0].contents[0] assert isinstance(content, Content) + async def test_failed_run_reports_the_reason_in_the_reply_text(self) -> None: + """A failure must not read as the agent having nothing to say. + + The entity absorbs exceptions so the session survives, but error content alone leaves + ``text`` empty, so a caller reading the reply the normal way sees silence and has to go + digging to find out that anything went wrong at all. + """ + mock_agent = Mock() + mock_agent.run = _create_mock_run(side_effect=ValueError("no such deployment")) + + entity = _make_entity(mock_agent) + + result = await entity.run({"message": "Message", "correlationId": "corr-entity-error-5"}) + + assert "no such deployment" in result.text + assert "ValueError" in result.text + # The typed error content is still first, so callers inspecting contents are unaffected. + assert result.messages[0].contents[0].type == "error" + + async def test_failed_turns_are_not_replayed_to_the_model(self) -> None: + """The error text is for the caller, not for the model's context.""" + mock_agent = Mock() + mock_agent.run = _create_mock_run(side_effect=ValueError("boom")) + + entity = _make_entity(mock_agent) + await entity.run({"message": "first", "correlationId": "corr-entity-error-6"}) + + replayed = [message.text for message in entity._replay_all_messages()] + + assert "first" in replayed + assert not any("boom" in text for text in replayed) + class TestConversationHistory: """Test suite for conversation history tracking.""" @@ -740,7 +783,7 @@ async def test_run_agent_with_run_request_object(self) -> None: async def test_run_agent_with_dict_request(self) -> None: """Test run_agent with a dictionary request.""" - mock_agent = Mock() + mock_agent = Mock(default_options={}) mock_agent.run = _create_mock_run(response=_agent_response("Response")) entity = _make_entity(mock_agent) @@ -810,7 +853,7 @@ async def test_run_agent_with_response_format(self) -> None: async def test_run_agent_disable_tool_calls(self) -> None: """Test run_agent with tool calls disabled.""" - mock_agent = Mock() + mock_agent = Mock(default_options={}) mock_agent.run = _create_mock_run(response=_agent_response("Response")) entity = _make_entity(mock_agent) diff --git a/python/packages/durabletask/tests/test_durable_history_autoswap.py b/python/packages/durabletask/tests/test_durable_history_autoswap.py new file mode 100644 index 0000000..c574d87 --- /dev/null +++ b/python/packages/durabletask/tests/test_durable_history_autoswap.py @@ -0,0 +1,1064 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Durable history substitution and ownership unit tests with recording doubles, without live services.""" + +import json +from collections.abc import AsyncIterable, Awaitable, Sequence +from copy import deepcopy +from typing import Any + +import pytest +from agent_framework import ( + Agent, + AgentSession, + ChatResponse, + ChatResponseUpdate, + Content, + ContextProvider, + HistoryProvider, + InMemoryHistoryProvider, + Message, + ResponseStream, + SessionContext, +) + +from agent_framework_durabletask import ( + AgentEntity, + AgentEntityStateProviderMixin, + DurableAgentState, + DurableHistoryProvider, + _entities, +) +from agent_framework_durabletask._history_provider import ensure_durable_history + + +class _StubClient: + """Chat client stand-in that stores history locally (the common case).""" + + STORES_BY_DEFAULT = False + + def __init__(self) -> None: + self.additional_properties: dict[str, Any] = {} + + +class _ServiceStoringClient(_StubClient): + """Chat client whose service keeps the conversation server-side.""" + + STORES_BY_DEFAULT = True + + +class _RecordingClient(_StubClient): + """Client that records the message list handed to it on each call. + + Needed to tell "the provider is attached" apart from "the provider is answering", which is the + distinction that keeps a service-backed agent from being sent its own transcript. + """ + + def __init__(self) -> None: + super().__init__() + self.received: list[list[Message]] = [] + self._counter = 0 + + def get_response( + self, + messages: str | Message | list[str] | list[Message], + *, + stream: bool = False, + options: dict[str, Any] | None = None, + **kwargs: Any, + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + options = options or {} + normalized = [m for m in messages if isinstance(m, Message)] if isinstance(messages, list) else [] + self.received.append(normalized) + + if stream: + return self._stream(options) + + async def _get() -> ChatResponse: + self._counter += 1 + return ChatResponse(messages=Message(role="assistant", contents=[f"reply-{self._counter}"])) + + return _get() + + def _stream(self, options: dict[str, Any]) -> ResponseStream[ChatResponseUpdate, ChatResponse]: + async def _updates() -> AsyncIterable[ChatResponseUpdate]: + self._counter += 1 + yield ChatResponseUpdate(contents=[Content.from_text(f"reply-{self._counter}")], role="assistant") + + def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse: + return ChatResponse.from_updates(updates, output_format_type=options.get("response_format")) + + return ResponseStream(_updates(), finalizer=_finalize) + + +class _RecordingServiceClient(_RecordingClient): + """The same, but its service keeps the conversation server-side.""" + + STORES_BY_DEFAULT = True + + +class _ConversationIdClient(_StubClient): + """Recording double that returns conversation IDs through core response types.""" + + def __init__(self, *, stores_by_default: bool, supports_streaming: bool) -> None: + super().__init__() + self.STORES_BY_DEFAULT = stores_by_default + self.supports_streaming = supports_streaming + self.calls: list[dict[str, Any]] = [] + self._counter = 0 + + def get_response( + self, + messages: str | Message | list[str] | list[Message], + *, + stream: bool = False, + options: dict[str, Any] | None = None, + **kwargs: Any, + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + options = options or {} + self.calls.append(deepcopy({"messages": messages, "stream": stream, "options": options, "kwargs": kwargs})) + if stream and not self.supports_streaming: + raise TypeError("stream is not supported") + + self._counter += 1 + text = f"reply-{self._counter}" + response_id = f"result-{self._counter}" + conversation_id = f"service-branch-{self._counter}" if options.get("store", self.STORES_BY_DEFAULT) else None + + if stream: + + async def _updates() -> AsyncIterable[ChatResponseUpdate]: + yield ChatResponseUpdate( + contents=[Content.from_text(text)], + role="assistant", + response_id=response_id, + conversation_id=conversation_id, + ) + + return ResponseStream(_updates(), finalizer=ChatResponse.from_updates) + + async def _get() -> ChatResponse: + return ChatResponse( + messages=Message(role="assistant", contents=[text]), + response_id=response_id, + conversation_id=conversation_id, + ) + + return _get() + + +class _ExternalHistoryProvider(HistoryProvider): + """Stand-in for Cosmos/Redis/file-backed history the user chose deliberately.""" + + def __init__(self) -> None: + super().__init__(source_id="external") + + async def get_messages(self, session_id: str | None, **kwargs: Any) -> list[Message]: + return [] + + async def save_messages(self, session_id: str | None, messages: Any, **kwargs: Any) -> None: + return None + + +class _InMemoryStateProvider(AgentEntityStateProviderMixin): + """JSON storage boundary without a durable backend.""" + + def __init__(self, *, session_id: str = "autoswap-session", raw: dict[str, Any] | None = None) -> None: + self._session_id = session_id + self._state_dict: dict[str, Any] = json.loads(json.dumps(raw or {})) + self.writes = 0 + + def _get_state_dict(self) -> dict[str, Any]: + return deepcopy(self._state_dict) + + def _set_state_dict(self, state: dict[str, Any]) -> None: + self._state_dict = json.loads(json.dumps(state)) + self.writes += 1 + + def _get_session_id_from_entity(self) -> str: + return self._session_id + + +class _PreviousResponseNotFound(Exception): + """Shaped like the provider's refusal of a conversation id it previously issued. + + Mirrors the real payload field for field, because the entity matches on the structured + ``code`` rather than on the message text. + """ + + def __init__(self) -> None: + super().__init__( + "Error code: 400 - {'error': {'message': \"Previous response with id 'resp_x' not " + "found.\", 'type': 'invalid_request_error', 'param': 'previous_response_id', " + "'code': 'previous_response_not_found'}}" + ) + self.status_code = 400 + self.code = "previous_response_not_found" + self.param = "previous_response_id" + self.body = { + "message": "Previous response with id 'resp_x' not found.", + "type": "invalid_request_error", + "param": "previous_response_id", + "code": "previous_response_not_found", + } + + +class _ContextLengthExceeded(Exception): + """A different 400, which must not be mistaken for a lost conversation.""" + + def __init__(self) -> None: + super().__init__("Error code: 400 - context_length_exceeded") + self.status_code = 400 + self.code = "context_length_exceeded" + + +def _agent(client: Any = None, **kwargs: Any) -> Agent: + """Build an agent with a stub client. + + The stubs cover the parts of the client protocol these tests exercise but not its full generic + signature, so the type is relaxed here rather than at every call site. + """ + chat_client: Any = client if client is not None else _StubClient() + return Agent(client=chat_client, name="a", **kwargs) + + +def _history_providers(agent: Any) -> list[Any]: + return [p for p in agent.context_providers if isinstance(p, HistoryProvider)] + + +class TestAutomaticDurableHistory: + """The durable runtime substitutes durable-backed history where appropriate.""" + + def test_agent_without_providers_gets_durable_history(self) -> None: + agent = _agent() + + prepared = ensure_durable_history(agent) + + providers = _history_providers(prepared) + assert len(providers) == 1 + assert isinstance(providers[0], DurableHistoryProvider) + # Uses the source id core's auto-injected provider would have, so a + # default-configured CompactionProvider still resolves it. + assert providers[0].source_id == InMemoryHistoryProvider.DEFAULT_SOURCE_ID + + def test_in_memory_history_is_replaced_preserving_source_id(self) -> None: + agent = _agent(context_providers=[InMemoryHistoryProvider(source_id="custom_slot", skip_excluded=True)]) + + prepared = ensure_durable_history(agent) + + providers = _history_providers(prepared) + assert len(providers) == 1 + replacement = providers[0] + assert isinstance(replacement, DurableHistoryProvider) + # Preserving these is what keeps an existing CompactionProvider wired up. + assert replacement.source_id == "custom_slot" + assert replacement.skip_excluded is True + + def test_external_history_provider_is_left_alone(self) -> None: + """The user deliberately chose their own storage; durable must not override it.""" + external = _ExternalHistoryProvider() + agent = _agent(context_providers=[external]) + + prepared = ensure_durable_history(agent) + + assert prepared is agent + assert _history_providers(prepared) == [external] + + def test_service_managed_history_still_gets_a_provider(self) -> None: + """The service owning the conversation is a per-run fact, not a per-registration one. + + Leaving a service-backed agent with no provider used to look right, because the service + holds the transcript. But ``store`` is an ordinary run option, so a single run can put the + conversation back in the client's hands, and core then injects a history provider of its + own. Its state is persisted along with the entity and retention cannot see it, so it grows + without bound. Claiming the slot up front is what keeps those turns reachable. + """ + agent = _agent(_ServiceStoringClient()) + + prepared = ensure_durable_history(agent) + + providers = _history_providers(prepared) + assert len(providers) == 1 + assert isinstance(providers[0], DurableHistoryProvider) + + def test_store_false_overrides_a_service_storing_client(self) -> None: + """``store=False`` puts history back in the client's hands, so durable must back it. + + Mirrors core's precedence: an explicit ``store`` wins over ``STORES_BY_DEFAULT``. Without + this, an agent using the Responses API with ``store=False`` would keep a plain in-memory + provider that the durable runtime never persists, silently losing the conversation. + """ + agent = _agent(_ServiceStoringClient(), default_options={"store": False}) + + prepared = ensure_durable_history(agent) + + providers = _history_providers(prepared) + assert len(providers) == 1 + assert isinstance(providers[0], DurableHistoryProvider) + + def test_store_true_still_gets_a_provider(self) -> None: + """Attached, but it yields nothing while the service owns the run. + + Attaching is about occupying the slot, not about taking over storage. What stops the model + being handed the transcript twice is the provider returning no history on a service-owned + run, which :class:`TestServiceManagedSessions` covers. + """ + agent = _agent(default_options={"store": True}) + + prepared = ensure_durable_history(agent) + + providers = _history_providers(prepared) + assert len(providers) == 1 + assert isinstance(providers[0], DurableHistoryProvider) + + def test_existing_durable_provider_is_untouched(self) -> None: + """Explicit configuration (for example to enable pruning) wins.""" + explicit = DurableHistoryProvider(prune_excluded=True) + agent = _agent(context_providers=[explicit]) + + prepared = ensure_durable_history(agent) + + assert prepared is agent + assert _history_providers(prepared) == [explicit] + + def test_agent_without_context_pipeline_is_left_alone(self) -> None: + """Custom agents that do not expose context_providers keep legacy replay.""" + + class _CustomAgent: + name = "custom" + + async def run(self, *args: Any, **kwargs: Any) -> Any: ... + + agent = _CustomAgent() + + assert ensure_durable_history(agent) is agent # type: ignore[arg-type] + + +class TestUserAgentIsNotMutated: + """Substitution must not change the object the caller handed us.""" + + def test_original_agent_keeps_its_providers(self) -> None: + original_provider = InMemoryHistoryProvider() + agent = _agent(context_providers=[original_provider]) + original_list = agent.context_providers + + prepared = ensure_durable_history(agent) + + assert prepared is not agent + assert agent.context_providers is original_list + assert agent.context_providers == [original_provider] + + def test_entity_construction_does_not_mutate_the_agent(self) -> None: + agent = _agent(context_providers=[InMemoryHistoryProvider()]) + + entity = AgentEntity(agent, state_provider=_InMemoryStateProvider()) + + assert isinstance(_history_providers(entity.agent)[0], DurableHistoryProvider) + assert isinstance(_history_providers(agent)[0], InMemoryHistoryProvider) + + +class TestFollowCompactionRetention: + """Follow-compaction retention physically deletes exclusions.""" + + def test_off_by_default(self) -> None: + entity = AgentEntity(_agent(), state_provider=_InMemoryStateProvider()) + + assert entity._retention == "keep_all" + assert entity._max_state_bytes is None + assert _history_providers(entity.agent)[0].prune_excluded is False + + def test_enabled_via_registration(self) -> None: + agent = _agent(context_providers=[InMemoryHistoryProvider()]) + + prepared = ensure_durable_history(agent, prune_excluded=True) + + assert _history_providers(prepared)[0].prune_excluded is True + + def test_entity_forwards_the_flag(self) -> None: + agent = _agent() + + entity = AgentEntity(agent, state_provider=_InMemoryStateProvider(), retention="follow_compaction") + + assert _history_providers(entity.agent)[0].prune_excluded is True + + @pytest.mark.parametrize("max_state_bytes", [None, 100_000]) + def test_keep_all_does_not_prune_on_write(self, max_state_bytes: int | None) -> None: + """A pressure budget does not enable eager pruning.""" + entity = AgentEntity( + _agent(), + state_provider=_InMemoryStateProvider(), + retention="keep_all", + max_state_bytes=max_state_bytes, + ) + + assert _history_providers(entity.agent)[0].prune_excluded is False + + def test_auto_is_not_a_retention_mode(self) -> None: + invalid_mode: Any = "auto" + with pytest.raises(ValueError, match="retention"): + AgentEntity(_agent(), state_provider=_InMemoryStateProvider(), retention=invalid_mode) + + def test_explicit_provider_configuration_wins(self) -> None: + """A hand-configured provider is never overridden by the registration flag.""" + explicit = DurableHistoryProvider(prune_excluded=False) + agent = _agent(context_providers=[explicit]) + + prepared = ensure_durable_history(agent, prune_excluded=True) + + assert _history_providers(prepared)[0] is explicit + assert explicit.prune_excluded is False + + def test_an_unset_provider_inherits_the_retention_mode(self) -> None: + """Constructing the provider by hand must not silently disable ``follow_compaction``. + + A caller who writes ``DurableHistoryProvider()`` has expressed no opinion about pruning, + so the entity's retention mode is the only instruction available. Treating the unset + default as a deliberate "no" made ``retention='follow_compaction'`` do nothing at all for + anyone who wired the provider themselves. + """ + unset = DurableHistoryProvider() + assert unset.prune_excluded is None + agent = _agent(context_providers=[unset]) + + prepared = ensure_durable_history(agent, prune_excluded=True) + + providers = _history_providers(prepared) + assert providers[0] is not unset + assert isinstance(providers[0], DurableHistoryProvider) + assert providers[0].prune_excluded is True + # The caller's own object is never mutated. + assert unset.prune_excluded is None + + def test_an_unset_provider_stays_unpruned_under_keep_all(self) -> None: + unset = DurableHistoryProvider() + agent = _agent(context_providers=[unset]) + + entity = AgentEntity(agent, state_provider=_InMemoryStateProvider(), retention="keep_all") + + providers = _history_providers(entity.agent) + assert providers[0].prune_excluded is False + + +class _StoringExternalProvider(HistoryProvider): + """External-store double with a blind append, not its own input deduplication.""" + + def __init__(self) -> None: + super().__init__(source_id="external-store") + self.saved: list[Message] = [] + self.saved_batches: list[list[Message]] = [] + + async def get_messages(self, session_id: str | None, **kwargs: Any) -> list[Message]: + return deepcopy(self.saved) + + async def save_messages(self, session_id: str | None, messages: Any, **kwargs: Any) -> None: + batch = deepcopy(list(messages)) + self.saved_batches.append(batch) + self.saved.extend(batch) + + +class _ServiceAwareExternalProvider(_StoringExternalProvider): + """Test provider whose hooks defer to a service ID on the active session.""" + + async def before_run( + self, *, agent: Any, session: AgentSession, context: SessionContext, state: dict[str, Any] + ) -> None: + if session.service_session_id is None: + await super().before_run(agent=agent, session=session, context=context, state=state) + + async def after_run( + self, *, agent: Any, session: AgentSession, context: SessionContext, state: dict[str, Any] + ) -> None: + if session.service_session_id is None: + await super().after_run(agent=agent, session=session, context=context, state=state) + + +class _SessionObserver(ContextProvider): + def __init__(self) -> None: + super().__init__("session-observer") + self.before: list[dict[str, Any]] = [] + self.after: list[dict[str, Any]] = [] + + @staticmethod + def _snapshot(session: AgentSession, context: SessionContext) -> dict[str, Any]: + return { + "service_session_id": session.service_session_id, + "context_service_session_id": context.service_session_id, + "texts": [message.text for message in context.get_messages(include_input=True)], + } + + async def before_run( + self, *, agent: Any, session: AgentSession, context: SessionContext, state: dict[str, Any] + ) -> None: + self.before.append(self._snapshot(session, context)) + + async def after_run( + self, *, agent: Any, session: AgentSession, context: SessionContext, state: dict[str, Any] + ) -> None: + self.after.append(self._snapshot(session, context)) + + +class TestWeDoNotKeepASecondCopyOfSomeoneElsesConversation: + """External history needs delivery and ingestion receipts, not a local message mirror.""" + + def _content_items(self, entity: AgentEntity, kind: str) -> int: + return sum( + len(m.contents) + for entry in entity.state.data.conversation_history + for m in entry.messages + if entry.json_type == kind + ) + + async def _run(self, providers: list[Any], turns: int = 4) -> AgentEntity: + agent = _agent(_RecordingClient(), context_providers=providers) + entity = AgentEntity(agent, state_provider=_InMemoryStateProvider()) + for index in range(turns): + await entity.run({"message": f"a reasonably long question number {index}", "correlationId": f"c{index}"}) + return entity + + async def test_requests_are_not_kept_twice(self) -> None: + external = _StoringExternalProvider() + + entity = await self._run([external]) + + assert len(external.saved) == 8 + assert entity.state.data.conversation_history == [] + + async def test_responses_are_kept_in_the_mailbox_for_delivery(self) -> None: + external = _StoringExternalProvider() + + entity = await self._run([external]) + + restored = DurableAgentState.from_json(entity.state.to_json()) + assert restored.data.conversation_history == [] + assert set(restored.data.response_mailbox) == {f"c{index}" for index in range(4)} + for index in range(4): + response = restored.try_get_agent_response(f"c{index}") + assert response is not None + assert response.text == f"reply-{index + 1}" + assert response.to_dict() == restored.data.response_mailbox[f"c{index}"]["response"] + + async def test_completion_is_recorded_separately_from_the_transcript(self) -> None: + external = _StoringExternalProvider() + + entity = await self._run([external]) + + data = json.loads(entity.state.to_json())["data"] + assert data["conversationHistory"] == [] + assert set(data["completedCorrelations"]) == {f"c{index}" for index in range(4)} + assert all(receipt["completedAt"] for receipt in data["completedCorrelations"].values()) + + @pytest.mark.parametrize("include_new_message", [False, True], ids=["repeated-only", "repeated-and-new"]) + async def test_custom_context_ids_are_deduplicated_after_json_cold_reload(self, include_new_message: bool) -> None: + external = _StoringExternalProvider() + client = _RecordingClient() + provider = _InMemoryStateProvider() + entity = AgentEntity(_agent(client, context_providers=[external]), state_provider=provider) + original = Message(role="user", contents=["upstream original"], message_id="custom-source-id") + fresh = Message(role="user", contents=["upstream new"], message_id="another-custom-id") + + first = await entity.run({ + "message": "upstream original", + "correlationId": "first-delivery", + "contextMessages": [original.to_dict()], + }) + raw = json.loads(json.dumps(provider._get_state_dict())) + assert raw["schemaVersion"] == "2.0.0" + assert raw["data"]["conversationHistory"] == [] + original_receipt = raw["data"]["ingestedMessages"]["custom-source-id"] + assert original_receipt + assert external.saved_batches[0][0].message_id == "custom-source-id" + + restarted_provider = _InMemoryStateProvider(raw=raw) + restarted = AgentEntity(_agent(client, context_providers=[external]), state_provider=restarted_provider) + follow_up = [original, fresh] if include_new_message else [original] + await restarted.run({ + "message": "logging-only input must not be replayed", + "correlationId": "new-delivery", + "contextMessages": [message.to_dict() for message in follow_up], + }) + + new_texts = [fresh.text] if include_new_message else [] + assert len(client.received) == 2, "a new correlation must run even when its projected input is already ingested" + assert [message.text for message in client.received[1]] == [original.text, "reply-1", *new_texts] + assert [message.text for message in external.saved_batches[1]] == [*new_texts, "reply-2"] + assert sum(message.message_id == original.message_id for message in external.saved) == 1 + + restored = DurableAgentState.from_json(json.dumps(restarted_provider._get_state_dict())) + assert restored.data.conversation_history == [] + assert restored.data.ingested_messages["custom-source-id"] == original_receipt + expected_ids = {"custom-source-id", "another-custom-id"} if include_new_message else {"custom-source-id"} + assert set(restored.data.ingested_messages) == expected_ids + assert set(restored.data.completed_correlations) == {"first-delivery", "new-delivery"} + delivered = restored.try_get_agent_response("first-delivery") + assert delivered is not None + assert delivered.to_dict() == first.to_dict() + + async def test_our_own_history_is_kept_in_full(self) -> None: + """Nothing else is holding it, so forgetting it would lose the conversation.""" + entity = await self._run([]) + + assert self._content_items(entity, "request") > 0 + assert self._content_items(entity, "response") > 0 + + +class TestServiceManagedSessions: + """Service-backed agents let the service own the conversation.""" + + @pytest.mark.parametrize("streaming", [False, True], ids=["nonstream-fallback", "streaming"]) + @pytest.mark.parametrize("external_history", [False, True], ids=["durable-primary", "external-primary"]) + @pytest.mark.parametrize( + ("stores_by_default", "default_options", "service_options", "local_options"), + [ + pytest.param(True, {}, {}, {"store": False}, id="client-default-true"), + pytest.param(False, {"store": True}, {}, {"store": False}, id="agent-default-true"), + pytest.param(False, {}, {"store": True}, {}, id="client-default-false"), + pytest.param(True, {"store": False}, {"store": True}, {}, id="agent-default-false"), + ], + ) + async def test_core_pipeline_isolates_service_and_local_branches_after_json_reload( + self, + streaming: bool, + external_history: bool, + stores_by_default: bool, + default_options: dict[str, Any], + service_options: dict[str, Any], + local_options: dict[str, Any], + ) -> None: + """True/False/False/True through core Agent, with recording doubles rather than live services.""" + client = _ConversationIdClient(stores_by_default=stores_by_default, supports_streaming=streaming) + observer = _SessionObserver() + external = _ServiceAwareExternalProvider() if external_history else None + providers: list[ContextProvider] = [external, observer] if external is not None else [observer] + prompts = ["service-first", "local-first", "local-second", "service-resumed"] + expected_inputs = [ + ["service-first"], + ["local-first"], + ["local-first", "reply-2", "local-second"], + ["service-resumed"], + ] + expected_local_history = [ + [], + ["local-first", "reply-2"], + ["local-first", "reply-2", "local-second", "reply-3"], + ["local-first", "reply-2", "local-second", "reply-3"], + ] + raw: dict[str, Any] = {} + originals: dict[str, dict[str, Any]] = {} + attempts = [True] if streaming else [True, False] + + for index, prompt in enumerate(prompts): + # Rebuild the agent, entity and state provider; only the recording doubles survive. + provider = _InMemoryStateProvider(raw=raw) + entity = AgentEntity( + _agent(client, context_providers=providers, default_options=default_options), + state_provider=provider, + ) + history_providers = _history_providers(entity.agent) + if external is not None: + assert history_providers == [external] + else: + assert len(history_providers) == 1 + assert isinstance(history_providers[0], DurableHistoryProvider) + + store = index in (0, 3) + options = dict(service_options if store else local_options) + start = len(client.calls) + response = await entity.run({"message": prompt, "correlationId": f"c{index}", "options": options}) + assert response.text == f"reply-{index + 1}" + assert response.response_id == f"result-{index + 1}" + originals[f"c{index}"] = json.loads(json.dumps(response.to_dict())) + + calls = client.calls[start:] + assert [call["stream"] for call in calls] == attempts + active_id = "service-branch-1" if index == 3 else None + for call in calls: + assert [message.text for message in call["messages"]] == expected_inputs[index] + assert call["options"].get("store", stores_by_default) is store + assert call["options"].get("conversation_id") == active_id + assert call["kwargs"].get("conversation_id") is None + assert call["kwargs"]["client_kwargs"].get("conversation_id") is None + forwarded_session = call["kwargs"]["client_kwargs"]["session"] + assert forwarded_session.service_session_id == active_id + assert forwarded_session.session_id == "autoswap-session" + + raw = json.loads(json.dumps(provider._get_state_dict())) + data = raw["data"] + assert provider.writes == 1 + assert data["session"]["service_session_id"] == ("service-branch-4" if index == 3 else "service-branch-1") + assert InMemoryHistoryProvider.DEFAULT_SOURCE_ID not in data["session"]["state"] + local_texts = [ + message.text for entry in entity.state.data.conversation_history for message in entry.messages + ] + assert local_texts == ([] if external is not None else expected_local_history[index]) + assert set(data["responseMailbox"]) == set(originals) + assert set(data["completedCorrelations"]) == set(originals) + assert {key: entry["response"] for key, entry in data["responseMailbox"].items()} == originals + + expected_active_ids = [None, None, None, "service-branch-1"] + assert [entry["service_session_id"] for entry in observer.before] == [ + value for value in expected_active_ids for _ in attempts + ] + assert [entry["context_service_session_id"] for entry in observer.before] == [ + value for value in expected_active_ids for _ in attempts + ] + # Implicit durable history is appended like Core's automatic provider, after the observer. + # Only this before-hook sees raw input; keep the full model-input checks above unchanged. + # An explicit external primary still runs before the observer and supplies its history. + expected_before_inputs = expected_inputs if external is not None else [[prompt] for prompt in prompts] + assert [entry["texts"] for entry in observer.before] == [ + batch for batch in expected_before_inputs for _ in attempts + ] + assert [entry["service_session_id"] for entry in observer.after] == [ + "service-branch-1", + None, + None, + "service-branch-4", + ] + assert [entry["context_service_session_id"] for entry in observer.after] == expected_active_ids + assert [entry["texts"] for entry in observer.after] == expected_inputs + if external is not None: + assert [message.text for message in external.saved] == expected_local_history[-1] + assert [[message.text for message in batch] for batch in external.saved_batches] == [ + ["local-first", "reply-2"], + ["local-second", "reply-3"], + ] + + reloaded = DurableAgentState.from_json(json.dumps(raw)) + for correlation_id, original in originals.items(): + delivered = reloaded.try_get_agent_response(correlation_id) + assert delivered is not None + assert delivered.to_dict() == original + + async def test_a_service_owned_run_is_not_sent_its_own_history(self) -> None: + """A service-owned run receives only new input, even before a service ID has been issued.""" + client = _RecordingServiceClient() + agent = _agent(client) + entity = AgentEntity(agent, state_provider=_InMemoryStateProvider()) + + for index in range(4): + await entity.run({"message": f"m{index}", "correlationId": f"c{index}"}) + + assert [len(batch) for batch in client.received] == [1, 1, 1, 1] + + async def test_a_client_side_run_does_get_its_history(self) -> None: + """The same provider, on runs the service is not holding, supplies the conversation.""" + client = _RecordingServiceClient() + agent = _agent(client) + entity = AgentEntity(agent, state_provider=_InMemoryStateProvider()) + + for index in range(4): + await entity.run({"message": f"m{index}", "correlationId": f"c{index}", "options": {"store": False}}) + + assert [len(batch) for batch in client.received] == [1, 3, 5, 7] + + async def test_a_client_side_run_does_not_grow_opaque_session_state(self) -> None: + """Client-owned turns stay in the local transcript, not a second history slice in the session bag.""" + provider = _InMemoryStateProvider() + entity = AgentEntity(_agent(_RecordingServiceClient()), state_provider=provider) + + sizes: list[int] = [] + for index in range(6): + await entity.run({"message": f"m{index}", "correlationId": f"c{index}", "options": {"store": False}}) + session_slice = provider._get_state_dict().get("data", {}).get("session", {}) + sizes.append(len(json.dumps(session_slice))) + + assert sizes[0] == sizes[-1], f"session state grew: {sizes}" + assert len(entity.state.data.conversation_history) == 12 + + async def test_only_new_messages_are_sent(self) -> None: + """History must not be replayed locally when the service already holds it.""" + recorded: list[list[Message]] = [] + + class _ServiceAgent: + name = "svc" + client = _ServiceStoringClient() + context_providers: list[Any] = [] + + def create_session(self, **kwargs: Any) -> Any: + from agent_framework import AgentSession + + return AgentSession() + + async def run(self, messages: Any = None, *, stream: bool = False, **kwargs: Any) -> Any: + from agent_framework import AgentResponse + + if stream: + raise TypeError("stream is not supported") + recorded.append(list(messages or [])) + return AgentResponse(messages=[Message(role="assistant", contents=["ok"])]) + + entity = AgentEntity(_ServiceAgent(), state_provider=_InMemoryStateProvider()) # type: ignore[arg-type] + + await entity.run({"message": "first", "correlationId": "c0"}) + await entity.run({"message": "second", "correlationId": "c1"}) + + # Each turn delivers only its own message; the service supplies the rest. + assert len(recorded[1]) == 1 + assert recorded[1][0].text == "second" + + async def test_service_conversation_id_is_persisted_and_restored(self) -> None: + """Without this the service would start a new thread on every turn.""" + seen_ids: list[str | None] = [] + + class _ThreadingAgent: + name = "svc" + client = _ServiceStoringClient() + context_providers: list[Any] = [] + + def create_session(self, **kwargs: Any) -> Any: + from agent_framework import AgentSession + + return AgentSession() + + async def run( + self, + messages: Any = None, + *, + stream: bool = False, + session: Any = None, + **kwargs: Any, + ) -> Any: + from agent_framework import AgentResponse + + if stream: + raise TypeError("stream is not supported") + seen_ids.append(getattr(session, "service_session_id", None)) + # The service issues (or confirms) the thread id on the session. + session.service_session_id = "svc-thread-1" + return AgentResponse(messages=[Message(role="assistant", contents=["ok"])]) + + provider = _InMemoryStateProvider() + entity = AgentEntity(_ThreadingAgent(), state_provider=provider) # type: ignore[arg-type] + + await entity.run({"message": "first", "correlationId": "c0"}) + await entity.run({"message": "second", "correlationId": "c1"}) + + assert seen_ids[0] is None # first turn has no thread yet + assert seen_ids[1] == "svc-thread-1" # second turn continues the same thread + assert provider._get_state_dict()["data"]["session"]["service_session_id"] == "svc-thread-1" + + +class TestRejectedConversationIdRecovery: + """Injected service refusals exercise bounded identical-request retries, not transcript recovery.""" + + @pytest.fixture(autouse=True) + def _no_backoff(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Remove the retry waits, which are about a real service catching up, not about tests.""" + monkeypatch.setattr(_entities, "_REJECTED_ID_BACKOFF_SECONDS", 0.0) + + async def test_a_late_id_is_recovered_without_resending_the_transcript(self) -> None: + """The common case: the id resolves a moment later, so nothing needs resending.""" + calls: list[dict[str, Any]] = [] + + class _SlowToCommitAgent: + name = "svc" + client = _ServiceStoringClient() + context_providers: list[Any] = [] + + def create_session(self, **kwargs: Any) -> Any: + from agent_framework import AgentSession + + return AgentSession() + + async def run( + self, + messages: Any = None, + *, + stream: bool = False, + session: Any = None, + **kwargs: Any, + ) -> Any: + from agent_framework import AgentResponse + + if stream: + raise TypeError("stream is not supported") + previous = getattr(session, "service_session_id", None) + calls.append({"previous": previous, "texts": [m.text for m in (messages or [])]}) + if previous is None: + session.service_session_id = "thread-1" + return AgentResponse(messages=[Message(role="assistant", contents=["ok"])]) + # Refused once, then the service catches up. + if len([c for c in calls if c["previous"] is not None]) == 1: + raise _PreviousResponseNotFound + return AgentResponse(messages=[Message(role="assistant", contents=["ok"])]) + + provider = _InMemoryStateProvider() + entity = AgentEntity(_SlowToCommitAgent(), state_provider=provider) # type: ignore[arg-type] + + await entity.run({"message": "first", "correlationId": "c0"}) + response = await entity.run({"message": "second", "correlationId": "c1"}) + + assert response.text == "ok" + # First turn, the refusal, then one retry that succeeded. No transcript replay. + assert len(calls) == 3 + assert calls[2]["previous"] == "thread-1" + assert calls[2]["texts"] == ["second"] + # The conversation continued on the same thread rather than starting a new one. + assert provider._get_state_dict()["data"]["session"]["service_session_id"] == "thread-1" + + async def test_an_id_that_never_resolves_fails_the_turn(self) -> None: + calls: list[dict[str, Any]] = [] + + class _ForgetfulAgent: + name = "svc" + client = _ServiceStoringClient() + context_providers: list[Any] = [] + + def create_session(self, **kwargs: Any) -> Any: + from agent_framework import AgentSession + + return AgentSession() + + async def run( + self, + messages: Any = None, + *, + stream: bool = False, + session: Any = None, + **kwargs: Any, + ) -> Any: + from agent_framework import AgentResponse + + if stream: + raise TypeError("stream is not supported") + previous = getattr(session, "service_session_id", None) + calls.append({"previous": previous, "texts": [m.text for m in (messages or [])]}) + # Any turn that arrives carrying a conversation id is refused. + if previous is not None: + raise _PreviousResponseNotFound + session.service_session_id = f"thread-{len(calls)}" + return AgentResponse(messages=[Message(role="assistant", contents=["ok"])]) + + provider = _InMemoryStateProvider() + entity = AgentEntity(_ForgetfulAgent(), state_provider=provider) # type: ignore[arg-type] + + await entity.run({"message": "first", "correlationId": "c0"}) + response = await entity.run({"message": "second", "correlationId": "c1"}) + + # Five calls: the first turn, the refused attempt, and three retries of the identical + # request. Nothing else is tried, because the only remaining recovery would be resending + # our own transcript, and that is only possible if the entity keeps a full second copy of + # a conversation the service is already holding. + assert len(calls) == 5 + # Every attempt after the first was the same request, unchanged, still chained on the id. + assert all(call["texts"] == ["second"] and call["previous"] == "thread-1" for call in calls[1:]) + # The turn is reported as failed rather than silently answered without its context. + assert any(content.type == "error" for content in response.messages[0].contents) + # The stored id is left alone, so a service that recovers later still works. + assert provider._get_state_dict()["data"]["session"]["service_session_id"] == "thread-1" + + async def test_streaming_rejection_does_not_add_a_nonstreamed_attempt(self) -> None: + """Falling back to a non-streamed call with the refused id only wastes a round trip.""" + attempts: list[tuple[str, str | None]] = [] + + class _StreamingForgetfulAgent: + name = "svc" + client = _ServiceStoringClient() + context_providers: list[Any] = [] + + def create_session(self, **kwargs: Any) -> Any: + from agent_framework import AgentSession + + return AgentSession() + + async def run( + self, + messages: Any = None, + *, + stream: bool = False, + session: Any = None, + **kwargs: Any, + ) -> Any: + from agent_framework import AgentResponse + + previous = getattr(session, "service_session_id", None) + attempts.append(("stream" if stream else "nonstream", previous)) + if previous is not None: + raise _PreviousResponseNotFound + if stream: + raise TypeError("stream is not supported") + session.service_session_id = "thread-1" + return AgentResponse(messages=[Message(role="assistant", contents=["ok"])]) + + entity = AgentEntity( + _StreamingForgetfulAgent(), # type: ignore[arg-type] + state_provider=_InMemoryStateProvider(), + ) + + await entity.run({"message": "first", "correlationId": "c0"}) + await entity.run({"message": "second", "correlationId": "c1"}) + + # Retry the streamed invocation at the entity boundary, without clearing the ID + # or adding a non-streamed attempt carrying the same refused ID. + assert ("stream", "thread-1") in attempts + assert ("nonstream", "thread-1") not in attempts + + async def test_unrelated_bad_request_is_not_replayed(self) -> None: + """Replaying on any 400 would answer without the context the caller asked for.""" + calls: list[str | None] = [] + + class _FailingAgent: + name = "svc" + client = _ServiceStoringClient() + context_providers: list[Any] = [] + + def create_session(self, **kwargs: Any) -> Any: + from agent_framework import AgentSession + + return AgentSession() + + async def run( + self, + messages: Any = None, + *, + stream: bool = False, + session: Any = None, + **kwargs: Any, + ) -> Any: + if stream: + raise TypeError("stream is not supported") + calls.append(getattr(session, "service_session_id", None)) + raise _ContextLengthExceeded + + entity = AgentEntity(_FailingAgent(), state_provider=_InMemoryStateProvider()) # type: ignore[arg-type] + + response = await entity.run({"message": "first", "correlationId": "c0"}) + + assert len(calls) == 1 # attempted once, not retried + assert any(content.type == "error" for content in response.messages[0].contents) + + async def test_retries_are_bounded(self) -> None: + """A retry loop against a service that keeps refusing would never terminate.""" + calls: list[str | None] = [] + + class _AlwaysRejectingAgent: + name = "svc" + client = _ServiceStoringClient() + context_providers: list[Any] = [] + + def create_session(self, **kwargs: Any) -> Any: + from agent_framework import AgentSession + + return AgentSession() + + async def run( + self, + messages: Any = None, + *, + stream: bool = False, + session: Any = None, + **kwargs: Any, + ) -> Any: + if stream: + raise TypeError("stream is not supported") + calls.append(getattr(session, "service_session_id", None)) + raise _PreviousResponseNotFound + + entity = AgentEntity(_AlwaysRejectingAgent(), state_provider=_InMemoryStateProvider()) # type: ignore[arg-type] + + response = await entity.run({"message": "first", "correlationId": "c0"}) + + # The original attempt plus a fixed number of retries, then the failure is reported rather + # than retried forever. + assert len(calls) == 4 + assert any(content.type == "error" for content in response.messages[0].contents) diff --git a/python/packages/durabletask/tests/test_durable_history_provider.py b/python/packages/durabletask/tests/test_durable_history_provider.py new file mode 100644 index 0000000..c1ffd2b --- /dev/null +++ b/python/packages/durabletask/tests/test_durable_history_provider.py @@ -0,0 +1,1011 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Core history-provider unit tests with recording clients and JSON state, without a live backend.""" + +import json +from collections.abc import AsyncIterable, Awaitable, Sequence +from copy import deepcopy +from typing import Any + +import pytest +from agent_framework import ( + Agent, + AgentSession, + ChatResponse, + ChatResponseUpdate, + CompactionProvider, + Content, + ContextProvider, + HistoryProvider, + InMemoryHistoryProvider, + Message, + ResponseStream, + SessionContext, +) + +from agent_framework_durabletask import ( + AgentEntity, + AgentEntityStateProviderMixin, + DurableAgentState, + DurableHistoryProvider, +) +from agent_framework_durabletask._history_provider import replayable_entries + +KEEP_LAST_MESSAGES = 2 + + +class RecordingChatClient: + """Minimal chat client that records the message list it receives per call.""" + + def __init__(self) -> None: + self.additional_properties: dict[str, Any] = {} + self.received_messages: list[list[Message]] = [] + self._counter = 0 + + def get_response( + self, + messages: str | Message | list[str] | list[Message], + *, + stream: bool = False, + options: dict[str, Any] | None = None, + **kwargs: Any, + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + options = options or {} + normalized = [m for m in messages if isinstance(m, Message)] if isinstance(messages, list) else [] + self.received_messages.append(normalized) + + if stream: + return self._stream(options) + + async def _get() -> ChatResponse: + self._counter += 1 + return ChatResponse(messages=Message(role="assistant", contents=[f"reply-{self._counter}"])) + + return _get() + + def _stream(self, options: dict[str, Any]) -> ResponseStream[ChatResponseUpdate, ChatResponse]: + async def _updates() -> AsyncIterable[ChatResponseUpdate]: + self._counter += 1 + yield ChatResponseUpdate(contents=[Content.from_text(f"reply-{self._counter}")], role="assistant") + + def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse: + return ChatResponse.from_updates(updates, output_format_type=options.get("response_format")) + + return ResponseStream(_updates(), finalizer=_finalize) + + +class _InMemoryStateProvider(AgentEntityStateProviderMixin): + """Test-only state provider that keeps the serialized entity state in memory.""" + + def __init__(self, *, session_id: str = "durable-history-session", raw: dict[str, Any] | None = None) -> None: + self._session_id = session_id + self._state_dict: dict[str, Any] = json.loads(json.dumps(raw or {})) + self.writes = 0 + + def _get_state_dict(self) -> dict[str, Any]: + return deepcopy(self._state_dict) + + def _set_state_dict(self, state: dict[str, Any]) -> None: + # Reject non-JSON state and avoid aliasing the staged operation snapshot. + self._state_dict = json.loads(json.dumps(state)) + self.writes += 1 + + def _get_session_id_from_entity(self) -> str: + return self._session_id + + +async def _keep_last_messages(messages: list[Message]) -> bool: + """Compaction strategy: mark everything except the most recent messages as excluded.""" + if len(messages) <= KEEP_LAST_MESSAGES: + return False + changed = False + for message in messages[:-KEEP_LAST_MESSAGES]: + if not message.additional_properties.get("_excluded"): + message.additional_properties["_excluded"] = True + changed = True + return changed + + +async def _summarize_oldest(messages: list[Message]) -> bool: + """Strategy that *inserts* a summary message, mimicking ToolResultCompactionStrategy. + + Uses a stable summary id derived from the messages it replaces, so re-running it must + not create duplicates. + """ + if len(messages) <= KEEP_LAST_MESSAGES: + return False + + older = [m for m in messages[:-KEEP_LAST_MESSAGES] if not m.additional_properties.get("_excluded")] + if not older: + return False + + summary_id = "summary_" + "_".join(sorted(m.message_id or "" for m in older)) + if any(m.message_id == summary_id for m in messages): + return False + + for message in older: + message.additional_properties["_excluded"] = True + message.additional_properties["_summarized_by_summary_id"] = summary_id + + summary = Message( + role="assistant", + contents=[f"[summary of {len(older)} messages]"], + message_id=summary_id, + additional_properties={"_summary_of_message_ids": [m.message_id for m in older]}, + ) + messages.insert(messages.index(older[-1]) + 1, summary) + return True + + +def _agent(providers: list[Any], client: RecordingChatClient | None = None) -> Agent: + """Build an agent with the given context providers. + + The stub client covers the parts of the client protocol these tests exercise but not its full + generic signature, so the type is relaxed here rather than at every call site. + """ + chat_client: Any = client or RecordingChatClient() + return Agent(client=chat_client, name="assistant", context_providers=providers) + + +def _providers_of(entity: AgentEntity) -> list[Any]: + """Return the context providers on the entity's (possibly substituted) agent.""" + return list(getattr(entity.agent, "context_providers", [])) + + +class _StubExternalProvider(HistoryProvider): + """Stand-in for a provider the user configured deliberately (Cosmos, Redis, file).""" + + def __init__(self) -> None: + super().__init__(source_id="external") + + async def get_messages(self, session_id: str | None, **kwargs: Any) -> list[Message]: + return [] + + async def save_messages(self, session_id: str | None, messages: Any, **kwargs: Any) -> None: + return None + + +def _build_agent( + client: RecordingChatClient, + *, + with_compaction: bool = False, + prune_excluded: bool = False, + strategy: Any = None, +) -> Agent: + history = DurableHistoryProvider(prune_excluded=prune_excluded) + providers: list[Any] = [history] + if with_compaction: + providers.append( + CompactionProvider( + after_strategy=strategy or _keep_last_messages, + history_source_id=history.source_id, + ) + ) + return _agent(providers, client) + + +def _make_entity(agent: Agent, provider: _InMemoryStateProvider) -> AgentEntity: + return AgentEntity(agent, state_provider=provider) + + +async def _run_turns(entity: AgentEntity, prompts: list[str]) -> None: + for index, prompt in enumerate(prompts): + await entity.run({"message": prompt, "correlationId": f"corr-{index}"}) + + +def _stored_messages(entity: AgentEntity) -> list[Any]: + return [m for entry in entity.state.data.conversation_history for m in entry.messages] + + +class TestDurableHistoryProvider: + """The local transcript backs core history independently of response delivery.""" + + async def test_state_is_written_once_per_turn(self) -> None: + """Each write serializes the whole conversation, so a spare one is not free. + + The provider used to persist at the end of ``flush``, which meant every turn serialized + the entire transcript twice: once mid-turn, before the response even existed, and again + when the entity finished. The mid-turn copy was always superseded, and with no compaction + configured it wrote back state nothing had touched. Cost grows with the conversation, so + this is pinned rather than left to drift back. + """ + for label, agent in ( + ("no compaction", _build_agent(RecordingChatClient())), + ("compaction", _build_agent(RecordingChatClient(), with_compaction=True)), + ( + "compaction and pruning", + _build_agent(RecordingChatClient(), with_compaction=True, prune_excluded=True), + ), + ): + provider = _InMemoryStateProvider() + entity = _make_entity(agent, provider) + + await _run_turns(entity, ["first", "second", "third"]) + + assert provider.writes == 3, f"{label}: expected one write per turn, got {provider.writes}" + + async def test_compaction_annotations_survive_the_turn(self) -> None: + """Removing the mid-turn write must not cost the annotations it used to persist.""" + provider = _InMemoryStateProvider() + entity = _make_entity(_build_agent(RecordingChatClient(), with_compaction=True), provider) + + await _run_turns(entity, ["first", "second", "third", "fourth"]) + + # Read from the serialized copy, not the in-memory objects, so this proves the + # annotations actually reached durable state. + persisted = provider._get_state_dict()["data"]["conversationHistory"] + annotated = [ + message + for entry in persisted + for message in entry.get("messages", []) + if (message.get("extensionData") or {}).get("_excluded") + ] + assert annotated, "compaction marked messages excluded but none of it was persisted" + + async def test_a_failed_turn_never_becomes_model_context(self) -> None: + """A failed result remains deliverable from the mailbox, but is never model history.""" + + class _FailingClient(RecordingChatClient): + def get_response(self, messages: Any, **kwargs: Any) -> Any: + raise RuntimeError("kaboom") + + provider = _InMemoryStateProvider() + entity = _make_entity(_build_agent(_FailingClient()), provider) # type: ignore[arg-type] + + failed = await entity.run({"message": "please fail", "correlationId": "corr-fail"}) + + reloaded = DurableAgentState.from_dict(provider._get_state_dict()) + replayed = [ + entry.messages[index].to_chat_message().text + for entry, index in replayable_entries(reloaded.data.conversation_history) + ] + + assert replayed == [] + delivered = reloaded.try_get_agent_response("corr-fail") + assert delivered is not None + assert delivered.to_dict() == failed.to_dict() + assert any(content.type == "error" for message in delivered.messages for content in message.contents) + client = RecordingChatClient() + restarted = _make_entity(_build_agent(client), _InMemoryStateProvider(raw=provider._get_state_dict())) + await restarted.run({"message": "next", "correlationId": "corr-next"}) + assert [[message.text for message in batch] for batch in client.received_messages] == [["next"]] + + @pytest.mark.parametrize("prune_excluded", [False, True], ids=["annotate", "prune"]) + async def test_a_summary_is_never_returned_as_an_answer(self, prune_excluded: bool) -> None: + """Original payloads and metadata survive compaction, transcript deletion and JSON reload.""" + + class _MetadataClient(RecordingChatClient): + def get_response(self, messages: Any, *, stream: bool = False, **kwargs: Any) -> Awaitable[ChatResponse]: + if stream: + raise TypeError("stream is not supported") + self.received_messages.append([message for message in messages if isinstance(message, Message)]) + + async def _get() -> ChatResponse: + self._counter += 1 + return ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_text( + f"reply-{self._counter}", additional_properties={"source": {"tags": ["original"]}} + ) + ], + message_id=f"answer-{self._counter}", + author_name="metadata-client", + additional_properties={"trace": {"tags": ["original"]}}, + ), + response_id=f"response-{self._counter}", + created_at="2026-09-08T12:00:00+00:00", + finish_reason="stop", + usage_details={"input_token_count": 7, "output_token_count": 11, "total_token_count": 18}, + additional_properties={"result_metadata": {"tags": ["original"]}}, + ) + + return _get() + + client = _MetadataClient() + provider = _InMemoryStateProvider() + entity = _make_entity( + _build_agent(client, with_compaction=True, prune_excluded=prune_excluded, strategy=_summarize_oldest), + provider, + ) + + originals: dict[str, dict[str, Any]] = {} + for index in range(6): + response = await entity.run({"message": f"t{index}", "correlationId": f"corr-{index}"}) + original = json.loads(json.dumps(response.to_dict())) + assert original["response_id"] == f"response-{index + 1}" + assert original["created_at"] == "2026-09-08T12:00:00+00:00" + assert original["finish_reason"] == "stop" + assert original["usage_details"] == { + "input_token_count": 7, + "output_token_count": 11, + "total_token_count": 18, + } + assert original["additional_properties"]["result_metadata"] == {"tags": ["original"]} + assert original["messages"][0]["message_id"] == f"answer-{index + 1}" + assert original["messages"][0]["author_name"] == "metadata-client" + assert original["messages"][0]["additional_properties"]["trace"] == {"tags": ["original"]} + assert original["messages"][0]["contents"][0]["text"] == f"reply-{index + 1}" + assert original["messages"][0]["contents"][0]["additional_properties"]["source"] == {"tags": ["original"]} + originals[f"corr-{index}"] = original + + # Mutating a returned result must not mutate its committed mailbox snapshot. + response.additional_properties["result_metadata"]["tags"].append("caller-mutation") + response.messages[0].additional_properties["trace"]["tags"].append("caller-mutation") + response.messages[0].contents[0].additional_properties["source"]["tags"].append("caller-mutation") + + summaries = [m for m in _stored_messages(entity) if "[summary of" in (m.to_chat_message().text or "")] + assert summaries, "compaction produced no summary, so this proves nothing" + first_transcript_answer = [message for message in _stored_messages(entity) if message.message_id == "answer-1"] + if prune_excluded: + assert not first_transcript_answer + else: + assert first_transcript_answer + assert (first_transcript_answer[0].extension_data or {}).get("_excluded") is True + + for correlation_id, original in originals.items(): + delivered = entity.state.try_get_agent_response(correlation_id) + assert delivered is not None + assert delivered.to_dict() == original + + mailbox = deepcopy(entity.state.data.response_mailbox) + completions = deepcopy(entity.state.data.completed_correlations) + entity.state.data.conversation_history.clear() + entity.persist_state() + restarted_provider = _InMemoryStateProvider(raw=provider._get_state_dict()) + restarted = _make_entity(_build_agent(client), restarted_provider) + assert restarted.state.data.conversation_history == [] + assert restarted.state.data.response_mailbox == mailbox + assert restarted.state.data.completed_correlations == completions + before_retry = len(client.received_messages) + for correlation_id, original in originals.items(): + delivered = await restarted.run({"message": "duplicate delivery", "correlationId": correlation_id}) + assert delivered.to_dict() == original + delivered.messages[0].contents[0].text = "caller-modified lookup" + polled_again = restarted.state.try_get_agent_response(correlation_id) + assert polled_again is not None + assert polled_again.to_dict() == original + assert len(client.received_messages) == before_retry + assert restarted_provider.writes == 0 + + async def test_transcript_is_not_duplicated_in_session_state(self) -> None: + """The local transcript and delivery mailbox do not add a third copy in the session bag.""" + client = RecordingChatClient() + provider = _InMemoryStateProvider() + entity = _make_entity(_build_agent(client), provider) + + await _run_turns(entity, ["first", "second"]) + + persisted = provider._get_state_dict()["data"] + assert "conversationHistory" in persisted + # The session is persisted for provider state, but the history provider's slice - the + # only place messages would appear - is excluded from it. + session_state = persisted["session"]["state"] + assert not any("messages" in slice_ for slice_ in session_state.values() if isinstance(slice_, dict)) + assert len(entity.state.data.conversation_history) == 4 + + async def test_provider_supplies_history_across_turns(self) -> None: + """Prior turns are loaded from durable state, not replayed by the entity.""" + client = RecordingChatClient() + entity = _make_entity(_build_agent(client), _InMemoryStateProvider()) + + await _run_turns(entity, ["first", "second", "third"]) + + assert len(client.received_messages[0]) == 1 + assert len(client.received_messages[1]) > len(client.received_messages[0]) + assert len(client.received_messages[2]) > len(client.received_messages[1]) + assert client.received_messages[1][0].text == "first" + + async def test_no_duplicate_of_in_flight_request(self) -> None: + """The in-flight request is delivered as input, not also loaded as history.""" + client = RecordingChatClient() + entity = _make_entity(_build_agent(client), _InMemoryStateProvider()) + + await _run_turns(entity, ["only-once"]) + + texts = [m.text for m in client.received_messages[0]] + assert texts.count("only-once") == 1 + + async def test_compaction_annotations_persist_in_durable_state(self) -> None: + """Core compaction plugs in and its annotations are stored with the messages.""" + client = RecordingChatClient() + entity = _make_entity(_build_agent(client, with_compaction=True), _InMemoryStateProvider()) + + await _run_turns(entity, ["t1", "t2", "t3", "t4", "t5"]) + + excluded = [m for m in _stored_messages(entity) if (m.extension_data or {}).get("_excluded")] + assert excluded, "expected compaction annotations persisted in conversation history" + + # Annotations survive a full serialize/deserialize round-trip of entity state. + restored = DurableAgentState.from_dict(entity.state.to_dict()) + restored_excluded = [ + m + for entry in restored.data.conversation_history + for m in entry.messages + if (m.extension_data or {}).get("_excluded") + ] + assert len(restored_excluded) == len(excluded) + + async def test_compaction_bounds_model_input(self) -> None: + """Excluded messages are withheld from the model, so context stops growing.""" + turns = ["t1", "t2", "t3", "t4", "t5", "t6"] + + plain_client = RecordingChatClient() + await _run_turns(_make_entity(_build_agent(plain_client), _InMemoryStateProvider()), turns) + + compacted_client = RecordingChatClient() + await _run_turns( + _make_entity(_build_agent(compacted_client, with_compaction=True), _InMemoryStateProvider()), + turns, + ) + + assert len(compacted_client.received_messages[-1]) < len(plain_client.received_messages[-1]) + + async def test_prune_excluded_bounds_persisted_state(self) -> None: + """Opt-in pruning physically shrinks durable storage (the lossy L2 step).""" + turns = ["t1", "t2", "t3", "t4", "t5", "t6"] + + kept_entity = _make_entity(_build_agent(RecordingChatClient(), with_compaction=True), _InMemoryStateProvider()) + await _run_turns(kept_entity, turns) + + pruned_entity = _make_entity( + _build_agent(RecordingChatClient(), with_compaction=True, prune_excluded=True), + _InMemoryStateProvider(), + ) + await _run_turns(pruned_entity, turns) + + assert len(_stored_messages(pruned_entity)) < len(_stored_messages(kept_entity)) + # Nothing marked excluded is left behind in storage. + assert not [m for m in _stored_messages(pruned_entity) if (m.extension_data or {}).get("_excluded")] + + async def test_summarizing_strategy_persists_inserted_messages(self) -> None: + """Strategies that insert a summary (not just annotate) are reconciled by message id.""" + client = RecordingChatClient() + entity = _make_entity( + _build_agent(client, with_compaction=True, strategy=_summarize_oldest), + _InMemoryStateProvider(), + ) + + await _run_turns(entity, ["t1", "t2", "t3", "t4"]) + + stored = _stored_messages(entity) + summaries = [m for m in stored if m.message_id and m.message_id.startswith("summary_")] + assert summaries, "expected the inserted summary message to be persisted" + + # Identity and annotations survive a durable state round-trip. + restored = DurableAgentState.from_dict(entity.state.to_dict()) + restored_ids = [ + m.message_id + for entry in restored.data.conversation_history + for m in entry.messages + if m.message_id and m.message_id.startswith("summary_") + ] + assert restored_ids == [m.message_id for m in summaries] + + async def test_summary_is_not_duplicated_across_turns(self) -> None: + """Re-running compaction with a stable summary id must not append duplicates.""" + client = RecordingChatClient() + entity = _make_entity( + _build_agent(client, with_compaction=True, strategy=_summarize_oldest), + _InMemoryStateProvider(), + ) + + await _run_turns(entity, ["t1", "t2", "t3", "t4", "t5", "t6"]) + + ids = [m.message_id for m in _stored_messages(entity) if m.message_id] + assert len(ids) == len(set(ids)), f"duplicate message ids persisted: {ids}" + + async def test_insertion_keeps_later_positions_valid(self) -> None: + """Inserting into an entry shifts its later messages, so recorded positions must follow. + + Entries normally hold a single message, which hides this. A workflow node receives the + upstream conversation as several messages in one request entry, so a summary inserted in + the middle of that entry invalidates the recorded index of everything after it, and the + annotation lands on the wrong stored message. + """ + + async def _insert_then_exclude_up3(messages: list[Message]) -> bool: + if any((m.additional_properties or {}).get("_marker") for m in messages): + return False + summary = Message( + role="assistant", + contents=["summary"], + message_id="summary_mid", + additional_properties={"_marker": True}, + ) + # Insert near the front, so messages later in the *same* durable entry shift. + messages.insert(1, summary) + for message in messages: + if message.message_id == "up-3": + message.additional_properties = dict(message.additional_properties or {}) | {"_excluded": True} + return True + + client = RecordingChatClient() + entity = _make_entity( + _build_agent(client, with_compaction=True, strategy=_insert_then_exclude_up3), + _InMemoryStateProvider(), + ) + + # An upstream conversation delivered as one multi-message request entry. + context = [ + Message(role="user", contents=["upstream one"], message_id="up-1").to_dict(), + Message(role="assistant", contents=["upstream two"], message_id="up-2").to_dict(), + Message(role="user", contents=["upstream three"], message_id="up-3").to_dict(), + ] + await entity.run({"message": "upstream three", "correlationId": "c0", "contextMessages": context}) + await entity.run({"message": "next", "correlationId": "c1"}) + + stored = {m.message_id: m for m in _stored_messages(entity) if m.message_id} + assert "up-3" in stored, f"expected the upstream messages to be persisted: {list(stored)}" + + # The annotation must land on up-3 itself, not on the neighbour that shifted when the + # summary was inserted earlier in the same entry. + assert (stored["up-3"].extension_data or {}).get("_excluded"), "annotation did not reach up-3" + assert not (stored["up-2"].extension_data or {}).get("_excluded"), "annotation shifted onto up-2" + + async def test_generated_ids_survive_a_cold_start(self) -> None: + """Ids synthesized for messages stored without one must derive from persisted state. + + A cold start or a retried flush rebuilds the entry objects at fresh addresses, so an id + taken from object identity would differ every run, and a recycled address could even + collide with an id an earlier run already persisted. + """ + provider = _InMemoryStateProvider() + entity = _make_entity(_build_agent(RecordingChatClient()), provider) + await _run_turns(entity, ["first", "second"]) + + # History as a producer that does not stamp message ids would have written it. + raw = deepcopy(provider._get_state_dict()) + for entry in raw["data"]["conversationHistory"]: + for message in entry["messages"]: + message.pop("messageId", None) + + async def _synthesized_ids() -> list[str]: + restarted_provider = _InMemoryStateProvider() + restarted_provider._set_state_dict(deepcopy(raw)) + restarted = _make_entity(_build_agent(RecordingChatClient()), restarted_provider) + await restarted.run({"message": "third", "correlationId": "corr-restart"}) + return [m.message_id for m in _stored_messages(restarted) if (m.message_id or "").startswith("durable_")] + + first = await _synthesized_ids() + second = await _synthesized_ids() + + assert first, "expected ids to be synthesized for the messages that had none" + assert len(first) == len(set(first)), f"synthesized ids collided within one run: {first}" + assert first == second, f"synthesized ids changed across a cold start: {first} != {second}" + + @pytest.mark.parametrize("service_session_id", [None, "svc-123"], ids=["no-service-id", "saved-service-id"]) + @pytest.mark.parametrize("service_owns_history", [False, True], ids=["client-owned", "service-owned"]) + async def test_history_hooks_use_binding_ownership_not_the_saved_service_id( + self, service_session_id: str | None, service_owns_history: bool + ) -> None: + from agent_framework_durabletask._history_provider import ( + DurableHistoryBinding, + bind_durable_history, + unbind_durable_history, + ) + + client = RecordingChatClient() + provider = _InMemoryStateProvider() + entity = _make_entity(_build_agent(client), provider) + await _run_turns(entity, ["first", "second"]) + + history = DurableHistoryProvider() + before = deepcopy(provider.state.to_dict()) + token = bind_durable_history( + DurableHistoryBinding(state_provider=provider, service_owns_history=service_owns_history) + ) + try: + state: dict[str, Any] = {} + session = AgentSession(session_id="s", service_session_id=service_session_id) + context = SessionContext(session_id="s", service_session_id=service_session_id, input_messages=[]) + await history.before_run(agent=entity.agent, session=session, context=context, state=state) + + if service_owns_history: + assert state == {} + assert context.get_messages() == [] + assert provider.state.to_dict() == before + stored = provider.state.data.conversation_history[0].messages[0] + changed = stored.to_chat_message() + changed.message_id = "must-not-flush" + state = { + "messages": [changed], + "_positions": {changed.message_id: (provider.state.data.conversation_history[0], 0)}, + } + else: + assert [message.text for message in context.get_messages()] == ["first", "reply-1", "second", "reply-2"] + assert len(state["messages"]) == 4 + assert len(state["_positions"]) == 4 + + state["messages"][0].additional_properties["hook-marker"] = {"kept": True} + await history.after_run(agent=entity.agent, session=session, context=context, state=state) + if service_owns_history: + assert provider.state.to_dict() == before + else: + stored = provider.state.data.conversation_history[0].messages[0] + assert (stored.extension_data or {})["hook-marker"] == {"kept": True} + assert provider.writes == 2, "history hooks must not commit an intermediate snapshot" + finally: + unbind_durable_history(token) + + async def test_core_configured_agent_gets_durable_history_automatically(self) -> None: + """An agent configured the ordinary core way runs durably with no changes.""" + client = RecordingChatClient() + agent = _agent([InMemoryHistoryProvider()], client) + entity = _make_entity(agent, _InMemoryStateProvider()) + + await _run_turns(entity, ["first", "second"]) + + # The entity swapped in durable-backed history without the user asking. + assert any(isinstance(p, DurableHistoryProvider) for p in _providers_of(entity)) + # The caller's agent is untouched. + assert any(isinstance(p, InMemoryHistoryProvider) for p in agent.context_providers) + # History is served from durable state, so turn 2 sees turn 1. + assert len(client.received_messages[1]) > len(client.received_messages[0]) + assert len(entity.state.data.conversation_history) == 4 + + +class TestExternalHistoryProviders: + """Providers that own their own storage (Cosmos, Redis, file) keep working durably.""" + + async def test_external_provider_receives_the_entity_session_id(self) -> None: + """Their storage is keyed by session id, so it must be the entity's stable id. + + The entity builds a fresh session per operation. If that session carried a generated id, + an external provider would read and write a different key every turn and never see prior + history - broken continuity with no error to show for it. + """ + seen: list[str | None] = [] + + class _RecordingExternalProvider(HistoryProvider): + def __init__(self) -> None: + super().__init__(source_id="external") + + async def get_messages(self, session_id: str | None, **kwargs: Any) -> list[Message]: + seen.append(session_id) + return [] + + async def save_messages(self, session_id: str | None, messages: Any, **kwargs: Any) -> None: + seen.append(session_id) + + agent = _agent([_RecordingExternalProvider()]) + entity = _make_entity(agent, _InMemoryStateProvider(session_id="stable-session")) + + await _run_turns(entity, ["first", "second"]) + + assert seen, "the external provider should have taken part in the run" + assert set(seen) == {"stable-session"} + + async def test_external_provider_is_not_replaced(self) -> None: + """The user chose their own storage; durable must not swap it out.""" + external = _StubExternalProvider() + agent = _agent([external]) + + entity = _make_entity(agent, _InMemoryStateProvider()) + + assert _providers_of(entity)[0] is external + + +class TestSessionStatePersistence: + """Provider state kept in the session bag survives across turns. + + Core documents the per-provider ``state`` dict as durable for the life of the session and + persists it through ``AgentSession.to_dict()``. The entity builds a fresh session per + operation, so it has to carry that state forward - otherwise providers silently start from + scratch every turn (tool approval rules and queued approval requests, todo lists, memory + extraction state). + """ + + async def test_a_failed_turn_still_carries_the_session_forward(self) -> None: + """The entity absorbs the failure, so the session has to survive it too. + + Providers run before the model call, so a turn that fails afterwards can still have queued + a tool approval or been handed a conversation id by the service. Capturing the session only + on success dropped both, and the next turn started from scratch while the service-side + conversation was left orphaned. + """ + + class _QueueingProvider(ContextProvider): + def __init__(self) -> None: + super().__init__("approvals") + + async def before_run(self, *, agent: Any, session: Any, context: Any, state: dict[str, Any]) -> None: + state["pending_approval"] = "delete-the-archive" + + class _FailingClient(RecordingChatClient): + def get_response(self, messages: Any, **kwargs: Any) -> Any: + raise RuntimeError("kaboom") + + provider = _InMemoryStateProvider() + agent = _agent([InMemoryHistoryProvider(), _QueueingProvider()], _FailingClient()) + entity = AgentEntity(agent, state_provider=provider) + + response = await entity.run({"message": "please fail", "correlationId": "boom"}) + + assert any(content.type == "error" for content in response.messages[0].contents) + stored_session = provider._get_state_dict()["data"].get("session") + assert stored_session, "a failed turn discarded the session" + assert stored_session["state"]["approvals"]["pending_approval"] == "delete-the-archive" + + async def test_a_failure_before_the_session_exists_reports_its_own_error(self) -> None: + """``session`` is referenced while handling the error, so it must always be bound. + + It used to be assigned only inside the ``try``. A ``create_session`` that raised would then + leave the name unbound, and the failure path would replace the agent's error with a + ``NameError`` while trying to persist the session. + """ + + class _NoSessionAgent: + name = "broken" + client = RecordingChatClient() + context_providers: list[Any] = [] + + def create_session(self, **kwargs: Any) -> Any: + raise TypeError("this agent cannot make a session") + + async def run(self, *args: Any, **kwargs: Any) -> Any: + raise AssertionError("should never be reached") + + entity = AgentEntity(_NoSessionAgent(), state_provider=_InMemoryStateProvider()) # type: ignore[arg-type] + + response = await entity.run({"message": "x", "correlationId": "c0"}) + + text = " ".join(content.text or "" for content in response.messages[0].contents) + assert "cannot make a session" in text + assert "NameError" not in text + + async def test_provider_state_survives_across_turns(self) -> None: + seen: list[dict[str, Any]] = [] + + class _CountingProvider(ContextProvider): + def __init__(self) -> None: + super().__init__("counter") + + async def before_run(self, *, agent: Any, session: Any, context: Any, state: dict[str, Any]) -> None: + seen.append(dict(state)) + + async def after_run(self, *, agent: Any, session: Any, context: Any, state: dict[str, Any]) -> None: + state["runs"] = state.get("runs", 0) + 1 + + agent = _agent([_CountingProvider()]) + entity = _make_entity(agent, _InMemoryStateProvider()) + + await _run_turns(entity, ["first", "second", "third"]) + + assert seen[0] == {} # nothing stored yet on the first turn + assert seen[1] == {"runs": 1} + assert seen[2] == {"runs": 2} + + async def test_state_is_persisted_as_plain_data(self) -> None: + """Values go through core's serialization, so entity state stays JSON-safe.""" + + class _StoringProvider(ContextProvider): + def __init__(self) -> None: + super().__init__("storer") + + async def after_run(self, *, agent: Any, session: Any, context: Any, state: dict[str, Any]) -> None: + state.setdefault("note", Message(role="user", contents=["remember me"])) + + provider = _InMemoryStateProvider() + agent = _agent([_StoringProvider()]) + await _run_turns(_make_entity(agent, provider), ["first"]) + + session_payload = provider._get_state_dict()["data"]["session"] + assert isinstance(session_payload["state"]["storer"]["note"], dict) + # ...and comes back as a Message, because core pre-registers that type. + restored = AgentSession.from_dict(dict(session_payload)) + assert isinstance(restored.state["storer"]["note"], Message) + + async def test_service_conversation_id_rides_along(self) -> None: + """It is part of the serialized session, so it needs no field of its own.""" + provider = _InMemoryStateProvider() + agent = _agent([InMemoryHistoryProvider()]) + entity = _make_entity(agent, provider) + + await _run_turns(entity, ["first"]) + assert "service_session_id" in provider._get_state_dict()["data"]["session"] + + async def test_tool_approval_state_survives_a_turn(self) -> None: + """The motivating case: standing approvals must outlive the turn that granted them. + + It also comes back as ``ToolApprovalState`` rather than a plain dict. Core seeds its state + type registry with only ``Message``, so the entity registers the serializable types loaded + in this process before restoring. + """ + # The harness is experimental; skip rather than fail if it moves. + tool_approval = pytest.importorskip("agent_framework._harness._tool_approval") + ToolApprovalRule = tool_approval.ToolApprovalRule + ToolApprovalState = tool_approval.ToolApprovalState + + seen: list[Any] = [] + approval_key = "_tool_approval" + + class _ApprovalCarryingProvider(ContextProvider): + def __init__(self) -> None: + super().__init__("approvals") + + async def before_run(self, *, agent: Any, session: Any, context: Any, state: dict[str, Any]) -> None: + seen.append(session.state.get(approval_key)) + + async def after_run(self, *, agent: Any, session: Any, context: Any, state: dict[str, Any]) -> None: + session.state.setdefault( + approval_key, + ToolApprovalState(rules=[ToolApprovalRule("delete_file")]), + ) + + agent = _agent([_ApprovalCarryingProvider()]) + await _run_turns(_make_entity(agent, _InMemoryStateProvider()), ["first", "second"]) + + assert seen[0] is None # nothing granted yet + restored = seen[1] + assert isinstance(restored, ToolApprovalState), f"approval state came back as {type(restored).__name__}" + assert restored.rules[0].tool_name == "delete_file" + + async def test_durable_history_slice_is_not_persisted(self) -> None: + """That slice is derived from conversation_history; storing it would duplicate it.""" + provider = _InMemoryStateProvider() + agent = _agent([InMemoryHistoryProvider()]) + entity = _make_entity(agent, provider) + + await _run_turns(entity, ["first", "second"]) + + durable_history = next(p for p in _providers_of(entity) if isinstance(p, DurableHistoryProvider)) + session_state = provider._get_state_dict()["data"]["session"]["state"] + assert durable_history.source_id not in session_state + + async def test_durable_history_slice_is_dropped_before_serializing(self) -> None: + """Not after. That slice holds the working buffer, so serializing it is wasted work. + + It also keeps a position index whose values reference durable state objects, so the less + of it that reaches core's serializer the better. + """ + serialized_keys: list[list[str]] = [] + + class _SpySession: + def __init__(self, state: dict[str, Any]) -> None: + self.state = state + self.service_session_id = None + + def to_dict(self) -> dict[str, Any]: + serialized_keys.append(sorted(self.state)) + return {"session_id": "spy", "state": dict(self.state)} + + entity = _make_entity(_build_agent(RecordingChatClient()), _InMemoryStateProvider()) + durable_history = next(p for p in _providers_of(entity) if isinstance(p, DurableHistoryProvider)) + session = _SpySession({durable_history.source_id: {"messages": ["transcript"]}, "other": {"keep": 1}}) + + entity._capture_session(session) + + assert serialized_keys == [["other"]], f"the durable slice was serialized: {serialized_keys}" + assert durable_history.source_id in session.state, "the caller's session was left modified" + + @pytest.mark.parametrize("prior_turn", [False, True], ids=["new-session", "existing-session"]) + @pytest.mark.filterwarnings("ignore:AgentSession state value .* has unsupported type:RuntimeWarning") + async def test_unserializable_provider_state_fails_without_committing(self, prior_turn: bool) -> None: + """A successful model call is not a committed outcome when session serialization fails.""" + + class _UnserializableProvider(ContextProvider): + def __init__(self) -> None: + super().__init__("unserializable") + self.poison = True + + async def after_run(self, *, agent: Any, session: Any, context: Any, state: dict[str, Any]) -> None: + state["handle"] = object() if self.poison else "serializable" + + provider = _InMemoryStateProvider() + client = RecordingChatClient() + if prior_turn: + await _make_entity(_build_agent(client), provider).run({"message": "first", "correlationId": "committed"}) + stateful = _UnserializableProvider() + agent = _agent([InMemoryHistoryProvider(), stateful], client) + entity = _make_entity(agent, provider) + before = json.loads(json.dumps(provider._get_state_dict())) + cached_before = json.loads(entity.state.to_json()) + writes_before = provider.writes + calls_before = len(client.received_messages) + request = { + "message": "uncommitted input", + "correlationId": "uncommitted", + "contextMessages": [ + Message(role="user", contents=["uncommitted input"], message_id="pending-id").to_dict() + ], + } + + with pytest.raises(ValueError, match="session state.*JSON-compatible.*cannot commit"): + await entity.run(request) + + assert len(client.received_messages) == calls_before + 1 + assert provider.writes == writes_before + assert provider._get_state_dict() == before + assert entity.state.to_dict() == cached_before + assert entity.state.try_get_agent_response("uncommitted") is None + assert "uncommitted" not in entity.state.data.response_mailbox + assert "uncommitted" not in entity.state.data.completed_correlations + assert "pending-id" not in entity.state.data.ingested_messages + + cold = _make_entity(_build_agent(client), _InMemoryStateProvider(raw=before)) + assert cold.state.to_dict() == cached_before + assert cold.state.try_get_agent_response("uncommitted") is None + + stateful.poison = False + response = await entity.run(request) + assert response.text == f"reply-{calls_before + 2}" + assert len(client.received_messages) == calls_before + 2 + assert [message.text for message in client.received_messages[-1]].count("uncommitted input") == 1 + assert provider.writes == writes_before + 1 + assert "uncommitted" in entity.state.data.completed_correlations + assert provider._get_state_dict()["data"]["session"]["state"]["unserializable"]["handle"] == "serializable" + + +class TestARequestIsAnsweredOnce: + """A committed correlation returns its recorded outcome; uncommitted effects may repeat.""" + + async def test_the_agent_does_not_run_twice(self) -> None: + client = RecordingChatClient() + entity = _make_entity(_build_agent(client), _InMemoryStateProvider()) + + await entity.run({"message": "what is the capital of Norway?", "correlationId": "dup"}) + await entity.run({"message": "what is the capital of Norway?", "correlationId": "dup"}) + + assert len(client.received_messages) == 1 + + async def test_the_same_answer_comes_back(self) -> None: + entity = _make_entity(_build_agent(RecordingChatClient()), _InMemoryStateProvider()) + + first = await entity.run({"message": "hello", "correlationId": "dup"}) + second = await entity.run({"message": "hello", "correlationId": "dup"}) + + assert first.text == second.text + + async def test_the_conversation_is_not_recorded_twice(self) -> None: + provider = _InMemoryStateProvider() + entity = _make_entity(_build_agent(RecordingChatClient()), provider) + + await entity.run({"message": "hello", "correlationId": "dup"}) + await entity.run({"message": "hello", "correlationId": "dup"}) + + entries = [e for e in entity.state.data.conversation_history if e.correlation_id == "dup"] + assert len(entries) == 2, "expected one request and one response, not a second pair" + + async def test_a_failed_turn_is_also_answered_once(self) -> None: + """The recorded failure comes back rather than the agent being run again. + + A caller retrying after a failure mints a new correlation id, so a repeat of this one is + still a duplicate delivery of the same request. + """ + + class _FailingClient(RecordingChatClient): + def get_response(self, messages: Any, **kwargs: Any) -> Any: + normalized = [m for m in messages if isinstance(m, Message)] if isinstance(messages, list) else [] + self.received_messages.append(normalized) + raise RuntimeError("kaboom") + + client = _FailingClient() + provider = _InMemoryStateProvider() + entity = _make_entity(_build_agent(client), provider) + + first = await entity.run({"message": "hello", "correlationId": "dup"}) + assert len(client.received_messages) == 1, "a failed stream must not trigger another agent execution" + raw = json.loads(json.dumps(provider._get_state_dict())) + assert raw["data"]["responseMailbox"]["dup"]["response"] == first.to_dict() + assert "dup" in raw["data"]["completedCorrelations"] + restarted_provider = _InMemoryStateProvider(raw=raw) + restarted = _make_entity(_build_agent(client), restarted_provider) + second = await restarted.run({"message": "hello", "correlationId": "dup"}) + + assert len(client.received_messages) == 1 + assert restarted_provider.writes == 0 + assert second.to_dict() == first.to_dict() + assert any(content.type == "error" for content in first.messages[0].contents) + assert any(content.type == "error" for content in second.messages[0].contents) + + async def test_a_different_request_still_runs(self) -> None: + """Only an exact repeat is short-circuited.""" + client = RecordingChatClient() + entity = _make_entity(_build_agent(client), _InMemoryStateProvider()) + + await entity.run({"message": "first", "correlationId": "c0"}) + await entity.run({"message": "second", "correlationId": "c1"}) + + assert len(client.received_messages) == 2 diff --git a/python/packages/durabletask/tests/test_durabletask_workflow_initial_input.py b/python/packages/durabletask/tests/test_durabletask_workflow_initial_input.py index a948d64..9034747 100644 --- a/python/packages/durabletask/tests/test_durabletask_workflow_initial_input.py +++ b/python/packages/durabletask/tests/test_durabletask_workflow_initial_input.py @@ -35,7 +35,14 @@ def supports_event_streaming(self) -> bool: def current_utc_datetime(self) -> datetime: return datetime.now(timezone.utc) - def prepare_agent_task(self, executor_id: str, message: str, orchestration_instance_id: str) -> Any: + def prepare_agent_task( + self, + executor_id: str, + message: str, + orchestration_instance_id: str, + context_messages: list[dict[str, Any]] | None = None, + context_message_ids: list[str] | None = None, + ) -> Any: raise AssertionError("This test workflow has no agent executors") def prepare_activity_task(self, activity_name: str, input_json: str) -> str: diff --git a/python/packages/durabletask/tests/test_execution_boundaries.py b/python/packages/durabletask/tests/test_execution_boundaries.py new file mode 100644 index 0000000..7b9aa10 --- /dev/null +++ b/python/packages/durabletask/tests/test_execution_boundaries.py @@ -0,0 +1,791 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Adversarial entity-operation boundaries using core agents and JSON storage, without live services.""" + +import json +from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence +from copy import deepcopy +from datetime import datetime, timedelta, timezone +from typing import Any + +import pytest +from agent_framework import ( + Agent, + AgentResponse, + AgentResponseUpdate, + AgentSession, + ChatOptions, + ChatResponse, + ChatResponseUpdate, + Content, + ContextProvider, + HistoryProvider, + InMemoryHistoryProvider, + Message, + ResponseStream, + SessionContext, + tool, +) +from test_durable_history_provider import RecordingChatClient +from test_history_pipeline_revision import CountingHistory, ToolChatClient +from test_revision_contract import JsonStateProvider + +from agent_framework_durabletask import AgentEntity, DurableAgentState, DurableHistoryProvider +from agent_framework_durabletask._callbacks import AgentCallbackContext +from agent_framework_durabletask._durable_agent_state import DurableAgentStateResponse +from agent_framework_durabletask._history_provider import current_durable_history_binding +from agent_framework_durabletask._message_identity import message_identity +from agent_framework_durabletask._state_migration import migrate_legacy_state, state_snapshot_digest + + +class _RecoverableExternalHistory(HistoryProvider): + """A primary whose reads fail until the test explicitly repairs the backing store.""" + + def __init__(self) -> None: + super().__init__("external") + self.fail_reads = True + self.read_sessions: list[str | None] = [] + self.saved: list[list[Message]] = [] + + async def get_messages(self, session_id: str | None, **kwargs: Any) -> list[Message]: + self.read_sessions.append(session_id) + if self.fail_reads: + raise OSError("temporary external history outage") + return deepcopy([message for batch in self.saved for message in batch]) + + async def save_messages(self, session_id: str | None, messages: Sequence[Message], **kwargs: Any) -> None: + self.saved.append(deepcopy(list(messages))) + + +class _ControlProvider(ContextProvider): + """Observe real core hooks without resolving the response's lazy structured value.""" + + def __init__(self) -> None: + super().__init__("control") + self.loaded: list[dict[str, Any]] = [] + self.inputs: list[list[Message]] = [] + self.sessions: list[AgentSession] = [] + self.agents: list[Any] = [] + self.responses: list[AgentResponse] = [] + + async def before_run( + self, + *, + agent: Any, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + self.loaded.append(deepcopy(state)) + self.inputs.append(deepcopy(context.input_messages)) + self.sessions.append(session) + self.agents.append(agent) + state["before_runs"] = state.get("before_runs", 0) + 1 + + async def after_run(self, *, context: SessionContext, state: dict[str, Any], **kwargs: Any) -> None: + assert isinstance(context.response, AgentResponse) + self.responses.append(context.response) + state["after_runs"] = state.get("after_runs", 0) + 1 + + +class _CountingAgent(Agent): + def __init__(self, *, client: Any, **kwargs: Any) -> None: + super().__init__(client=client, **kwargs) + self.run_modes: list[bool] = [] + + def run(self, *args: Any, **kwargs: Any) -> Any: + self.run_modes.append(bool(kwargs.get("stream", False))) + return super().run(*args, **kwargs) + + +class _FinalFlushFailureHistory(DurableHistoryProvider): + def __init__(self) -> None: + # Pin the policy so registration retains this provider rather than replacing it. + super().__init__(prune_excluded=False) + self.fail_final_flush = False + self.after_run_finished = False + self.failed_snapshot: dict[str, Any] | None = None + self.failures = 0 + + async def after_run(self, **kwargs: Any) -> None: + self.after_run_finished = False + await super().after_run(**kwargs) + self.after_run_finished = True + + def flush(self, state: dict[str, Any]) -> None: + if self.fail_final_flush and self.after_run_finished: + binding = current_durable_history_binding() + assert binding is not None + self.failed_snapshot = deepcopy(binding.state_provider.state.to_dict()) + self.failures += 1 + raise OSError("final durable history flush failed") + super().flush(state) + + +class _FailAfterFirstServiceCall(ToolChatClient): + def _inner_get_response( + self, *, messages: Sequence[Message], stream: bool, options: Mapping[str, Any], **kwargs: Any + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + # The first call requests a real tool. Every subsequent model attempt fails, + # including any inappropriate non-streaming fallback made by the entity. + self.fail = bool(self.received_messages) + return super()._inner_get_response(messages=messages, stream=stream, options=options, **kwargs) + + +class _InterruptedStreamClient(ToolChatClient): + """Yield a model update, then fail after an optional real core tool invocation.""" + + def __init__(self, *, use_tool: bool) -> None: + super().__init__(tool_calls=False) + self.use_tool = use_tool + self.stream_modes: list[bool] = [] + + def _inner_get_response( + self, *, messages: Sequence[Message], stream: bool, options: Mapping[str, Any], **kwargs: Any + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + self.received_messages.append(deepcopy(list(messages))) + self.received_options.append(dict(options)) + self.stream_modes.append(stream) + has_result = any( + content.type == "function_result" and content.call_id == "boundary-lookup" + for message in messages + for content in message.contents + ) + calls_tool = self.use_tool and not has_result + contents = ( + [Content.from_function_call("boundary-lookup", "lookup", arguments={"key": "durable"})] + if calls_tool + else [Content.from_text("partial answer")] + ) + response = ChatResponse( + messages=[Message("assistant", contents)], + response_id=f"boundary-response-{len(self.received_messages)}", + finish_reason="tool_calls" if calls_tool else "stop", + ) + if stream: + + async def updates() -> AsyncIterable[ChatResponseUpdate]: + yield ChatResponseUpdate( + role="assistant", + contents=contents, + response_id=response.response_id, + finish_reason=response.finish_reason, + ) + if not calls_tool: + raise RuntimeError("model stream interrupted after an update") + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) + + async def get() -> ChatResponse: + # A second invocation can succeed, but it must never hide a failed stream + # or repeat the tool requested by the first invocation. + return response + + return get() + + +class _NonStreamingClient(RecordingChatClient): + def get_response(self, messages: Any, *, stream: bool = False, **kwargs: Any) -> Any: + if stream: + raise TypeError("stream is not supported") + return super().get_response(messages, stream=False, **kwargs) + + +class _RecordingCallback: + def __init__(self) -> None: + self.updates: list[AgentResponseUpdate] = [] + self.responses: list[AgentResponse] = [] + + async def on_streaming_response_update(self, update: AgentResponseUpdate, context: AgentCallbackContext) -> None: + self.updates.append(deepcopy(update)) + + async def on_agent_response(self, response: AgentResponse, context: AgentCallbackContext) -> None: + self.responses.append(response) + + +def _committed(provider: JsonStateProvider) -> dict[str, Any]: + # Neither a cached state object nor a to_dict() alias proves a receipt was committed. + return json.loads(json.dumps(provider.raw)) + + +def _request(correlation_id: str, message: Message) -> dict[str, Any]: + return { + "message": message.text, + "correlationId": correlation_id, + "contextMessages": [deepcopy(message.to_dict())], + } + + +def _assert_committed_error( + provider: JsonStateProvider, + correlation_id: str, + response: AgentResponse, + *, + error_code: str, + detail: str, +) -> dict[str, Any]: + raw = _committed(provider) + assert raw["schemaVersion"] == "2.0.0" + data = raw["data"] + mailbox = data["responseMailbox"][correlation_id] + assert mailbox["response"] == json.loads(json.dumps(response.to_dict())) + assert data["completedCorrelations"][correlation_id] == {"completedAt": mailbox["createdAt"], "outcome": "failed"} + assert datetime.fromisoformat(mailbox["expiresAt"]) > datetime.fromisoformat(mailbox["createdAt"]) + delivered = DurableAgentState.from_json(json.dumps(raw)).try_get_agent_response(correlation_id) + assert isinstance(delivered, AgentResponse) + errors = [content for message in delivered.messages for content in message.contents if content.type == "error"] + assert len(errors) == 1 + assert errors[0].error_code == error_code + assert detail in (errors[0].message or "") + assert detail in delivered.text + assert delivered.value is None + return data + + +@pytest.mark.parametrize("per_call", [False, True], ids=["per-run", "per-service-call"]) +async def test_external_load_failure_commits_error_without_consuming_projected_input(per_call: bool) -> None: + external = _RecoverableExternalHistory() + client = ToolChatClient(tool_calls=False) + agent = Agent( + client=client, + context_providers=[external], + require_per_service_call_history_persistence=per_call, + ) + provider = JsonStateProvider() + entity = AgentEntity(agent, state_provider=provider) + message = Message("user", ["same projected payload"], message_id="upstream-0") + request = _request("external-failed", message) + + failed = await entity.run(request) + + data = _assert_committed_error( + provider, "external-failed", failed, error_code="OSError", detail="temporary external history outage" + ) + assert provider.writes == 1 + assert external.read_sessions and set(external.read_sessions) == {provider.core_session_id} + assert external.saved == [] and client.received_messages == [] + assert data["conversationHistory"] == [] + assert "upstream-0" not in data.get("ingestedMessages", {}) + original_failure = deepcopy(data["responseMailbox"]["external-failed"]) + original_receipt = deepcopy(data["completedCorrelations"]["external-failed"]) + + external.fail_reads = False + cold_provider = JsonStateProvider(_committed(provider)) + cold = AgentEntity(agent, state_provider=cold_provider) + reads_before_duplicate = len(external.read_sessions) + duplicate = await cold.run(request) + assert duplicate.to_dict() == failed.to_dict() + assert len(external.read_sessions) == reads_before_duplicate + assert client.received_messages == [] and cold_provider.writes == 0 + + # Same identity AND payload, but a new execution. The failed read never delivered it. + recovered = await cold.run(_request("external-recovered", message)) + assert recovered.text == "answer-1" + assert [[item.to_dict() for item in batch] for batch in client.received_messages] == [[message.to_dict()]] + assert [[item.text for item in batch] for batch in external.saved] == [[message.text, "answer-1"]] + assert set(external.read_sessions) == {cold_provider.core_session_id} + saved = _committed(cold_provider)["data"] + assert saved["conversationHistory"] == [] + assert saved["ingestedMessages"] == {"upstream-0": [message_identity(message)]} + assert saved["responseMailbox"]["external-failed"] == original_failure + assert saved["completedCorrelations"]["external-failed"] == original_receipt + assert cold_provider.writes == 1 + + before = _committed(cold_provider) + reads_before_duplicate = len(external.read_sessions) + assert (await cold.run(request)).to_dict() == failed.to_dict() + assert len(external.read_sessions) == reads_before_duplicate + assert len(client.received_messages) == 1 and len(external.saved) == 1 + assert _committed(cold_provider) == before and cold_provider.writes == 1 + + +async def test_lazy_invalid_structured_value_is_a_committed_error_with_provider_control_state() -> None: + session = AgentSession(session_id="revision-session") + control_state = {"approval": {"call_id": "pending-approval", "approved": False}, "cursor": [1, 3]} + session.state = {"control": deepcopy(control_state), "foreign-provider": {"pending": ["keep"]}} + initial = DurableAgentState() + initial.data.session = session.to_dict() + provider = JsonStateProvider(json.loads(initial.to_json())) + client: Any = RecordingChatClient() + control = _ControlProvider() + agent = Agent(client=client, context_providers=[control]) + entity = AgentEntity(agent, state_provider=provider) + request = { + "message": "return structured output", + "correlationId": "invalid-json", + "options": {"response_format": {"type": "object", "properties": {"answer": {"type": "integer"}}}}, + } + + response = await entity.run(request) + + data = _assert_committed_error( + provider, "invalid-json", response, error_code="ValueError", detail="Response text is not valid JSON" + ) + assert provider.writes == 1 and len(client.received_messages) == 1 + assert control.loaded == [control_state] + assert len(control.responses) == 1 and control.responses[0].text == "reply-1" + # This is core's actual lazy parser, not a fabricated exception from a response mock. + with pytest.raises(ValueError, match="not valid JSON"): + _ = control.responses[0].value + expected_control = {**control_state, "before_runs": 1, "after_runs": 1} + assert data["session"]["state"]["control"] == expected_control + assert data["session"]["state"]["foreign-provider"] == {"pending": ["keep"]} + + cold_control = _ControlProvider() + cold_provider = JsonStateProvider(_committed(provider)) + cold = AgentEntity(Agent(client=client, context_providers=[cold_control]), state_provider=cold_provider) + assert (await cold.run(request)).to_dict() == response.to_dict() + assert cold_control.loaded == [] and cold_control.responses == [] + assert len(client.received_messages) == 1 and cold_provider.writes == 0 + await cold.run({"message": "continue without a schema", "correlationId": "after-invalid-json"}) + assert cold_control.loaded == [expected_control] + assert _committed(cold_provider)["data"]["session"]["state"]["foreign-provider"] == {"pending": ["keep"]} + + +@pytest.mark.parametrize("per_call", [False, True], ids=["per-run", "per-service-call"]) +@pytest.mark.parametrize("explicit_run_id", [False, True], ids=["default-id", "default-and-run-id"]) +async def test_store_false_removes_default_and_saved_conversation_ids_without_mutating_caller( + per_call: bool, explicit_run_id: bool +) -> None: + class ServiceClient(ToolChatClient): + STORES_BY_DEFAULT = True + + client = ServiceClient(tool_calls=False) + defaults: ChatOptions = {"conversation_id": "stale", "store": True, "metadata": {"labels": ["caller"]}} + caller_defaults = deepcopy(defaults) + agent = Agent( + client=client, + default_options=defaults, + require_per_service_call_history_persistence=per_call, + ) + original_options = agent.default_options + original_options_value = deepcopy(original_options) + original_providers = agent.context_providers + provider = JsonStateProvider() + await AgentEntity(agent, state_provider=provider).run({"message": "service turn", "correlationId": "service"}) + first = _committed(provider) + assert first["data"]["session"]["service_session_id"] == "service-thread" + assert first["data"]["conversationHistory"] == [] + assert len(client.received_messages) == 1 + + cold_provider = JsonStateProvider(first) + cold = AgentEntity(agent, state_provider=cold_provider) + prepared_agent = cold.agent + options: dict[str, Any] = {"store": False} + if explicit_run_id: + options["conversation_id"] = "stale-per-run" + original_run_options = deepcopy(options) + response = await cold.run({"message": "client-owned turn", "correlationId": "local", "options": options}) + + assert response.text == "answer-2" + assert len(client.received_options) == 2 + assert client.received_options[-1]["store"] is False + assert "conversation_id" not in client.received_options[-1] + assert [message.text for message in client.received_messages[-1]] == ["client-owned turn"] + assert _committed(cold_provider)["data"]["session"]["service_session_id"] == "service-thread" + assert cold.agent is prepared_agent + assert agent.default_options is original_options and agent.default_options == original_options_value + assert agent.context_providers is original_providers + assert defaults == caller_defaults and options == original_run_options + assert current_durable_history_binding() is None + + +async def test_final_flush_failure_unbinds_restores_agent_and_rolls_back_all_local_state() -> None: + history = _FinalFlushFailureHistory() + control = _ControlProvider() + client = ToolChatClient(tool_calls=False) + agent = Agent( + client=client, + default_options={"conversation_id": "stale", "store": False}, + context_providers=[history, control], + ) + initial = DurableAgentState() + initial.data.session = AgentSession(session_id="revision-session", service_session_id="saved-service-id").to_dict() + provider = JsonStateProvider(json.loads(initial.to_json())) + entity = AgentEntity(agent, state_provider=provider) + await entity.run(_request("previous", Message("user", ["previous input"], message_id="previous-input"))) + before = _committed(provider) + original_state = entity.state + original_agent = entity.agent + original_options = agent.default_options + original_options_value = deepcopy(original_options) + history.fail_final_flush = True + request = _request("flush-failed", Message("user", ["uncommitted input"], message_id="uncommitted-input")) + assert current_durable_history_binding() is None + + with pytest.raises(OSError, match="final durable history flush failed"): + await entity.run(request) + + assert history.failures == 1 + assert history.failed_snapshot is not None + assert any( + entry.get("correlationId") == "flush-failed" for entry in history.failed_snapshot["data"]["conversationHistory"] + ), "the failure must occur after real history hooks staged this turn" + assert len(client.received_messages) == 2 and len(control.responses) == 2 + assert current_durable_history_binding() is None + assert control.agents[-1] is not original_agent, "exercise the temporary default-options clone" + assert control.sessions[-1].service_session_id == "saved-service-id" + assert entity.agent is original_agent + assert agent.default_options is original_options and agent.default_options == original_options_value + assert entity.state is original_state and entity.state.to_dict() == before + assert _committed(provider) == before and provider.writes == 1 + assert entity.state.try_get_agent_response("flush-failed") is None + assert "uncommitted-input" not in _committed(provider)["data"]["ingestedMessages"] + + history.fail_final_flush = False + response = await entity.run(request) + assert response.text == "answer-3" + assert len(client.received_messages) == 3 and provider.writes == 2 + assert control.loaded[-1] == before["data"]["session"]["state"]["control"] + assert "flush-failed" in _committed(provider)["data"]["completedCorrelations"] + assert current_durable_history_binding() is None and entity.agent is original_agent + + +@pytest.mark.parametrize("per_call", [False, True], ids=["per-run", "per-service-call"]) +@pytest.mark.parametrize("after_first_response", [False, True], ids=["before-first-response", "after-first-response"]) +async def test_failed_model_turn_consumes_only_inputs_actually_saved_by_history( + per_call: bool, after_first_response: bool +) -> None: + prior = Message("user", ["previously consumed"], message_id="prior-input") + initial = DurableAgentState() + initial.data.ingested_messages = {"prior-input": [message_identity(prior)]} + provider = JsonStateProvider(json.loads(initial.to_json())) + history = CountingHistory([]) + client = _FailAfterFirstServiceCall() if after_first_response else ToolChatClient(fail=True) + tool_calls: list[str] = [] + + @tool(name="lookup", approval_mode="never_require") + def lookup(key: str) -> str: + tool_calls.append(key) + return f"value:{key}" + + agent = Agent( + client=client, + tools=[lookup], + context_providers=[history], + require_per_service_call_history_persistence=per_call, + ) + entity = AgentEntity(agent, state_provider=provider) + message = Message("user", ["Use lookup for durable."], message_id="current-input") + request = _request("failed-model", message) + + response = await entity.run(request) + + data = _assert_committed_error( + provider, "failed-model", response, error_code="RuntimeError", detail="model failed before history persistence" + ) + was_saved = per_call and after_first_response + assert history.after_calls == int(was_saved) + assert tool_calls == (["durable"] if after_first_response else []) + expected_receipts = {"prior-input": [message_identity(prior)]} + if was_saved: + expected_receipts["current-input"] = [message_identity(message)] + assert data["ingestedMessages"] == expected_receipts + stored_inputs = [ + item + for entry in data["conversationHistory"] + if entry["$type"] == "request" and entry.get("correlationId") == "failed-model" + for item in entry["messages"] + ] + if was_saved: + assert len(stored_inputs) == 2 + assert stored_inputs[0]["messageId"] == "current-input" + assert stored_inputs[1]["role"] == "tool" + assert stored_inputs[1]["contents"][0]["$type"] == "functionResult" + else: + assert stored_inputs == [] + assert provider.writes == 1 + calls_before_duplicate = len(client.received_messages) + assert (await entity.run(request)).to_dict() == response.to_dict() + assert len(client.received_messages) == calls_before_duplicate and provider.writes == 1 + + healthy_client = ToolChatClient(tool_calls=False) + probe = _ControlProvider() + cold = AgentEntity( + Agent( + client=healthy_client, + context_providers=[probe], + require_per_service_call_history_persistence=per_call, + ), + state_provider=JsonStateProvider(_committed(provider)), + ) + recovered = await cold.run(_request("model-recovered", message)) + assert recovered.text == "answer-1" + expected_input_ids = [[]] if was_saved else [["current-input"]] + assert [[item.message_id for item in batch] for batch in probe.inputs] == expected_input_ids + assert [item.message_id for item in healthy_client.received_messages[0]].count("current-input") == 1 + + +@pytest.mark.parametrize("per_call", [False, True], ids=["per-run", "per-service-call"]) +async def test_partial_external_save_does_not_claim_a_local_ingestion_receipt(per_call: bool) -> None: + external = _RecoverableExternalHistory() + external.fail_reads = False + client = _FailAfterFirstServiceCall() + tool_calls: list[str] = [] + + @tool(name="lookup", approval_mode="never_require") + def lookup(key: str) -> str: + tool_calls.append(key) + return f"value:{key}" + + agent = Agent( + client=client, + tools=[lookup], + context_providers=[external], + require_per_service_call_history_persistence=per_call, + ) + provider = JsonStateProvider() + entity = AgentEntity(agent, state_provider=provider) + message = Message("user", ["Use lookup for durable."], message_id="external-partial-input") + request = _request("external-partial", message) + + response = await entity.run(request) + + data = _assert_committed_error( + provider, + "external-partial", + response, + error_code="RuntimeError", + detail="model failed before history persistence", + ) + assert tool_calls == ["durable"] and provider.writes == 1 + assert len(external.saved) == int(per_call) + if per_call: + assert external.saved[0][0].message_id == message.message_id + assert any(content.type == "function_call" for item in external.saved[0] for content in item.contents) + assert data["conversationHistory"] == [] + # External appends are outside the entity transaction. Even a saved first call + # cannot establish a portable local receipt for the interrupted whole run. + assert "external-partial-input" not in data.get("ingestedMessages", {}) + calls_before_duplicate = len(client.received_messages) + saved_before_duplicate = deepcopy(external.saved) + assert (await entity.run(request)).to_dict() == response.to_dict() + assert len(client.received_messages) == calls_before_duplicate + assert [[item.to_dict() for item in batch] for batch in external.saved] == [ + [item.to_dict() for item in batch] for batch in saved_before_duplicate + ] + assert tool_calls == ["durable"] and provider.writes == 1 + + +@pytest.mark.parametrize("use_tool", [False, True], ids=["model-stream", "tool-then-model-stream"]) +async def test_started_stream_failure_is_not_reexecuted_as_a_non_streaming_run(use_tool: bool) -> None: + client = _InterruptedStreamClient(use_tool=use_tool) + tool_calls: list[str] = [] + + @tool(name="lookup", approval_mode="never_require") + def lookup(key: str) -> str: + tool_calls.append(key) + return f"value:{key}" + + agent = _CountingAgent(client=client, tools=[lookup] if use_tool else []) + callback = _RecordingCallback() + provider = JsonStateProvider() + entity = AgentEntity(agent, callback=callback, state_provider=provider) + request = {"message": "start the operation", "correlationId": "interrupted-stream"} + + response = await entity.run(request) + + assert any(update.text == "partial answer" for update in callback.updates), "the model stream must actually start" + if use_tool: + assert any(content.type == "function_result" for update in callback.updates for content in update.contents) + assert agent.run_modes == [True], "a runtime stream failure must not start another agent/tool execution" + assert client.stream_modes == [True] * (2 if use_tool else 1) + assert tool_calls == (["durable"] if use_tool else []) + assert callback.responses == [] + _assert_committed_error( + provider, + "interrupted-stream", + response, + error_code="RuntimeError", + detail="model stream interrupted after an update", + ) + assert provider.writes == 1 + cold_provider = JsonStateProvider(_committed(provider)) + cold = AgentEntity(agent, state_provider=cold_provider) + assert (await cold.run(request)).to_dict() == response.to_dict() + assert agent.run_modes == [True] and cold_provider.writes == 0 + assert tool_calls == (["durable"] if use_tool else []) + + +async def test_unsupported_stream_type_error_still_allows_one_non_streaming_invocation() -> None: + client = _NonStreamingClient() + agent = _CountingAgent(client=client) + callback = _RecordingCallback() + provider = JsonStateProvider() + entity = AgentEntity(agent, callback=callback, state_provider=provider) + request = {"message": "non-streaming client", "correlationId": "unsupported-stream"} + + response = await entity.run(request) + + assert agent.run_modes == [True, False] + assert len(client.received_messages) == 1 and response.text == "reply-1" + assert callback.updates == [] and len(callback.responses) == 1 + assert callback.responses[0] is not response + assert callback.responses[0].to_dict() == response.to_dict() + assert callback.responses[0].messages[0] is not response.messages[0] + data = _committed(provider)["data"] + assert data["responseMailbox"]["unsupported-stream"]["response"] == response.to_dict() + assert "unsupported-stream" in data["completedCorrelations"] and provider.writes == 1 + assert (await entity.run(request)).to_dict() == response.to_dict() + assert agent.run_modes == [True, False] and provider.writes == 1 + + +def test_registration_fails_if_public_provider_list_cannot_be_replaced() -> None: + class ReadOnlyProvidersAgent(_CountingAgent): + @property + def context_providers(self) -> list[ContextProvider]: + return self._context_providers + + @context_providers.setter + def context_providers(self, providers: list[ContextProvider]) -> None: + if hasattr(self, "_context_providers"): + raise AttributeError("context_providers cannot be replaced after construction") + self._context_providers = providers + + agent = ReadOnlyProvidersAgent( + client=RecordingChatClient(), + name="read-only-providers", + context_providers=[InMemoryHistoryProvider("registered-history")], + ) + original = agent.context_providers + provider = JsonStateProvider() + with pytest.raises(ValueError, match="attach durable history"): + AgentEntity(agent, state_provider=provider) + assert agent.context_providers is original and agent.run_modes == [] + assert isinstance(original[0], InMemoryHistoryProvider) + assert _committed(provider) == {} and provider.writes == 0 + + +def test_registration_fails_if_core_agent_cannot_be_copied() -> None: + class UncopyableAgent(_CountingAgent): + def __copy__(self) -> Any: + raise TypeError("agent cannot be copied") + + agent = UncopyableAgent(client=RecordingChatClient(), context_providers=[InMemoryHistoryProvider()]) + original = agent.context_providers + provider = JsonStateProvider() + with pytest.raises(ValueError, match="attach durable history"): + AgentEntity(agent, state_provider=provider) + assert agent.context_providers is original and agent.run_modes == [] + assert _committed(provider) == {} and provider.writes == 0 + + +async def test_reset_clears_local_context_but_keeps_delivery_and_ingestion_receipts() -> None: + initial = DurableAgentState().to_dict() + initial["futureRoot"] = {"opaque": [1, 2]} + initial["data"]["futureData"] = {"opaque": [3, 4]} + initial["data"]["conversationHistory"] = [{"$type": "future-kind", "opaque": ["old local history"]}] + provider = JsonStateProvider(initial) + control = _ControlProvider() + initial_client: Any = RecordingChatClient() + entity = AgentEntity(Agent(client=initial_client, context_providers=[control]), state_provider=provider) + message = Message("user", ["before reset"], message_id="before-reset-input") + request = _request("before-reset", message) + original_response = await entity.run(request) + before = _committed(provider) + assert len(before["data"]["conversationHistory"]) > 1 + assert before["data"]["session"]["state"]["control"] == {"before_runs": 1, "after_runs": 1} + + entity.reset() + + reset = _committed(provider) + assert reset["data"]["conversationHistory"] == [] + assert "session" not in reset["data"] + # Explicit reset may delete even opaque local history; unrelated data is not history. + assert reset["futureRoot"] == before["futureRoot"] + assert reset["data"]["futureData"] == before["data"]["futureData"] + for field in ("responseMailbox", "completedCorrelations", "ingestedMessages"): + assert reset["data"][field] == before["data"][field] + assert provider.writes == 2 + + client: Any = RecordingChatClient() + cold_control = _ControlProvider() + cold_provider = JsonStateProvider(reset) + cold = AgentEntity(Agent(client=client, context_providers=[cold_control]), state_provider=cold_provider) + assert (await cold.run(request)).to_dict() == original_response.to_dict() + assert client.received_messages == [] and cold_control.loaded == [] and cold_provider.writes == 0 + await cold.run(_request("after-reset", Message("user", ["after reset"], message_id="after-reset-input"))) + assert cold_control.loaded == [{}] + assert [[item.text for item in batch] for batch in client.received_messages] == [["after reset"]] + saved = _committed(cold_provider) + assert saved["data"]["session"]["state"]["control"] == {"before_runs": 1, "after_runs": 1} + assert saved["data"]["responseMailbox"]["before-reset"] == before["data"]["responseMailbox"]["before-reset"] + assert ( + saved["data"]["completedCorrelations"]["before-reset"] + == before["data"]["completedCorrelations"]["before-reset"] + ) + assert saved["data"]["ingestedMessages"]["before-reset-input"] == [message_identity(message)] + assert (await cold.run(request)).to_dict() == original_response.to_dict() + assert len(client.received_messages) == 1 and cold_provider.writes == 1 + + +@pytest.mark.parametrize("legacy", [False, True], ids=["writer", "legacy-migration"]) +def test_response_writer_and_migration_keep_completion_evidence_after_payload_expiry(legacy: bool) -> None: + response = AgentResponse(messages=[Message("assistant", ["original result"])]) + state = DurableAgentState("1.1.0" if legacy else "2.0.0") + if legacy: + state.data.conversation_history.append(DurableAgentStateResponse.from_run_response("completed", response)) + source = state.to_dict() + state = migrate_legacy_state( + source, + source_digest=state_snapshot_digest(source), + source_session_id="source-session", + migration_id="expiry-migration", + ownership_transfer_id="quiesced-owner", + delivery_window_seconds=3600, + ) + else: + state.record_response("completed", response, delivery_window_seconds=3600) + raw = json.loads(state.to_json()) + assert raw["schemaVersion"] == "2.0.0" + mailbox = raw["data"]["responseMailbox"]["completed"] + receipt = raw["data"]["completedCorrelations"]["completed"] + assert receipt["completedAt"] == mailbox["createdAt"] + assert receipt.get("legacy", False) is legacy + + restored = DurableAgentState.from_json(json.dumps(raw)) + restored.expire_responses(now=datetime.fromisoformat(mailbox["expiresAt"])) + expired = DurableAgentState.from_json(restored.to_json()) + assert expired.data.response_mailbox == {} + assert expired.data.completed_correlations["completed"] == receipt + delivered = expired.try_get_agent_response("completed") + assert isinstance(delivered, AgentResponse) + assert delivered.additional_properties["durable_status"] == "already_completed" + assert delivered.messages[0].contents[0].error_code == "response_expired" + before = expired.to_json() + expired.record_response( + "completed", AgentResponse(messages=[Message("assistant", ["replacement"])]), delivery_window_seconds=3600 + ) + assert expired.to_json() == before + + +@pytest.mark.parametrize("expired", [False, True], ids=["live-mailbox", "expired-mailbox"]) +@pytest.mark.parametrize("receipt_shape", ["missing-container", "empty-container", "unrelated-receipt"]) +def test_version_two_mailbox_without_matching_receipt_is_rejected_on_initial_read( + expired: bool, receipt_shape: str +) -> None: + state = DurableAgentState() + now = datetime.now(timezone.utc) - (timedelta(days=2) if expired else timedelta()) + state.record_response( + "completed", + AgentResponse(messages=[Message("assistant", ["original result"])]), + delivery_window_seconds=3600, + now=now, + ) + raw = json.loads(state.to_json()) + # Positive control: this is an otherwise valid writer-produced mailbox and receipt. + assert DurableAgentState.from_json(json.dumps(raw)).to_dict() == raw + receipt = raw["data"]["completedCorrelations"].pop("completed") + if receipt_shape == "missing-container": + del raw["data"]["completedCorrelations"] + elif receipt_shape == "unrelated-receipt": + raw["data"]["completedCorrelations"]["different-correlation"] = receipt + original = deepcopy(raw) + + # An orphan mailbox must not be the last completion evidence: expiry could delete + # it and reopen execution. Reject corruption before polling or writing another turn. + with pytest.raises(ValueError): + DurableAgentState.from_dict(raw) + with pytest.raises(ValueError): + DurableAgentState.from_json(json.dumps(original)) + assert raw == original diff --git a/python/packages/durabletask/tests/test_execution_followup_review.py b/python/packages/durabletask/tests/test_execution_followup_review.py new file mode 100644 index 0000000..3e90465 --- /dev/null +++ b/python/packages/durabletask/tests/test_execution_followup_review.py @@ -0,0 +1,823 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Execution follow-ups through real core invocation and detached JSON entity storage.""" + +from __future__ import annotations + +import json +from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass +from typing import Any, cast + +import pytest +from agent_framework import ( + Agent, + AgentExecutor, + AgentResponse, + AgentResponseUpdate, + AgentSession, + ChatResponse, + ChatResponseUpdate, + Content, + ContextProvider, + FunctionInvocationLayer, + Message, + ResponseStream, + SessionContext, + SupportsAgentRun, + WorkflowBuilder, + tool, +) +from pydantic import BaseModel, ValidationError +from test_history_pipeline_revision import ToolChatClient +from test_revision_contract import JsonStateProvider + +from agent_framework_durabletask import AgentEntity, DurableAgentState, DurableHistoryProvider, RunRequest +from agent_framework_durabletask import _entities as entities_module +from agent_framework_durabletask._callbacks import AgentCallbackContext +from agent_framework_durabletask._history_provider import ( + POSITIONS_KEY, + WORKING_BUFFER_KEY, + current_durable_history_binding, +) +from agent_framework_durabletask._invocation_safety import DurableToolGuard +from agent_framework_durabletask._message_identity import message_identity +from agent_framework_durabletask._response_utils import ensure_response_format, serialize_agent_response +from agent_framework_durabletask._state_migration import migrate_legacy_state, state_snapshot_digest +from agent_framework_durabletask._workflows.naming import workflow_message_id + + +def _object(value: object) -> dict[str, object]: + assert isinstance(value, dict) + assert all(isinstance(key, str) for key in value) + return cast(dict[str, object], value) + + +def _array(value: object) -> list[object]: + assert isinstance(value, list) + return cast(list[object], value) + + +def _wire(value: object) -> dict[str, object]: + return _object(json.loads(json.dumps(value, allow_nan=False))) + + +def _data(provider: JsonStateProvider) -> dict[str, object]: + return _object(_wire(provider.raw)["data"]) + + +def _mailbox(provider: JsonStateProvider, correlation: str) -> dict[str, object]: + return _object(_object(_object(_data(provider)["responseMailbox"])[correlation])["response"]) + + +def _delivered(provider: JsonStateProvider, correlation: str) -> AgentResponse[Any]: + response = DurableAgentState.from_json(json.dumps(provider.raw)).try_get_agent_response(correlation) + assert isinstance(response, AgentResponse) + return response + + +class _ObservedAgent(Agent): + """Observe the Agent.run boundary without replacing core execution.""" + + def __init__(self, *, client: Any, streaming: bool = True, **kwargs: Any) -> None: + super().__init__(client=client, **kwargs) + self.streaming = streaming + self.run_modes: list[bool] = [] + self.run_client_kwargs: list[dict[str, object]] = [] + self.sessions: list[AgentSession] = [] + + def run(self, *args: Any, **kwargs: Any) -> Any: + if kwargs.get("stream") and not self.streaming: + raise TypeError("stream is not supported") + self.run_modes.append(bool(kwargs.get("stream"))) + self.run_client_kwargs.append(dict(kwargs.get("client_kwargs") or {})) + session = kwargs.get("session") + assert isinstance(session, AgentSession) + self.sessions.append(session) + return super().run(*args, **kwargs) + + +class _DelegatingClient: + """A non-invocation wrapper whose declared configuration belongs to its inner client.""" + + def __init__(self, inner: ToolChatClient) -> None: + self.inner = inner + self.forwarded: list[dict[str, object]] = [] + self.inner_configurations: list[dict[str, object]] = [] + + def __getattr__(self, name: str) -> object: + # Do not delegate copy/pickle special methods or fabricate arbitrary attributes. + if name in {"function_invocation_configuration", "additional_properties"}: + return getattr(self.inner, name) + raise AttributeError(name) + + def get_response( + self, messages: Sequence[Message], *, stream: bool = False, **kwargs: Any + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + self.forwarded.append(dict(kwargs.get("client_kwargs") or {})) + self.inner_configurations.append(dict(self.inner.function_invocation_configuration)) + # Crucially, this uses inner's configuration, not a shadow assigned to a wrapper copy. + return cast(Callable[..., Any], self.inner.get_response)(messages=messages, stream=stream, **kwargs) + + +@pytest.mark.parametrize("per_call", [False, True], ids=["per-run", "per-service-call"]) +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +@pytest.mark.parametrize("enabled", [False, True], ids=["disabled", "enabled-control"]) +async def test_tool_guard_reaches_delegated_core_invocation_without_mutating_configuration( + per_call: bool, stream: bool, enabled: bool +) -> None: + calls: list[str] = [] + + @tool(name="lookup", approval_mode="never_require") + def lookup(key: str) -> str: + calls.append(key) + return f"value:{key}" + + class ProviderTools(ContextProvider): + async def before_run(self, *, context: SessionContext, **kwargs: Any) -> None: + context.tools.append(lookup) + + inner = ToolChatClient() + wrapper = _DelegatingClient(inner) + assert not isinstance(wrapper, FunctionInvocationLayer) + configuration = inner.function_invocation_configuration + original_configuration = deepcopy(configuration) + assert wrapper.function_invocation_configuration is configuration + agent = _ObservedAgent( + client=wrapper, + streaming=stream, + context_providers=[ProviderTools("provider-tools")], + require_per_service_call_history_persistence=per_call, + ) + defaults, providers = agent.default_options, agent.context_providers + original_defaults = deepcopy(defaults) + initial_session = AgentSession(session_id="revision-session", service_session_id="parked-service-id") + initial_session.state["foreign"] = {"pending": ["keep"]} + session_snapshot = deepcopy(initial_session.to_dict()) + initial = DurableAgentState() + initial.data.session = initial_session.to_dict() + provider = JsonStateProvider(_wire(initial.to_dict())) + entity = AgentEntity(agent, state_provider=provider) + registered = entity.agent + request = {"message": "use lookup", "correlationId": "guard", "enable_tool_calls": enabled} + before_request = deepcopy(request) + + response = await entity.run(request) + + assert response.text == "answer-2" + assert len(inner.received_messages) == 2 and agent.run_modes == [stream] + assert calls == (["durable"] if enabled else []) + assert len(wrapper.forwarded) == 1 + guards = [item for item in _array(agent.run_client_kwargs[0]["middleware"]) if isinstance(item, DurableToolGuard)] + assert len(guards) == 1 and guards[0].enabled is enabled + assert guards[0] in _array(wrapper.forwarded[0]["middleware"]) + assert guards[0].progress.function_started is enabled + assert wrapper.forwarded[0]["session"] is agent.sessions[0] + assert wrapper.inner_configurations == [original_configuration] + results = [ + content + for message in inner.received_messages[1] + for content in message.contents + if content.type == "function_result" + ] + assert len(results) == 1 and results[0].call_id == "call-1" + assert results[0].result == ("value:durable" if enabled else "Tool execution is disabled for this invocation.") + if not enabled: + assert inner.received_options[0]["tool_choice"] == "none" + assert inner.function_invocation_configuration is configuration + assert configuration == original_configuration + assert wrapper.function_invocation_configuration is configuration + assert agent.default_options is defaults and defaults == original_defaults + assert agent.context_providers is providers and agent.client is wrapper + assert entity.agent is registered and request == before_request + assert initial_session.to_dict() == session_snapshot + assert agent.sessions[0] is not initial_session + saved_session = _object(_data(provider)["session"]) + assert saved_session["service_session_id"] == "parked-service-id" + assert _object(saved_session["state"])["foreign"] == {"pending": ["keep"]} + assert _delivered(provider, "guard").to_dict() == response.to_dict() + assert current_durable_history_binding() is None + + +class _PreviousResponseMissing(RuntimeError): + code = "previous_response_not_found" + + +class _VisibilityClient(ToolChatClient): + """Fail a real model boundary, with a successful third call available to expose restarts.""" + + def __init__(self, *, tool_followup: bool = False, partial: bool = False) -> None: + super().__init__(tool_calls=tool_followup) + self.tool_followup = tool_followup + self.partial = partial + self.failure = _PreviousResponseMissing("service parent is not visible") + + def _inner_get_response( + self, *, messages: Sequence[Message], stream: bool, options: Mapping[str, Any], **kwargs: Any + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + call = len(self.received_messages) + 1 + failing_call = 2 if self.tool_followup else 1 + if call != failing_call: + return super()._inner_get_response(messages=messages, stream=stream, options=options, **kwargs) + self.received_messages.append(deepcopy(list(messages))) + self.received_options.append(dict(options)) + if stream: + + async def updates() -> AsyncIterable[ChatResponseUpdate]: + if self.partial: + # No conversation ID: this case must be stopped by output progress alone. + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("partial visible answer")]) + raise self.failure + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) + + async def get() -> ChatResponse: + raise self.failure + + return get() + + +class _CallbackRecorder: + def __init__(self) -> None: + self.updates: list[AgentResponseUpdate] = [] + self.responses: list[AgentResponse[Any]] = [] + + async def on_streaming_response_update(self, update: AgentResponseUpdate, context: AgentCallbackContext) -> None: + self.updates.append(update) + + async def on_agent_response(self, response: AgentResponse[Any], context: AgentCallbackContext) -> None: + self.responses.append(response) + + +def _service_state() -> dict[str, object]: + state = DurableAgentState() + state.data.session = AgentSession( + session_id="revision-session", service_session_id="original-service-parent" + ).to_dict() + return _wire(state.to_dict()) + + +def _assert_missing_error(provider: JsonStateProvider, response: AgentResponse[Any], correlation: str) -> None: + assert response.additional_properties["durable_status"] == "error" + assert response.text == "_PreviousResponseMissing: service parent is not visible" + errors = [content for message in response.messages for content in message.contents if content.type == "error"] + assert len(errors) == 1 and errors[0].error_code == "_PreviousResponseMissing" + assert _delivered(provider, correlation).to_dict() == response.to_dict() + assert correlation in _object(_data(provider)["completedCorrelations"]) + assert provider.writes == 1 + + +@pytest.mark.parametrize("per_call", [False, True], ids=["per-run", "per-service-call"]) +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +async def test_missing_parent_on_tool_followup_does_not_restart_the_agent( + monkeypatch: pytest.MonkeyPatch, per_call: bool, stream: bool +) -> None: + monkeypatch.setattr(entities_module, "_REJECTED_ID_BACKOFF_SECONDS", 0) + calls: list[str] = [] + + @tool(name="lookup", approval_mode="never_require") + def lookup(key: str) -> str: + calls.append(key) + return f"value:{key}" + + client = _VisibilityClient(tool_followup=True) + agent = _ObservedAgent( + client=_DelegatingClient(client), + streaming=stream, + tools=[lookup], + default_options={"store": True}, + require_per_service_call_history_persistence=per_call, + ) + provider = JsonStateProvider(_service_state()) + request = {"message": "use lookup", "correlationId": "failed-followup"} + + response = await AgentEntity(agent, state_provider=provider).run(request) + + assert len(client.received_messages) == 2, "a successful third model call must not hide the follow-up error" + assert agent.run_modes == [stream], "retrying the whole Agent.run would restart the tool loop" + assert calls == ["durable"] + assert [options.get("conversation_id") for options in client.received_options] == [ + "original-service-parent", + "service-thread", + ] + results = [ + content + for message in client.received_messages[1] + for content in message.contents + if content.type == "function_result" + ] + assert len(results) == 1 and results[0].call_id == "call-1" and results[0].result == "value:durable" + assert agent.sessions[0].service_session_id == "service-thread", "the first response advanced the session" + assert _object(_data(provider)["session"])["service_session_id"] == "service-thread" + _assert_missing_error(provider, response, "failed-followup") + cold_provider = JsonStateProvider(_wire(provider.raw)) + duplicate = await AgentEntity(agent, state_provider=cold_provider).run(request) + assert duplicate.to_dict() == response.to_dict() and cold_provider.writes == 0 + assert len(client.received_messages) == 2 and calls == ["durable"] and agent.run_modes == [stream] + + +async def test_partial_stream_missing_parent_is_not_retryable(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(entities_module, "_REJECTED_ID_BACKOFF_SECONDS", 0) + client = _VisibilityClient(partial=True) + agent = _ObservedAgent(client=client, default_options={"store": True}) + callback = _CallbackRecorder() + provider = JsonStateProvider(_service_state()) + + response = await AgentEntity(agent, callback=callback, state_provider=provider).run({ + "message": "continue", + "correlationId": "partial-stream", + }) + + assert [update.text for update in callback.updates] == ["partial visible answer"] + assert callback.responses == [] + assert agent.sessions[0].service_session_id == "original-service-parent" + assert len(client.received_messages) == 1 and agent.run_modes == [True] + _assert_missing_error(provider, response, "partial-stream") + + +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +async def test_zero_output_first_request_missing_parent_still_retries_identically( + monkeypatch: pytest.MonkeyPatch, stream: bool +) -> None: + monkeypatch.setattr(entities_module, "_REJECTED_ID_BACKOFF_SECONDS", 0) + client = _VisibilityClient() + agent = _ObservedAgent(client=client, streaming=stream, default_options={"store": True}) + provider = JsonStateProvider(_service_state()) + callback = _CallbackRecorder() + + response = await AgentEntity(agent, callback=callback, state_provider=provider).run({ + "message": "continue", + "correlationId": "zero-output", + }) + + assert response.text == "answer-2" and response.additional_properties.get("durable_status") != "error" + assert len(client.received_messages) == 2 and agent.run_modes == [stream, stream] + assert client.received_options[0] == client.received_options[1] + assert client.received_options[0]["conversation_id"] == "original-service-parent" + assert [[message.to_dict() for message in batch] for batch in client.received_messages] == [ + [message.to_dict() for message in client.received_messages[0]] + ] * 2 + assert agent.sessions[0] is agent.sessions[1] + assert len(callback.responses) == 1 + assert len(callback.updates) == int(stream) + assert all(update.text == "answer-2" for update in callback.updates) + assert _delivered(provider, "zero-output").text == "answer-2" and provider.writes == 1 + + +class _NestedValue(BaseModel): + values: list[int] + + +class ReviewValue(BaseModel): + nested: _NestedValue + + +@dataclass +class _SDKPayload: + labels: list[str] + + +class _UncopyableSDK: + def __deepcopy__(self, memo: dict[int, object]) -> _UncopyableSDK: + raise TypeError("opaque SDK handle cannot be copied") + + +class _ReplyAgent: + """Custom non-pipeline agent for responses that should not be interpreted as core runs.""" + + name = "custom-response" + id = "custom-response" + description = None + + def __init__(self, response: AgentResponse[Any]) -> None: + self.response = response + self.inputs: list[list[Message]] = [] + + def create_session(self, **kwargs: Any) -> AgentSession: + return AgentSession(**kwargs) + + async def run(self, messages: list[Message], **kwargs: Any) -> AgentResponse[Any]: + self.inputs.append(deepcopy(messages)) + return self.response + + +class _TypedMutatingCallback(_CallbackRecorder): + def __init__(self) -> None: + super().__init__() + self.mutations: list[str] = [] + + async def on_agent_response(self, response: AgentResponse[Any], context: AgentCallbackContext) -> None: + self.responses.append(response) + value = response.value + if isinstance(value, ReviewValue): + value.nested.values.append(99) + self.mutations.append("typed-value") + response.messages[0].contents[0].text = "callback changed text" + response.messages[0].contents[0].additional_properties["source"]["labels"].append("callback") + self.mutations.append("content") + if isinstance(response.raw_representation, _SDKPayload): + response.raw_representation.labels.append("callback") + self.mutations.append("raw") + + +@pytest.mark.parametrize("lazy", [False, True], ids=["already-typed", "lazy-typed"]) +@pytest.mark.parametrize("opaque", [False, True], ids=["copyable-sdk", "uncopyable-sdk"]) +async def test_final_callback_keeps_model_format_and_detaches_value_content_and_sdk(lazy: bool, opaque: bool) -> None: + model = ReviewValue(nested=_NestedValue(values=[1, 2])) + sdk = _UncopyableSDK() if opaque else _SDKPayload(["original"]) + original: AgentResponse[Any] = AgentResponse( + messages=[ + Message( + "assistant", + [ + Content.from_text( + model.model_dump_json(), additional_properties={"source": {"labels": ["original"]}} + ) + ], + ) + ], + value=None if lazy else model, + response_format=ReviewValue, + raw_representation=sdk, + ) + expected = _wire(serialize_agent_response(original)) + agent = _ReplyAgent(original) + callback = _TypedMutatingCallback() + provider = JsonStateProvider() + request = RunRequest("return typed output", "typed-callback", response_format=ReviewValue) + + response = await AgentEntity(cast(SupportsAgentRun, agent), callback=callback, state_provider=provider).run(request) + + # Callback exceptions are swallowed by the host, so verify completed mutations outside it. + assert callback.mutations == ["typed-value", "content"] + ([] if opaque else ["raw"]) + assert len(callback.responses) == 1 and callback.updates == [] + snapshot = callback.responses[0] + assert snapshot is not original and response is original + assert isinstance(snapshot.value, ReviewValue) and isinstance(original.value, ReviewValue) + assert snapshot.value is not original.value and snapshot.value.nested is not original.value.nested + assert snapshot.value.nested.values == [1, 2, 99] + assert snapshot._response_format is ReviewValue + assert original._response_format is ReviewValue + assert original.value.nested.values == [1, 2] + assert snapshot.messages[0].contents[0] is not original.messages[0].contents[0] + assert original.messages[0].contents[0].additional_properties == {"source": {"labels": ["original"]}} + assert original.raw_representation is sdk + if opaque: + # Only an uncopyable opaque SDK field may be omitted, never the rest of the callback response. + assert snapshot.raw_representation is None + else: + assert isinstance(sdk, _SDKPayload) and sdk.labels == ["original"] + assert isinstance(snapshot.raw_representation, _SDKPayload) + assert snapshot.raw_representation is not sdk and snapshot.raw_representation.labels == ["original", "callback"] + assert _wire(serialize_agent_response(original)) == expected + assert _mailbox(provider, "typed-callback") == expected + assert "raw_representation" not in expected and "response_format" not in expected + delivered = _delivered(provider, "typed-callback") + ensure_response_format(ReviewValue, "typed-callback", delivered) + assert delivered.value == ReviewValue(nested=_NestedValue(values=[1, 2])) + before = _wire(provider.raw) + snapshot.value.nested.values.append(101) + snapshot.messages[0].contents[0].text = "late callback mutation" + assert _wire(provider.raw) == before and original.value.nested.values == [1, 2] + cold_provider = JsonStateProvider(before) + duplicate = await AgentEntity(cast(SupportsAgentRun, agent), callback=callback, state_provider=cold_provider).run( + request + ) + assert _wire(serialize_agent_response(duplicate)) == expected + assert len(agent.inputs) == 1 and len(callback.responses) == 1 and cold_provider.writes == 0 + + +async def test_custom_terminal_text_skips_typed_validation_and_is_not_replayed_next_turn() -> None: + original = AgentResponse( + messages=[Message("assistant", ["original terminal text, not JSON"])], + response_format=ReviewValue, + additional_properties={"durable_status": "error", "provider_detail": {"labels": ["keep"]}}, + ) + # This is a genuinely invalid lazy value, not an inert format marker. + with pytest.raises(ValidationError): + _ = deepcopy(original).value + agent = _ReplyAgent(original) + provider = JsonStateProvider() + request = RunRequest("first input", "terminal", response_format=ReviewValue) + + response = await AgentEntity(cast(SupportsAgentRun, agent), state_provider=provider).run(request) + + assert response is original and response.text == "original terminal text, not JSON" + delivered = _delivered(provider, "terminal") + assert delivered.text == original.text and delivered.additional_properties == original.additional_properties + assert delivered.value is None + assert all(content.type == "text" for message in delivered.messages for content in message.contents) + assert "value" not in _mailbox(provider, "terminal") + assert all(_object(entry)["$type"] == "request" for entry in _array(_data(provider)["conversationHistory"])) + cold_provider = JsonStateProvider(_wire(provider.raw)) + cold = AgentEntity(cast(SupportsAgentRun, agent), state_provider=cold_provider) + assert (await cold.run(request)).to_dict() == delivered.to_dict() + assert len(agent.inputs) == 1 and cold_provider.writes == 0 + agent.response = AgentResponse(messages=[Message("assistant", ["next answer"])]) + assert (await cold.run({"message": "second input", "correlationId": "next"})).text == "next answer" + assert [message.text for message in agent.inputs[1]] == ["first input", "second input"] + assert _mailbox(cold_provider, "terminal") == _mailbox(provider, "terminal") + + +class _MessageClient(ToolChatClient): + def __init__(self, response_message: Message) -> None: + super().__init__(tool_calls=False) + self.response_message = response_message + + def _inner_get_response( + self, *, messages: Sequence[Message], stream: bool, options: Mapping[str, Any], **kwargs: Any + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + self.received_messages.append(deepcopy(list(messages))) + self.received_options.append(dict(options)) + message = deepcopy(self.response_message) + if stream: + + async def updates() -> AsyncIterable[ChatResponseUpdate]: + yield ChatResponseUpdate( + role=cast(Any, message.role), + contents=message.contents, + message_id=message.message_id, + additional_properties=message.additional_properties, + ) + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) + + async def get() -> ChatResponse: + return ChatResponse(messages=[message]) + + return get() + + +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +async def test_core_user_input_request_precedes_lazy_typed_parsing(stream: bool) -> None: + client = _MessageClient( + Message( + "assistant", + [ + Content.from_text("Approval required, not a typed JSON answer"), + Content.from_function_approval_request( + "approval-1", Content.from_function_call("call-1", "lookup", arguments={"key": "durable"}) + ), + ], + ) + ) + agent = _ObservedAgent(client=client, streaming=stream) + provider = JsonStateProvider() + request = RunRequest("request approval", "approval", response_format=ReviewValue) + + response = await AgentEntity(agent, state_provider=provider).run(request) + + assert len(client.received_messages) == 1 + assert response.additional_properties.get("durable_status") != "error" + assert response.text == "Approval required, not a typed JSON answer" + assert len(response.user_input_requests) == 1 and response.user_input_requests[0].id == "approval-1" + with pytest.raises(ValidationError): + _ = deepcopy(response).value + assert "value" not in _mailbox(provider, "approval") + delivered = _delivered(provider, "approval") + ensure_response_format(ReviewValue, "approval", delivered) + assert delivered.value is None and delivered.text == response.text + assert delivered.user_input_requests[0].to_dict() == response.user_input_requests[0].to_dict() + cold_provider = JsonStateProvider(_wire(provider.raw)) + assert (await AgentEntity(agent, state_provider=cold_provider).run(request)).to_dict() == delivered.to_dict() + assert len(client.received_messages) == 1 and cold_provider.writes == 0 + + +class _Inputs(ContextProvider): + def __init__(self) -> None: + super().__init__("input-probe") + self.inputs: list[list[Message]] = [] + + async def before_run(self, *, context: SessionContext, **kwargs: Any) -> None: + self.inputs.append(deepcopy(context.input_messages)) + + +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +async def test_rich_image_context_and_paired_ids_cross_entity_without_decoding_opaque_outputs(stream: bool) -> None: + outputs = [{"type": "text", "provider_only": {"type": "error", "pixels": [0, False, None, "雪"]}}] + message = Message( + "assistant", + [Content.from_image_generation_tool_result(image_id="image-1", outputs=deepcopy(outputs))], + message_id="shared-app-id", + additional_properties={"origin": {"labels": ["keep"]}}, + ) + raw_message = _wire(message.to_dict()) + raw_message["future_message"] = {"opaque": [1]} + _object(_array(raw_message["contents"])[0])["future_content"] = {"opaque": [2]} + request = { + "message": "logging only", + "correlationId": "rich-input", + "contextMessages": [deepcopy(raw_message), deepcopy(raw_message)], + "contextMessageIds": ["image-occurrence-1", "image-occurrence-2"], + } + before = deepcopy(request) + client = _MessageClient(deepcopy(message)) + probe = _Inputs() + provider = JsonStateProvider() + agent = _ObservedAgent(client=client, streaming=stream, context_providers=[probe]) + + response = await AgentEntity(agent, state_provider=provider).run(request) + + assert response.additional_properties.get("durable_status") != "error" + assert response.messages[0].contents[0].outputs == outputs and len(client.received_messages) == 1 + assert [item.to_dict() for item in probe.inputs[0]] == [message.to_dict()] * 2 + assert len(client.received_messages[0]) == 2 + for item in client.received_messages[0]: + assert item.message_id == "shared-app-id" + assert item.contents[0].type == "image_generation_tool_result" + assert item.contents[0].outputs == outputs + assert isinstance(item.contents[0].outputs[0], dict) + assert request == before and message.contents[0].outputs == outputs + assert _data(provider)["ingestedMessages"] == { + identity: [message_identity(message)] for identity in ("image-occurrence-1", "image-occurrence-2") + } + raw = _wire(provider.raw) + mailbox = _object(_object(_object(_object(raw["data"])["responseMailbox"])["rich-input"])["response"]) + assert mailbox == _wire(serialize_agent_response(response)) + mailbox["future_response"] = {"opaque": [3]} + # Simulate a newer writer adding optional fields to the actual committed model response. + saved_message = _object(_array(mailbox["messages"])[0]) + saved_message["future_message"] = {"opaque": [1]} + _object(_array(saved_message["contents"])[0])["future_content"] = {"opaque": [2]} + expected_raw = deepcopy(mailbox) + cold_provider = JsonStateProvider(raw) + cold = AgentEntity(agent, state_provider=cold_provider) + delivered = await cold.run(request) + assert delivered.messages[0].contents[0].outputs == outputs + delivered_output = delivered.messages[0].contents[0].outputs[0] + assert isinstance(delivered_output, dict) + delivered_output["provider_only"]["pixels"].append("consumer edit") + assert _mailbox(cold_provider, "rich-input") == expected_raw + assert len(client.received_messages) == 1 and cold_provider.writes == 0 + followup = await cold.run({**request, "correlationId": "rich-next"}) + assert followup.additional_properties.get("durable_status") != "error" + assert probe.inputs[-1] == [] + assert len(client.received_messages) == 2 + assert len(client.received_messages[1]) == 3, "two ingested occurrences plus the actual model response" + assert all(item.contents[0].outputs == outputs for item in client.received_messages[1]) + assert _mailbox(cold_provider, "rich-input") == expected_raw and cold_provider.writes == 1 + assert request == before + + +async def test_contentless_migrated_workflow_id_never_backfills_an_ingestion_receipt() -> None: + identity = workflow_message_id("upstream", 3) + source = { + "schemaVersion": "1.1.0", + "data": { + "conversationHistory": [ + { + "$type": "request", + "correlationId": "legacy", + "createdAt": "2024-01-01T00:00:00+00:00", + "messages": [{"role": "user", "messageId": identity, "contents": []}], + } + ] + }, + } + before_source = deepcopy(source) + migrated = migrate_legacy_state( + source, + source_digest=state_snapshot_digest(source), + source_session_id="legacy-session", + migration_id="execution-followup", + ownership_transfer_id="quiesced-owner", + delivery_window_seconds=3600, + ) + assert migrated.data.ingested_messages == {} + probe = _Inputs() + client = ToolChatClient(tool_calls=False) + agent = Agent(client=client, context_providers=[probe]) + provider = JsonStateProvider(_wire(migrated.to_dict())) + await AgentEntity(agent, state_provider=provider).run({"message": "unrelated turn", "correlationId": "unrelated"}) + assert _data(provider).get("ingestedMessages", {}) == {}, "loading old history is not proof of ingestion" + cold_provider = JsonStateProvider(_wire(provider.raw)) + message = Message("user", ["complete incoming payload"], message_id=identity) + request = {"message": "logging only", "correlationId": "incoming", "contextMessages": [message.to_dict()]} + + response = await AgentEntity(agent, state_provider=cold_provider).run(request) + + assert response.text == "answer-2" + assert [item.to_dict() for item in probe.inputs[-1]] == [message.to_dict()] + assert [item.text for item in client.received_messages[-1]].count(message.text) == 1 + assert _data(cold_provider)["ingestedMessages"] == {identity: [message_identity(message)]} + stored = _array(_data(cold_provider)["conversationHistory"]) + legacy = next(_object(entry) for entry in stored if _object(entry).get("correlationId") == "legacy") + assert _object(_array(legacy["messages"])[0])["contents"] == [] + assert source == before_source + + +async def test_new_direct_context_same_application_id_uses_actual_ingestion_receipt_after_cold_reload() -> None: + probe = _Inputs() + client = ToolChatClient(tool_calls=False) + agent = Agent(client=client, context_providers=[probe]) + provider = JsonStateProvider() + message = Message("user", ["direct projected input"], message_id="application-id") + request = {"message": "logging only", "correlationId": "direct-first", "contextMessages": [message.to_dict()]} + assert (await AgentEntity(agent, state_provider=provider).run(request)).text == "answer-1" + assert [item.to_dict() for item in probe.inputs[0]] == [message.to_dict()] + expected_receipts = {"application-id": [message_identity(message)]} + assert _data(provider)["ingestedMessages"] == expected_receipts + cold_provider = JsonStateProvider(_wire(provider.raw)) + + response = await AgentEntity(agent, state_provider=cold_provider).run({**request, "correlationId": "direct-next"}) + + assert response.text == "answer-2" and probe.inputs[-1] == [] + assert [item.text for item in client.received_messages[-1]].count(message.text) == 1 + assert _data(cold_provider)["ingestedMessages"] == expected_receipts + assert request["contextMessages"] == [message.to_dict()] + + +class _StatefulHistory(DurableHistoryProvider): + def __init__(self) -> None: + super().__init__(source_id="stateful-history", prune_excluded=False) + self.loaded_counters: list[int] = [] + self.live_states: list[dict[str, object]] = [] + + async def before_run( + self, *, agent: SupportsAgentRun, session: AgentSession, context: SessionContext, state: dict[str, Any] + ) -> None: + counter = state.get("counter", 0) + assert isinstance(counter, int) + self.loaded_counters.append(counter) + state["counter"] = counter + 1 + self.live_states.append(state) + await super().before_run(agent=agent, session=session, context=context, state=state) + + +async def test_custom_history_state_and_unknown_session_envelope_survive_two_json_cold_runs() -> None: + initial = DurableAgentState() + session = AgentSession(session_id="revision-session") + pending = {"approval": {"ids": ["pending-1"], "approved": False}, "additional": {"cursor": [1, 3]}} + session.state["stateful-history"] = {"counter": 0, "pending": deepcopy(pending)} + session.state["foreign"] = {"pending": ["untouched"]} + initial.data.session = session.to_dict() + future = {"type": "future_session_metadata", "opaque": [None, False, {"labels": ["keep"]}]} + initial.data.session["future_session"] = deepcopy(future) + raw = _wire(initial.to_dict()) + original = deepcopy(raw) + histories: list[_StatefulHistory] = [] + clients: list[ToolChatClient] = [] + for turn in (1, 2): + history = _StatefulHistory() + histories.append(history) + client = ToolChatClient(tool_calls=False) + clients.append(client) + provider = JsonStateProvider(_wire(raw)) + agent = Agent(client=client, context_providers=[history]) + + response = await AgentEntity(agent, state_provider=provider).run({ + "message": f"input-{turn}", + "correlationId": f"state-{turn}", + }) + + assert response.text == "answer-1" and provider.writes == 1 + assert history.loaded_counters == [turn - 1] + saved_session = _object(_data(provider)["session"]) + saved_state = _object(saved_session["state"]) + assert saved_state[history.source_id] == {"counter": turn, "pending": pending} + assert saved_state["foreign"] == {"pending": ["untouched"]} + assert saved_session["future_session"] == future + assert WORKING_BUFFER_KEY in history.live_states[0] and POSITIONS_KEY in history.live_states[0] + assert len(_array(history.live_states[0][WORKING_BUFFER_KEY])) == turn * 2 + raw = _wire(provider.raw) + assert histories[0] is not histories[1] and histories[0].live_states[0] is not histories[1].live_states[0] + assert [message.text for message in clients[1].received_messages[0]] == ["input-1", "answer-1", "input-2"] + assert _wire(initial.to_dict()) == original + + +def test_optional_af_unnamed_agent_registers_under_explicit_workflow_executor_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + app_module = pytest.importorskip("agent_framework_azurefunctions._app") + original_factory = app_module.create_agent_entity + created: list[SupportsAgentRun] = [] + + def capture_factory(agent: SupportsAgentRun, *args: Any, **kwargs: Any) -> Any: + created.append(agent) + return original_factory(agent, *args, **kwargs) + + monkeypatch.setattr(app_module, "create_agent_entity", capture_factory) + client = ToolChatClient(tool_calls=False) + agent = Agent(client=client) + assert agent.name is None + executor = AgentExecutor(agent, id="reviewer") + workflow = WorkflowBuilder(name="execution_followup", start_executor=executor, output_from=[executor]).build() + app = app_module.AgentFunctionApp( + workflow=workflow, + deployment_mode="isolated_v2", + enable_health_check=False, + enable_http_endpoints=False, + enable_mcp_tool_trigger=False, + ) + assert app.agents == {"execution_followup-reviewer": agent} + assert created == [agent] and agent.name is None + functions = {function.get_function_name(): function for function in app.get_functions()} + registered = functions["dafx-execution_followup-reviewer"] + assert registered.get_trigger().get_dict_repr()["type"] == "entityTrigger" + assert callable(registered.get_user_function()) + assert client.received_messages == [] + # A missing standalone identity is still rejected rather than silently inventing a name. + with pytest.raises(ValueError, match="name"): + app.add_agent(agent) + assert created == [agent] diff --git a/python/packages/durabletask/tests/test_execution_review.py b/python/packages/durabletask/tests/test_execution_review.py new file mode 100644 index 0000000..e1ed389 --- /dev/null +++ b/python/packages/durabletask/tests/test_execution_review.py @@ -0,0 +1,816 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Execution regressions using real core pipelines and detached JSON storage.""" + +import json +from collections.abc import AsyncIterable, Sequence +from copy import deepcopy +from typing import Any + +import pytest +from agent_framework import ( + Agent, + AgentResponse, + AgentResponseUpdate, + AgentSession, + ChatResponse, + ChatResponseUpdate, + Content, + ContextProvider, + HistoryProvider, + Message, + ResponseStream, + SessionContext, + tool, +) +from test_durable_history_provider import RecordingChatClient +from test_history_pipeline_revision import CountingHistory, NonStreamingAgent, ToolChatClient +from test_revision_contract import JsonStateProvider + +from agent_framework_durabletask import AgentEntity, DurableAgentState, DurableHistoryProvider, RunRequest +from agent_framework_durabletask import _entities as entities_module +from agent_framework_durabletask._durable_agent_state import DurableAgentStateRequest +from agent_framework_durabletask._history_provider import current_durable_history_binding +from agent_framework_durabletask._message_identity import message_identity +from agent_framework_durabletask._retention import enforce_budget + + +def _wire(value: Any) -> Any: + return json.loads(json.dumps(value, allow_nan=False)) + + +def _committed(provider: JsonStateProvider) -> dict[str, Any]: + return _wire(provider.raw) + + +def _make_agent(client: Any, **kwargs: Any) -> Agent: + return Agent(client=client, **kwargs) + + +def _projection(correlation: str, messages: list[Message], occurrences: list[str]) -> dict[str, Any]: + return { + "message": "logging-only input must not become model context", + "correlationId": correlation, + "contextMessages": _wire([message.to_dict() for message in messages]), + "contextMessageIds": list(occurrences), + } + + +def _delivered(provider: JsonStateProvider, correlation: str) -> AgentResponse[Any]: + response = DurableAgentState.from_json(json.dumps(_committed(provider))).try_get_agent_response(correlation) + assert isinstance(response, AgentResponse) + return response + + +class _Probe(ContextProvider): + def __init__(self) -> None: + super().__init__("execution-probe") + self.inputs: list[list[Message]] = [] + self.agents: list[Any] = [] + self.sessions: list[AgentSession] = [] + self.responses: list[AgentResponse[Any]] = [] + self.response_snapshots: list[dict[str, Any]] = [] + self.fail_at: str | None = None + + async def before_run(self, *, agent: Any, session: AgentSession, context: SessionContext, **kwargs: Any) -> None: + self.inputs.append(deepcopy(context.input_messages)) + self.agents.append(agent) + self.sessions.append(session) + if self.fail_at == "before": + raise RuntimeError("probe before-run failure") + + async def after_run(self, *, context: SessionContext, **kwargs: Any) -> None: + assert isinstance(context.response, AgentResponse) + self.responses.append(context.response) + self.response_snapshots.append(_wire(context.response.to_dict())) + if self.fail_at == "after": + raise RuntimeError("probe after-run failure") + + +class _RichClient(RecordingChatClient): + text = '{"nested":{"items":[1,2]}}' + + @staticmethod + def _content() -> Content: + return Content.from_text(_RichClient.text, additional_properties={"source": {"tags": ["original"]}}) + + def get_response(self, messages: Any, *, stream: bool = False, **kwargs: Any) -> Any: + if stream: + return super().get_response(messages, stream=True, **kwargs) + self.received_messages.append(list(messages)) + + async def get() -> ChatResponse[Any]: + return ChatResponse( + messages=[Message("assistant", [self._content()], message_id="rich-answer")], + response_id="rich-response", + additional_properties={"result": {"tags": ["original"]}}, + value={"nested": {"items": [1, 2]}}, + ) + + return get() + + def _stream(self, options: dict[str, Any]) -> ResponseStream[ChatResponseUpdate, ChatResponse]: + async def updates() -> AsyncIterable[ChatResponseUpdate]: + yield ChatResponseUpdate( + role="assistant", + contents=[self._content()], + message_id="rich-answer", + response_id="rich-response", + additional_properties={"result": {"tags": ["original"]}}, + ) + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) + + +class _MutatingCallback: + def __init__(self, *, mutate_updates: bool, mutate_final: bool) -> None: + self.mutate_updates = mutate_updates + self.mutate_final = mutate_final + self.updates: list[AgentResponseUpdate] = [] + self.responses: list[AgentResponse[Any]] = [] + self.mutations: list[str] = [] + + async def on_streaming_response_update(self, update: AgentResponseUpdate, context: Any) -> None: + self.updates.append(update) + if self.mutate_updates: + update.contents[0].text = '{"nested":{"items":[99]}}' + update.contents[0].additional_properties["source"]["tags"].append("callback-update") + assert update.additional_properties is not None + update.additional_properties["result"]["tags"].append("callback-update") + self.mutations.append("update") + + async def on_agent_response(self, response: AgentResponse[Any], context: Any) -> None: + self.responses.append(response) + if self.mutate_final: + value = response.value + assert value is not None + value["nested"]["items"].append(99) + response.messages[0].contents[0].text = "callback final text" + response.messages[0].contents[0].additional_properties["source"]["tags"].append("callback-final") + response.additional_properties["result"]["tags"].append("callback-final") + self.mutations.append("final") + + +@pytest.mark.parametrize( + ("stream", "mutate_updates", "mutate_final"), + [(True, True, False), (True, False, True), (True, True, True), (False, False, True)], + ids=["stream-update", "stream-final", "stream-both", "non-streaming-final"], +) +async def test_callbacks_cannot_change_core_response_history_or_cold_mailbox( + stream: bool, mutate_updates: bool, mutate_final: bool +) -> None: + client = _RichClient() + probe = _Probe() + agent = (Agent if stream else NonStreamingAgent)(client=client, context_providers=[probe]) + callback = _MutatingCallback(mutate_updates=mutate_updates, mutate_final=mutate_final) + provider = JsonStateProvider() + entity = AgentEntity(agent, callback=callback, state_provider=provider) + request = { + "message": "return a nested value", + "correlationId": "callback-copy", + "options": {"response_format": {"type": "object"}}, + } + + response = await entity.run(request) + + assert callback.mutations == (["update"] if mutate_updates else []) + (["final"] if mutate_final else []) + assert len(callback.updates) == int(stream) and len(callback.responses) == 1 + assert len(probe.responses) == 1 and probe.responses[0] is response + assert callback.responses[0] is not response + assert callback.responses[0].messages[0].contents[0] is not response.messages[0].contents[0] + assert response.text == _RichClient.text + assert response.value == {"nested": {"items": [1, 2]}} + assert response.additional_properties["result"] == {"tags": ["original"]} + assert response.messages[0].contents[0].additional_properties == {"source": {"tags": ["original"]}} + assert probe.response_snapshots[0]["messages"][0]["contents"][0]["text"] == _RichClient.text + stored = DurableAgentState.from_json(json.dumps(_committed(provider))) + answers = [ + message.to_chat_message() + for entry in stored.data.conversation_history + for message in entry.messages + if message.role == "assistant" + ] + assert len(answers) == 1 and answers[0].text == _RichClient.text + assert answers[0].contents[0].additional_properties == {"source": {"tags": ["original"]}} + delivered = _delivered(provider, "callback-copy") + assert delivered.text == response.text and delivered.value == response.value + assert delivered.additional_properties == response.additional_properties + assert delivered.messages[0].contents[0].to_dict() == response.messages[0].contents[0].to_dict() + before = _committed(provider) + # A callback may retain its objects and mutate them after the entity has committed. + retained_value = callback.responses[0].value + assert retained_value is not None + retained_value["nested"]["items"].append(101) + callback.responses[0].messages[0].contents[0].additional_properties["source"]["tags"].append("late") + assert response.value == {"nested": {"items": [1, 2]}} + assert _committed(provider) == before + cold_provider = JsonStateProvider(before) + duplicate = await AgentEntity(agent, callback=callback, state_provider=cold_provider).run(request) + assert duplicate.to_dict() == delivered.to_dict() + assert len(client.received_messages) == 1 and cold_provider.writes == 0 + assert len(callback.responses) == 1 + + +class _StructuredFailure(RuntimeError): + def __init__(self, code: str, *, body: bool = False) -> None: + super().__init__(code) + if body: + self.body = {"code": code} + else: + self.code = code + + +class _RetryClient(RecordingChatClient): + STORES_BY_DEFAULT = True + + def __init__(self, second: str, *, body: bool = False) -> None: + super().__init__() + self.second = second + self.body = body + self.errors: list[BaseException] = [] + self.options: list[dict[str, Any]] = [] + + def get_response(self, messages: Any, *, options: Any = None, **kwargs: Any) -> Any: + self.options.append(deepcopy(dict(options or {}))) + return super().get_response(messages, options=options, **kwargs) + + def _stream(self, options: dict[str, Any]) -> ResponseStream[ChatResponseUpdate, ChatResponse]: + async def updates() -> AsyncIterable[ChatResponseUpdate]: + attempt = len(self.received_messages) + try: + if attempt == 1: + raise _StructuredFailure("previous_response_not_found", body=self.body) + if attempt == 2: + if self.second == "implicit": + raise RuntimeError("unrelated second failure") + code = "previous_response_not_found" if self.second == "wrapped-missing" else "invalid_api_key" + current = _StructuredFailure(code, body=self.body) + if self.second.startswith("wrapped"): + try: + raise current + except _StructuredFailure as cause: + raise RuntimeError(f"current wrapper: {code}") from cause + raise current + except Exception as exc: + self.errors.append(exc) + raise + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("retry recovered")]) + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) + + +@pytest.mark.parametrize("second", ["different-code", "implicit", "wrapped-different"]) +@pytest.mark.parametrize("body", [False, True], ids=["code-attribute", "body-code"]) +async def test_retry_delivers_current_failure_without_following_stale_implicit_context( + monkeypatch: pytest.MonkeyPatch, second: str, body: bool +) -> None: + monkeypatch.setattr(entities_module, "_REJECTED_ID_BACKOFF_SECONDS", 0) + client = _RetryClient(second, body=body) + provider = JsonStateProvider() + entity = AgentEntity(Agent(client=client), state_provider=provider) + request = {"message": "continue", "correlationId": "retry-current"} + + response = await entity.run(request) + + assert len(client.received_messages) == 2, "only the missing-response failure authorizes another attempt" + assert len(client.errors) == 2 + current = client.errors[1].__cause__ if second == "wrapped-different" else client.errors[1] + assert current is not None + assert current.__context__ is client.errors[0], "exercise Python's implicit prior-exception chain" + assert response.additional_properties["durable_status"] == "error" + expected = "unrelated second failure" if second == "implicit" else "invalid_api_key" + assert expected in response.text and "previous_response_not_found" not in response.text + assert client.options[0] == client.options[1] + assert [[message.to_dict() for message in batch] for batch in client.received_messages] == [ + [message.to_dict() for message in client.received_messages[0]] + ] * 2 + assert _delivered(provider, "retry-current").to_dict() == response.to_dict() + cold_provider = JsonStateProvider(_committed(provider)) + assert (await AgentEntity(Agent(client=client), state_provider=cold_provider).run(request)).text == response.text + assert len(client.received_messages) == 2 and cold_provider.writes == 0 + + +@pytest.mark.parametrize("body", [False, True], ids=["code-attribute", "body-code"]) +async def test_explicit_current_missing_cause_remains_retryable(monkeypatch: pytest.MonkeyPatch, body: bool) -> None: + monkeypatch.setattr(entities_module, "_REJECTED_ID_BACKOFF_SECONDS", 0) + client = _RetryClient("wrapped-missing", body=body) + provider = JsonStateProvider() + + response = await AgentEntity(Agent(client=client), state_provider=provider).run({ + "message": "continue", + "correlationId": "wrapped-retry", + }) + + assert response.text == "retry recovered" and len(client.received_messages) == 3 + assert len(client.errors) == 2 + assert isinstance(client.errors[1].__cause__, _StructuredFailure) + assert client.errors[1].__cause__ is not client.errors[0] + assert client.options == [client.options[0]] * 3 + assert _delivered(provider, "wrapped-retry").text == "retry recovered" + assert provider.writes == 1 + + +@pytest.mark.parametrize("per_call", [False, True], ids=["per-run", "per-service-call"]) +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +@pytest.mark.parametrize("tool_source", ["default", "context-provider"]) +async def test_disabled_tools_never_execute_even_when_the_model_returns_a_function_call( + per_call: bool, stream: bool, tool_source: str +) -> None: + calls: list[str] = [] + + @tool(name="lookup", approval_mode="never_require") + def lookup(key: str) -> str: + calls.append(key) + return f"value:{key}" + + class ToolProvider(ContextProvider): + async def before_run(self, *, context: SessionContext, **kwargs: Any) -> None: + context.tools.append(lookup) + + def make(client: ToolChatClient) -> Agent: + return (Agent if stream else NonStreamingAgent)( + client=client, + tools=[lookup] if tool_source == "default" else [], + context_providers=[ToolProvider("tools")] if tool_source == "context-provider" else [], + require_per_service_call_history_persistence=per_call, + ) + + # Positive control proves that the exact helper/model request reaches real core invocation. + enabled_client = ToolChatClient() + enabled = await AgentEntity(make(enabled_client), state_provider=JsonStateProvider()).run({ + "message": "use lookup", + "correlationId": "tools-enabled", + }) + assert enabled.text == "answer-2" and calls == ["durable"] + assert len(enabled_client.received_messages) == 2 + calls.clear() + client = ToolChatClient() + agent = make(client) + original_options = agent.default_options + original_tools = original_options["tools"] + original_tool_items = list(original_tools) + original_providers = agent.context_providers + config = deepcopy(client.function_invocation_configuration) + provider = JsonStateProvider() + entity = AgentEntity(agent, state_provider=provider) + registered = entity.agent + request = {"message": "use lookup", "correlationId": "tools-disabled", "enable_tool_calls": False} + + response = await entity.run(request) + + assert client.received_messages, "reach the model that requests lookup despite tool_choice=none" + assert client.received_options[0].get("tool_choice") == "none" + assert calls == [], "forwarding tool_choice is insufficient if the invocation layer still has a callable" + assert not any( + content.type == "function_result" and content.result == "value:durable" + for batch in client.received_messages + for message in batch + for content in message.contents + ) + assert entity.agent is registered + assert agent.default_options is original_options and original_options["tools"] is original_tools + assert original_tools == original_tool_items and agent.context_providers is original_providers + assert client.function_invocation_configuration == config + assert current_durable_history_binding() is None + # An error response or an unexecuted function call are both safe outcomes. + assert _delivered(provider, "tools-disabled").to_dict() == response.to_dict() + count = len(client.received_messages) + await AgentEntity(agent, state_provider=JsonStateProvider(_committed(provider))).run(request) + assert len(client.received_messages) == count and calls == [] + + +class _ExternalHistory(HistoryProvider): + """Ordinary blind-append storage with no awareness of store options or durable bindings.""" + + def __init__(self, source_id: str, **kwargs: Any) -> None: + super().__init__(source_id, **kwargs) + self.messages: list[Message] = [] + self.loads: list[str | None] = [] + self.saves: list[tuple[str | None, list[Message]]] = [] + + async def get_messages(self, session_id: str | None, **kwargs: Any) -> list[Message]: + self.loads.append(session_id) + return deepcopy(self.messages) + + async def save_messages(self, session_id: str | None, messages: Sequence[Message], **kwargs: Any) -> None: + self.saves.append((session_id, deepcopy(list(messages)))) + self.messages.extend(deepcopy(list(messages))) + + +@pytest.mark.parametrize("per_call", [False, True], ids=["per-run", "per-service-call"]) +async def test_store_true_false_true_parks_external_primary_but_not_store_only_sinks(per_call: bool) -> None: + primary = _ExternalHistory("external") + primary.messages = [Message("user", ["external seed"], message_id="seed")] + outputs = _ExternalHistory("audit-outputs", load_messages=False, store_inputs=False) + inputs = _ExternalHistory("audit-inputs", load_messages=False, store_outputs=False) + probe = _Probe() + client = ToolChatClient(tool_calls=False) + agent = Agent( + client=client, + context_providers=[primary, outputs, inputs, probe], + require_per_service_call_history_persistence=per_call, + ) + providers = agent.context_providers + defaults = agent.default_options + original_defaults = deepcopy(defaults) + raw: dict[str, Any] = {} + for index, (store, text) in enumerate(((True, "service-first"), (False, "local"), (True, "service-again")), 1): + provider = JsonStateProvider(_wire(raw)) + entity = AgentEntity(agent, state_provider=provider) + registered = entity.agent + request = {"message": text, "correlationId": f"ownership-{index}", "options": {"store": store}} + original_request = deepcopy(request) + + response = await entity.run(request) + + assert response.text == f"answer-{index}" + assert entity.agent is registered and agent.context_providers is providers + assert agent.default_options is defaults and defaults == original_defaults and request == original_request + invocation = probe.agents[-1] + if store: + assert invocation is not registered + assert invocation.context_providers[0].__wrapped__ is primary + else: + assert invocation.context_providers[0] is primary + assert invocation.context_providers[1] is outputs and invocation.context_providers[2] is inputs + assert outputs.load_messages is False and outputs.store_inputs is False and outputs.store_outputs is True + assert inputs.load_messages is False and inputs.store_inputs is True and inputs.store_outputs is False + assert outputs.loads == inputs.loads == [] + assert [batch[0].text for _, batch in outputs.saves] == [f"answer-{i}" for i in range(1, index + 1)] + assert all(len(batch) == 1 for _, batch in outputs.saves + inputs.saves) + raw = _committed(provider) + assert raw["data"]["conversationHistory"] == [] + assert raw["data"]["session"]["service_session_id"] == "service-thread" + assert provider.writes == 1 and current_durable_history_binding() is None + assert primary.loads == ["revision-session"] and len(primary.saves) == 1 + assert [message.text for message in primary.messages] == ["external seed", "local", "answer-2"] + assert [batch[0].text for _, batch in inputs.saves] == ["service-first", "local", "service-again"] + assert {session_id for session_id, _ in primary.saves + inputs.saves + outputs.saves} == {"revision-session"} + assert [[message.text for message in batch] for batch in client.received_messages] == [ + ["service-first"], + ["external seed", "local"], + ["service-again"], + ] + assert [options.get("conversation_id") for options in client.received_options] == [None, None, "service-thread"] + + +@pytest.mark.parametrize("per_call", [False, True], ids=["per-run", "per-service-call"]) +@pytest.mark.parametrize("failure", ["before", "model", "after", "commit"]) +async def test_service_owner_view_is_restored_on_provider_model_and_commit_errors(per_call: bool, failure: str) -> None: + primary = _ExternalHistory("external") + sink = _ExternalHistory("audit", load_messages=False) + probe = _Probe() + probe.fail_at = failure + client = ToolChatClient(tool_calls=False, fail=failure == "model") + agent = Agent( + client=client, + context_providers=[primary, sink, probe], + require_per_service_call_history_persistence=per_call, + ) + providers = agent.context_providers + provider = JsonStateProvider() + provider.fail_writes = failure == "commit" + entity = AgentEntity(agent, state_provider=provider) + registered = entity.agent + original_state = entity.state + request = {"message": "service failure", "correlationId": "owner-error", "options": {"store": True}} + + if failure == "commit": + with pytest.raises(OSError, match="commit failure"): + await entity.run(request) + assert entity.state is original_state and _committed(provider) == {} and provider.writes == 0 + assert entity.state.try_get_agent_response("owner-error") is None + else: + response = await entity.run(request) + assert response.additional_properties["durable_status"] == "error" + assert _delivered(provider, "owner-error").to_dict() == response.to_dict() + assert probe.agents and probe.agents[0] is not registered + assert probe.agents[0].context_providers[0].__wrapped__ is primary + assert probe.agents[0].context_providers[1] is sink + assert entity.agent is registered and agent.context_providers is providers and providers[0] is primary + assert primary.loads == [] and primary.saves == [] + assert current_durable_history_binding() is None + probe.fail_at = None + client.fail = False + provider.fail_writes = False + recovered = await entity.run({ + "message": "local recovery", + "correlationId": "owner-recovered", + "options": {"store": False}, + }) + assert recovered.additional_properties.get("durable_status") != "error" + assert primary.loads == [provider.core_session_id] and len(primary.saves) == 1 + assert entity.agent is registered and agent.context_providers is providers + + +@pytest.mark.parametrize("invalid", [17, 1.5, None, "", " \t", True, False, [], ["unsafe"]]) +@pytest.mark.parametrize("boundary", ["constructor", "from-dict", "entity-dict", "entity-json"]) +async def test_correlation_id_requires_a_nonblank_string_before_any_client_or_state_write( + invalid: Any, boundary: str +) -> None: + provider = JsonStateProvider() + client = RecordingChatClient() + entity = AgentEntity(_make_agent(client), state_provider=provider) + original = entity.state + payload = {"message": "must not execute", "correlationId": invalid} + with pytest.raises(ValueError, match="correlationId"): + if boundary == "constructor": + request: Any = RunRequest(message="must not execute", correlation_id=invalid) + elif boundary == "from-dict": + request = RunRequest.from_dict(payload) + else: + request = json.dumps(payload) if boundary == "entity-json" else payload + await entity.run(request) + assert client.received_messages == [] and provider.writes == 0 and _committed(provider) == {} + assert entity.state is original + + +@pytest.mark.parametrize("payload", [None, [], [1], 17, 1.5, True, False, "input"]) +async def test_nonobject_json_requests_are_rejected_before_execution(payload: Any) -> None: + provider = JsonStateProvider() + client = RecordingChatClient() + with pytest.raises(ValueError, match="object"): + await AgentEntity(_make_agent(client), state_provider=provider).run(json.dumps(payload)) + assert client.received_messages == [] and provider.writes == 0 and provider.raw == {} + + +async def test_string_numeric_correlation_survives_cold_reload_without_reexecution() -> None: + provider = JsonStateProvider() + client = RecordingChatClient() + agent = _make_agent(client) + request = RunRequest(message="valid string ID", correlation_id="17") + response = await AgentEntity(agent, state_provider=provider).run(request) + cold_provider = JsonStateProvider(_committed(provider)) + duplicate = await AgentEntity(agent, state_provider=cold_provider).run(request.to_dict()) + assert duplicate.to_dict() == response.to_dict() + assert set(provider.raw["data"]["completedCorrelations"]) == {"17"} + assert len(client.received_messages) == 1 and cold_provider.writes == 0 + + +@pytest.mark.parametrize("anonymous", [False, True], ids=["application-id", "anonymous"]) +async def test_paired_occurrences_preserve_exact_raw_messages_and_canonical_content_metadata(anonymous: bool) -> None: + message = Message( + "user", + [Content.from_text("identical payload", additional_properties={"source": {"tags": ["keep"]}})], + message_id=None if anonymous else "same-application-id", + additional_properties={"application": {"labels": [1, 2]}}, + ) + request = _projection("paired", [message, deepcopy(message)], ["o1", "o2"]) + original = deepcopy(request) + client = RecordingChatClient() + provider = JsonStateProvider() + + response = await AgentEntity(_make_agent(client), state_provider=provider).run(request) + + assert response.text == "reply-1" + assert len(client.received_messages) == 1 + assert [item.to_dict() for item in client.received_messages[0]] == original["contextMessages"] + assert [item.message_id for item in client.received_messages[0]] == [message.message_id] * 2 + assert request == original and message.to_dict() == original["contextMessages"][0] + assert provider.raw["data"]["ingestedMessages"] == {key: [message_identity(message)] for key in ("o1", "o2")} + restored = DurableAgentState.from_json(json.dumps(_committed(provider))) + inputs = [ + stored.to_chat_message() + for entry in restored.data.conversation_history + if isinstance(entry, DurableAgentStateRequest) + for stored in entry.messages + ] + assert len(inputs) == 2 + assert [item.contents[0].to_dict() for item in inputs] == [message.contents[0].to_dict()] * 2 + + +@pytest.mark.parametrize("evict", [False, True], ids=["retained-history", "pressure-evicted-history"]) +async def test_occurrence_receipts_survive_cold_reload_and_eviction_without_blocking_new_runs(evict: bool) -> None: + message = Message("user", ["projected input " * 1000], message_id="application-id") + client = RecordingChatClient() + probe = _Probe() + agent = _make_agent(client, context_providers=[probe]) + provider = JsonStateProvider() + entity = AgentEntity(agent, state_provider=provider) + original_request = _projection("first", [message], ["o1"]) + original_response = await entity.run(original_request) + await entity.run({"message": "newest exchange", "correlationId": "anchor"}) + if evict: + removed = await enforce_budget(entity.state, max_state_bytes=6000) + assert removed > 0 + entity.persist_state() + assert not any( + stored.text == message.text for entry in entity.state.data.conversation_history for stored in entry.messages + ) + assert provider.raw["data"]["ingestedMessages"] == {"o1": [message_identity(message)]} + cold_provider = JsonStateProvider(_committed(provider)) + cold = AgentEntity(agent, state_provider=cold_provider) + calls_before = len(client.received_messages) + inputs_before = len(probe.inputs) + duplicate = await cold.run(original_request) + assert duplicate.to_dict() == original_response.to_dict() + assert len(client.received_messages) == calls_before and len(probe.inputs) == inputs_before + assert cold_provider.writes == 0 + repeated = await cold.run(_projection("new-correlation-same-occurrence", [message], ["o1"])) + assert repeated.text == "reply-3" and probe.inputs[-1] == [] + assert len(client.received_messages) == calls_before + 1 and cold_provider.writes == 1 + assert "new-correlation-same-occurrence" in cold_provider.raw["data"]["completedCorrelations"] + await cold.run(_projection("new-occurrence", [message], ["o2"])) + assert [item.to_dict() for item in probe.inputs[-1]] == [message.to_dict()] + revised = deepcopy(message) + revised.contents[0].text = "revised payload" + await cold.run(_projection("revised-occurrence", [revised], ["o1"])) + assert [item.to_dict() for item in probe.inputs[-1]] == [revised.to_dict()] + assert cold_provider.raw["data"]["ingestedMessages"] == { + "o1": [message_identity(message), message_identity(revised)], + "o2": [message_identity(message)], + } + + +async def test_paired_empty_projection_roundtrips_and_never_falls_back_to_logging_text() -> None: + request = RunRequest( + message="must not reach the model", correlation_id="empty-pair", context_messages=[], context_message_ids=[] + ) + for restored in (RunRequest.from_dict(request.to_dict()), RunRequest.from_json(json.dumps(request.to_dict()))): + assert restored.context_messages == restored.context_message_ids == [] + assert restored.to_dict()["contextMessages"] == restored.to_dict()["contextMessageIds"] == [] + client = RecordingChatClient() + provider = JsonStateProvider() + response = await AgentEntity(_make_agent(client), state_provider=provider).run(request) + assert response.text == "reply-1" and client.received_messages == [[]] + assert provider.raw["data"].get("ingestedMessages", {}) == {} + + +@pytest.mark.parametrize("occurrences", [[], ["o1", "o2"], "o1", [None], [17], [""]]) +@pytest.mark.parametrize("direct", [False, True], ids=["wire", "constructor"]) +async def test_malformed_paired_occurrence_ids_are_rejected_before_client_calls(occurrences: Any, direct: bool) -> None: + client = RecordingChatClient() + provider = JsonStateProvider() + message = Message("user", ["input"], message_id="raw-id") + with pytest.raises(ValueError, match="contextMessageIds"): + request: Any = ( + RunRequest( + message="input", + correlation_id="bad-pair", + context_messages=[message.to_dict()], + context_message_ids=occurrences, + ) + if direct + else { + "message": "input", + "correlationId": "bad-pair", + "contextMessages": [message.to_dict()], + "contextMessageIds": occurrences, + } + ) + await AgentEntity(_make_agent(client), state_provider=provider).run(request) + assert client.received_messages == [] and provider.writes == 0 and provider.raw == {} + + +@pytest.mark.parametrize("per_call", [False, True], ids=["per-run", "per-service-call"]) +async def test_failed_tool_followup_retains_receipts_by_occurrence_not_application_id(per_call: bool) -> None: + calls: list[str] = [] + + @tool(name="lookup", approval_mode="never_require") + def lookup(key: str) -> str: + calls.append(key) + return f"value:{key}" + + history = CountingHistory([]) + client = ToolChatClient(fail_on_call=2) + provider = JsonStateProvider() + agent = Agent( + client=client, + tools=[lookup], + context_providers=[history], + require_per_service_call_history_persistence=per_call, + ) + message = Message("user", ["use lookup"], message_id="shared-application-id") + request = _projection("partial", [message, deepcopy(message)], ["o1", "o2"]) + response = await AgentEntity(agent, state_provider=provider).run(request) + assert response.additional_properties["durable_status"] == "error" + assert "model failed before history persistence" in response.text + assert calls == ["durable"] and len(client.received_messages) == 2 + assert history.after_calls == int(per_call) + raw = _committed(provider) + expected = {identity: [message_identity(message)] for identity in ("o1", "o2")} if per_call else {} + assert raw["data"].get("ingestedMessages", {}) == expected + saved_inputs = [ + stored + for entry in raw["data"]["conversationHistory"] + if entry["$type"] == "request" + for stored in entry["messages"] + if stored["role"] == "user" + ] + assert len(saved_inputs) == (2 if per_call else 0) + probe = _Probe() + cold_client = RecordingChatClient() + cold_provider = JsonStateProvider(raw) + cold = AgentEntity(_make_agent(cold_client, context_providers=[probe]), state_provider=cold_provider) + assert (await cold.run(request)).to_dict() == response.to_dict() + assert cold_client.received_messages == [] and cold_provider.writes == 0 + await cold.run(_projection("after-partial", [message, deepcopy(message)], ["o1", "o3"])) + assert [item.to_dict() for item in probe.inputs[-1]] == [message.to_dict()] * (1 if per_call else 2) + expected.update({identity: [message_identity(message)] for identity in ("o1", "o3")}) + assert cold_provider.raw["data"]["ingestedMessages"] == expected + + +@pytest.mark.parametrize("saved_count", [0, 1, 2], ids=["no-append", "partial-append", "full-append"]) +async def test_partial_append_does_not_consume_an_unsaved_equal_occurrence(saved_count: int) -> None: + class InterruptedHistory(DurableHistoryProvider): + def __init__(self) -> None: + super().__init__(prune_excluded=False) + self.appended = 0 + + async def after_run( + self, *, session: AgentSession, context: SessionContext, state: dict[str, Any], **kwargs: Any + ) -> None: + batch = context.input_messages[:saved_count] + await self.save_messages(session.session_id, batch, state=state) + self.appended = len(batch) + raise OSError("history interrupted after selected inputs") + + history = InterruptedHistory() + client = RecordingChatClient() + provider = JsonStateProvider() + message = Message("user", ["equal input"], message_id="same-application-id") + request = _projection("append-failed", [message, deepcopy(message)], ["o1", "o2"]) + entity = AgentEntity(_make_agent(client, context_providers=[history]), state_provider=provider) + + response = await entity.run(request) + + assert response.additional_properties["durable_status"] == "error" + assert "history interrupted after selected inputs" in response.text + assert len(client.received_messages) == 1 and history.appended == saved_count + raw = _committed(provider) + saved = [item for entry in raw["data"]["conversationHistory"] for item in entry["messages"]] + assert len(saved) == saved_count, "the failure must occur after the selected real durable appends" + assert raw["data"].get("ingestedMessages", {}) == { + occurrence: [message_identity(message)] for occurrence in ["o1", "o2"][:saved_count] + } + assert provider.writes == 1 + probe = _Probe() + cold_client = RecordingChatClient() + cold_provider = JsonStateProvider(raw) + cold = AgentEntity(_make_agent(cold_client, context_providers=[probe]), state_provider=cold_provider) + assert (await cold.run(request)).to_dict() == response.to_dict() + assert cold_client.received_messages == [] and cold_provider.writes == 0 + await cold.run(_projection("append-recovered", [message, deepcopy(message)], ["o1", "o2"])) + assert [item.to_dict() for item in probe.inputs[-1]] == [message.to_dict()] * (2 - saved_count) + assert cold_provider.raw["data"]["ingestedMessages"] == { + occurrence: [message_identity(message)] for occurrence in ("o1", "o2") + } + + +def test_reset_cannot_commit_a_retained_floor_above_the_configured_budget() -> None: + initial = DurableAgentState() + initial.data.session = AgentSession(session_id="revision-session", service_session_id="keep-on-rollback").to_dict() + initial.data.conversation_history.append(DurableAgentStateRequest.from_run_request(RunRequest("history", "old"))) + initial.record_response( + "protected", AgentResponse(messages=[Message("assistant", ["x" * 8000])]), delivery_window_seconds=3600 + ) + provider = JsonStateProvider(_wire(initial.to_dict())) + client = RecordingChatClient() + entity = AgentEntity(_make_agent(client), state_provider=provider, max_state_bytes=512) + original = entity.state + before = _committed(provider) + + with pytest.raises(ValueError, match="[Bb]udget|max_state_bytes|[Cc]apacity|floor"): + entity.reset() + + assert entity.state is original and entity.state.to_dict() == before + assert _committed(provider) == before and provider.writes == 0 and client.received_messages == [] + assert _delivered(provider, "protected").text == "x" * 8000 + + +async def test_nan_session_state_aborts_commit_and_does_not_cache_completion() -> None: + class NonFiniteProvider(ContextProvider): + poison = True + + async def after_run(self, *, state: dict[str, Any], **kwargs: Any) -> None: + state["nested"] = {"value": float("nan") if self.poison else 1} + + control = NonFiniteProvider("nonfinite") + initial = DurableAgentState() + initial.data.session = AgentSession(session_id="revision-session").to_dict() + initial.data.session["state"] = {"foreign": {"pending": ["keep"]}} + initial.record_response( + "prior", AgentResponse(messages=[Message("assistant", ["prior answer"])]), delivery_window_seconds=3600 + ) + provider = JsonStateProvider(_wire(initial.to_dict())) + client = RecordingChatClient() + entity = AgentEntity(_make_agent(client, context_providers=[control]), state_provider=provider) + original = entity.state + before = _committed(provider) + request = _projection("nan-run", [Message("user", ["input"], message_id="raw-id")], ["nan-occurrence"]) + + with pytest.raises(ValueError, match="JSON|finite|NaN"): + await entity.run(request) + + assert len(client.received_messages) == 1 + assert entity.state is original and entity.state.to_dict() == before + assert _committed(provider) == before and provider.writes == 0 + assert entity.state.try_get_agent_response("nan-run") is None + assert "nan-run" not in entity.state.data.completed_correlations + assert "nan-occurrence" not in entity.state.data.ingested_messages + assert current_durable_history_binding() is None + control.poison = False + response = await entity.run(request) + assert response.text == "reply-2" and len(client.received_messages) == 2 and provider.writes == 1 + assert provider.raw["data"]["session"]["state"]["foreign"] == {"pending": ["keep"]} + assert _delivered(provider, "prior").text == "prior answer" diff --git a/python/packages/durabletask/tests/test_history_pipeline_revision.py b/python/packages/durabletask/tests/test_history_pipeline_revision.py new file mode 100644 index 0000000..98f0072 --- /dev/null +++ b/python/packages/durabletask/tests/test_history_pipeline_revision.py @@ -0,0 +1,1550 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Core-pipeline cadence, detached transcript appends and compaction reconciliation.""" + +import json +from collections.abc import AsyncIterable, Awaitable, Iterator, Mapping, Sequence +from contextlib import contextmanager +from copy import deepcopy +from datetime import datetime, timezone +from typing import Any + +import pytest +from agent_framework import ( + Agent, + AgentResponse, + AgentSession, + BaseChatClient, + ChatMiddlewareLayer, + ChatResponse, + ChatResponseUpdate, + CompactionProvider, + Content, + ContextProvider, + FunctionInvocationLayer, + HistoryProvider, + InMemoryHistoryProvider, + Message, + ResponseStream, + SessionContext, + SummarizationStrategy, + annotate_message_groups, + tool, +) +from test_durable_history_provider import RecordingChatClient, _InMemoryStateProvider + +from agent_framework_durabletask import AgentEntity, DurableAgentState, DurableHistoryProvider +from agent_framework_durabletask._durable_agent_state import ( + DurableAgentStateCompaction, + DurableAgentStateEntry, + DurableAgentStateEntryJsonType, + DurableAgentStateMessage, + DurableAgentStateRequest, + DurableAgentStateResponse, + DurableAgentStateUnknownEntry, + DurableAgentStateUsage, +) +from agent_framework_durabletask._history_provider import ( + POSITIONS_KEY, + WORKING_BUFFER_KEY, + DurableHistoryBinding, + bind_durable_history, + current_durable_history_binding, + ensure_durable_history, + prune_messages, + unbind_durable_history, +) +from agent_framework_durabletask._message_identity import message_identity + +OLD = datetime(2026, 1, 1, tzinfo=timezone.utc) +PROMPT = "Use lookup for durable." + + +class ToolChatClient(FunctionInvocationLayer, ChatMiddlewareLayer, BaseChatClient): + """Exercise real core middleware and function invocation, not just the client protocol.""" + + def __init__( + self, + *, + tool_calls: bool = True, + response_message_id: str | None = None, + fail: bool = False, + fail_on_call: int | None = None, + events: list[str] | None = None, + ) -> None: + super().__init__(middleware=[]) + self.tool_calls = tool_calls + self.response_message_id = response_message_id + self.fail = fail + self.fail_on_call = fail_on_call + self.events = events if events is not None else [] + self.received_messages: list[list[Message]] = [] + self.received_options: list[dict[str, Any]] = [] + + def _inner_get_response( + self, + *, + messages: Sequence[Message], + stream: bool, + options: Mapping[str, Any], + **kwargs: Any, + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + self.received_messages.append(deepcopy(list(messages))) + self.received_options.append(dict(options)) + call = len(self.received_messages) + self.events.append(f"model-{call}") + if self.fail or call == self.fail_on_call: + raise RuntimeError("model failed before history persistence") + calls_tool = self.tool_calls and call == 1 + contents = ( + [Content.from_function_call(call_id="call-1", name="lookup", arguments='{"key":"durable"}')] + if calls_tool + else [Content.from_text(f"answer-{call}")] + ) + response = ChatResponse( + messages=[ + Message( + "assistant", + contents, + message_id=self.response_message_id, + additional_properties={"model_metadata": {"tags": ["original"]}}, + ) + ], + response_id=f"response-{call}", + conversation_id="service-thread" if options.get("store") else None, + finish_reason="tool_calls" if calls_tool else "stop", + ) + if stream: + + async def updates() -> AsyncIterable[ChatResponseUpdate]: + for message in response.messages: + yield ChatResponseUpdate( + role="assistant", + contents=message.contents, + message_id=message.message_id, + additional_properties=deepcopy(message.additional_properties), + response_id=response.response_id, + conversation_id=response.conversation_id, + finish_reason=response.finish_reason, + ) + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) + + async def get() -> ChatResponse: + return response + + return get() + + +class NonStreamingAgent(Agent): + """Negotiate non-streaming before any core model or tool execution.""" + + def run(self, *args: Any, **kwargs: Any) -> Any: + if kwargs.get("stream"): + raise TypeError("stream is not supported") + return super().run(*args, **kwargs) + + +class CountingHistory(DurableHistoryProvider): + def __init__(self, events: list[str]) -> None: + super().__init__(prune_excluded=False) + self.events = events + self.before_calls = 0 + self.after_calls = 0 + + async def before_run(self, **kwargs: Any) -> None: + self.before_calls += 1 + self.events.append("history-before") + await super().before_run(**kwargs) + + async def after_run(self, **kwargs: Any) -> None: + self.after_calls += 1 + self.events.append("history-after") + await super().after_run(**kwargs) + + +class CaptureHistory(HistoryProvider): + def __init__(self, *, load_messages: bool = False) -> None: + super().__init__("audit", load_messages=load_messages) + self.saved: list[list[Message]] = [] + + async def get_messages(self, session_id: str | None, **kwargs: Any) -> list[Message]: + return [deepcopy(message) for batch in self.saved for message in batch] + + async def save_messages(self, session_id: str | None, messages: Sequence[Message], **kwargs: Any) -> None: + self.saved.append(deepcopy(list(messages))) + + +class AddContext(ContextProvider): + async def before_run(self, *, context: SessionContext, **kwargs: Any) -> None: + context.extend_messages(self, [Message("user", [f"context-{self.source_id}"])]) + + +class SummarizeSeed: + def __init__(self, events: list[str]) -> None: + self.events = events + self.calls = 0 + + async def __call__(self, messages: list[Message]) -> bool: + self.calls += 1 + self.events.append("compaction") + seed = next(message for message in messages if message.message_id == "seed-user") + seed.additional_properties.update({"_excluded": True, "after_hook": {"tags": ["kept"]}}) + if any(message.message_id == "seed-summary" for message in messages): + return False + messages.insert( + messages.index(seed) + 1, + Message( + "assistant", + ["seed summary"], + message_id="seed-summary", + additional_properties={"_summary_of_message_ids": ["seed-user"]}, + ), + ) + return True + + +@contextmanager +def bound( + provider: _InMemoryStateProvider, + correlation_id: str | None = "current", + *, + service_owns_history: bool = False, +) -> Iterator[DurableHistoryBinding]: + binding = DurableHistoryBinding(provider, correlation_id, service_owns_history) + token = bind_durable_history(binding) + try: + yield binding + finally: + unbind_durable_history(token) + + +def stored(message_id: str | None, text: str, role: str = "user") -> DurableAgentStateMessage: + return DurableAgentStateMessage.from_chat_message(Message(role, [text], message_id=message_id)) + + +def seed(provider: _InMemoryStateProvider) -> None: + provider.state.data.conversation_history.extend([ + DurableAgentStateRequest("seed", OLD, [stored("seed-user", "seed question")]), + DurableAgentStateResponse("seed", OLD, [stored("seed-assistant", "seed answer", "assistant")]), + ]) + + +def transcript(provider: _InMemoryStateProvider) -> list[DurableAgentStateMessage]: + return [message for entry in provider.state.data.conversation_history for message in entry.messages] + + +def ids(provider: _InMemoryStateProvider) -> list[str | None]: + return [message.message_id for message in transcript(provider)] + + +def assert_current_positions(provider: _InMemoryStateProvider, state: dict[str, Any]) -> None: + history = provider.state.data.conversation_history + positions = state[POSITIONS_KEY] + for message in state[WORKING_BUFFER_KEY]: + entry, index = positions[message.message_id] + assert any(candidate is entry for candidate in history) + assert entry.messages[index].message_id == message.message_id + assert len(positions) == len({message.message_id for message in transcript(provider)}) + + +def assert_tool_follow_up(client: ToolChatClient) -> None: + assert len(client.received_messages) == 2 + second = client.received_messages[1] + assert [message.text for message in second].count(PROMPT) == 1 + calls = [content for message in second for content in message.contents if content.type == "function_call"] + results = [content for message in second for content in message.contents if content.type == "function_result"] + assert len(calls) == len(results) == 1 + assert calls[0].call_id == results[0].call_id == "call-1" + assert results[0].result == "value:durable" + assert not client.received_options[1].get("conversation_id"), "the core sentinel must not reach the model" + + +@tool(name="lookup", approval_mode="never_require") +def lookup(key: str) -> str: + return f"value:{key}" + + +@pytest.mark.parametrize("per_call", [False, True], ids=["per-run", "per-service-call"]) +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +async def test_real_core_tool_loop_and_final_flush(per_call: bool, stream: bool) -> None: + events: list[str] = [] + history = CountingHistory(events) + strategy = SummarizeSeed(events) + provider = _InMemoryStateProvider() + seed(provider) + client = ToolChatClient(events=events) + agent = Agent( + client=client, + name="tool-agent", + tools=[lookup], + context_providers=[history, CompactionProvider(after_strategy=strategy, history_source_id=history.source_id)], + require_per_service_call_history_persistence=per_call, + ) + session = agent.create_session(session_id="pipeline-session") + with bound(provider) as binding: + if stream: + response = await agent.run(PROMPT, session=session, stream=True).get_final_response() + else: + response = await agent.run(PROMPT, session=session) + + assert_tool_follow_up(client) + assert response.text == "answer-2" + assert history.before_calls == history.after_calls == (2 if per_call else 1) + assert binding.pending_inputs == [] + assert strategy.calls == 1 + assert events[-1] == ("compaction" if per_call else "history-after") + if per_call: + assert "seed-summary" not in ids(provider), "run-end compaction still needs the entity's final flush" + original_response = deepcopy(response.to_dict()) + state = session.state[history.source_id] + history.flush(state) + assert current_durable_history_binding() is binding + assert [message.text for message in transcript(provider)] == [ + "seed question", + "seed summary", + "seed answer", + PROMPT, + "", + "", + "answer-2", + ] + assert all(ids(provider)) and len(ids(provider)) == len(set(ids(provider))) + assert (transcript(provider)[0].extension_data or {})["after_hook"] == {"tags": ["kept"]} + current = [entry for entry in provider.state.data.conversation_history if entry.correlation_id == "current"] + assert len(current) == (4 if per_call else 2) + assert [entry.json_type for entry in current] == ( + [DurableAgentStateEntryJsonType.REQUEST, DurableAgentStateEntryJsonType.RESPONSE] * (2 if per_call else 1) + ) + assert_current_positions(provider, state) + snapshot = deepcopy(provider.state.to_dict()) + ordinal = binding.append_ordinal + history.finalize_failed_run(state) + history.flush(state) + history.flush(state) + assert provider.state.to_dict() == snapshot + assert binding.append_ordinal == ordinal + assert strategy.calls == 1 and len(client.received_messages) == 2 + assert response.to_dict() == original_response + assert provider.writes == 0 + + +@pytest.mark.parametrize("per_call", [False, True]) +async def test_entity_flushes_after_all_core_providers_and_cold_reload(per_call: bool) -> None: + events: list[str] = [] + history = CountingHistory(events) + strategy = SummarizeSeed(events) + bindings: list[DurableHistoryBinding] = [] + + class LastAfterProvider(ContextProvider): + async def after_run(self, *, session: AgentSession, **kwargs: Any) -> None: + binding = current_durable_history_binding() + assert binding is not None + bindings.append(binding) + session.state[history.source_id][WORKING_BUFFER_KEY][-1].additional_properties["last_after"] = True + + provider = _InMemoryStateProvider() + seed(provider) + client = ToolChatClient(events=events) + agent = Agent( + client=client, + tools=[lookup], + context_providers=[ + LastAfterProvider("last-after"), + history, + CompactionProvider(after_strategy=strategy, history_source_id=history.source_id), + ], + require_per_service_call_history_persistence=per_call, + ) + entity = AgentEntity(agent, state_provider=provider) + response = await entity.run({"message": PROMPT, "correlationId": "tool-turn"}) + + assert_tool_follow_up(client) + assert history.after_calls == (2 if per_call else 1) and strategy.calls == 1 + assert len(bindings) == 1 and bindings[0].correlation_id == "tool-turn" + assert provider.writes == 1 + assert len(transcript(provider)) == 7 + assert ids(provider).count("seed-summary") == 1 + assert (transcript(provider)[-1].extension_data or {})["last_after"] is True + assert not response.messages[-1].additional_properties.get("last_after") + assert history.source_id not in provider._get_state_dict()["data"]["session"]["state"] + original = deepcopy(response.to_dict()) + mailbox = deepcopy(provider.state.data.response_mailbox) + receipts = deepcopy(provider.state.data.completed_correlations) + + cold_provider = _InMemoryStateProvider(raw=provider._get_state_dict()) + cold_client = ToolChatClient(tool_calls=False) + cold_history = DurableHistoryProvider(prune_excluded=False) + cold = AgentEntity( + Agent( + client=cold_client, + context_providers=[cold_history], + require_per_service_call_history_persistence=per_call, + ), + state_provider=cold_provider, + ) + repeated = await cold.run({"message": "must not execute", "correlationId": "tool-turn"}) + assert repeated.to_dict() == original + assert cold_client.received_messages == [] and cold_provider.writes == 0 + assert cold_provider.state.data.response_mailbox == mailbox + assert cold_provider.state.data.completed_correlations == receipts + await cold.run({"message": "continue", "correlationId": "next"}) + assert [message.text for message in cold_client.received_messages[0]] == [ + "seed summary", + "seed answer", + PROMPT, + "", + "", + "answer-2", + "continue", + ] + assert cold_provider.writes == 1 + assert len(ids(cold_provider)) == len(set(ids(cold_provider))) + + +@pytest.mark.parametrize("provider_kind", ["in-memory", "explicit-durable"]) +@pytest.mark.parametrize("per_call", [False, True]) +@pytest.mark.parametrize("store_inputs", [False, True]) +@pytest.mark.parametrize("store_outputs", [False, True]) +@pytest.mark.parametrize("store_context_messages", [False, True]) +@pytest.mark.parametrize("store_context_from", [None, set(), {"selected"}]) +async def test_all_store_flags_survive_substitution_and_control_real_core_hooks( + provider_kind: str, + per_call: bool, + store_inputs: bool, + store_outputs: bool, + store_context_messages: bool, + store_context_from: set[str] | None, +) -> None: + factory = InMemoryHistoryProvider if provider_kind == "in-memory" else DurableHistoryProvider + original = factory( + "custom-history", + store_inputs=store_inputs, + store_outputs=store_outputs, + store_context_messages=store_context_messages, + store_context_from=store_context_from, + skip_excluded=False, + ) + client = ToolChatClient(tool_calls=False) + agent = Agent( + client=client, + context_providers=[original, AddContext("selected"), AddContext("other")], + require_per_service_call_history_persistence=per_call, + ) + prepared: Any = ensure_durable_history(agent, prune_excluded=True) + history = prepared.context_providers[0] + assert isinstance(history, DurableHistoryProvider) + assert history is not original and agent.context_providers[0] is original + assert history.source_id == "custom-history" + assert history.skip_excluded is False and history.prune_excluded is True + assert history.store_inputs is store_inputs and history.store_outputs is store_outputs + assert history.store_context_messages is store_context_messages + assert history.store_context_from == store_context_from + if store_context_from is not None: + assert history.store_context_from is not original.store_context_from + if isinstance(original, DurableHistoryProvider): + assert original.prune_excluded is None + provider = _InMemoryStateProvider() + session = prepared.create_session() + with bound(provider): + await prepared.run("input", session=session) + history.flush(session.state[history.source_id]) + expected_context = [ + f"context-{source}" + for source in ("selected", "other") + if store_context_messages and (store_context_from is None or source in store_context_from) + ] + expected_inputs = [*expected_context, *(["input"] if store_inputs else [])] + expected_outputs = ["answer-1"] if store_outputs else [] + assert [message.text for message in transcript(provider)] == expected_inputs + expected_outputs + entries = provider.state.data.conversation_history + assert [entry.json_type for entry in entries] == ( + ([DurableAgentStateEntryJsonType.REQUEST] if expected_inputs else []) + + ([DurableAgentStateEntryJsonType.RESPONSE] if expected_outputs else []) + ) + assert [message.text for message in client.received_messages[0]] == ["context-selected", "context-other", "input"] + assert provider.writes == 0 + + +@pytest.mark.parametrize("prune_excluded", [False, True]) +def test_explicit_pruning_and_store_only_sinks_are_not_reconfigured(prune_excluded: bool) -> None: + history = DurableHistoryProvider( + store_inputs=False, + store_outputs=False, + store_context_messages=True, + store_context_from={"selected"}, + prune_excluded=prune_excluded, + ) + sink = InMemoryHistoryProvider("sink", load_messages=False, store_inputs=False) + client: Any = RecordingChatClient() + agent = Agent(client=client, context_providers=[history, sink]) + assert ensure_durable_history(agent, prune_excluded=not prune_excluded) is agent + assert agent.context_providers == [history, sink] + external = CaptureHistory(load_messages=True) + external_client: Any = RecordingChatClient() + external_agent = Agent(client=external_client, context_providers=[external, sink]) + assert ensure_durable_history(external_agent) is external_agent + with pytest.raises(ValueError, match="primary"): + conflicting_client: Any = RecordingChatClient() + ensure_durable_history(Agent(client=conflicting_client, context_providers=[history, external, sink])) + + +@pytest.mark.parametrize("per_call", [False, True]) +async def test_service_ownership_suppresses_only_durable_history(per_call: bool) -> None: + provider = _InMemoryStateProvider() + seed(provider) + before = deepcopy(provider.state.to_dict()) + history = DurableHistoryProvider(prune_excluded=True) + sink = CaptureHistory() + client = ToolChatClient(tool_calls=False) + agent = Agent( + client=client, + default_options={"store": True}, + context_providers=[history, sink], + require_per_service_call_history_persistence=per_call, + ) + session = agent.create_session() + with bound(provider, service_owns_history=True) as binding: + await agent.run("service input", session=session) + assert await history.get_messages(session.session_id) == [] + await history.save_messages(session.session_id, [Message("user", ["do not store"])]) + history.flush({WORKING_BUFFER_KEY: [Message("assistant", ["do not insert"])]}) + assert binding.append_ordinal == 0 + assert provider.state.to_dict() == before and provider.writes == 0 + assert [message.text for message in client.received_messages[0]] == ["service input"] + assert [[message.text for message in batch] for batch in sink.saved] == [["service input", "answer-1"]] + + +@pytest.mark.parametrize("per_call", [False, True]) +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +async def test_failed_core_call_does_not_persist_inputs_before_the_history_hook(per_call: bool, stream: bool) -> None: + provider = _InMemoryStateProvider() + seed(provider) + before = deepcopy(provider.state.to_dict()) + history = CountingHistory([]) + agent = Agent( + client=ToolChatClient(fail=True), + context_providers=[history], + require_per_service_call_history_persistence=per_call, + ) + session = agent.create_session() + with bound(provider) as binding: + with pytest.raises(RuntimeError, match="model failed"): + if stream: + await agent.run("failed input", session=session, stream=True).get_final_response() + else: + await agent.run("failed input", session=session) + history.finalize_failed_run(session.state[history.source_id]) + history.flush(session.state[history.source_id]) + assert binding.pending_inputs == [] + assert history.after_calls == 0 + assert provider.state.to_dict() == before + assert provider.writes == 0 + + +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +async def test_failed_second_service_call_commits_actual_tool_result_and_cold_replays_it(stream: bool) -> None: + tool_calls: list[str] = [] + + @tool(name="lookup", approval_mode="never_require") + def counted_lookup(key: str) -> str: + tool_calls.append(key) + return f"value:{key}" + + history = CountingHistory([]) + client = ToolChatClient(fail_on_call=2) + agent_type = Agent if stream else NonStreamingAgent + agent = agent_type( + client=client, + tools=[counted_lookup], + context_providers=[history], + require_per_service_call_history_persistence=True, + ) + provider = _InMemoryStateProvider() + session = agent.create_session() + foreign_state = {"approval": {"pending": ["keep"]}} + session.state["foreign-provider"] = deepcopy(foreign_state) + provider.state.data.session = session.to_dict() + entity = AgentEntity(agent, state_provider=provider) + message = Message( + "user", [PROMPT], message_id="projected-input", additional_properties={"trace": {"tags": ["original"]}} + ) + original_input = deepcopy(message.to_dict()) + request = {"message": PROMPT, "correlationId": "failed-tool", "contextMessages": [deepcopy(original_input)]} + + failed = await entity.run(request) + + assert_tool_follow_up(client) + assert history.before_calls == 2 and history.after_calls == 1 + assert tool_calls == ["durable"] + assert failed.additional_properties["durable_status"] == "error" + assert "model failed" in failed.text and provider.writes == 1 + assert message.to_dict() == original_input and request["contextMessages"] == [original_input] + assert [entry.json_type for entry in provider.state.data.conversation_history] == [ + DurableAgentStateEntryJsonType.REQUEST, + DurableAgentStateEntryJsonType.RESPONSE, + DurableAgentStateEntryJsonType.REQUEST, + ] + assert all(entry.correlation_id == "failed-tool" for entry in provider.state.data.conversation_history) + assert [message.text for message in transcript(provider)] == [PROMPT, "", ""] + actual_result = next( + content + for message in client.received_messages[1] + for content in message.contents + if content.type == "function_result" + ) + saved_result = transcript(provider)[-1].to_chat_message() + assert len(saved_result.contents) == 1 + assert saved_result.contents[0].call_id == actual_result.call_id == "call-1" + assert saved_result.contents[0].result == actual_result.result == "value:durable" + + raw = provider._get_state_dict() + data = raw["data"] + assert data["responseMailbox"]["failed-tool"]["response"] == failed.to_dict() + assert "failed-tool" in data["completedCorrelations"] + assert data["ingestedMessages"] == {"projected-input": [message_identity(message)]} + assert history.source_id not in data["session"]["state"] + assert data["session"]["state"]["foreign-provider"] == foreign_state + cold_provider = _InMemoryStateProvider(raw=raw) + cold_client = ToolChatClient(tool_calls=False) + cold = AgentEntity( + agent_type( + client=cold_client, + tools=[counted_lookup], + context_providers=[DurableHistoryProvider(prune_excluded=False)], + require_per_service_call_history_persistence=True, + ), + state_provider=cold_provider, + ) + repeated = await cold.run(request) + assert repeated.to_dict() == failed.to_dict() + assert cold_client.received_messages == [] and cold_provider.writes == 0 + await cold.run({"message": "continue", "correlationId": "next"}) + replayed = cold_client.received_messages[0] + assert [message.text for message in replayed] == [PROMPT, "", "", "continue"] + assert [message.role for message in replayed] == ["user", "assistant", "tool", "user"] + calls = [content for message in replayed for content in message.contents if content.type == "function_call"] + results = [content for message in replayed for content in message.contents if content.type == "function_result"] + assert len(calls) == len(results) == 1 + assert calls[0].call_id == results[0].call_id == actual_result.call_id + assert results[0].result == actual_result.result + assert tool_calls == ["durable"] and cold_provider.writes == 1 + assert cold_provider.state.data.response_mailbox["failed-tool"] == data["responseMailbox"]["failed-tool"] + assert ( + cold_provider.state.data.completed_correlations["failed-tool"] == (data["completedCorrelations"]["failed-tool"]) + ) + + +@pytest.mark.parametrize("per_call", [False, True], ids=["per-run", "per-service-call"]) +@pytest.mark.parametrize("store_inputs", [False, True]) +@pytest.mark.parametrize("store_outputs", [False, True]) +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +async def test_failure_finalization_respects_real_core_cadence_and_store_flags( + per_call: bool, store_inputs: bool, store_outputs: bool, stream: bool +) -> None: + tool_calls: list[str] = [] + + @tool(name="lookup", approval_mode="never_require") + def counted_lookup(key: str) -> str: + tool_calls.append(key) + return f"value:{key}" + + provider = _InMemoryStateProvider() + history = DurableHistoryProvider(store_inputs=store_inputs, store_outputs=store_outputs, prune_excluded=False) + client = ToolChatClient(fail_on_call=2) + agent = Agent( + client=client, + tools=[counted_lookup], + context_providers=[history], + require_per_service_call_history_persistence=per_call, + ) + session = agent.create_session() + with bound(provider) as binding: + with pytest.raises(RuntimeError, match="model failed"): + if stream: + await agent.run(PROMPT, session=session, stream=True).get_final_response() + else: + await agent.run(PROMPT, session=session) + state = session.state[history.source_id] + before = deepcopy(provider.state.to_dict()) + assert bool(binding.pending_inputs) is (per_call and store_inputs) + expected = ([PROMPT] if per_call and store_inputs else []) + ([""] if per_call and store_outputs else []) + assert [message.text for message in transcript(provider)] == expected + history.finalize_failed_run(state) + history.flush(state) + assert binding.pending_inputs == [] + expected_result = per_call and store_inputs and store_outputs + assert [message.text for message in transcript(provider)] == expected + ([""] if expected_result else []) + results = [ + content + for message in transcript(provider) + for content in message.to_chat_message().contents + if content.type == "function_result" + ] + assert len(results) == int(expected_result) + if expected_result: + assert results[0].result == "value:durable" + assert_tool_follow_up(client) + else: + assert provider.state.to_dict() == before + snapshot = deepcopy(provider.state.to_dict()) + ordinal = binding.append_ordinal + history.finalize_failed_run(state) + history.flush(state) + assert provider.state.to_dict() == snapshot and binding.append_ordinal == ordinal + assert len(client.received_messages) == 2 and tool_calls == ["durable"] + assert provider.writes == 0 + + +@pytest.mark.parametrize("message_id", [None, "shared"], ids=["anonymous", "reused-id"]) +async def test_failed_inputs_preserve_matched_groups_metadata_and_original_ingestion_hash( + message_id: str | None, +) -> None: + provider = _InMemoryStateProvider() + history = DurableHistoryProvider(prune_excluded=False) + agent = Agent(client=ToolChatClient(), require_per_service_call_history_persistence=True) + session = agent.create_session() + state: dict[str, Any] = {} + session.state[history.source_id] = state + old_call = Message( + "assistant", + [Content.from_function_call("historical", "lookup", arguments={})], + message_id="old-call", + ) + provider.state.data.conversation_history.append( + DurableAgentStateResponse("previous", OLD, [DurableAgentStateMessage.from_chat_message(old_call)]) + ) + with bound(provider) as binding: + await history.before_run( + agent=agent, + session=session, + context=SessionContext(input_messages=[Message("user", ["discard this snapshot"])]), + state=state, + ) + assert binding.pending_inputs + context = SessionContext(input_messages=[Message("user", [PROMPT], message_id="shared")]) + context._response = AgentResponse( + messages=[ + Message( + "assistant", + [ + Content.from_function_call(call_id, "lookup", arguments={}) + for call_id in ("call-1", "call-2", "completed") + ], + ) + ] + ) + await history.after_run(agent=agent, session=session, context=context, state=state) + assert binding.pending_inputs == [], "a standalone after_run must also clear the pending snapshot" + await history.save_messages( + session.session_id, + [Message("tool", [Content.from_function_result("completed", result="already stored")])], + state=state, + ) + result = Message( + "tool", + [ + Content("function_result", call_id=call_id, result={"values": [call_id]}) + for call_id in ("call-1", "call-2") + ], + message_id=message_id, + additional_properties={"trace": {"tags": ["original"]}}, + ) + original = deepcopy(result.to_dict()) + inputs = [ + Message("user", ["unrelated fresh request"]), + Message("tool", [Content.from_function_result("historical", result="wrong correlation")]), + Message("tool", [Content.from_function_result("unknown", result="no stored call")]), + Message("tool", [Content.from_function_result("completed", result="duplicate result")]), + Message("tool", []), + Message("tool", [Content("function_result", result="no call id")]), + Message("user", [Content.from_function_result("call-1", result="wrong role")]), + Message( + "tool", [Content.from_function_result("call-1", result="mixed input"), Content.from_text("fresh input")] + ), + Message( + "tool", + [ + Content.from_function_result("call-1", result="mixed ids"), + Content.from_function_result("unknown", result="no"), + ], + ), + Message("tool", [Content.from_function_result("call-1", result="duplicate id")] * 2), + result, + deepcopy(result), + ] + before = deepcopy(provider.state.to_dict()) + await history.before_run( + agent=agent, + session=session, + context=SessionContext( + input_messages=[Message("tool", [Content.from_function_result("call-1", result="superseded snapshot")])] + ), + state=state, + ) + await history.before_run( + agent=agent, session=session, context=SessionContext(input_messages=inputs), state=state + ) + assert provider.state.to_dict() == before, "before_run must capture, not append" + assert binding.pending_inputs[-2] is not result + assert binding.pending_inputs[-2].additional_properties["trace"] is not result.additional_properties["trace"] + assert result.to_dict() == original + result.contents[0].result["values"].append("caller mutation") + result.additional_properties["trace"]["tags"].append("caller mutation") + history.finalize_failed_run(state) + assert binding.pending_inputs == [] + assert set(state) == {WORKING_BUFFER_KEY, POSITIONS_KEY} + assert len(transcript(provider)) == 5 + saved = transcript(provider)[-1] + assert saved.ingestion_identity == message_identity(Message.from_dict(original)) + assert saved.message_id and saved.message_id != message_id + assert result.message_id == message_id + assert [content.to_dict()["result"] for content in saved.contents] == [ + {"values": ["call-1"]}, + {"values": ["call-2"]}, + ] + assert saved.extension_data == {"trace": {"tags": ["original"]}} + working = state[WORKING_BUFFER_KEY][-1] + working.additional_properties["trace"]["tags"].append("compaction") + working.contents[0].result["values"].append("working mutation") + assert saved.extension_data == {"trace": {"tags": ["original"]}} + history.flush(state) + assert saved.extension_data == {"trace": {"tags": ["original", "compaction"]}} + assert saved.contents[0].to_dict()["result"] == {"values": ["call-1"]} + assert_current_positions(provider, state) + snapshot = deepcopy(provider.state.to_dict()) + ordinal = binding.append_ordinal + history.finalize_failed_run(state) + history.flush(state) + assert provider.state.to_dict() == snapshot and binding.append_ordinal == ordinal + assert provider.writes == 0 + + +@pytest.mark.parametrize("disabled_by", ["service", "store-inputs", "no-correlation"]) +async def test_finalizing_pending_inputs_rechecks_binding_and_input_storage(disabled_by: str) -> None: + provider = _InMemoryStateProvider() + history = DurableHistoryProvider(prune_excluded=False) + agent = Agent(client=ToolChatClient(), require_per_service_call_history_persistence=True) + session = agent.create_session() + state: dict[str, Any] = {} + with bound(provider) as binding: + context = SessionContext(input_messages=[]) + context._response = AgentResponse( + messages=[Message("assistant", [Content.from_function_call("call-1", "lookup", arguments={})])] + ) + await history.after_run(agent=agent, session=session, context=context, state=state) + await history.before_run( + agent=agent, + session=session, + context=SessionContext( + input_messages=[Message("tool", [Content.from_function_result("call-1", result="actual result")])] + ), + state=state, + ) + assert binding.pending_inputs + before = deepcopy(provider.state.to_dict()) + ordinal = binding.append_ordinal + if disabled_by == "service": + binding.service_owns_history = True + elif disabled_by == "store-inputs": + history.store_inputs = False + else: + binding.correlation_id = None + history.finalize_failed_run(state) + assert binding.pending_inputs == [] + assert provider.state.to_dict() == before and binding.append_ordinal == ordinal + history.finalize_failed_run(state) + assert provider.state.to_dict() == before and provider.writes == 0 + + +async def test_failed_load_does_not_store_a_raw_historical_tool_continuation() -> None: + class FailingLoadHistory(DurableHistoryProvider): + async def get_messages(self, session_id: str | None, **kwargs: Any) -> list[Message]: + raise OSError("history load failed") + + provider = _InMemoryStateProvider() + old_call = Message( + "assistant", [Content.from_function_call("call-1", "lookup", arguments={})], message_id="old-call" + ) + provider.state.data.conversation_history.append( + DurableAgentStateResponse("previous", OLD, [DurableAgentStateMessage.from_chat_message(old_call)]) + ) + before = deepcopy(provider.state.to_dict()) + history = FailingLoadHistory(prune_excluded=False) + client = ToolChatClient() + agent = Agent(client=client, context_providers=[history], require_per_service_call_history_persistence=True) + session = agent.create_session() + messages = [ + Message("tool", [Content.from_function_result("call-1", result="caller result")]), + Message("user", ["new input"]), + ] + with bound(provider) as binding: + with pytest.raises(OSError, match="history load failed"): + await agent.run(messages, session=session) + assert binding.pending_inputs + history.finalize_failed_run(session.state[history.source_id]) + history.flush(session.state[history.source_id]) + assert binding.pending_inputs == [] + assert provider.state.to_dict() == before + assert client.received_messages == [] and provider.writes == 0 + + +@pytest.mark.parametrize("with_state", [False, True]) +async def test_generic_save_appends_anonymous_messages_with_stable_write_time_ids(with_state: bool) -> None: + provider = _InMemoryStateProvider() + history = DurableHistoryProvider() + state: dict[str, Any] | None = {} if with_state else None + messages = [Message("user", ["repeat"]), Message("user", ["repeat"])] + with bound(provider, "append") as binding: + await history.save_messages("session", messages, state=state) + await history.save_messages("session", messages[:1], state=state) + await history.save_messages("session", [], state=state) + assert binding.append_ordinal == 2 + assert ids(provider) == [ + "durable_request_append_0_0", + "durable_request_append_0_1", + "durable_request_append_1_0", + ] + assert [message.message_id for message in messages] == [None, None] + assert [message.text for message in transcript(provider)] == ["repeat"] * 3 + if state is not None: + assert_current_positions(provider, state) + assert provider.writes == 0 + cold = _InMemoryStateProvider(raw=json.loads(provider.state.to_json())) + with bound(cold, "append"): + loaded = await history.get_messages("session") + assert [message.message_id for message in loaded] == ids(provider) + assert [message.text for message in loaded] == ["repeat"] * 3 + + +async def test_reused_ids_get_internal_revisions_without_changing_external_ids() -> None: + provider = _InMemoryStateProvider() + history = DurableHistoryProvider(prune_excluded=False) + sink = CaptureHistory() + agent = Agent( + client=ToolChatClient(tool_calls=False, response_message_id="shared"), + context_providers=[history, sink], + ) + session = agent.create_session() + inputs = [ + Message("user", ["version one"], message_id="shared"), + Message("user", ["version two"], message_id="shared"), + ] + responses: list[AgentResponse] = [] + for index, message in enumerate(inputs): + with bound(provider, f"revision-{index}"): + responses.append(await agent.run(message, session=session)) + history.flush(session.state[history.source_id]) + assert [message.message_id for message in inputs] == ["shared", "shared"] + assert [response.messages[0].message_id for response in responses] == ["shared", "shared"] + assert [[message.message_id for message in batch] for batch in sink.saved] == [["shared", "shared"]] * 2 + assert len(ids(provider)) == len(set(ids(provider))) == 4 + assert ids(provider)[0] == "shared" + assert all(message_id and message_id.startswith("durable_revision_") for message_id in ids(provider)[1:]) + assert [message.text for message in transcript(provider)] == ["version one", "answer-1", "version two", "answer-2"] + before = deepcopy(provider.state.to_dict()) + inputs[0].contents[0].text = "caller changed input" + responses[-1].messages[0].additional_properties["model_metadata"]["tags"].append("caller changed output") + assert provider.state.to_dict() == before + cold = _InMemoryStateProvider(raw=json.loads(provider.state.to_json())) + state: dict[str, Any] = {} + with bound(cold): + loaded = await history.get_messages("session", state=state) + loaded[2].additional_properties["revision_marker"] = True + history.flush(state) + marked = [message.text for message in transcript(cold) if (message.extension_data or {}).get("revision_marker")] + assert marked == ["version two"] + assert_current_positions(cold, state) + assert [message.message_id for message in loaded] == ids(provider) + assert [message.text for message in loaded] == ["version one", "answer-1", "version two", "answer-2"] + + +async def test_generated_ids_reserve_supplied_ids_in_the_same_batch() -> None: + provider = _InMemoryStateProvider() + history = DurableHistoryProvider() + reserved_id = "durable_request_current_0_0" + messages = [Message("user", ["anonymous"]), Message("user", ["supplied"], message_id=reserved_id)] + with bound(provider): + await history.save_messages("session", messages) + assert ids(provider) == [f"{reserved_id}_1", reserved_id] + assert [message.message_id for message in messages] == [None, reserved_id] + + +async def test_append_and_flush_never_alias_tool_payloads_or_response_annotations() -> None: + provider = _InMemoryStateProvider() + history = DurableHistoryProvider(prune_excluded=False) + session = AgentSession() + result = Message( + "tool", + [Content("function_result", call_id="call-1", result={"values": [1]})], + additional_properties={"trace": {"tags": ["original"]}}, + ) + response = AgentResponse(messages=[result], additional_properties={"receipt": {"tags": ["original"]}}) + before = deepcopy(response.to_dict()) + context = SessionContext(input_messages=[Message("user", ["input"])]) + context._response = response + state: dict[str, Any] = {} + with bound(provider): + await history.after_run(agent=None, session=session, context=context, state=state) + assert ids(provider) == ["durable_request_current_0_0", "durable_response_current_1_0"] + assert context.input_messages[0].message_id is None and response.messages[0].message_id is None + working = state[WORKING_BUFFER_KEY][-1] + working.additional_properties["trace"]["tags"].append("compaction") + working.contents[0].result["values"].append(2) + assert transcript(provider)[-1].contents[0].to_dict()["result"] == {"values": [1]} + assert (transcript(provider)[-1].extension_data or {})["trace"] == {"tags": ["original"]} + history.flush(state) + assert (transcript(provider)[-1].extension_data or {})["trace"] == {"tags": ["original", "compaction"]} + assert transcript(provider)[-1].contents[0].to_dict()["result"] == {"values": [1]} + assert response.to_dict() == before + working.additional_properties["trace"]["tags"].append("not flushed") + assert (transcript(provider)[-1].extension_data or {})["trace"] == {"tags": ["original", "compaction"]} + assert provider.writes == 0 + + +async def test_repeated_core_summaries_survive_pruning_id_reuse_and_cold_reload() -> None: + provider = _InMemoryStateProvider() + seed(provider) + history = DurableHistoryProvider(prune_excluded=True) + summary_client = ToolChatClient(tool_calls=False) + compaction = CompactionProvider( + after_strategy=SummarizationStrategy( + client=summary_client, target_count=2, threshold=0, max_summary_input_tokens=None + ), + history_source_id=history.source_id, + ) + session = AgentSession() + state: dict[str, Any] = {} + session.state[history.source_id] = state + generated_ids: list[str | None] = [] + for turn in range(3): + with bound(provider, f"turn-{turn}") as binding: + await history.get_messages(session.session_id, state=state) + await history.save_messages( + session.session_id, + [ + Message("user", [f"question-{turn}"], message_id=f"user-{turn}"), + Message("assistant", [f"response-{turn}"], message_id=f"assistant-{turn}"), + ], + state=state, + ) + await compaction.after_run(agent=None, session=session, context=None, state={}) + buffer = state[WORKING_BUFFER_KEY] + summary = buffer[0] + assert summary.text == f"answer-{turn + 1}" + generated_id = summary.message_id + generated_ids.append(generated_id) + original_links = deepcopy(summary.additional_properties["_group"]) + originals = [ + message for message in buffer[1:] if message.message_id in original_links["_summary_of_message_ids"] + ] + # Core may re-include the older summary while grouping duplicate IDs. Preserve its + # actual inclusion decisions, without hiding the new summary's text or identity. + expected_texts = [message.text for message in buffer if not message.additional_properties.get("_excluded")] + assert originals + if turn == 2: + older_summary = next(message for message in originals if message.message_id == generated_id) + assert older_summary.text == "answer-2" + assert sum(message.message_id == generated_id for message in buffer) == 2 + + history.flush(state) + if turn == 2: + assert summary.message_id != generated_id + assert summary.message_id.startswith("durable_revision_compaction_turn-2_") + assert [message.text for message in transcript(provider)] == expected_texts + assert len(ids(provider)) == len(set(ids(provider))) == len(expected_texts) + assert ( + summary.additional_properties["_group"]["_summary_of_message_ids"] + == (original_links["_summary_of_message_ids"]) + ) + assert ( + summary.additional_properties["_group"]["_summary_of_group_ids"] + == (original_links["_summary_of_group_ids"]) + ) + assert summary.additional_properties["_group"]["id"] == f"group_{summary.message_id}" + assert all( + message.additional_properties["_group"]["_summarized_by_summary_id"] == summary.message_id + for message in originals + ) + assert all( + (message.extension_data or {})["_group"]["_summarized_by_summary_id"] == summary.message_id + for message in transcript(provider) + if message.message_id in original_links["_summary_of_message_ids"] + ) + assert_current_positions(provider, state) + snapshot = deepcopy(provider.state.to_dict()) + ordinal = binding.append_ordinal + history.flush(state) + assert provider.state.to_dict() == snapshot and binding.append_ordinal == ordinal + assert generated_ids == ["summary_4", "summary_5", "summary_5"] + assert len(summary_client.received_messages) == 3 and provider.writes == 0 + + cold = _InMemoryStateProvider(raw=json.loads(provider.state.to_json())) + cold_history = DurableHistoryProvider(prune_excluded=True) + cold_state: dict[str, Any] = {} + with bound(cold, "cold"): + loaded = await cold_history.get_messages(session.session_id, state=cold_state) + assert loaded[0].text == "answer-3" + assert [message.text for message in loaded] == expected_texts + assert [message.message_id for message in loaded] == ids(provider) + assert loaded[0].additional_properties == transcript(provider)[0].extension_data + cold_history.flush(cold_state) + assert cold.state.to_dict() == provider.state.to_dict() + assert_current_positions(cold, cold_state) + assert cold.writes == 0 + + +@pytest.mark.parametrize("nested_links", [False, True], ids=["top-level", "core-group"]) +@pytest.mark.parametrize("remove_old_summary", [False, True], ids=["old-in-buffer", "old-removed"]) +async def test_reused_summary_ids_keep_older_links_and_original_contents( + nested_links: bool, remove_old_summary: bool +) -> None: + provider = _InMemoryStateProvider() + seed(provider) + provider.state.data.conversation_history.append( + DurableAgentStateRequest("current", OLD, [stored("current", "current")]) + ) + history = DurableHistoryProvider(prune_excluded=False) + session = AgentSession() + state: dict[str, Any] = {} + session.state[history.source_id] = state + calls = 0 + + def links(message: Message) -> dict[str, Any]: + if nested_links: + return message.additional_properties.setdefault("_group", {}) + return message.additional_properties + + async def summarize(messages: list[Message]) -> bool: + nonlocal calls + source_id = "seed-user" if calls == 0 else "seed-assistant" + source = next(message for message in messages if message.message_id == source_id) + calls += 1 + summary_id = "repeated-summary" + source.additional_properties["_excluded"] = True + links(source)["_summarized_by_summary_id"] = summary_id + if calls == 2: + # Ordinary source edits must not turn annotation reconciliation into content replacement. + source.contents[0].text = "working-only text" + if remove_old_summary: + messages[:] = [message for message in messages if message.message_id != summary_id] + summary = Message( + "assistant", + [Content.from_text(f"summary version {calls}", additional_properties={"trace": {"call": calls}})], + message_id=summary_id, + ) + links(summary)["_summary_of_message_ids"] = [source_id] + insertion_index = messages.index(source) + 1 + messages.insert(insertion_index, summary) + annotate_message_groups(messages, from_index=insertion_index) + return True + + compaction = CompactionProvider(after_strategy=summarize, history_source_id=history.source_id) + with bound(provider) as binding: + await history.get_messages(session.session_id, state=state) + await compaction.after_run(agent=None, session=session, context=None, state={}) + history.flush(state) + await compaction.after_run(agent=None, session=session, context=None, state={}) + new_summary = next(message for message in state[WORKING_BUFFER_KEY] if message.text == "summary version 2") + assert new_summary.message_id == "repeated-summary" + history.flush(state) + assert new_summary.message_id != "repeated-summary" + # Removing a summary from the logical buffer does not erase its body or lineage under keep_all. + assert [message.text for message in transcript(provider)] == [ + "seed question", + "summary version 1", + "seed answer", + "summary version 2", + "current", + ] + assert len(ids(provider)) == len(set(ids(provider))) == 5 + stored_messages = {message.message_id: message.to_chat_message() for message in transcript(provider)} + assert links(stored_messages["seed-user"])["_summarized_by_summary_id"] == "repeated-summary" + assert links(stored_messages["seed-assistant"])["_summarized_by_summary_id"] == new_summary.message_id + assert links(stored_messages["repeated-summary"])["_summary_of_message_ids"] == ["seed-user"] + assert links(stored_messages[new_summary.message_id])["_summary_of_message_ids"] == ["seed-assistant"] + assert_current_positions(provider, state) + snapshot = deepcopy(provider.state.to_dict()) + ordinal = binding.append_ordinal + history.flush(state) + assert provider.state.to_dict() == snapshot and binding.append_ordinal == ordinal + assert calls == 2 and provider.writes == 0 + + cold = _InMemoryStateProvider(raw=json.loads(provider.state.to_json())) + cold_history = DurableHistoryProvider(prune_excluded=False) + cold_state: dict[str, Any] = {} + with bound(cold): + loaded = await cold_history.get_messages(session.session_id, state=cold_state) + assert [message.text for message in loaded] == [ + *([] if remove_old_summary else ["summary version 1"]), + "summary version 2", + "current", + ] + assert [message.message_id for message in loaded] == [ + *([] if remove_old_summary else ["repeated-summary"]), + new_summary.message_id, + "current", + ] + cold_history.flush(cold_state) + assert cold.state.to_dict() == provider.state.to_dict() + assert_current_positions(cold, cold_state) + assert cold.writes == 0 + + +@pytest.mark.parametrize("entry_type", [DurableAgentStateRequest, DurableAgentStateResponse]) +@pytest.mark.parametrize("message_count", [2, 4]) +async def test_multiple_mid_entry_summaries_keep_exact_order_metadata_and_receipts( + entry_type: type[DurableAgentStateRequest] | type[DurableAgentStateResponse], message_count: int +) -> None: + provider = _InMemoryStateProvider() + source_ids = [f"item-{index}" for index in range(message_count)] + owner = entry_type("old", OLD, [stored(message_id, message_id) for message_id in source_ids]) + owner.extension_data = {"envelope": {"tags": ["original"]}} + owner.unknown_fields = {"futureField": {"keep": [1]}} + if isinstance(owner, DurableAgentStateRequest): + owner.orchestration_id = "workflow" + owner.response_schema = {"properties": {"value": {"type": "string"}}} + else: + owner.usage = DurableAgentStateUsage(input_token_count=7) + unknown = DurableAgentStateUnknownEntry({"$type": "futureKind", "future": {"opaque": [1, 2]}}) + current = DurableAgentStateRequest("current", OLD, [stored("current", "current")]) + provider.state.data.conversation_history.extend([unknown, owner, current]) + provider.state.record_response( + "old", + AgentResponse(messages=[Message("assistant", ["original answer"])]), + delivery_window_seconds=3600, + ) + mailbox = deepcopy(provider.state.data.response_mailbox) + receipts = deepcopy(provider.state.data.completed_correlations) + history = DurableHistoryProvider(prune_excluded=False) + state: dict[str, Any] = {} + with bound(provider): + loaded = await history.get_messages("session", state=state) + buffer: list[Message] = [] + expected: list[str] = [] + for index, message in enumerate(loaded[:-1]): + buffer.append(message) + expected.append(source_ids[index]) + if index + 1 < message_count: + summary_id = f"summary-{index}" + buffer.append(Message("assistant", [summary_id], message_id=summary_id)) + expected.append(summary_id) + buffer[-1].additional_properties["target"] = "last original" + buffer.append(loaded[-1]) + expected.append("current") + state[WORKING_BUFFER_KEY] = buffer + history.flush(state) + assert ids(provider) == expected + assert (transcript(provider)[-2].extension_data or {})["target"] == "last original" + assert_current_positions(provider, state) + envelopes: list[Any] = [ + entry + for entry in provider.state.data.conversation_history + if isinstance(entry, entry_type) and entry.correlation_id == "old" + ] + assert len(envelopes) == message_count + for entry in envelopes: + assert entry.correlation_id == "old" and entry.created_at == OLD + assert entry.extension_data == owner.extension_data and entry.unknown_fields == owner.unknown_fields + assert all( + message.message_id and not message.message_id.startswith("summary-") for message in entry.messages + ) + assert envelopes[0].extension_data is not envelopes[-1].extension_data + assert envelopes[0].unknown_fields is not envelopes[-1].unknown_fields + if isinstance(owner, DurableAgentStateRequest): + assert all(entry.response_schema == owner.response_schema for entry in envelopes) + assert all(entry.orchestration_id == "workflow" for entry in envelopes) + assert envelopes[0].response_schema is not envelopes[-1].response_schema + else: + assert owner.usage is not None + expected_usage = owner.usage.to_dict() + assert all(entry.usage.to_dict() == expected_usage for entry in envelopes) + assert envelopes[0].usage is not envelopes[-1].usage + snapshot = deepcopy(provider.state.to_dict()) + history.flush(state) + assert provider.state.to_dict() == snapshot + assert provider.state.data.response_mailbox == mailbox + assert provider.state.data.completed_correlations == receipts + + cold = _InMemoryStateProvider(raw=json.loads(provider.state.to_json())) + history = DurableHistoryProvider(prune_excluded=True) + with bound(cold): + loaded = await history.get_messages("session", state=state) + assert [message.message_id for message in loaded] == expected + next(message for message in loaded if message.message_id == "item-1").additional_properties["_excluded"] = True + state[WORKING_BUFFER_KEY].insert(1, Message("assistant", ["later"], message_id="later-summary")) + loaded[0].additional_properties["target"] = "first original" + history.flush(state) + assert ids(cold) == [ + expected[0], + "later-summary", + *[message_id for message_id in expected[1:] if message_id != "item-1"], + ] + assert (transcript(cold)[0].extension_data or {})["target"] == "first original" + assert (cold.state.data.truncation or {})["evictedMessageCount"] == 1 + assert_current_positions(cold, state) + snapshot = deepcopy(cold.state.to_dict()) + history.flush(state) + assert cold.state.to_dict() == snapshot + assert cold.state.data.conversation_history[0].to_dict() == unknown.to_dict() + assert cold.state.data.response_mailbox == mailbox + assert cold.state.data.completed_correlations == receipts + + +@pytest.mark.parametrize("remove_entry", [False, True]) +async def test_flush_rebuilds_positions_after_detached_pruning_without_resurrection(remove_entry: bool) -> None: + provider = _InMemoryStateProvider() + owner = DurableAgentStateRequest("old", OLD, [stored("a", "a"), stored("b", "b")]) + other = DurableAgentStateResponse("older", OLD, [stored("c", "c", "assistant")]) + provider.state.data.conversation_history.extend([owner, other]) + history = DurableHistoryProvider(prune_excluded=False) + state: dict[str, Any] = {} + with bound(provider): + await history.get_messages("session", state=state) + replacement = deepcopy(provider.state.data.conversation_history) + if remove_entry: + replacement.pop(0) + else: + replacement[0].messages.pop(0) + provider.state.data.conversation_history = replacement + state[WORKING_BUFFER_KEY][-1].additional_properties["current_owner"] = True + history.flush(state) + assert ids(provider) == (["c"] if remove_entry else ["b", "c"]) + assert (transcript(provider)[-1].extension_data or {})["current_owner"] is True + assert not (other.messages[0].extension_data or {}).get("current_owner") + assert_current_positions(provider, state) + snapshot = deepcopy(provider.state.to_dict()) + history.flush(state) + assert provider.state.to_dict() == snapshot + + +@pytest.mark.parametrize("kind", list(DurableAgentStateEntryJsonType)) +def test_prune_only_drops_changed_bare_known_envelopes(kind: DurableAgentStateEntryJsonType) -> None: + bare = DurableAgentStateEntry(kind, "bare", OLD, [stored("bare", "remove")]) + metadata = DurableAgentStateEntry(kind, "metadata", OLD, [stored("metadata", "remove")], extension_data={}) + metadata.unknown_fields = {"future": {"keep": True}} + empty = DurableAgentStateRequest("already-empty", OLD, []) + opaque = DurableAgentStateUnknownEntry({"$type": "futureKind", "payload": {"keep": [1]}}) + unknown = DurableAgentStateEntry("futureKind", "future", OLD, [stored("future", "opaque")]) + history = [opaque, empty, bare, metadata, unknown] + unknown_before = deepcopy(unknown.to_dict()) + metadata_before = deepcopy(metadata.to_dict()) + message = bare.messages[0] + prune_messages( + history, + [(bare, message), (bare, message), (metadata, metadata.messages[0]), (unknown, unknown.messages[0])], + ) + assert history == [opaque, empty, metadata, unknown] + metadata_before["messages"] = [] + assert metadata.to_dict() == metadata_before + assert unknown.to_dict() == unknown_before + snapshot = [deepcopy(entry.to_dict()) for entry in history] + prune_messages(history, [(bare, message)]) + assert [entry.to_dict() for entry in history] == snapshot + + +@pytest.mark.parametrize("correlation_id", [None, "current"]) +async def test_eager_pruning_protects_system_and_current_exchange_with_exact_count(correlation_id: str | None) -> None: + provider = _InMemoryStateProvider() + old = DurableAgentStateRequest("old", OLD, [stored("system", "keep instructions", "system"), stored("old", "drop")]) + metadata = DurableAgentStateResponse( + "old", + OLD, + [stored("old-answer", "drop", "assistant")], + extension_data={"usage-note": {"keep": [1]}}, + ) + current = DurableAgentStateRequest("current", OLD, [stored("current-input", "keep")]) + answer = DurableAgentStateResponse("current", OLD, [stored("current-answer", "keep", "assistant")]) + opaque = DurableAgentStateUnknownEntry({"$type": "futureKind", "payload": {"keep": [1]}}) + empty = DurableAgentStateRequest("already-empty", OLD, []) + provider.state.data.conversation_history.extend([opaque, empty, old, metadata, current, answer]) + provider.state.record_response( + "old", + AgentResponse(messages=[Message("assistant", ["mailbox original"])]), + delivery_window_seconds=3600, + ) + mailbox = deepcopy(provider.state.data.response_mailbox) + receipts = deepcopy(provider.state.data.completed_correlations) + provider.state.data.truncation = {"evictedMessageCount": 7, "firstEvictedAt": OLD.isoformat(), "future": [1]} + history = DurableHistoryProvider(prune_excluded=True) + state: dict[str, Any] = {} + with bound(provider, correlation_id) as binding: + await history.get_messages("session", state=state) + old_message = old.messages[-1] + for message in state[WORKING_BUFFER_KEY]: + message.additional_properties["_excluded"] = True + history.flush(state) + assert ids(provider) == ["system", "current-input", "current-answer"] + assert metadata.messages == [] and metadata in provider.state.data.conversation_history + assert empty in provider.state.data.conversation_history and opaque in provider.state.data.conversation_history + assert (provider.state.data.truncation or {})["evictedMessageCount"] == 9 + assert (provider.state.data.truncation or {})["firstEvictedAt"] == OLD.isoformat() + assert (provider.state.data.truncation or {})["future"] == [1] + snapshot = deepcopy(provider.state.to_dict()) + history._prune(binding, [(old, old_message), (old, old_message)]) + history.flush(state) + assert provider.state.to_dict() == snapshot + assert_current_positions(provider, state) + assert provider.state.data.response_mailbox == mailbox + assert provider.state.data.completed_correlations == receipts + assert provider.writes == 0 + + +@pytest.mark.parametrize("protected_by", ["system", "current", "newest"]) +async def test_eager_pruning_protects_atomic_groups_intersecting_the_floor(protected_by: str) -> None: + provider = _InMemoryStateProvider() + call = DurableAgentStateMessage.from_chat_message( + Message( + "assistant", + [Content.from_function_call(call_id="lookup", name="lookup", arguments="{}")], + message_id="call", + ) + ) + result = DurableAgentStateMessage.from_chat_message( + Message("tool", [Content.from_function_result(call_id="lookup", result="keep")], message_id="result") + ) + policy = stored("policy", "keep instructions", "system") + if protected_by == "system": + policy.extension_data = {"_group": {"id": "saved-policy"}} + call.extension_data = {"_group": {"id": "saved-policy"}} + result_owner = "active" if protected_by == "current" else "newest" if protected_by == "newest" else "old" + provider.state.data.conversation_history.extend([ + DurableAgentStateRequest("policy", OLD, [policy]), + DurableAgentStateResponse("old", OLD, [call]), + DurableAgentStateRequest("gap", OLD, [stored("gap", "drop")]), + DurableAgentStateResponse(result_owner, OLD, [result]), + DurableAgentStateRequest("newest", OLD, [stored("newest-user", "keep")]), + DurableAgentStateResponse("newest", OLD, [stored("newest-answer", "keep", "assistant")]), + ]) + history = DurableHistoryProvider(skip_excluded=False, prune_excluded=True) + state: dict[str, Any] = {} + with bound(provider, "active" if protected_by == "current" else None): + await history.get_messages("session", state=state) + for message in state[WORKING_BUFFER_KEY]: + message.additional_properties["_excluded"] = True + history.flush(state) + assert ids(provider) == ["policy", "call", "result", "newest-user", "newest-answer"] + assert (provider.state.data.truncation or {})["evictedMessageCount"] == 1 + if protected_by == "system": + assert (call.extension_data or {})["_group"] == {"id": "saved-policy"} + assert (policy.extension_data or {})["_group"] == {"id": "saved-policy"} + assert_current_positions(provider, state) + snapshot = deepcopy(provider.state.to_dict()) + history.flush(state) + assert provider.state.to_dict() == snapshot + assert provider.writes == 0 + + cold = _InMemoryStateProvider(raw=json.loads(provider.state.to_json())) + with bound(cold, "active" if protected_by == "current" else None): + cold_state: dict[str, Any] = {} + loaded = await history.get_messages("session", state=cold_state) + assert [message.message_id for message in loaded] == ids(provider) + history.flush(cold_state) + assert cold.state.to_dict() == provider.state.to_dict() + assert_current_positions(cold, cold_state) + assert cold.writes == 0 + + +async def test_legacy_repeated_and_missing_ids_survive_multiple_cold_loads() -> None: + provider = _InMemoryStateProvider() + provider.state.data.conversation_history.extend([ + DurableAgentStateRequest("same", OLD, [stored(None, "first")]), + DurableAgentStateRequest("same", OLD, [stored(None, "second")]), + DurableAgentStateResponse("same", OLD, [stored("reused", "version one", "assistant")]), + DurableAgentStateResponse("same", OLD, [stored("reused", "version two", "assistant")]), + ]) + raw = json.loads(provider.state.to_json()) + history = DurableHistoryProvider() + snapshots: list[dict[str, Any]] = [] + for _ in range(2): + cold = _InMemoryStateProvider(raw=raw) + with bound(cold, "same"): + loaded = await history.get_messages("session") + assert [message.text for message in loaded] == ["first", "second", "version one", "version two"] + assert all(ids(cold)) and len(ids(cold)) == len(set(ids(cold))) == 4 + snapshots.append(cold.state.to_dict()) + assert snapshots[0] == snapshots[1] + assert DurableAgentState.from_json(json.dumps(snapshots[0])).to_dict() == snapshots[0] + + +async def test_anonymous_summary_can_be_inserted_into_an_empty_history_once() -> None: + provider = _InMemoryStateProvider() + history = DurableHistoryProvider(prune_excluded=False) + state: dict[str, Any] = {WORKING_BUFFER_KEY: [Message("assistant", ["summary"])], POSITIONS_KEY: {}} + with bound(provider): + history.flush(state) + assert len(provider.state.data.conversation_history) == 1 + assert isinstance(provider.state.data.conversation_history[0], DurableAgentStateCompaction) + assert ids(provider) == ["durable_compaction_current_0_0"] + history.flush(state) + assert len(provider.state.data.conversation_history) == 1 + assert_current_positions(provider, state) + assert provider.writes == 0 + + +async def test_newest_exchange_is_protected_before_current_inputs_are_appended() -> None: + provider = _InMemoryStateProvider() + seed(provider) + history = DurableHistoryProvider(prune_excluded=True) + state: dict[str, Any] = {} + with bound(provider, "not-yet-appended"): + await history.get_messages("session", state=state) + for message in state[WORKING_BUFFER_KEY]: + message.additional_properties["_excluded"] = True + history.flush(state) + assert ids(provider) == ["seed-user", "seed-assistant"] + assert provider.state.data.truncation is None + + +async def test_direct_save_initializes_the_complete_working_buffer() -> None: + provider = _InMemoryStateProvider() + seed(provider) + history = DurableHistoryProvider() + state: dict[str, Any] = {} + with bound(provider): + await history.save_messages("session", [Message("user", ["next"])], state=state) + assert [message.text for message in state[WORKING_BUFFER_KEY]] == ["seed question", "seed answer", "next"] + assert_current_positions(provider, state) + + +@pytest.mark.parametrize("provider_type", [InMemoryHistoryProvider, DurableHistoryProvider]) +@pytest.mark.parametrize("store_inputs", [False, True]) +@pytest.mark.parametrize("store_outputs", [False, True]) +async def test_entity_does_not_bypass_provider_store_choices( + provider_type: type[InMemoryHistoryProvider] | type[DurableHistoryProvider], + store_inputs: bool, + store_outputs: bool, +) -> None: + original = provider_type( + store_inputs=store_inputs, + store_outputs=store_outputs, + store_context_messages=True, + store_context_from={"selected"}, + ) + provider = _InMemoryStateProvider() + client: Any = RecordingChatClient() + entity = AgentEntity( + Agent(client=client, context_providers=[original, AddContext("selected"), AddContext("other")]), + state_provider=provider, + ) + response = await entity.run({"message": "input", "correlationId": "choices"}) + assert [message.text for message in transcript(provider)] == [ + "context-selected", + *(["input"] if store_inputs else []), + *(["reply-1"] if store_outputs else []), + ] + assert provider.state.data.response_mailbox["choices"]["response"] == response.to_dict() + assert len(client.received_messages) == provider.writes == 1 + + +async def test_entity_failure_before_history_hook_is_mailbox_only() -> None: + provider = _InMemoryStateProvider() + entity = AgentEntity( + Agent(client=ToolChatClient(fail=True), context_providers=[DurableHistoryProvider()]), + state_provider=provider, + ) + response = await entity.run({"message": "not saved", "correlationId": "failed"}) + assert provider.state.data.conversation_history == [] + assert provider.state.data.response_mailbox["failed"]["response"] == response.to_dict() + assert any(content.type == "error" for message in response.messages for content in message.contents) + assert provider.writes == 1 diff --git a/python/packages/durabletask/tests/test_hosting_review_dt.py b/python/packages/durabletask/tests/test_hosting_review_dt.py new file mode 100644 index 0000000..6781b76 --- /dev/null +++ b/python/packages/durabletask/tests/test_hosting_review_dt.py @@ -0,0 +1,324 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Derived-name ownership and registration failure boundaries for the worker host.""" + +from collections.abc import Callable +from dataclasses import fields +from typing import Any +from unittest.mock import Mock + +import pytest +from agent_framework import Agent, AgentExecutor, Executor, InMemoryHistoryProvider, WorkflowExecutor +from durabletask.azuremanaged.worker import DurableTaskSchedulerWorker +from durabletask.worker import TaskHubGrpcWorker, _Registry + +from agent_framework_durabletask import DTS_MAX_STATE_BYTES, DurableAIAgentWorker, DurableHistoryProvider +from agent_framework_durabletask._configuration import AgentRegistrationSettings, validate_agent_configuration + + +class RecordingWorker: + def __init__(self) -> None: + self.calls: list[tuple[str, str]] = [] + self.entities: dict[str, Any] = {} + self.fail_at: int | None = None + + def _record(self, kind: str, name: str) -> str: + self.calls.append((kind, name)) + if len(self.calls) == self.fail_at: + raise RuntimeError("injected native registration failure") + return name + + def add_entity(self, entity: Any) -> str: + self.entities[entity.__name__] = entity + return self._record("entity", entity.__name__) + + def add_activity(self, activity: Callable[..., Any]) -> str: + return self._record("activity", activity.__name__) + + def add_orchestrator(self, orchestrator: Callable[..., Any]) -> str: + return self._record("orchestration", orchestrator.__name__) + + def start(self) -> None: + self._record("start", "worker") + + +def _worker(native: Any, **kwargs: Any) -> DurableAIAgentWorker: + return DurableAIAgentWorker(native, **kwargs) + + +def _agent(name: str = "assistant") -> Agent: + client: Any = Mock(additional_properties={}, STORES_BY_DEFAULT=False) + return Agent(client=client, name=name, context_providers=[InMemoryHistoryProvider("primary")]) + + +def _workflow(name: str, executor_id: str = "node", *, agent: Any = None, children: tuple[Any, ...] = ()) -> Any: + executor = Mock(spec=Executor if agent is None else AgentExecutor) + executor.id = executor_id + if agent is not None: + executor.agent = agent + executors = {executor_id: executor} + for index, child in enumerate(children): + nested = Mock(spec=WorkflowExecutor) + nested.id = f"child{index}" + nested.workflow = child + executors[nested.id] = nested + workflow = Mock() + workflow.name = name + workflow.executors = executors + return workflow + + +@pytest.mark.parametrize("kinds", [(False, False), (True, True)]) +@pytest.mark.parametrize("nested", [False, True]) +def test_ambiguous_concatenations_fail_before_backend_or_metadata_changes( + kinds: tuple[bool, bool], nested: bool +) -> None: + native = RecordingWorker() + host = _worker(native) + left = _workflow("alpha-beta", "gamma", agent=_agent("left") if kinds[0] else None) + right = _workflow("alpha", "beta-gamma", agent=_agent("right") if kinds[1] else None) + if nested: + candidate = _workflow("root", children=(left, right)) + else: + host.configure_workflow(left) + candidate = right + calls = list(native.calls) + agents, workflows = host.registered_agent_names, host.registered_workflow_names + with pytest.raises(ValueError, match="Derived name.*collides"): + host.configure_workflow(candidate) + assert native.calls == calls + assert host.registered_agent_names == agents + assert host.registered_workflow_names == workflows + host.configure_workflow(_workflow("corrected")) + + +@pytest.mark.parametrize("standalone_first", [False, True]) +@pytest.mark.parametrize("same_agent", [False, True]) +def test_standalone_and_workflow_owners_cannot_share_an_entity(standalone_first: bool, same_agent: bool) -> None: + native = RecordingWorker() + host = _worker(native) + standalone = _agent("alpha-beta-node") + workflow = _workflow("alpha-beta", agent=standalone if same_agent else _agent()) + if standalone_first: + host.add_agent(standalone) + else: + host.configure_workflow(workflow) + calls = list(native.calls) + with pytest.raises(ValueError): + if standalone_first: + host.configure_workflow(workflow) + else: + host.add_agent(standalone) + assert native.calls == calls + + +@pytest.mark.parametrize("reverse", [False, True]) +@pytest.mark.parametrize("nested", [False, True]) +def test_native_registry_allows_the_same_name_for_different_artifact_kinds(reverse: bool, nested: bool) -> None: + registry = _Registry() + native = Mock(spec=TaskHubGrpcWorker) + native.add_entity.side_effect = registry.add_entity + native.add_activity.side_effect = registry.add_activity + native.add_orchestrator.side_effect = registry.add_orchestrator + host = _worker(native) + workflows = [ + _workflow("alpha-beta", "gamma", agent=_agent()), + _workflow("alpha", "beta-gamma"), + _workflow("alpha-beta-gamma"), + ] + if reverse: + workflows.reverse() + if nested: + host.configure_workflow(_workflow("root", children=tuple(workflows))) + else: + for workflow in workflows: + host.configure_workflow(workflow) + + name = "dafx-alpha-beta-gamma" + assert name in registry.entities + assert name in registry.activities + assert name in registry.orchestrators + assert {namespace for namespace, registered_name in host._registration_identities if registered_name == name} == { + "entity-name", + "activity-name", + "orchestrator-name", + } + assert set(host.registered_workflow_names) == ({"root"} if nested else {workflow.name for workflow in workflows}) + + +def test_case_folded_agent_identity_is_checked_before_native_registration() -> None: + native = RecordingWorker() + host = _worker(native) + host.add_agent(_agent("Assistant")) + with pytest.raises(ValueError, match="case-insensitively"): + host.add_agent(_agent("assistant")) + assert native.calls == [("entity", "dafx-Assistant")] + + +_CHANGED_SETTINGS: dict[str, Any] = { + "retention": "follow_compaction", + "max_state_bytes": 8192, + "high_watermark": 0.99, + "low_watermark": 0.1, + "response_delivery_window_seconds": 17, + "callback": Mock(), +} + + +def test_shared_configuration_cases_cover_every_setting_field() -> None: + assert set(_CHANGED_SETTINGS) == {field.name for field in fields(AgentRegistrationSettings)} + + +@pytest.mark.parametrize("setting", _CHANGED_SETTINGS) +def test_shared_workflow_reuse_requires_identical_resolved_settings(setting: str) -> None: + native = RecordingWorker() + host = _worker(native) + child = _workflow("shared", agent=_agent()) + host.configure_workflow(_workflow("first", children=(child,))) + calls = list(native.calls) + with pytest.raises(ValueError, match="different settings"): + host.configure_workflow(_workflow("second", children=(child,)), **{setting: _CHANGED_SETTINGS[setting]}) + assert native.calls == calls + assert host.registered_workflow_names == ["first"] + host.configure_workflow(_workflow("second", children=(child,))) + assert native.calls.count(("entity", "dafx-shared-node")) == 1 + + +def test_repeated_identical_workflow_is_benign_and_names_are_unchanged() -> None: + native = RecordingWorker() + host = _worker(native) + workflow = _workflow("Orders", "review", agent=_agent()) + host.configure_workflow(workflow) + host.configure_workflow(workflow) + assert native.calls == [("entity", "dafx-Orders-review"), ("orchestration", "dafx-Orders")] + + +class UncopyableAgent: + name = "uncopyable" + context_providers = [InMemoryHistoryProvider("history")] + + def __copy__(self) -> Any: + raise TypeError("cannot copy") + + +class ReadOnlyProviders: + name = "readonly" + + @property + def context_providers(self) -> list[Any]: + return [InMemoryHistoryProvider("history")] + + +@pytest.mark.parametrize("factory", [UncopyableAgent, ReadOnlyProviders]) +@pytest.mark.parametrize("surface", ["agent", "nested"]) +def test_actual_adapter_preparation_fails_during_registration(factory: Any, surface: str) -> None: + native = RecordingWorker() + host = _worker(native) + agent = factory() + with pytest.raises(ValueError, match="attach durable history"): + if surface == "agent": + host.add_agent(agent) + else: + host.configure_workflow(_workflow("root", agent=_agent(), children=(_workflow("child", agent=agent),))) + assert native.calls == [] + assert host.registered_agent_names == [] + assert host.registered_workflow_names == [] + host.add_agent(_agent("valid")) + + +def test_dry_preparation_preserves_original_agent_and_provider_configuration() -> None: + native = RecordingWorker() + host = _worker(native, retention="follow_compaction") + agent = _agent() + providers = agent.context_providers + provider = providers[0] + validate_agent_configuration(agent, retention="follow_compaction") + host.add_agent(agent) + assert host._registered_agents["assistant"] is agent + assert agent.context_providers is providers and providers[0] is provider + assert isinstance(provider, InMemoryHistoryProvider) + + +def test_uncopyable_unresolved_durable_provider_fails_at_registration() -> None: + class UncopyableHistory(DurableHistoryProvider): + def __copy__(self) -> Any: + raise TypeError("provider cannot copy") + + native = RecordingWorker() + agent = _agent() + provider = UncopyableHistory() + agent.context_providers = [provider] + with pytest.raises(ValueError, match="prepare.*durable history"): + _worker(native).add_agent(agent, retention="follow_compaction") + assert native.calls == [] + assert agent.context_providers == [provider] and provider.prune_excluded is None + + +def test_standalone_native_failure_also_blocks_reuse_and_start() -> None: + native = RecordingWorker() + native.fail_at = 1 + host = _worker(native) + with pytest.raises(RuntimeError, match="injected native"): + host.add_agent(_agent()) + assert host.registered_agent_names == [] + for action in (lambda: host.add_agent(_agent()), host.start): + with pytest.raises(RuntimeError, match="partially registered"): + action() + assert len(native.calls) == 1 + + +@pytest.mark.parametrize("fail_at", [1, 2, 3]) +def test_partial_native_failure_blocks_start_and_retry_without_false_metadata(fail_at: int) -> None: + native = RecordingWorker() + host = _worker(native) + host.add_agent(_agent("existing")) + previous_calls = len(native.calls) + native.fail_at = previous_calls + fail_at + workflow = _workflow("root", agent=_agent(), children=(_workflow("child"),)) + with pytest.raises(RuntimeError, match="injected native"): + host.configure_workflow(workflow) + assert host.registered_agent_names == ["existing"] + assert host.registered_workflow_names == [] + assert host._registered_orchestrations == {} + calls = list(native.calls) + for action in (lambda: host.add_agent(_agent("retry")), lambda: host.configure_workflow(workflow), host.start): + with pytest.raises(RuntimeError, match="partially registered"): + action() + assert native.calls == calls + + +@pytest.mark.parametrize("surface", ["host", "agent", "workflow"]) +def test_generic_grpc_worker_does_not_imply_a_dts_budget(surface: str) -> None: + native = Mock(spec=TaskHubGrpcWorker) + with pytest.raises(ValueError, match="known backend_limit"): + if surface == "host": + _worker(native, max_state_bytes="backend_limit") + else: + host = _worker(native) + if surface == "agent": + host.add_agent(_agent(), max_state_bytes="backend_limit") + else: + host.configure_workflow(_workflow("flow", agent=_agent()), max_state_bytes="backend_limit") + assert native.mock_calls == [] + + +@pytest.mark.parametrize("surface", ["host", "agent", "workflow"]) +def test_dts_worker_resolves_backend_budget_on_every_surface(surface: str) -> None: + native = Mock(spec=DurableTaskSchedulerWorker) + host = _worker(native, **({"max_state_bytes": "backend_limit"} if surface == "host" else {})) + if surface == "workflow": + host.configure_workflow(_workflow("flow", agent=_agent()), max_state_bytes="backend_limit") + elif surface == "agent": + host.add_agent(_agent(), max_state_bytes="backend_limit") + else: + host.add_agent(_agent()) + entity = native.add_entity.call_args.args[0]() + assert entity._agent_entity._max_state_bytes == DTS_MAX_STATE_BYTES + + +def test_explicit_budget_works_with_an_unknown_backend() -> None: + native = RecordingWorker() + host = _worker(native, max_state_bytes=8192) + host.add_agent(_agent()) + entity = native.entities["dafx-assistant"]() + assert entity._agent_entity._max_state_bytes == 8192 diff --git a/python/packages/durabletask/tests/test_integration_environment.py b/python/packages/durabletask/tests/test_integration_environment.py new file mode 100644 index 0000000..1a4b467 --- /dev/null +++ b/python/packages/durabletask/tests/test_integration_environment.py @@ -0,0 +1,111 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Regression tests for the integration worker's subprocess environment.""" + +import os +import subprocess +import sys +import uuid +from importlib.util import module_from_spec, spec_from_file_location +from pathlib import Path +from types import ModuleType, SimpleNamespace +from typing import Any +from unittest.mock import Mock + +import pytest + + +@pytest.fixture +def _dt_harness(monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest) -> ModuleType: + path = Path(__file__).parent / "integration_tests" / "conftest.py" + spec = spec_from_file_location(f"_dt_integration_environment_{uuid.uuid4().hex}", path) + assert spec is not None and spec.loader is not None + module = module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + # Import as an ordinary module, without loading local secrets or registering pytest hooks. + with monkeypatch.context() as import_patch: + import_patch.setattr("dotenv.load_dotenv", Mock(return_value=False)) + import_patch.setattr("logging.basicConfig", Mock()) + spec.loader.exec_module(module) + assert not request.config.pluginmanager.is_registered(module) + return module + + +@pytest.mark.parametrize("parent_mode", [None, "legacy"], ids=["missing-mode", "invalid-mode"]) +@pytest.mark.parametrize("platform", ["win32", "linux"], ids=["windows", "unix"]) +def test_worker_subprocess_opts_into_isolated_mode( + _dt_harness: ModuleType, + monkeypatch: pytest.MonkeyPatch, + parent_mode: str | None, + platform: str, +) -> None: + harness = _dt_harness + # Set the case after importing the harness so neither dotenv nor the root fixture can mask it. + if parent_mode is None: + monkeypatch.delenv("DURABLE_AGENTS_DEPLOYMENT_MODE", raising=False) + else: + monkeypatch.setenv("DURABLE_AGENTS_DEPLOYMENT_MODE", parent_mode) + monkeypatch.setenv("TASKHUB", "parent-hub") + monkeypatch.setenv("ENDPOINT", "http://parent.invalid:8080") + parent_env = dict(os.environ) + + process = Mock(spec=subprocess.Popen) + process.poll.return_value = None + process.wait.return_value = 0 + + def start_worker(*_args: object, **kwargs: Any) -> Mock: + assert dict(os.environ) == parent_env + assert kwargs["env"].get("DURABLE_AGENTS_DEPLOYMENT_MODE") == "isolated_v2" + return process + + popen = Mock(side_effect=start_worker) + monkeypatch.setattr(harness, "sys", SimpleNamespace(platform=platform, executable=sys.executable)) + monkeypatch.setattr( + harness, + "subprocess", + SimpleNamespace(Popen=popen, CREATE_NEW_PROCESS_GROUP=512, TimeoutExpired=subprocess.TimeoutExpired), + ) + monkeypatch.setattr(harness, "time", SimpleNamespace(sleep=Mock())) + for probe in ("_check_dts_available", "_check_redis_available"): + monkeypatch.setattr(harness, probe, Mock(side_effect=AssertionError("Unexpected infrastructure probe"))) + + sample_name = "12_subworkflow_hitl" + assert harness.__file__ is not None + sample_path = Path(harness.__file__).parents[4] / "samples" / sample_name + request_stub = Mock(spec=pytest.FixtureRequest) + request_stub.node.get_closest_marker.return_value = SimpleNamespace(args=(sample_name,)) + taskhub = harness.unique_taskhub.__wrapped__() + assert taskhub.startswith("test-") and taskhub != parent_env["TASKHUB"] + endpoint = "http://localhost:8080" + lifecycle = harness.worker_process.__wrapped__( + dts_available=True, + check_sample_env=None, + dts_endpoint=endpoint, + unique_taskhub=taskhub, + request=request_stub, + ) + try: + worker_info = next(lifecycle) + expected_options: dict[str, Any] = { + "cwd": str(sample_path), + "env": { + **parent_env, + "ENDPOINT": endpoint, + "TASKHUB": taskhub, + "DURABLE_AGENTS_DEPLOYMENT_MODE": "isolated_v2", + }, + "text": True, + } + if platform == "win32": + expected_options.update(creationflags=512, shell=True) + popen.assert_called_once_with([sys.executable, str(sample_path / "worker.py")], **expected_options) + assert popen.call_args.kwargs["env"] is not os.environ + assert worker_info == {"process": process, "endpoint": endpoint, "taskhub": taskhub} + request_stub.node.get_closest_marker.assert_called_once_with("sample") + assert dict(os.environ) == parent_env + finally: + lifecycle.close() + + process.terminate.assert_called_once_with() + process.wait.assert_called_once_with(timeout=5) + assert dict(os.environ) == parent_env diff --git a/python/packages/durabletask/tests/test_maintenance_review.py b/python/packages/durabletask/tests/test_maintenance_review.py new file mode 100644 index 0000000..bc93d78 --- /dev/null +++ b/python/packages/durabletask/tests/test_maintenance_review.py @@ -0,0 +1,780 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Maintenance contracts through real entities and registered Durable Task methods.""" + +import hashlib +import json +from copy import deepcopy +from datetime import datetime, timedelta, timezone +from typing import Any +from unittest.mock import Mock + +import pytest +from agent_framework import Agent, AgentResponse, Content, ContextProvider, HistoryProvider, Message +from durabletask.entities import EntityInstanceId +from test_durable_history_provider import RecordingChatClient +from test_revision_contract import JsonStateProvider +from typing_extensions import Self + +from agent_framework_durabletask import AgentEntity, DurableAgentState, DurableAIAgentWorker, serialize_agent_response +from agent_framework_durabletask import _durable_agent_state as state_module +from agent_framework_durabletask import _entities as entities_module +from agent_framework_durabletask import _retention as retention_module +from agent_framework_durabletask import _state_migration as migration_module +from agent_framework_durabletask._message_identity import message_identity + +NOW = datetime(2040, 1, 1, 12, tzinfo=timezone.utc) +WINDOW = 60 +SOURCE_ID = "@dafx-maintenance@legacy-source" +DESTINATION_ID = "@dafx-maintenance@destination" + + +class _ClockType(type): + def __instancecheck__(cls, instance: Any) -> bool: + # Parsed timestamps remain real datetime objects, not instances of the test subclass. + return isinstance(instance, datetime) + + +class Clock(datetime, metaclass=_ClockType): + current = NOW + + @classmethod + def now(cls, tz: Any = None) -> Self: + return cls.fromtimestamp(cls.current.timestamp(), tz=tz) + + +@pytest.fixture +def clock(monkeypatch: pytest.MonkeyPatch) -> type[Clock]: + monkeypatch.setattr(Clock, "current", NOW) + for module in (state_module, entities_module, retention_module, migration_module): + monkeypatch.setattr(module, "datetime", Clock) + return Clock + + +class Store(JsonStateProvider): + def __init__(self, raw: dict[str, Any] | None = None) -> None: + super().__init__(raw) + self.attempts = 0 + + def _set_state_dict(self, state: dict[str, Any]) -> None: + self.attempts += 1 + super()._set_state_dict(state) + + def _get_session_id_from_entity(self) -> str: + return "destination" + + def _get_entity_name_from_entity(self) -> str: + return "dafx-maintenance" + + +class Hooks(ContextProvider): + def __init__(self) -> None: + super().__init__("maintenance-probe") + self.calls: list[str] = [] + + async def before_run(self, **kwargs: Any) -> None: + self.calls.append("before") + + async def after_run(self, **kwargs: Any) -> None: + self.calls.append("after") + + +class ExternalHistory(HistoryProvider): + def __init__(self) -> None: + super().__init__("external") + self.calls: list[tuple[str, str | None]] = [] + self.rows: dict[str | None, list[Message]] = { + SOURCE_ID: [Message("user", ["already in the external store"], message_id="external-old")] + } + + async def get_messages(self, session_id: str | None, **kwargs: Any) -> list[Message]: + self.calls.append(("get", session_id)) + return deepcopy(self.rows.get(session_id, [])) + + async def save_messages(self, session_id: str | None, messages: Any, **kwargs: Any) -> None: + self.calls.append(("save", session_id)) + self.rows.setdefault(session_id, []).extend(deepcopy(list(messages))) + + +def _agent() -> tuple[Agent, RecordingChatClient, Hooks, Mock]: + client: Any = RecordingChatClient() + hooks = Hooks() + callback = Mock(spec=["on_streaming_response_update", "on_agent_response"]) + return Agent(client=client, name="maintenance", context_providers=[hooks]), client, hooks, callback + + +def _quiet(client: RecordingChatClient, hooks: Hooks, callback: Mock) -> None: + assert client.received_messages == [] + assert hooks.calls == [] + assert callback.mock_calls == [] + + +def _digest(raw: dict[str, Any]) -> str: + encoded = json.dumps(raw, sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +def _size(raw: dict[str, Any]) -> int: + return len(json.dumps(raw, allow_nan=False)) + + +def _source() -> dict[str, Any]: + # Z is intentional: digest the export, not its timestamp-normalized reader projection. + return { + "schemaVersion": "1.1.0", + "futureRoot": {"keep": ["雪", None]}, + "data": { + "conversationHistory": [ + { + "$type": kind, + "correlationId": "legacy-done", + "createdAt": "2024-01-01T00:00:00Z", + "messages": [ + { + "role": role, + "messageId": identity, + "contents": [{"$type": "text", "text": text}], + } + ], + } + for kind, role, identity, text in ( + ("request", "user", "legacy-input", "retained legacy input"), + ("response", "assistant", "legacy-answer", "retained legacy answer"), + ) + ], + "session": {"session_id": SOURCE_ID, "state": {"opaque": {"keep": [1, 3]}}}, + "futureData": {"keep": [False, 0]}, + }, + } + + +def _request(source: dict[str, Any] | None = None, *, evidence: bool = False) -> dict[str, Any]: + source = _source() if source is None else source + messages = [Message("user", [f"accepted {position}"], message_id=f"wf_upstream_{position}") for position in (1, 3)] + if evidence: + source["data"]["ingestedPositions"] = {"upstream": 3} + source["data"]["conversationHistory"][0]["messages"][0]["messageId"] = "wf_upstream_3" + digest = _digest(source) + request: dict[str, Any] = { + "source": source, + "sourceDigest": digest, + "sourceSessionId": SOURCE_ID, + "destinationSessionId": DESTINATION_ID, + "migrationId": "migration-1", + "ownershipTransferId": "operator-transfer-1", + } + if evidence: + request["deliveryEvidence"] = { + "sourceDigest": digest, + "evidenceId": "operator-journal-1", + "complete": True, + "messages": [message.to_dict() for message in messages], + } + return request + + +def _mailboxes() -> dict[str, Any]: + raw = _source() + raw["schemaVersion"] = "2.0.0" + state = DurableAgentState.from_dict(raw) + state.data.ingested_messages = {"old-input": ["a" * 64]} + state.data.completed_correlations["long-gone"] = {"completedAt": "2020-01-01T00:00:00+00:00"} + for correlation, error in (("expired-success", False), ("expired-error", True), ("live", False)): + response = AgentResponse[Any]( + messages=[ + Message( + "assistant", + [Content.from_error(message="original failure", error_code="OriginalError")] + if error + else [Content.from_text("original result", additional_properties={"tags": ["雪"]})], + message_id=f"answer-{correlation}", + ) + ], + response_id=f"response-{correlation}", + additional_properties={"durable_status": "error" if error else "success", "nested": {"keep": [1]}}, + value={"original": [1, 2]} if not error else None, + ) + state.record_response( + correlation, + response, + delivery_window_seconds=WINDOW, + now=NOW if correlation == "live" else NOW - timedelta(minutes=2), + ) + assert state.data.response_mailbox[correlation]["response"] == serialize_agent_response(response) + return json.loads(state.to_json()) + + +def _without_expired(raw: dict[str, Any]) -> dict[str, Any]: + expected = deepcopy(raw) + for correlation in ("expired-success", "expired-error"): + del expected["data"]["responseMailbox"][correlation] + return expected + + +@pytest.mark.parametrize("operation", ["new-run", "duplicate-run", "reset", "expire_responses"]) +async def test_legacy_entity_is_readable_but_every_writer_fails_before_execution(operation: str) -> None: + raw = _source() + store = Store(raw) + agent, client, hooks, callback = _agent() + entity = AgentEntity(agent, state_provider=store, callback=callback) + cached = entity.state + before = cached.to_dict() + legacy = cached.try_get_agent_response("legacy-done") + assert legacy is not None and legacy.text == "retained legacy answer" + + with pytest.raises(ValueError, match="[Ll]egacy.*read-only"): + if operation.endswith("run"): + correlation = "legacy-done" if operation == "duplicate-run" else "new" + await entity.run({"message": "must not execute", "correlationId": correlation}) + else: + getattr(entity, operation)() + + assert entity.state is cached and entity.state.to_dict() == before + assert store.raw == raw and store.attempts == store.writes == 0 + _quiet(client, hooks, callback) + + +@pytest.mark.parametrize("evidence", [False, True]) +def test_migrate_empty_destination_binds_raw_export_and_optional_sparse_evidence( + clock: type[Clock], evidence: bool +) -> None: + request = _request(evidence=evidence) + before = deepcopy(request) + assert request["sourceDigest"] != _digest(DurableAgentState.from_dict(request["source"]).to_dict()) + store = Store() + agent, client, hooks, callback = _agent() + entity = AgentEntity(agent, callback=callback, state_provider=store, response_delivery_window_seconds=WINDOW) + + result = entity.migrate(request) + + assert result == {"status": "migrated", "migrationId": "migration-1", "sessionId": store.core_session_id} + assert store.core_session_id == DESTINATION_ID != SOURCE_ID + assert store.writes == store.attempts == 1 + cold = Store(store.raw).state + assert cold.schema_version == "2.0.0" + metadata = cold.data.unknown_fields["migration"] + assert metadata == { + "id": "migration-1", + "sourceDigest": request["sourceDigest"], + "sourceSessionId": SOURCE_ID, + "destinationSessionId": DESTINATION_ID, + "ownershipTransferId": "operator-transfer-1", + "requestDigest": _digest(request), + "createdAt": NOW.isoformat(), + **({"evidenceId": "operator-journal-1"} if evidence else {}), + } + assert cold.data.session == request["source"]["data"]["session"] + expected_history = DurableAgentState.from_dict(request["source"]).to_dict()["data"]["conversationHistory"] + assert cold.to_dict()["data"]["conversationHistory"] == expected_history + assert cold.data.response_mailbox["legacy-done"]["expiresAt"] == (NOW + timedelta(seconds=WINDOW)).isoformat() + assert cold.data.completed_correlations["legacy-done"]["legacy"] is True + assert cold.to_dict()["futureRoot"] == request["source"]["futureRoot"] + assert cold.data.unknown_fields["futureData"] == request["source"]["data"]["futureData"] + if evidence: + expected = { + message["message_id"]: [message_identity(Message.from_dict(deepcopy(message)))] + for message in request["deliveryEvidence"]["messages"] + } + assert cold.data.ingested_messages == expected + assert "wf_upstream_2" not in cold.data.ingested_messages + assert request == before + _quiet(client, hooks, callback) + + +@pytest.mark.parametrize( + "field", + ["sourceDigest", "sourceSessionId", "destinationSessionId", "migrationId", "ownershipTransferId"], +) +def test_migration_requires_nonblank_explicit_identifiers(field: str) -> None: + store = Store() + agent, client, hooks, callback = _agent() + entity = AgentEntity(agent, state_provider=store, callback=callback) + request = _request() + request[field] = " \t" + with pytest.raises(ValueError, match="nonblank"): + entity.migrate(request) + assert store.raw == {} and store.attempts == 0 + _quiet(client, hooks, callback) + + +@pytest.mark.parametrize("invalid", ["raw-digest", "wrong-destination", "same-source", "missing-id", "non-json"]) +def test_migration_rejects_invalid_export_or_address_without_mutating_destination(invalid: str) -> None: + store = Store() + agent, client, hooks, callback = _agent() + entity = AgentEntity(agent, state_provider=store, callback=callback) + request = _request() + if invalid == "raw-digest": + request["sourceDigest"] = _digest(DurableAgentState.from_dict(request["source"]).to_dict()) + elif invalid == "wrong-destination": + request["destinationSessionId"] = "@dafx-another-agent@destination" + elif invalid == "same-source": + request["sourceSessionId"] = DESTINATION_ID + elif invalid == "missing-id": + del request["ownershipTransferId"] + else: + request["source"]["futureRoot"] = {"bad": (1, 2)} + before = deepcopy(request) + cached = entity.state + with pytest.raises(ValueError): + entity.migrate(request) + assert entity.state is cached and store.raw == {} and store.attempts == 0 + assert request == before + _quiet(client, hooks, callback) + + +async def test_cold_exact_retry_after_a_v2_run_never_writes_or_refreshes_grace(clock: type[Clock]) -> None: + request = _request() + before_source = deepcopy(request) + store = Store() + agent, client, hooks, callback = _agent() + first = AgentEntity(agent, state_provider=store, callback=callback) + result = first.migrate(request) + expiry = store.raw["data"]["responseMailbox"]["legacy-done"]["expiresAt"] + clock.current = NOW + timedelta(seconds=10) + response = await first.run({"message": "v2 turn", "correlationId": "v2-done"}) + assert response.text == "reply-1" and len(client.received_messages) == 1 + assert hooks.calls == ["before", "after"] + before = deepcopy(store.raw) + calls = list(callback.mock_calls) + + clock.current = NOW + timedelta(days=1) + cold_store = Store(before) + cold = AgentEntity(agent, state_provider=cold_store, callback=callback) + assert cold.migrate(deepcopy(request)) == result + assert cold.migrate(deepcopy(request)) == result + assert cold_store.raw == before and cold.state.to_dict() == before + assert cold_store.writes == cold_store.attempts == 0 + assert cold.state.data.response_mailbox["legacy-done"]["expiresAt"] == expiry + expired = cold.state.try_get_agent_response("legacy-done") + assert expired is not None and expired.additional_properties["durable_status"] == "already_completed" + assert len(client.received_messages) == 1 and hooks.calls == ["before", "after"] + assert callback.mock_calls == calls and request == before_source + + +@pytest.mark.parametrize( + "change", ["migrationId", "source", "sourceSessionId", "ownershipTransferId", "deliveryEvidence"] +) +def test_existing_migration_rejects_reused_identity_with_any_changed_request(clock: type[Clock], change: str) -> None: + request = _request() + store = Store() + agent, client, hooks, callback = _agent() + entity = AgentEntity(agent, state_provider=store, callback=callback) + entity.migrate(request) + changed = deepcopy(request) + if change == "source": + changed["source"]["futureRoot"]["keep"].append("different source") + changed["sourceDigest"] = _digest(changed["source"]) + elif change == "sourceSessionId": + changed[change] = "@dafx-maintenance@other-source" + changed["source"]["data"]["session"]["session_id"] = changed[change] + changed["sourceDigest"] = _digest(changed["source"]) + elif change == "deliveryEvidence": + changed[change] = { + "sourceDigest": changed["sourceDigest"], + "evidenceId": "new-journal", + "complete": True, + "messages": [], + } + else: + changed[change] += "-different" + before = deepcopy(store.raw) + cold_store = Store(before) + cold = AgentEntity(agent, state_provider=cold_store, callback=callback) + with pytest.raises(ValueError, match="empty|different migration"): + cold.migrate(changed) + assert cold_store.raw == before and cold_store.attempts == 0 and cold.state.to_dict() == before + _quiet(client, hooks, callback) + + +def test_nonempty_destination_without_migration_is_never_overwritten(clock: type[Clock]) -> None: + store = Store(_mailboxes()) + before = deepcopy(store.raw) + agent, client, hooks, callback = _agent() + with pytest.raises(ValueError, match="empty"): + AgentEntity(agent, state_provider=store, callback=callback).migrate(_request()) + assert store.raw == before and store.attempts == 0 + _quiet(client, hooks, callback) + + +def test_failed_migration_commit_restores_warm_cache_and_retry_can_commit(clock: type[Clock]) -> None: + store = Store() + agent, client, hooks, callback = _agent() + entity = AgentEntity(agent, state_provider=store, callback=callback) + original = entity.state + request = _request(evidence=True) + before = deepcopy(request) + store.fail_writes = True + with pytest.raises(OSError, match="commit failure"): + entity.migrate(request) + assert entity.state is original and original.to_dict() == DurableAgentState().to_dict() + assert store.raw == {} and store.writes == 0 and store.attempts == 1 + store.fail_writes = False + assert entity.migrate(request)["status"] == "migrated" + assert store.writes == 1 and store.attempts == 2 and request == before + _quiet(client, hooks, callback) + + +def test_migration_budget_counts_full_destination_metadata_and_never_prunes(clock: type[Clock]) -> None: + request = _request() + request["source"]["data"]["conversationHistory"][0]["messages"][0]["contents"][0]["text"] = "雪" * 500 + request["sourceDigest"] = _digest(request["source"]) + before = deepcopy(request) + agent, client, hooks, callback = _agent() + sizing_store = Store() + AgentEntity(agent, state_provider=sizing_store).migrate(request) + full_size = _size(sizing_store.raw) + assert full_size > _size(request["source"]) + expected_history = DurableAgentState.from_dict(request["source"]).to_dict()["data"]["conversationHistory"] + assert sizing_store.raw["data"]["conversationHistory"] == expected_history + + rejected = Store() + entity = AgentEntity(agent, state_provider=rejected, callback=callback, max_state_bytes=full_size - 1) + cached = entity.state + with pytest.raises(ValueError, match="max_state_bytes|capacity|budget"): + entity.migrate(request) + assert entity.state is cached and rejected.raw == {} and rejected.attempts == 0 + + accepted = Store() + AgentEntity(agent, state_provider=accepted, callback=callback, max_state_bytes=full_size).migrate(request) + assert accepted.raw == sizing_store.raw and accepted.writes == 1 + assert "truncation" not in accepted.raw["data"] and request == before + _quiet(client, hooks, callback) + + +async def test_external_get_and_save_keep_source_logical_identity_after_cold_run_and_retry(clock: type[Clock]) -> None: + external = ExternalHistory() + client: Any = RecordingChatClient() + agent = Agent(client=client, name="maintenance", context_providers=[external]) + request = _request() + before_request = deepcopy(request) + external_before = {key: [message.to_dict() for message in messages] for key, messages in external.rows.items()} + store = Store() + AgentEntity(agent, state_provider=store).migrate(request) + # The operator authorizes transfer outside this API; migration must not copy provider history. + assert external.calls == [] + assert { + key: [message.to_dict() for message in messages] for key, messages in external.rows.items() + } == external_before + for index in range(2): + store = Store(store.raw) + entity = AgentEntity(agent, state_provider=store) + if index: + before = deepcopy(store.raw) + clock.current = NOW + timedelta(seconds=20) + assert entity.migrate(request)["status"] == "migrated" + assert store.raw == before and store.attempts == 0 + response = await entity.run({"message": f"destination turn {index}", "correlationId": f"new-{index}"}) + assert response.text == f"reply-{index + 1}" + assert store.raw["data"]["session"]["session_id"] == SOURCE_ID + assert external.calls == [(phase, SOURCE_ID) for _ in range(2) for phase in ("get", "save")] + assert set(external.rows) == {SOURCE_ID} and store.core_session_id == DESTINATION_ID + assert [message.text for message in external.rows[SOURCE_ID]] == [ + "already in the external store", + "destination turn 0", + "reply-1", + "destination turn 1", + "reply-2", + ] + assert "retained legacy input" not in [message.text for batch in client.received_messages for message in batch] + assert request == before_request + + +@pytest.mark.parametrize("correlation", ["expired-success", "expired-error"]) +async def test_expired_duplicate_run_removes_physical_payloads_without_model_or_hooks( + clock: type[Clock], correlation: str +) -> None: + raw = _mailboxes() + before = deepcopy(raw) + read_only = DurableAgentState.from_dict(raw) + lookup = read_only.try_get_agent_response(correlation) + assert lookup is not None and lookup.additional_properties["durable_status"] == "already_completed" + assert read_only.to_dict() == before + assert before["data"]["responseMailbox"][correlation]["response"]["response_id"] == f"response-{correlation}" + store = Store(raw) + agent, client, hooks, callback = _agent() + entity = AgentEntity(agent, state_provider=store, callback=callback) + + response = await entity.run({"message": "duplicate", "correlationId": correlation}) + + assert response.additional_properties == { + "durable_status": "already_completed", + "correlation_id": correlation, + "durable_outcome": "failed" if correlation == "expired-error" else "succeeded", + } + assert response.messages[0].contents[0].error_code == "response_expired" + assert store.raw == _without_expired(before) and store.writes == store.attempts == 1 + assert entity.state.to_dict() == store.raw and raw == before + _quiet(client, hooks, callback) + + +def test_idle_expiry_only_writes_for_removal_and_keeps_history_and_receipts_indefinitely(clock: type[Clock]) -> None: + raw = _mailboxes() + store = Store(raw) + agent, client, hooks, callback = _agent() + entity = AgentEntity(agent, state_provider=store, callback=callback) + assert entity.expire_responses() == 2 + assert store.raw == _without_expired(raw) + assert entity.expire_responses() == 0 and store.writes == store.attempts == 1 + live = entity.state.try_get_agent_response("live") + assert live is not None and serialize_agent_response(live) == raw["data"]["responseMailbox"]["live"]["response"] + + clock.current = NOW + timedelta(days=36500) + cold_store = Store(store.raw) + cold = AgentEntity(agent, state_provider=cold_store, callback=callback) + assert cold.expire_responses() == 1 + expected = _without_expired(raw) + del expected["data"]["responseMailbox"] + assert cold_store.raw == expected + assert cold.expire_responses() == 0 and cold_store.writes == cold_store.attempts == 1 + for correlation in raw["data"]["completedCorrelations"]: + result = cold.state.try_get_agent_response(correlation) + assert result is not None and result.additional_properties["durable_status"] == "already_completed" + _quiet(client, hooks, callback) + + +def test_expiry_commit_failure_restores_cache_and_payloads_before_successful_retry(clock: type[Clock]) -> None: + raw = _mailboxes() + store = Store(raw) + agent, client, hooks, callback = _agent() + entity = AgentEntity(agent, state_provider=store, callback=callback) + cached = entity.state + store.fail_writes = True + with pytest.raises(OSError, match="commit failure"): + entity.expire_responses() + assert entity.state is cached and cached.to_dict() == raw + assert store.raw == raw and store.writes == 0 and store.attempts == 1 + store.fail_writes = False + assert entity.expire_responses() == 2 + assert store.raw == _without_expired(raw) and store.writes == 1 + _quiet(client, hooks, callback) + + +def test_expiry_budget_uses_whole_retained_floor_without_pruning_live_result_or_history(clock: type[Clock]) -> None: + raw = _mailboxes() + expected = _without_expired(raw) + full_size = _size(expected) + assert full_size > _size(expected["data"]["responseMailbox"]) + agent, client, hooks, callback = _agent() + rejected = Store(raw) + entity = AgentEntity(agent, state_provider=rejected, callback=callback, max_state_bytes=full_size - 1) + cached = entity.state + with pytest.raises(ValueError, match="max_state_bytes|capacity|budget"): + entity.expire_responses() + assert entity.state is cached and cached.to_dict() == raw + assert rejected.raw == raw and rejected.attempts == 0 + + accepted = Store(raw) + entity = AgentEntity(agent, state_provider=accepted, callback=callback, max_state_bytes=full_size) + assert entity.expire_responses() == 2 + assert accepted.raw == expected and accepted.writes == accepted.attempts == 1 + _quiet(client, hooks, callback) + + +def _registered(agent: Agent, callback: Mock, **settings: Any) -> Any: + native = Mock() + host = DurableAIAgentWorker(native, deployment_mode="isolated_v2", callback=callback, **settings) + host.add_agent(agent) + entity_type = native.add_entity.call_args.args[0] + assert entity_type.__name__ == "dafx-maintenance" + return entity_type + + +def _host_entity(entity_type: Any, store: Store) -> Any: + context = Mock() + context.entity_id = EntityInstanceId("dafx-maintenance", "destination") + context.get_state.side_effect = lambda *args, **kwargs: store._get_state_dict() + context.set_state.side_effect = store._set_state_dict + entity = entity_type() + entity._initialize_entity_context(context) + assert isinstance(entity._agent_entity, AgentEntity) + assert entity.core_session_id == DESTINATION_ID + return entity + + +@pytest.mark.parametrize("operation", ["new-run", "duplicate-run", "reset", "expire_responses"]) +def test_registered_dt_legacy_writer_guards_execute_actual_entity(operation: str) -> None: + raw = _source() + store = Store(raw) + agent, client, hooks, callback = _agent() + hosted = _host_entity(_registered(agent, callback), store) + assert hosted.state.try_get_agent_response("legacy-done").text == "retained legacy answer" + with pytest.raises(ValueError, match="[Ll]egacy.*read-only"): + if operation.endswith("run"): + hosted.run({ + "message": "blocked", + "correlationId": "legacy-done" if operation == "duplicate-run" else "new", + }) + else: + getattr(hosted, operation)() + assert store.raw == raw and store.attempts == 0 + _quiet(client, hooks, callback) + + +def test_registered_dt_migrate_cold_retry_and_expiry_use_configured_window(clock: type[Clock]) -> None: + agent, client, hooks, callback = _agent() + entity_type = _registered(agent, callback, response_delivery_window_seconds=17) + store = Store() + hosted = _host_entity(entity_type, store) + request = _request(evidence=True) + before_request = deepcopy(request) + result = hosted.migrate(request) + assert result == {"status": "migrated", "migrationId": "migration-1", "sessionId": DESTINATION_ID} + assert store.raw["data"]["responseMailbox"]["legacy-done"]["expiresAt"] == (NOW + timedelta(seconds=17)).isoformat() + _quiet(client, hooks, callback) + assert hosted.run({"message": "v2 turn", "correlationId": "v2-done"})["type"] == "agent_response" + before = deepcopy(store.raw) + model_calls, hook_calls, callbacks = len(client.received_messages), list(hooks.calls), list(callback.mock_calls) + assert model_calls == 1 + clock.current = NOW + timedelta(seconds=18) + cold_store = Store(before) + cold = _host_entity(entity_type, cold_store) + assert cold.migrate(request) == result and cold_store.raw == before and cold_store.attempts == 0 + assert cold.expire_responses() == 2 + expected = deepcopy(before) + del expected["data"]["responseMailbox"] + assert cold_store.raw == expected + assert cold.expire_responses() == 0 and cold_store.writes == cold_store.attempts == 1 + assert len(client.received_messages) == model_calls + assert hooks.calls == hook_calls and callback.mock_calls == callbacks + assert request == before_request + + +@pytest.mark.parametrize("operation", ["migrate", "expire_responses"]) +def test_registered_dt_maintenance_commit_failure_rolls_back_actual_entity_cache( + clock: type[Clock], operation: str +) -> None: + raw = {} if operation == "migrate" else _mailboxes() + store = Store(raw) + agent, client, hooks, callback = _agent() + hosted = _host_entity(_registered(agent, callback), store) + cached = hosted.state + request = _request() + store.fail_writes = True + with pytest.raises(OSError, match="commit failure"): + hosted.migrate(request) if operation == "migrate" else hosted.expire_responses() + assert hosted.state is cached + assert cached.to_dict() == (raw or DurableAgentState().to_dict()) + assert store.raw == raw and store.writes == 0 and store.attempts == 1 + store.fail_writes = False + result = hosted.migrate(request) if operation == "migrate" else hosted.expire_responses() + expected = ( + {"status": "migrated", "migrationId": "migration-1", "sessionId": DESTINATION_ID} + if operation == "migrate" + else 2 + ) + assert result == expected + assert store.writes == 1 + _quiet(client, hooks, callback) + + +@pytest.mark.parametrize("case", ["digest", "destination", "same-source", "blank-id", "nonempty", "changed-retry"]) +def test_registered_dt_migration_rejections_leave_storage_and_execution_untouched( + clock: type[Clock], case: str +) -> None: + agent, client, hooks, callback = _agent() + entity_type = _registered(agent, callback) + store = Store(_mailboxes() if case == "nonempty" else None) + hosted = _host_entity(entity_type, store) + request = _request() + if case == "changed-retry": + hosted.migrate(request) + hosted = _host_entity(entity_type, store) + request["ownershipTransferId"] = "different-operator-transfer" + elif case == "digest": + request["sourceDigest"] = _digest(DurableAgentState.from_dict(request["source"]).to_dict()) + elif case == "destination": + request["destinationSessionId"] = "@dafx-other@destination" + elif case == "same-source": + request["sourceSessionId"] = DESTINATION_ID + elif case == "blank-id": + request["migrationId"] = " \t" + before, before_request, attempts = deepcopy(store.raw), deepcopy(request), store.attempts + with pytest.raises(ValueError): + hosted.migrate(request) + assert store.raw == before and request == before_request and store.attempts == attempts + _quiet(client, hooks, callback) + + +@pytest.mark.parametrize("operation", ["migrate", "expire_responses"]) +def test_registered_dt_maintenance_budget_rejects_one_byte_short_accepts_full_floor( + clock: type[Clock], operation: str +) -> None: + raw = {} if operation == "migrate" else _mailboxes() + request = _request() + agent, client, hooks, callback = _agent() + sizing_store = Store(raw) + sizing = _host_entity(_registered(agent, callback), sizing_store) + result = sizing.migrate(request) if operation == "migrate" else sizing.expire_responses() + expected = deepcopy(sizing_store.raw) + full_size = _size(expected) + if operation == "migrate": + normalized = DurableAgentState.from_dict(request["source"]).to_dict() + assert expected["data"]["conversationHistory"] == normalized["data"]["conversationHistory"] + else: + assert expected == _without_expired(raw) + rejected_store = Store(raw) + rejected = _host_entity(_registered(agent, callback, max_state_bytes=full_size - 1), rejected_store) + cached = rejected.state + with pytest.raises(ValueError, match="max_state_bytes|capacity|budget"): + rejected.migrate(request) if operation == "migrate" else rejected.expire_responses() + assert rejected.state is cached and rejected_store.raw == raw and rejected_store.attempts == 0 + accepted_store = Store(raw) + accepted = _host_entity(_registered(agent, callback, max_state_bytes=full_size), accepted_store) + actual = accepted.migrate(request) if operation == "migrate" else accepted.expire_responses() + assert actual == result and accepted_store.raw == expected and accepted_store.writes == 1 + _quiet(client, hooks, callback) + + +@pytest.mark.parametrize("correlation", ["expired-success", "expired-error"]) +def test_registered_dt_expired_duplicate_removes_mailbox_without_execution( + clock: type[Clock], correlation: str +) -> None: + raw = _mailboxes() + agent, client, hooks, callback = _agent() + store = Store(raw) + hosted = _host_entity(_registered(agent, callback), store) + result = hosted.run({"message": "duplicate", "correlationId": correlation}) + assert result["additional_properties"] == { + "durable_status": "already_completed", + "correlation_id": correlation, + "durable_outcome": "failed" if correlation == "expired-error" else "succeeded", + } + assert result["messages"][0]["contents"][0]["error_code"] == "response_expired" + assert store.raw == _without_expired(raw) and store.writes == 1 + _quiet(client, hooks, callback) + + +def test_registered_dt_external_identity_survives_cold_destination_and_migration_retry(clock: type[Clock]) -> None: + external = ExternalHistory() + agent, client, hooks, callback = _agent() + agent.context_providers = [external, hooks] + entity_type = _registered(agent, callback) + store = Store() + request = _request() + before_request = deepcopy(request) + external_before = {key: [message.to_dict() for message in messages] for key, messages in external.rows.items()} + _host_entity(entity_type, store).migrate(request) + assert external.calls == [] + assert { + key: [message.to_dict() for message in messages] for key, messages in external.rows.items() + } == external_before + for index in range(2): + store = Store(store.raw) + hosted = _host_entity(entity_type, store) + if index: + before = deepcopy(store.raw) + assert hosted.migrate(request)["status"] == "migrated" + assert store.raw == before and store.attempts == 0 + result = hosted.run({"message": f"new turn {index}", "correlationId": f"new-{index}"}) + assert result["type"] == "agent_response" + assert store.raw["data"]["session"]["session_id"] == SOURCE_ID + assert external.calls == [(phase, SOURCE_ID) for _ in range(2) for phase in ("get", "save")] + assert [message.text for message in external.rows[SOURCE_ID]] == [ + "already in the external store", + "new turn 0", + "reply-1", + "new turn 1", + "reply-2", + ] + assert set(external.rows) == {SOURCE_ID} and request == before_request + assert "retained legacy input" not in [message.text for batch in client.received_messages for message in batch] diff --git a/python/packages/durabletask/tests/test_media_retention_boundaries.py b/python/packages/durabletask/tests/test_media_retention_boundaries.py new file mode 100644 index 0000000..c3ff747 --- /dev/null +++ b/python/packages/durabletask/tests/test_media_retention_boundaries.py @@ -0,0 +1,330 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Media retention through core hooks and JSON cold starts, not live model services.""" + +import hashlib +import json +import struct +import zlib +from collections.abc import Awaitable, Iterator +from copy import deepcopy +from typing import Any + +import pytest +from agent_framework import GROUP_ANNOTATION_KEY, ChatResponse, CompactionProvider, Content, Message +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import InMemoryMetricReader, Sum +from test_durable_history_provider import RecordingChatClient +from test_history_pipeline_revision import NonStreamingAgent +from test_revision_contract import JsonStateProvider + +from agent_framework_durabletask import AgentEntity, DurableAgentState, DurableHistoryProvider +from agent_framework_durabletask import _retention_telemetry as telemetry +from agent_framework_durabletask._retention import RetentionMode, StateCapacityError + +MEDIA_CASES = ("inline-png", "inline-text", "uri-image", "hosted-file", "mixed-tool", "large-tool") + + +def _png() -> bytes: + def chunk(kind: bytes, data: bytes) -> bytes: + return struct.pack(">I", len(data)) + kind + data + struct.pack(">I", zlib.crc32(kind + data)) + + # A valid grayscale PNG with incompressible pixels. The binary data, not just + # text padding, must contribute materially to pressure and protected-floor checks. + width, height = 128, 64 + pixels = hashlib.shake_256(b"durable-media-pressure").digest(width * height) + rows = b"".join(b"\x00" + pixels[row * width : (row + 1) * width] for row in range(height)) + return ( + b"\x89PNG\r\n\x1a\n" + + chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 0, 0, 0, 0)) + + chunk(b"IDAT", zlib.compress(rows)) + + chunk(b"IEND", b"") + ) + + +PNG = _png() + + +@pytest.fixture +def media_metrics(monkeypatch: pytest.MonkeyPatch) -> Iterator[InMemoryMetricReader]: + reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[reader], shutdown_on_exit=False) + monkeypatch.setattr(telemetry, "get_meter", meter_provider.get_meter) + telemetry._instruments.cache_clear() + try: + yield reader + finally: + telemetry._instruments.cache_clear() + meter_provider.shutdown() + + +def _removed_counter(reader: InMemoryMetricReader) -> int: + data = reader.get_metrics_data() + total = 0 + if data is not None: + for resource in data.resource_metrics: + for scope in resource.scope_metrics: + for metric in scope.metrics: + if metric.name != "durable.retention.removed_messages": + continue + assert isinstance(metric.data, Sum) and metric.data.is_monotonic + for point in metric.data.data_points: + assert point.attributes is not None + assert point.attributes["outcome"] == "staged" + assert point.attributes["commit_status"] == "not_attempted" + assert isinstance(point.value, int) + total += point.value + return total + + +def _payload_messages(kind: str, turn: str) -> list[Message]: + application = {"type": "text", "nested": {"type": "error", "labels": [turn, "界", 0, False, None]}} + media = { + "inline-png": Content.from_data(PNG, "image/png"), + "inline-text": Content.from_data(("inline document 界\n" * 400).encode(), "text/plain"), + "uri-image": Content.from_uri("https://example.test/image.png?version=1", media_type="image/png"), + "hosted-file": Content("hosted_file", file_id=f"file-{turn}", additional_properties=deepcopy(application)), + }.get(kind, Content.from_text("Use the tool payload")) + arguments: Any = {"query": turn, "metadata": deepcopy(application)} + result: Any = {"records": [turn], "metadata": deepcopy(application)} + if kind == "mixed-tool": + result = [ + Content.from_text("tool text 界", additional_properties=deepcopy(application)), + Content.from_data(PNG, "image/png"), + Content.from_data(b"inline tool document", "text/plain"), + ] + elif kind == "large-tool": + # Preserve the original JSON string, including its whitespace, not a reparsed equivalent. + arguments = json.dumps({"query": "界🚀" * 700, "metadata": application}, ensure_ascii=False, indent=2) + result = {"records": ["result 界🚀" * 700], "metadata": deepcopy(application)} + return [ + Message( + "user", + [Content.from_text(f"{turn}: " + "context " * 400), media], + message_id=f"{turn}-input", + author_name="media-user", + additional_properties=deepcopy(application), + ), + Message( + "assistant", + [ + Content.from_text_reasoning( + id=f"{turn}-reasoning", + text="Retain this reasoning summary", + protected_data=f"opaque-{turn}", + additional_properties=deepcopy(application), + ), + Content.from_function_call(f"{turn}-call", "lookup", arguments=arguments), + ], + message_id=f"{turn}-call-message", + author_name="planner", + additional_properties=deepcopy(application), + ), + Message( + "tool", + [Content.from_function_result(f"{turn}-call", result=result, additional_properties=deepcopy(application))], + message_id=f"{turn}-result-message", + author_name="lookup", + additional_properties=deepcopy(application), + ), + ] + + +class _MediaClient(RecordingChatClient): + """Core Agent still owns history hooks, only the model response is deterministic.""" + + def get_response(self, messages: Any, *, stream: bool = False, **kwargs: Any) -> Awaitable[ChatResponse]: + if stream: + raise TypeError("stream is not supported") + self.received_messages.append(deepcopy(list(messages))) + self._counter += 1 + counter = self._counter + + async def get() -> ChatResponse: + return ChatResponse( + messages=[ + Message( + "assistant", + [Content.from_text(f"answer-{counter}", additional_properties={"json": {"keep": [0, False]}})], + message_id=f"answer-{counter}", + author_name="media-client", + additional_properties={"json": {"type": "text", "keep": [counter, None]}}, + ) + ], + response_id=f"response-{counter}", + ) + + return get() + + +def _agent(client: _MediaClient, **kwargs: Any) -> NonStreamingAgent: + chat_client: Any = client + return NonStreamingAgent(client=chat_client, **kwargs) + + +def _request(correlation: str, messages: list[Message]) -> dict[str, Any]: + return { + "message": "projected media turn", + "correlationId": correlation, + "contextMessages": [deepcopy(message.to_dict()) for message in messages], + } + + +def _messages(raw: dict[str, Any]) -> list[Message]: + cold = DurableAgentState.from_json(json.dumps(raw)) + return [message.to_chat_message() for entry in cold.data.conversation_history for message in entry.messages] + + +@pytest.mark.parametrize("kind", MEDIA_CASES) +@pytest.mark.parametrize("retention", ["keep_all", "follow_compaction"]) +@pytest.mark.parametrize("pressure", [False, True], ids=["no-budget", "pressure"]) +async def test_media_payloads_survive_policy_matrix_json_reload_and_next_model_call( + kind: str, retention: RetentionMode, pressure: bool, media_metrics: InMemoryMetricReader +) -> None: + provider = JsonStateProvider() + client = _MediaClient() + seed_agent: Any = _agent(client=client, name="media") + seed = AgentEntity(seed_agent, state_provider=provider) + originals: dict[str, dict[str, Any]] = {} + inputs: dict[str, list[Message]] = {} + atomic_pairs: list[set[str]] = [] + for index in range(8): + turn = f"seed-{index}" + inputs[turn] = _payload_messages(kind, turn) + response = await seed.run(_request(turn, inputs[turn])) + for message in [*inputs[turn], *response.messages]: + assert message.message_id is not None + originals[message.message_id] = deepcopy(message.to_dict()) + atomic_pairs.append({f"{turn}-call-message", f"{turn}-result-message"}) + + persisted_seed = json.loads(json.dumps(provider.raw)) + assert [message.to_dict() for message in _messages(persisted_seed)] == list(originals.values()) + assert provider.writes == 8 + # The request payloads dominate the small answer mailboxes, leaving a reachable floor. + # Use the actual complete JSON size so all media forms exercise pressure without a guessed cap. + budget = int(len(json.dumps(persisted_seed)) * 0.8) if pressure else None + excluded = {message.message_id for message in inputs["seed-0"]} | {"answer-1", "seed-1-call-message"} + strategy_calls: list[int] = [] + core_groups: dict[str, Any] = {} + + async def exclude_old(messages: list[Message]) -> bool: + strategy_calls.append(len(messages)) + changed = False + for message in messages: + assert message.message_id is not None + # Core adds grouping annotations before invoking a custom strategy. Preserve that + # generated metadata separately from the independent original-payload oracle. + core_groups[message.message_id] = deepcopy(message.additional_properties[GROUP_ANNOTATION_KEY]) + if message.message_id in excluded and not message.additional_properties.get("_excluded"): + message.additional_properties["_excluded"] = True + changed = True + return changed + + history = DurableHistoryProvider() + current_agent: Any = _agent( + client=client, + name="media", + context_providers=[ + history, + CompactionProvider(after_strategy=exclude_old, history_source_id=history.source_id), + ], + ) + current_provider = JsonStateProvider(persisted_seed) + entity = AgentEntity(current_agent, state_provider=current_provider, retention=retention, max_state_bytes=budget) + current_inputs = _payload_messages(kind, "current") + response = await entity.run(_request("current", current_inputs)) + for message in [*current_inputs, *response.messages]: + assert message.message_id is not None + originals[message.message_id] = deepcopy(message.to_dict()) + for message_id in excluded: + assert message_id is not None + originals[message_id]["additional_properties"]["_excluded"] = True + for message_id, group in core_groups.items(): + originals[message_id]["additional_properties"][GROUP_ANNOTATION_KEY] = group + # Core compaction explicitly materializes False on included messages. + originals[message_id]["additional_properties"].setdefault("_excluded", False) + assert strategy_calls, "the fixture must execute real core compaction hooks" + assert current_provider.writes == 1 + + raw = json.loads(json.dumps(current_provider.raw)) + retained = _messages(raw) + retained_ids = {message.message_id for message in retained} + removed = set(originals) - retained_ids + assert [message.to_dict() for message in retained] == [ + payload for message_id, payload in originals.items() if message_id in retained_ids + ] + newest_ids = {message.message_id for message in [*current_inputs, *response.messages]} + assert newest_ids <= retained_ids, "the entire newest exchange is protected, including non-text payloads" + for pair in atomic_pairs: + assert pair <= retained_ids or pair.isdisjoint(retained_ids), "no half tool-call/result group may be deleted" + if pressure: + assert len(removed) > 4, "pressure must remove more than the one fully excluded exchange" + assert budget is not None and len(json.dumps(raw)) < budget * 0.85 + elif retention == "follow_compaction": + assert removed == {message.message_id for message in inputs["seed-0"]} | {"answer-1"} + assert {"seed-1-call-message", "seed-1-result-message"} <= retained_ids + else: + assert removed == set() + truncation = raw["data"].get("truncation") or {} + assert truncation.get("evictedMessageCount", 0) == len(removed) + assert _removed_counter(media_metrics) == len(removed) + if removed: + assert truncation["firstEvictedAt"] and truncation["lastEvictedAt"] + + # A new provider and core Agent must reconstruct only committed, included payloads. + # Resending an old projected input also checks that eviction did not erase its receipt. + cold_client = _MediaClient() + cold_client._counter = client._counter + cold_agent: Any = _agent(client=cold_client, name="media", context_providers=[DurableHistoryProvider()]) + cold_provider = JsonStateProvider(raw) + cold = AgentEntity(cold_agent, state_provider=cold_provider, retention=retention, max_state_bytes=budget) + next_input = Message("user", ["next model call"], message_id="next-input", additional_properties={"json": [1]}) + duplicate = await cold.run(_request("current", current_inputs)) + assert duplicate.to_dict() == response.to_dict() + assert cold_client.received_messages == [] and cold_provider.writes == 0 + await cold.run(_request("next", [*inputs["seed-0"], next_input])) + expected = [message.to_dict() for message in retained if not message.additional_properties.get("_excluded")] + for payload in expected: + # HistoryProvider.before_run contributes source attribution to model copies only. + payload["additional_properties"]["_attribution"] = { + "source_id": DurableHistoryProvider.DEFAULT_SOURCE_ID, + "source_type": "DurableHistoryProvider", + } + assert len(cold_client.received_messages) == 1 + assert [message.to_dict() for message in cold_client.received_messages[0]] == [*expected, next_input.to_dict()] + cold_ids = {message.message_id for message in _messages(cold_provider.raw)} + assert removed.isdisjoint(cold_ids), "a cold flush or replayed transport input must not resurrect deleted payloads" + for pair in atomic_pairs: + assert pair <= cold_ids or pair.isdisjoint(cold_ids) + assert cold_provider.raw["data"]["ingestedMessages"] == { + **raw["data"]["ingestedMessages"], + "next-input": cold_provider.raw["data"]["ingestedMessages"]["next-input"], + } + assert cold_provider.writes == 1 + final_removed = len(originals) + 2 - len(_messages(cold_provider.raw)) + assert (cold_provider.raw["data"].get("truncation") or {}).get("evictedMessageCount", 0) == final_removed + assert _removed_counter(media_metrics) == final_removed + + +@pytest.mark.parametrize("kind", MEDIA_CASES) +async def test_newest_media_floor_cannot_be_deleted_to_make_a_commit_fit(kind: str) -> None: + client = _MediaClient() + agent: Any = _agent(client=client, name="protected-media") + probe_provider = JsonStateProvider() + probe = AgentEntity(agent, state_provider=probe_provider) + request = _request("protected", _payload_messages(kind, "protected")) + await probe.run(request) + full_size = len(json.dumps(probe_provider.raw)) + # The same turn cannot fit at this budget unless its newest protected payload is deleted. + budget = int(full_size * 0.8) + provider = JsonStateProvider() + entity = AgentEntity(agent, state_provider=provider, max_state_bytes=budget) + before = entity.state.to_dict() + + with pytest.raises(StateCapacityError) as error: + await entity.run(request) + + assert error.value.floor_bytes >= budget * 0.85 + assert len(client.received_messages) == 2, "the model succeeded before the commit was rejected" + assert entity.state.to_dict() == before and provider.raw == {} and provider.writes == 0 + assert entity.state.try_get_agent_response("protected") is None diff --git a/python/packages/durabletask/tests/test_provider_composition_review.py b/python/packages/durabletask/tests/test_provider_composition_review.py new file mode 100644 index 0000000..db89b13 --- /dev/null +++ b/python/packages/durabletask/tests/test_provider_composition_review.py @@ -0,0 +1,361 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Provider ownership, namespace and atomic reconciliation regressions against Core.""" + +import json +from copy import deepcopy +from itertools import combinations +from typing import Any + +import pytest +from agent_framework import ( + Agent, + AgentSession, + Content, + ContextProvider, + HistoryProvider, + InMemoryHistoryProvider, + Message, + SessionContext, +) +from test_durable_history_provider import _InMemoryStateProvider +from test_history_pipeline_revision import OLD, AddContext, ToolChatClient, bound, ids, seed, stored, transcript + +from agent_framework_durabletask import DurableHistoryProvider +from agent_framework_durabletask import _history_provider as history_module +from agent_framework_durabletask._durable_agent_state import ( + DurableAgentStateMessage, + DurableAgentStateRequest, + DurableAgentStateResponse, +) +from agent_framework_durabletask._history_provider import WORKING_BUFFER_KEY, ensure_durable_history + + +class OrdinaryExternalHistory(HistoryProvider): + """Blind append storage, deliberately unaware of service ownership.""" + + def __init__(self, source_id: str = "external", **kwargs: Any) -> None: + super().__init__(source_id, **kwargs) + self.saved: list[Message] = [] + self.calls: list[tuple[str, str | None]] = [] + self.resource = object() + self.lifecycle: list[str] = [] + + async def get_messages(self, session_id: str | None, **kwargs: Any) -> list[Message]: + self.calls.append(("load", session_id)) + return deepcopy(self.saved) + + async def save_messages(self, session_id: str | None, messages: Any, **kwargs: Any) -> None: + self.calls.append(("save", session_id)) + self.saved.extend(deepcopy(list(messages))) + + async def __aenter__(self) -> "OrdinaryExternalHistory": + self.lifecycle.append("enter") + return self + + async def __aexit__(self, *args: Any) -> None: + self.lifecycle.append("exit") + + +def prepare_owner(agent: Any, service_owned: bool) -> Any: + prepare = getattr(history_module, "prepare_history_owner", None) + assert callable(prepare), "Durable must provide per-run ownership for ordinary external providers" + return prepare(agent, service_owns_history=service_owned) + + +@pytest.mark.parametrize("per_call", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_external_service_branches_and_sink_choices_survive_cold_session_reload( + per_call: bool, stream: bool +) -> None: + primary = OrdinaryExternalHistory(store_context_messages=True, store_context_from={"selected"}) + sink = InMemoryHistoryProvider("audit", load_messages=False, store_inputs=False, store_outputs=True) + client = ToolChatClient(tool_calls=False) + agent = Agent( + client=client, + context_providers=[primary, AddContext("selected"), sink], + require_per_service_call_history_persistence=per_call, + ) + providers = agent.context_providers + mask = primary.store_context_from + assert ensure_durable_history(agent) is agent + session = agent.create_session(session_id="external-session") + saved_service_id = None + for turn, service_owned in enumerate((True, False, True, False), start=1): + # The parent owns service-ID parking. Exercise its contract without editing entities. + session.service_session_id = saved_service_id if service_owned else None + prepared = prepare_owner(agent, service_owned) + if service_owned: + assert prepared is not agent and prepared.context_providers is not providers + wrapper = prepared.context_providers[0] + assert isinstance(wrapper, HistoryProvider) + adapter: Any = wrapper + assert adapter.__wrapped__ is primary + assert wrapper.source_id == primary.source_id + assert wrapper.load_messages is primary.load_messages + assert wrapper.store_inputs is primary.store_inputs + assert wrapper.store_outputs is primary.store_outputs + assert wrapper.store_context_messages is primary.store_context_messages + assert wrapper.store_context_from == mask + assert adapter.resource is primary.resource + assert prepare_owner(prepared, True) is prepared + local_view = prepare_owner(prepared, False) + assert local_view.context_providers[0] is primary + assert prepared.client is agent.client + else: + assert prepared is agent and prepared.context_providers[0] is primary + assert prepared.context_providers[-1] is sink + options = {"store": service_owned} + if stream: + await prepared.run(f"turn-{turn}", session=session, options=options, stream=True).get_final_response() + else: + await prepared.run(f"turn-{turn}", session=session, options=options) + if service_owned: + saved_service_id = session.service_session_id + session = AgentSession.from_dict(json.loads(json.dumps(session.to_dict()))) + assert [message.text for message in session.state["audit"]["messages"]] == [ + f"answer-{index}" for index in range(1, turn + 1) + ] + assert agent.context_providers is providers and providers[0] is primary + assert primary.store_context_from is mask + assert sink.source_id == "audit" and sink.store_inputs is False and sink.load_messages is False + assert [message.text for message in primary.saved] == [ + "context-selected", + "turn-2", + "answer-2", + "context-selected", + "turn-4", + "answer-4", + ] + assert primary.calls == [(phase, "external-session") for _ in range(2) for phase in ("load", "save")] + assert primary.lifecycle == [] + assert all("turn-2" not in [message.text for message in client.received_messages[index]] for index in (0, 2)) + assert [message.text for message in client.received_messages[3]].count("turn-2") == 1 + assert not {"turn-1", "turn-3"} & {message.text for message in client.received_messages[3]} + loaded = next(message for message in client.received_messages[3] if message.text == "turn-2") + assert loaded.additional_properties["_attribution"] == { + "source_id": "external", + "source_type": "OrdinaryExternalHistory", + } + + +async def test_service_view_never_calls_custom_primary_hooks_or_direct_storage_methods() -> None: + class CustomExternal(OrdinaryExternalHistory): + async def before_run(self, **kwargs: Any) -> None: + self.calls.append(("before", None)) + await super().before_run(**kwargs) + + async def after_run(self, **kwargs: Any) -> None: + self.calls.append(("after", None)) + await super().after_run(**kwargs) + + primary = CustomExternal() + agent = Agent(client=ToolChatClient(tool_calls=False), context_providers=[primary]) + service = prepare_owner(agent, True) + wrapper = service.context_providers[0] + state = {"cursor": {"keep": [1, 3]}} + before = deepcopy(state) + assert await wrapper.get_messages("session", state=state) == [] + await wrapper.save_messages("session", [Message("user", ["must not save"])], state=state) + await service.run("service", session=service.create_session(), options={"store": True}) + assert state == before and primary.calls == [] and primary.lifecycle == [] + await prepare_owner(agent, False).run("local", session=agent.create_session(), options={"store": False}) + assert [phase for phase, _ in primary.calls] == ["before", "load", "after", "save"] + + +def test_default_sink_collision_fails_without_reconfiguring_the_caller() -> None: + sink = InMemoryHistoryProvider(load_messages=False) + agent = Agent(client=ToolChatClient(tool_calls=False), context_providers=[sink]) + providers = agent.context_providers + with pytest.raises(ValueError, match="in_memory.*source_id"): + ensure_durable_history(agent) + assert agent.context_providers is providers and providers == [sink] + assert sink.source_id == "in_memory" and sink.load_messages is False + + +@pytest.mark.parametrize("other", [ContextProvider("same"), InMemoryHistoryProvider("same", load_messages=False)]) +def test_duplicate_source_ids_fail_before_substitution(other: ContextProvider) -> None: + primary = InMemoryHistoryProvider("same") + agent = Agent(client=ToolChatClient(tool_calls=False), context_providers=[primary, other]) + with pytest.raises(ValueError, match="source_id.*same"): + ensure_durable_history(agent) + assert agent.context_providers == [primary, other] + + +async def test_uniquely_named_sink_only_keeps_separate_durable_and_sink_history() -> None: + sink = InMemoryHistoryProvider("audit", load_messages=False) + client = ToolChatClient(tool_calls=False) + agent: Any = ensure_durable_history(Agent(client=client, context_providers=[sink])) + # Like Core's automatic history, the implicit durable primary follows the caller's sink. + assert len(agent.context_providers) == 2 + assert agent.context_providers[0] is sink + primary = agent.context_providers[-1] + assert isinstance(primary, DurableHistoryProvider) and primary.source_id == "in_memory" + provider = _InMemoryStateProvider() + session = agent.create_session() + for turn in range(3): + with bound(provider, f"turn-{turn}"): + await agent.run(f"turn-{turn}", session=session) + primary.flush(session.state[primary.source_id]) + assert len(transcript(provider)) == len(session.state["audit"]["messages"]) == (turn + 1) * 2 + session.state.pop(primary.source_id) + session = AgentSession.from_dict(json.loads(json.dumps(session.to_dict()))) + assert [len(messages) for messages in client.received_messages] == [1, 3, 5] + + +@pytest.mark.parametrize("per_call", [False, True]) +@pytest.mark.parametrize("factory", [InMemoryHistoryProvider, DurableHistoryProvider]) +async def test_self_context_mask_does_not_reappend_history(factory: Any, per_call: bool) -> None: + original = factory("history", store_context_messages=True, store_context_from={"history"}) + client = ToolChatClient(tool_calls=False) + agent: Any = ensure_durable_history( + Agent(client=client, context_providers=[original], require_per_service_call_history_persistence=per_call) + ) + provider = _InMemoryStateProvider() + session = agent.create_session() + history = agent.context_providers[0] + counts = [] + for turn in range(3): + with bound(provider, f"turn-{turn}"): + await agent.run(f"turn-{turn}", session=session) + history.flush(session.state[history.source_id]) + counts.append(len(transcript(provider))) + core_history = InMemoryHistoryProvider("history", store_context_messages=True, store_context_from={"history"}) + core = Agent(client=ToolChatClient(tool_calls=False), context_providers=[core_history]) + core_session = core.create_session() + core_counts = [] + for turn in range(3): + await core.run(f"turn-{turn}", session=core_session) + core_counts.append(len(core_session.state["history"]["messages"])) + assert counts == [2, 4, 6] + # Core 1.13 still re-appends its own contribution; later core releases fix it. + # Durable must not copy that historical duplication bug into its append path. + assert core_counts in ([2, 4, 6], [2, 6, 14]) + assert [len(messages) for messages in client.received_messages] == [1, 3, 5] + assert original.store_context_from == {"history"} + + +@pytest.mark.parametrize("mask", [None, set(), {"history"}, {"selected"}, {"history", "selected"}]) +def test_context_mask_excludes_only_self_not_selected_sources(mask: set[str] | None) -> None: + history = DurableHistoryProvider("history", store_context_messages=True, store_context_from=mask) + context = SessionContext(input_messages=[]) + for source in ("history", "selected", "other"): + context.extend_messages(source, [Message("user", [source])]) + assert [message.text for message in history._get_context_messages_to_store(context)] == [ + source for source in ("selected", "other") if mask is None or source in mask + ] + + +@pytest.mark.parametrize("prune", [False, True]) +async def test_unset_durable_subclass_keeps_overrides_and_resources(prune: bool) -> None: + class CustomDurable(DurableHistoryProvider): + def __init__(self, resource: object) -> None: + super().__init__("custom", store_context_messages=True, store_context_from={"selected"}) + self.resource = resource + self.events: list[str] = [] + + async def before_run(self, **kwargs: Any) -> None: + self.events.append("before") + await super().before_run(**kwargs) + + async def after_run(self, **kwargs: Any) -> None: + self.events.append("after") + await super().after_run(**kwargs) + + original = CustomDurable(object()) + agent = Agent(client=ToolChatClient(tool_calls=False), context_providers=[original]) + prepared: Any = ensure_durable_history(agent, prune_excluded=prune) + replacement = prepared.context_providers[0] + assert type(replacement) is CustomDurable + assert replacement is not original and replacement.resource is original.resource + assert replacement.prune_excluded is prune and original.prune_excluded is None + assert replacement.store_context_from == original.store_context_from + assert replacement.store_context_from is not original.store_context_from + with bound(_InMemoryStateProvider()): + await prepared.run("input", session=prepared.create_session()) + assert replacement.events == ["before", "after"] + + +@pytest.mark.parametrize( + "excluded", + [set(members) for size in range(4) for members in combinations(("reason", "call", "result"), size)], +) +@pytest.mark.parametrize("non_contiguous", [False, True]) +async def test_eager_pruning_requires_the_entire_old_atomic_group_to_be_excluded( + excluded: set[str], non_contiguous: bool +) -> None: + provider = _InMemoryStateProvider() + messages = [ + Message("assistant", [Content.from_text_reasoning(text="reason")], message_id="reason"), + Message("assistant", [Content.from_function_call(call_id="t", name="tool", arguments="{}")], message_id="call"), + Message("tool", [Content.from_function_result(call_id="t", result="result")], message_id="result"), + ] + if non_contiguous: + messages.insert(2, Message("user", ["gap"], message_id="gap")) + provider.state.data.conversation_history.extend([ + DurableAgentStateResponse("old", OLD, [DurableAgentStateMessage.from_chat_message(m) for m in messages]), + DurableAgentStateRequest("current", OLD, [stored("current", "current")]), + ]) + history = DurableHistoryProvider(prune_excluded=True) + state: dict[str, Any] = {} + with bound(provider): + await history.get_messages("session", state=state) + for message in state[WORKING_BUFFER_KEY]: + if message.message_id in excluded: + message.additional_properties["_excluded"] = True + history.flush(state) + removed = 3 if len(excluded) == 3 else 0 + assert len(transcript(provider)) == len(messages) + 1 - removed + assert {"reason", "call", "result"} & set(ids(provider)) == (set() if removed else {"reason", "call", "result"}) + assert (provider.state.data.truncation or {}).get("evictedMessageCount", 0) == removed + snapshot = provider.state.to_dict() + history.flush(state) + assert provider.state.to_dict() == snapshot + + +@pytest.mark.parametrize("nested", [False, True]) +@pytest.mark.parametrize("remove_old", [False, True]) +async def test_same_body_summary_id_reuse_keeps_distinct_lineage(nested: bool, remove_old: bool) -> None: + def links(message: Message) -> dict[str, Any]: + return message.additional_properties.setdefault("_group", {}) if nested else message.additional_properties + + provider = _InMemoryStateProvider() + seed(provider) + history = DurableHistoryProvider(prune_excluded=False) + state: dict[str, Any] = {} + with bound(provider) as binding: + await history.get_messages("session", state=state) + for source_id in ("seed-user", "seed-assistant"): + buffer = state[WORKING_BUFFER_KEY] + source = next(message for message in buffer if message.message_id == source_id) + source.additional_properties["_excluded"] = True + links(source)["_summarized_by_summary_id"] = "summary" + if remove_old and source_id == "seed-assistant": + buffer[:] = [message for message in buffer if message.message_id != "summary"] + summary = Message("assistant", ["identical summary"], message_id="summary") + links(summary)["_summary_of_message_ids"] = [source_id] + buffer.insert(buffer.index(source) + 1, summary) + history.flush(state) + assert summary.message_id != "summary" + assert ids(provider) == ["seed-user", "summary", "seed-assistant", summary.message_id] + saved = {message.message_id: message.to_chat_message() for message in transcript(provider)} + assert links(saved["summary"])["_summary_of_message_ids"] == ["seed-user"] + assert links(saved[summary.message_id])["_summary_of_message_ids"] == ["seed-assistant"] + assert links(saved["seed-user"])["_summarized_by_summary_id"] == "summary" + assert links(saved["seed-assistant"])["_summarized_by_summary_id"] == summary.message_id + snapshot = provider.state.to_dict() + ordinal = binding.append_ordinal + history.flush(state) + assert provider.state.to_dict() == snapshot and binding.append_ordinal == ordinal + cold = _InMemoryStateProvider(raw=json.loads(provider.state.to_json())) + with bound(cold): + cold_state: dict[str, Any] = {} + loaded = await history.get_messages("session", state=cold_state) + # keep_all retains both bodies and their distinct lineage, but a removed summary is not replayed. + assert [message.text for message in loaded] == ["identical summary"] * (1 if remove_old else 2) + assert [message.message_id for message in loaded] == [ + *([] if remove_old else ["summary"]), + summary.message_id, + ] + history.flush(cold_state) + assert cold.state.to_dict() == provider.state.to_dict() diff --git a/python/packages/durabletask/tests/test_provider_hook_followup.py b/python/packages/durabletask/tests/test_provider_hook_followup.py new file mode 100644 index 0000000..b6e568f --- /dev/null +++ b/python/packages/durabletask/tests/test_provider_hook_followup.py @@ -0,0 +1,476 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Core hook ordering, custom history preservation and list-reducing compaction.""" + +import json +from copy import deepcopy +from itertools import combinations +from typing import Any + +import pytest +from agent_framework import Agent, AgentSession, CompactionProvider, Content, InMemoryHistoryProvider, Message +from test_durable_history_provider import _InMemoryStateProvider +from test_history_pipeline_revision import OLD, ToolChatClient, bound, ids, seed, stored, transcript + +from agent_framework_durabletask import AgentEntity, DurableHistoryProvider +from agent_framework_durabletask._durable_agent_state import ( + DurableAgentStateCompaction, + DurableAgentStateMessage, + DurableAgentStateRequest, + DurableAgentStateResponse, + DurableAgentStateUnknownEntry, +) +from agent_framework_durabletask._history_provider import ( + WORKING_BUFFER_KEY, + ensure_durable_history, + prepare_history_owner, + prune_messages, +) + + +class RecordingStrategy: + def __init__(self) -> None: + self.seen: list[list[tuple[str, str]]] = [] + + async def __call__(self, messages: list[Message]) -> bool: + self.seen.append([(message.role, message.text) for message in messages]) + return False + + +class CustomMemory(InMemoryHistoryProvider): + after_run_once_per_turn = True + + def __init__(self, source_id: str = "in_memory") -> None: + super().__init__(source_id) + self.events: list[tuple[str, int]] = [] + self.resource = object() + + async def before_run(self, *, state: dict[str, Any], **kwargs: Any) -> None: + self.events.append(("before", state.get("hook_runs", 0))) + await super().before_run(state=state, **kwargs) + + async def after_run(self, *, state: dict[str, Any], **kwargs: Any) -> None: + await super().after_run(state=state, **kwargs) + state["hook_runs"] = state.get("hook_runs", 0) + 1 + state["custom"] = {"safe": [None, False, 3, {"nested": "kept"}]} + self.events.append(("after", state["hook_runs"])) + + +class CustomDurable(DurableHistoryProvider): + after_run_once_per_turn = True + + def __init__(self) -> None: + super().__init__("custom-durable", store_context_from={"selected"}) + self.resource = object() + self.events: list[str] = [] + + async def after_run(self, **kwargs: Any) -> None: + self.events.append("after") + await super().after_run(**kwargs) + + +@pytest.mark.parametrize("per_call", [False, True]) +async def test_implicit_history_matches_core_first_turn_after_strategy(per_call: bool) -> None: + core_strategy = RecordingStrategy() + core_compaction = CompactionProvider(after_strategy=core_strategy) + core = Agent( + client=ToolChatClient(tool_calls=False), + name="assistant", + context_providers=[core_compaction], + require_per_service_call_history_persistence=per_call, + ) + core_session = core.create_session() + await core.run("first", session=core_session) + assert core.context_providers[0] is core_compaction + assert type(core.context_providers[-1]) is InMemoryHistoryProvider + + strategy = RecordingStrategy() + compaction = CompactionProvider(after_strategy=strategy) + original = Agent( + client=ToolChatClient(tool_calls=False), + name="assistant", + context_providers=[compaction], + require_per_service_call_history_persistence=per_call, + ) + entity = AgentEntity(original, state_provider=_InMemoryStateProvider()) + await entity.run({"message": "first", "correlationId": "first"}) + prepared: Any = entity.agent + assert prepared.context_providers[0] is compaction + history = prepared.context_providers[-1] + assert isinstance(history, DurableHistoryProvider) + assert history.source_id == core.context_providers[-1].source_id == compaction.history_source_id == "in_memory" + assert original.context_providers == [compaction] + assert strategy.seen == core_strategy.seen == [[("user", "first"), ("assistant", "answer-1")]] + + +@pytest.mark.parametrize("history_first", [False, True]) +async def test_explicit_history_registration_order_is_not_rewritten(history_first: bool) -> None: + core_strategy = RecordingStrategy() + core_history = InMemoryHistoryProvider("chosen") + core_compaction = CompactionProvider(after_strategy=core_strategy, history_source_id="chosen") + core_providers = [core_history, core_compaction] if history_first else [core_compaction, core_history] + core = Agent(client=ToolChatClient(tool_calls=False), context_providers=core_providers) + await core.run("first", session=core.create_session()) + + strategy = RecordingStrategy() + history = InMemoryHistoryProvider("chosen") + compaction = CompactionProvider(after_strategy=strategy, history_source_id="chosen") + providers = [history, compaction] if history_first else [compaction, history] + original = Agent(client=ToolChatClient(tool_calls=False), context_providers=providers) + entity = AgentEntity(original, state_provider=_InMemoryStateProvider()) + await entity.run({"message": "first", "correlationId": "first"}) + prepared: Any = entity.agent + assert [provider.source_id for provider in prepared.context_providers] == [p.source_id for p in providers] + assert original.context_providers == providers + assert strategy.seen == core_strategy.seen + assert bool(strategy.seen) is not history_first + + +@pytest.mark.parametrize("per_call", [False, True]) +@pytest.mark.parametrize("source_id", ["in_memory", "custom-history"]) +async def test_custom_memory_hooks_and_transcript_survive_entity_cold_reload(per_call: bool, source_id: str) -> None: + core_history = CustomMemory(source_id) + core = Agent( + client=ToolChatClient(tool_calls=False), + name="assistant", + context_providers=[core_history], + require_per_service_call_history_persistence=per_call, + ) + session = core.create_session() + history = CustomMemory(source_id) + client = ToolChatClient(tool_calls=False) + provider = _InMemoryStateProvider() + for turn in (1, 2): + original = Agent( + client=client, + name="assistant", + context_providers=[history], + require_per_service_call_history_persistence=per_call, + ) + assert ensure_durable_history(original, prune_excluded=True) is original + entity = AgentEntity(original, state_provider=provider, retention="follow_compaction") + prepared: Any = entity.agent + assert prepared.context_providers == [history] + assert prepared.require_per_service_call_history_persistence is per_call + await core.run(f"turn-{turn}", session=session, stream=True).get_final_response() + await entity.run({"message": f"turn-{turn}", "correlationId": f"turn-{turn}"}) + + raw = provider._get_state_dict() + restored = AgentSession.from_dict(raw["data"]["session"]) + assert restored.to_dict()["state"][source_id] == session.to_dict()["state"][source_id] + assert restored.state[source_id]["hook_runs"] == turn + assert all(isinstance(message, Message) for message in restored.state[source_id]["messages"]) + assert provider.state.data.conversation_history == [] + assert ( + history.events + == core_history.events + == [event for index in range(1, turn + 1) for event in (("before", index - 1), ("after", index))] + ) + session = AgentSession.from_dict(json.loads(json.dumps(session.to_dict()))) + provider = _InMemoryStateProvider(raw=raw) + assert [message.text for message in client.received_messages[-1]] == ["turn-1", "answer-1", "turn-2"] + assert client.received_messages[-1][0].additional_properties["_attribution"]["source_type"] == "CustomMemory" + + +@pytest.mark.parametrize("factory", [InMemoryHistoryProvider, CustomMemory, CustomDurable]) +@pytest.mark.parametrize("once_per_turn", [False, True]) +def test_substitution_preserves_once_per_turn_metadata(factory: Any, once_per_turn: bool) -> None: + original = factory() + original.after_run_once_per_turn = once_per_turn + agent = Agent(client=ToolChatClient(tool_calls=False), context_providers=[original]) + prepared: Any = ensure_durable_history(agent, prune_excluded=True) + replacement = prepared.context_providers[0] + assert replacement.after_run_once_per_turn is once_per_turn + assert agent.context_providers == [original] + if type(original) is InMemoryHistoryProvider: + assert type(replacement) is DurableHistoryProvider + elif isinstance(original, CustomMemory): + assert prepared is agent and replacement is original + else: + assert type(replacement) is CustomDurable and replacement is not original + assert replacement.resource is original.resource and replacement.events is original.events + assert replacement.store_context_from == original.store_context_from + assert replacement.store_context_from is not original.store_context_from + assert replacement.prune_excluded is True and original.prune_excluded is None + + +@pytest.mark.parametrize("factory", [InMemoryHistoryProvider, CustomMemory, CustomDurable]) +@pytest.mark.parametrize("once_per_turn", [False, True]) +async def test_preserved_metadata_controls_real_core_loop_iteration_hooks(factory: Any, once_per_turn: bool) -> None: + supports_once_per_turn = hasattr(InMemoryHistoryProvider(), "after_run_once_per_turn") + original = factory() + original.after_run_once_per_turn = once_per_turn + prepared: Any = ensure_durable_history(Agent(client=ToolChatClient(tool_calls=False), context_providers=[original])) + provider = _InMemoryStateProvider() + session = prepared.create_session() + history = prepared.context_providers[0] + with bound(provider): + await prepared.run("iteration", session=session, options={"_agent_loop_iteration": "turn"}) + if isinstance(history, DurableHistoryProvider): + history.flush(session.state[history.source_id]) + saved = [m.to_chat_message().text for m in transcript(provider)] + else: + saved = [m.text for m in session.state[history.source_id].get("messages", [])] + deferred = once_per_turn and supports_once_per_turn + assert saved == ([] if deferred else ["iteration", "answer-1"]) + if isinstance(history, CustomDurable): + assert history.events == ([] if deferred else ["after"]) + if isinstance(history, CustomMemory): + assert history.events == ([("before", 0)] if deferred else [("before", 0), ("after", 1)]) + + +@pytest.mark.parametrize("per_call", [False, True]) +async def test_custom_memory_uses_existing_inactive_primary_service_adapter(per_call: bool) -> None: + history = CustomMemory() + sink = InMemoryHistoryProvider("audit", load_messages=False) + client = ToolChatClient(tool_calls=False) + agent = Agent( + client=client, + context_providers=[history, sink], + require_per_service_call_history_persistence=per_call, + ) + assert ensure_durable_history(agent) is agent + session = agent.create_session() + service_id = None + for turn, service_owned in enumerate((True, False, True), start=1): + # Ownership intentionally isolates branches, rather than mirroring Core's service-call saves. + session.service_session_id = service_id if service_owned else None + prepared: Any = prepare_history_owner(agent, service_owned) + if service_owned: + wrapper = prepared.context_providers[0] + assert wrapper.__wrapped__ is history and wrapper.resource is history.resource + assert wrapper.after_run_once_per_turn is True + local: Any = prepare_history_owner(prepared, False) + assert local.context_providers[0] is history + else: + assert prepared is agent + assert prepared.context_providers[1] is sink + await prepared.run(f"turn-{turn}", session=session, options={"store": service_owned}) + if service_owned: + service_id = session.service_session_id + session = AgentSession.from_dict(json.loads(json.dumps(session.to_dict()))) + assert history.events == [("before", 0), ("after", 1)] + assert session.state[history.source_id]["hook_runs"] == 1 + assert [m.text for m in session.state[history.source_id]["messages"]] == ["turn-2", "answer-2"] + assert len(session.state["audit"]["messages"]) == 6 + assert [m.text for m in client.received_messages[2]] == ["turn-3"] + assert agent.context_providers == [history, sink] + + +class SliceOldExchange: + def __init__(self) -> None: + self.current: list[tuple[str, str]] = [] + + async def __call__(self, messages: list[Message]) -> bool: + self.current = [(m.role, m.text) for m in messages[-2:]] + start = next(index for index, message in enumerate(messages) if message.message_id == "seed-user") + assert messages[start + 1].message_id == "seed-assistant" + del messages[start : start + 2] + return True + + +@pytest.mark.parametrize("retention", ["keep_all", "follow_compaction"]) +async def test_after_strategy_list_removal_survives_cold_next_turn(retention: Any) -> None: + provider = _InMemoryStateProvider() + seed(provider) + originals = transcript(provider) + for message in originals: + message.extension_data = { + "future": {"safe": [None, False, 3, {"nested": "kept"}]}, + "_group": {"_summarized_by_summary_id": "kept-summary"}, + } + summary = Message( + "assistant", + ["old summary"], + message_id="kept-summary", + additional_properties={"_group": {"_summary_of_message_ids": ["seed-user", "seed-assistant"]}}, + ) + unknown = DurableAgentStateUnknownEntry({"$type": "futureKind", "payload": {"keep": [None, False, {"x": 1}]}}) + provider.state.data.conversation_history.insert(0, unknown) + provider.state.data.conversation_history.insert( + 1, DurableAgentStateRequest("system", OLD, [stored("system", "instructions", "system")]) + ) + provider.state.data.conversation_history.append( + DurableAgentStateCompaction(OLD, [DurableAgentStateMessage.from_chat_message(summary)]) + ) + unknown_before = deepcopy(unknown.to_dict()) + source_metadata = deepcopy(originals[0].extension_data) + assert source_metadata is not None + core_strategy = SliceOldExchange() + core_client = ToolChatClient(tool_calls=False) + core = Agent(client=core_client, context_providers=[CompactionProvider(after_strategy=core_strategy)]) + session = core.create_session() + session.state["in_memory"] = {"messages": [deepcopy(m).to_chat_message() for m in transcript(provider)]} + strategy = SliceOldExchange() + entity = AgentEntity( + Agent(client=ToolChatClient(tool_calls=False), context_providers=[CompactionProvider(after_strategy=strategy)]), + state_provider=provider, + retention=retention, + ) + await core.run("current", session=session) + response = await entity.run({"message": "current", "correlationId": "current"}) + assert response.text == "answer-1" + assert strategy.current == core_strategy.current == [("user", "current"), ("assistant", "answer-1")] + saved = {message.message_id: message for message in transcript(provider)} + old_ids = {"seed-user", "seed-assistant"} + if retention == "keep_all": + assert old_ids <= saved.keys() + assert all(saved[message_id].extension_data == {**source_metadata, "_excluded": True} for message_id in old_ids) + assert provider.state.data.truncation is None + else: + assert not old_ids & saved.keys() + assert (provider.state.data.truncation or {})["evictedMessageCount"] == 2 + assert "system" in saved + assert [ + m.to_chat_message().text + for entry in provider.state.data.conversation_history + if entry.correlation_id == "current" + for m in entry.messages + ] == ["current", "answer-1"] + assert (saved["kept-summary"].extension_data or {})["_group"]["_summary_of_message_ids"] == [ + "seed-user", + "seed-assistant", + ] + assert unknown.to_dict() == unknown_before + delivered = provider.state.try_get_agent_response("current") + assert delivered is not None and delivered.to_dict() == response.to_dict() + cold_provider = _InMemoryStateProvider(raw=provider._get_state_dict()) + cold_client = ToolChatClient(tool_calls=False) + cold = AgentEntity(Agent(client=cold_client), state_provider=cold_provider, retention=retention) + core.context_providers = [core.context_providers[-1]] + await core.run("next", session=AgentSession.from_dict(json.loads(json.dumps(session.to_dict())))) + await cold.run({"message": "next", "correlationId": "next"}) + core_input = [m.text for m in core_client.received_messages[-1]] + assert ( + [m.text for m in cold_client.received_messages[0]] + == core_input + == [ + "instructions", + "old summary", + "current", + "answer-1", + "next", + ] + ) + assert cold_provider.state.data.conversation_history[0].to_dict() == unknown_before + delivered = cold_provider.state.try_get_agent_response("current") + assert delivered is not None and delivered.to_dict() == response.to_dict() + + +@pytest.mark.parametrize( + "removed", [set(group) for size in range(3) for group in combinations(("call", "result"), size)] +) +async def test_list_removal_prunes_only_complete_atomic_groups_and_keeps_floor(removed: set[str]) -> None: + provider = _InMemoryStateProvider() + call = Message( + "assistant", [Content.from_function_call(call_id="t", name="lookup", arguments="{}")], message_id="call" + ) + result = Message("tool", [Content.from_function_result(call_id="t", result="value")], message_id="result") + provider.state.data.conversation_history.extend([ + DurableAgentStateRequest("system", OLD, [stored("system", "instructions", "system")]), + DurableAgentStateResponse("old", OLD, [DurableAgentStateMessage.from_chat_message(call)]), + DurableAgentStateRequest("old", OLD, [DurableAgentStateMessage.from_chat_message(result)]), + DurableAgentStateRequest("current", OLD, [stored("current-input", "current")]), + DurableAgentStateResponse("current", OLD, [stored("current-answer", "answer", "assistant")]), + ]) + history = DurableHistoryProvider(prune_excluded=True) + state: dict[str, Any] = {} + with bound(provider): + await history.get_messages("session", state=state) + # Even removing protected messages from the buffer cannot authorize their physical deletion. + missing = removed | {"system", "current-input", "current-answer"} + state[WORKING_BUFFER_KEY][:] = [m for m in state[WORKING_BUFFER_KEY] if m.message_id not in missing] + history.flush(state) + assert {"system", "current-input", "current-answer"} <= set(ids(provider)) + assert {"call", "result"} & set(ids(provider)) == (set() if len(removed) == 2 else {"call", "result"}) + assert (provider.state.data.truncation or {}).get("evictedMessageCount", 0) == (2 if len(removed) == 2 else 0) + snapshot = deepcopy(provider.state.to_dict()) + history.flush(state) + assert provider.state.to_dict() == snapshot + + +@pytest.mark.parametrize("nested", [False, True]) +@pytest.mark.parametrize("prune", [False, True]) +async def test_removed_summary_revision_stays_excluded_without_rewriting_backlinks(nested: bool, prune: bool) -> None: + def links(message: Message) -> dict[str, Any]: + return message.additional_properties.setdefault("_group", {}) if nested else message.additional_properties + + provider = _InMemoryStateProvider() + seed(provider) + provider.state.data.conversation_history.append( + DurableAgentStateRequest("current", OLD, [stored("current", "current")]) + ) + history = DurableHistoryProvider(prune_excluded=prune) + state: dict[str, Any] = {} + with bound(provider) as binding: + await history.get_messages("session", state=state) + buffer = state[WORKING_BUFFER_KEY] + links(buffer[0])["_summarized_by_summary_id"] = "summary" + first = Message( + "assistant", ["first summary"], message_id="summary", additional_properties={"future": {"keep": [1]}} + ) + links(first)["_summary_of_message_ids"] = ["seed-user"] + buffer.insert(1, first) + history.flush(state) + buffer.remove(first) + source = next(m for m in buffer if m.message_id == "seed-assistant") + links(source)["_summarized_by_summary_id"] = "summary" + second = Message("assistant", ["second summary"], message_id="summary") + links(second)["_summary_of_message_ids"] = ["seed-assistant"] + buffer.insert(buffer.index(source) + 1, second) + history.flush(state) + assert second.message_id != "summary" + saved = {m.message_id: deepcopy(m).to_chat_message() for m in transcript(provider)} + assert links(saved["seed-user"])["_summarized_by_summary_id"] == "summary" + assert links(saved["seed-assistant"])["_summarized_by_summary_id"] == second.message_id + assert links(saved[second.message_id])["_summary_of_message_ids"] == ["seed-assistant"] + if prune: + assert "summary" not in saved + else: + assert saved["summary"].additional_properties["_excluded"] is True + assert saved["summary"].additional_properties["future"] == {"keep": [1]} + assert links(saved["summary"])["_summary_of_message_ids"] == ["seed-user"] + snapshot = deepcopy(provider.state.to_dict()) + ordinal = binding.append_ordinal + history.flush(state) + assert provider.state.to_dict() == snapshot and binding.append_ordinal == ordinal + cold = _InMemoryStateProvider(raw=json.loads(provider.state.to_json())) + with bound(cold): + loaded = await history.get_messages("session", state={}) + assert "first summary" not in [m.text for m in loaded] + assert [m.text for m in loaded].count("second summary") == 1 + + +async def test_stale_summary_removed_from_storage_is_not_reinserted() -> None: + provider = _InMemoryStateProvider() + summary = DurableAgentStateCompaction(OLD, [stored("summary", "old summary", "assistant")]) + provider.state.data.conversation_history.extend([ + summary, + DurableAgentStateRequest("current", OLD, [stored("current", "current")]), + ]) + history = DurableHistoryProvider(prune_excluded=False) + state: dict[str, Any] = {} + with bound(provider) as binding: + await history.get_messages("session", state=state) + prune_messages(provider.state.data.conversation_history, [(summary, summary.messages[0])]) + history.flush(state) + assert ids(provider) == ["current"] and binding.append_ordinal == 0 + assert [m.message_id for m in state[WORKING_BUFFER_KEY]] == ["current"] + + +@pytest.mark.parametrize("prune", [False, True]) +async def test_unloaded_empty_payload_is_not_mistaken_for_a_strategy_removal(prune: bool) -> None: + provider = _InMemoryStateProvider() + empty = DurableAgentStateMessage("user", [], message_id="empty", extension_data={"future": {"keep": [1]}}) + provider.state.data.conversation_history.extend([ + DurableAgentStateRequest("old", OLD, [empty]), + DurableAgentStateRequest("current", OLD, [stored("current", "current")]), + ]) + history = DurableHistoryProvider(prune_excluded=prune) + state: dict[str, Any] = {} + with bound(provider): + loaded = await history.get_messages("session", state=state) + assert [m.message_id for m in loaded] == ["current"] + snapshot = deepcopy(provider.state.to_dict()) + history.flush(state) + assert provider.state.to_dict() == snapshot diff --git a/python/packages/durabletask/tests/test_response_fidelity_review.py b/python/packages/durabletask/tests/test_response_fidelity_review.py new file mode 100644 index 0000000..cbf4c96 --- /dev/null +++ b/python/packages/durabletask/tests/test_response_fidelity_review.py @@ -0,0 +1,609 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Cold-delivery fidelity tests against public core constructors, without host mocks.""" + +import builtins +import importlib +import json +from copy import deepcopy +from datetime import date +from inspect import Parameter, signature +from typing import Any, cast, get_args, get_type_hints +from unittest.mock import Mock + +import pytest +from agent_framework import AgentResponse, Content, ContinuationToken, Message +from pydantic import BaseModel, ConfigDict, Field, Json, RootModel, ValidationError + +from agent_framework_durabletask._response_utils import ( + ensure_response_format, + is_terminal_agent_response, + load_agent_response, + serialize_agent_response, +) + +CORRELATION_ID = "fidelity-review" + + +class AliasCount(BaseModel): + count: int = Field(alias="aliasCount") + + +class JsonValue(BaseModel): + document: Json[list[int]] + + +class NullValue(RootModel[None]): + pass + + +def _wire(response: AgentResponse[Any]) -> dict[str, Any]: + return json.loads(json.dumps(serialize_agent_response(response), allow_nan=False)) + + +def _response(value: Any) -> AgentResponse[Any]: + return AgentResponse(messages=[Message("assistant", ["Not the structured result"])], value=value) + + +@pytest.mark.parametrize( + "value", + [AliasCount(aliasCount=7), JsonValue(document="[1,2]"), NullValue(root=None)], + ids=["validation-alias", "json-round-trip", "explicit-null"], +) +def test_structured_models_survive_cold_delivery_without_using_text(value: BaseModel) -> None: + payload = _wire(_response(value)) + assert payload["value"] == value.model_dump(mode="json", by_alias=True, round_trip=True) + loaded = load_agent_response(payload) + + ensure_response_format(type(value), CORRELATION_ID, loaded) + + assert type(loaded.value) is type(value) + assert loaded.value == value + assert loaded.text == "Not the structured result" + + +def test_nested_aliases_and_json_fields_round_trip_together() -> None: + class NestedValue(BaseModel): + child: AliasCount = Field(alias="nestedChild") + document: Json[list[int]] = Field(alias="nestedDocument") + + value = NestedValue(nestedChild={"aliasCount": 4}, nestedDocument="[2,3]") + payload = _wire(_response(value)) + assert payload["value"] == {"nestedChild": {"aliasCount": 4}, "nestedDocument": "[2,3]"} + loaded = load_agent_response(payload) + + ensure_response_format(NestedValue, CORRELATION_ID, loaded) + + assert loaded.value == value + + +@pytest.mark.parametrize("defaulted", [False, True]) +def test_distinct_serialization_aliases_use_checked_field_name_input(defaulted: bool) -> None: + default: Any = 0 if defaulted else ... + + class SeparateAliases(BaseModel): + count: int = Field( + default=default, + validation_alias="inputCount", + serialization_alias="outputCount", + ) + + class NestedValue(BaseModel): + child: SeparateAliases = Field(alias="nestedChild") + values: list[SeparateAliases] + + value = NestedValue(nestedChild={"inputCount": 7}, values=[SeparateAliases(inputCount=8)]) + payload = _wire(_response(value)) + assert payload["value"] == {"child": {"count": 7}, "values": [{"count": 8}]} + assert payload["_durable_value_by_name"] is True + loaded = load_agent_response(payload) + # An untyped delivery can cross another JSON boundary before a caller requests its model. + loaded = load_agent_response(_wire(loaded)) + + ensure_response_format(NestedValue, CORRELATION_ID, loaded) + + assert loaded.value == value + assert loaded.additional_properties == {} + + +def test_strict_json_types_are_validated_as_json_not_python_values() -> None: + class StrictValue(BaseModel): + model_config = ConfigDict(strict=True) + day: date + coordinates: tuple[int, int] + + value = StrictValue(day=date(2026, 9, 9), coordinates=(1, 2)) + loaded = load_agent_response(_wire(_response(value))) + + ensure_response_format(StrictValue, CORRELATION_ID, loaded) + + assert loaded.value == value + + +@pytest.mark.parametrize("value", [None, False, 0, "", [], {}, {"type": "text", "custom": [1]}]) +def test_retained_value_presence_is_not_truthiness(value: Any) -> None: + payload = {"type": "agent_response", "messages": [Message("assistant", ["42"]).to_dict()], "value": value} + loaded = load_agent_response(payload) + assert "value" in _wire(loaded) + assert _wire(loaded)["value"] == value + + ensure_response_format(RootModel[Any], CORRELATION_ID, loaded) + + assert isinstance(loaded.value, RootModel) + assert loaded.value.root == value + assert type(loaded.value.root) is type(value) + + +def test_explicit_null_is_not_replaced_by_valid_conflicting_text() -> None: + loaded = load_agent_response({"messages": [Message("assistant", ['{"aliasCount":7}']).to_dict()], "value": None}) + + with pytest.raises(ValidationError): + ensure_response_format(AliasCount, CORRELATION_ID, loaded) + + +def test_absent_value_uses_requested_format_instead_of_original_lazy_format() -> None: + response = AgentResponse( + messages=[Message("assistant", ['{"aliasCount":7}'])], + response_format=JsonValue, + ) + + ensure_response_format(AliasCount, CORRELATION_ID, response) + + assert response.value == AliasCount(aliasCount=7) + assert "value" not in _wire(AgentResponse()) + + +def test_matching_model_value_is_not_replaced() -> None: + value = AliasCount(aliasCount=7) + response = _response(value) + + ensure_response_format(AliasCount, CORRELATION_ID, response) + + assert response.value is value + + +def test_serializing_a_lazy_value_does_not_mutate_the_original_response() -> None: + response = AgentResponse(messages=[Message("assistant", ['{"aliasCount":7}'])], response_format=AliasCount) + before = dict(vars(response)) + before_fields = deepcopy(response.to_dict()) + + payload = _wire(response) + + assert payload["value"] == {"aliasCount": 7} + assert vars(response) == before + assert response.to_dict() == before_fields + response.messages[0].contents[0].text = '{"aliasCount":9}' + assert response.value == AliasCount(aliasCount=9) + + +def test_lazy_schema_null_is_present_without_changing_the_original_cache() -> None: + response = AgentResponse(messages=[Message("assistant", ["null"])], response_format={"type": "null"}) + before = dict(vars(response)) + + payload = _wire(response) + + assert "value" in payload and payload["value"] is None + assert vars(response) == before + loaded = load_agent_response(payload) + ensure_response_format(NullValue, CORRELATION_ID, loaded) + assert loaded.value == NullValue(root=None) + + +def test_different_model_types_use_alias_json_when_validating_a_retained_model() -> None: + class OtherCount(BaseModel): + count: int = Field(alias="aliasCount") + + response = _response(AliasCount(aliasCount=7)) + + ensure_response_format(OtherCount, CORRELATION_ID, response) + + assert type(response.value) is OtherCount + assert response.value == OtherCount(aliasCount=7) + + +def test_subclass_snapshot_has_canonical_base_fields_and_retains_raw_extras() -> None: + class CustomResponse(AgentResponse[Any]): + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.custom_payload = {"type": "text", "application_field": [1]} + + def to_dict(self, **kwargs: Any) -> dict[str, Any]: + return {"type": "custom_response", "response_id": "not-the-public-id"} + + response = CustomResponse( + messages=[Message("assistant", ["answer"])], + response_id="public-id", + agent_id="public-agent", + value=AliasCount(aliasCount=7), + additional_properties={"type": "provider", "opaque": {"answer": 42}}, + ) + payload = _wire(response) + + assert payload["type"] == "agent_response" + assert payload["response_id"] == "public-id" + assert payload["agent_id"] == "public-agent" + assert payload["custom_payload"] == response.custom_payload + assert response.to_dict()["type"] == "custom_response" + loaded = load_agent_response(payload) + assert type(loaded) is AgentResponse + assert not hasattr(loaded, "custom_payload") + assert loaded.response_id == "public-id" + assert loaded.additional_properties == response.additional_properties + assert loaded.value == {"aliasCount": 7} + + +def test_ordinary_response_payload_is_canonical_and_accepted_by_core_loader() -> None: + response = AgentResponse( + messages=[Message("assistant", ["answer"])], + response_id="public-id", + additional_properties={"future": {"opaque": [1]}}, + ) + payload = _wire(response) + + assert payload["type"] == "agent_response" + assert AgentResponse.from_dict(deepcopy(payload)).to_dict() == response.to_dict() + assert load_agent_response(payload).to_dict() == response.to_dict() + + +def test_unknown_envelope_fields_are_ignored_without_changing_raw_snapshot() -> None: + payload = { + "type": "custom_response", + "future_response": {"type": "future_type", "keep": [1]}, + "response_format": {"type": "object", "required": ["not_in_value"]}, + "messages": [ + { + "type": "custom_message", + "role": "assistant", + "future_message": [2], + "contents": [{"type": "text", "text": "answer", "future_content": [3]}], + } + ], + "value": {"type": "agent_response", "arbitrary": {"type": "text", "future": [4]}}, + "additional_properties": {"type": "content", "keep": [5]}, + } + before = deepcopy(payload) + + loaded = load_agent_response(payload) + + assert loaded.text == "answer" + assert loaded.value == before["value"] + assert not hasattr(loaded, "future_response") + assert not hasattr(loaded.messages[0], "future_message") + assert not hasattr(loaded.messages[0].contents[0], "future_content") + loaded.value["arbitrary"]["future"].append(9) + loaded.additional_properties["keep"].append(9) + loaded.messages[0].contents[0].text = "changed" + assert payload == before + + +def test_stored_type_names_never_select_or_import_python_classes(monkeypatch: pytest.MonkeyPatch) -> None: + forbidden_import = Mock(side_effect=AssertionError("Stored type names must not trigger imports")) + original_import = builtins.__import__ + + def guarded_import(name: str, *args: Any, **kwargs: Any) -> Any: + if name.startswith("untrusted"): + return forbidden_import(name) + return original_import(name, *args, **kwargs) + + payload = { + "type": "untrusted.provider.CustomResponse", + "response_format": {"type": "untrusted.provider.CustomModel"}, + "future_response": {"class": "untrusted.provider.Future"}, + "messages": [ + { + "type": "untrusted.provider.CustomMessage", + "role": "assistant", + "contents": [ + { + "type": "untrusted.provider.FutureContent", + "future_content": [1], + "additional_properties": {"opaque": [2]}, + } + ], + } + ], + } + before = deepcopy(payload) + with monkeypatch.context() as patch: + patch.setattr(builtins, "__import__", guarded_import) + patch.setattr(importlib, "import_module", forbidden_import) + loaded = load_agent_response(payload) + + forbidden_import.assert_not_called() + assert type(loaded) is AgentResponse + assert type(loaded.messages[0]) is Message + content = loaded.messages[0].contents[0] + assert type(content) is Content and content.type == "untrusted.provider.FutureContent" + assert content.additional_properties == {"opaque": [2]} + assert not hasattr(content, "future_content") + assert payload == before + + +@pytest.mark.parametrize("content_type", get_args(get_type_hints(Content.__init__)["type"])) +def test_all_public_content_kinds_tolerate_unknown_optional_envelope_fields(content_type: Any) -> None: + original = Content(content_type, additional_properties={"type": "opaque", "unknown": [1]}) + data = original.to_dict() + data["future_content"] = {"custom": True} + loaded = load_agent_response({"messages": [{"role": "assistant", "contents": [data]}]}) + + assert type(loaded.messages[0].contents[0]) is Content + assert loaded.messages[0].contents[0].to_dict() == original.to_dict() + assert data["future_content"] == {"custom": True} + + +@pytest.mark.parametrize( + ("kind", "field", "sequence"), + [ + ("function_result", "items", True), + ("search_tool_result", "items", True), + ("code_interpreter_tool_call", "inputs", True), + ("code_interpreter_tool_result", "outputs", True), + ("shell_tool_result", "outputs", True), + ("function_approval_request", "function_call", False), + ("function_approval_response", "function_call", False), + ], +) +def test_nested_framework_content_is_reconstructed_at_known_edges(kind: str, field: str, sequence: bool) -> None: + inner = {"type": "text", "text": "result", "future_inner": {"keep": 1}} + middle = {"type": "function_result", "items": [inner], "future_middle": [2]} + content = {"type": kind, field: [middle] if sequence else middle, "future_outer": [3]} + raw = {"messages": [{"role": "tool", "contents": [content]}]} + before = deepcopy(raw) + + loaded = load_agent_response(raw) + + nested = getattr(loaded.messages[0].contents[0], field) + nested = nested[0] if sequence else nested + assert type(nested) is Content + assert nested.items is not None + assert type(nested.items[0]) is Content + assert nested.items[0].text == "result" + assert raw == before + + +@pytest.mark.parametrize("field", ["arguments", "result", "output", "outputs", "additional_properties"]) +def test_application_payloads_with_framework_type_names_are_not_reconstructed(field: str) -> None: + application = {"type": "text", "contents": [{"type": "error", "custom": [1]}], "not_a_core_field": [2]} + value: Any = [application] if field == "outputs" else application + raw = {"messages": [{"role": "assistant", "contents": [{"type": "image_generation_tool_result", field: value}]}]} + + loaded = load_agent_response(raw) + + assert getattr(loaded.messages[0].contents[0], field) == value + assert not is_terminal_agent_response(loaded) + + +def test_rich_response_metadata_and_value_round_trip_independently() -> None: + application = {"type": "text", "provider_field": {"type": "error", "unknown": [1]}} + citation: Any = { + "type": "citation", + "title": "Source", + "url": "https://example.test/source", + "annotated_regions": [{"type": "text_span", "start_index": 0, "end_index": 6, "future": [1]}], + "additional_properties": application, + "future_annotation": [2], + } + response = AgentResponse( + messages=[ + Message( + "assistant", + [ + Content.from_text("answer", annotations=[citation], additional_properties=deepcopy(application)), + Content.from_text_reasoning(id="reason", text="summary", protected_data="opaque"), + Content.from_function_call("call", "lookup", arguments=deepcopy(application)), + ], + author_name="writer", + message_id="message", + additional_properties=deepcopy(application), + raw_representation=object(), + ), + Message( + "tool", + [ + Content.from_function_result( + "call", result=[Content.from_text("tool result"), Content.from_data(b"data", "image/png")] + ) + ], + ), + ], + response_id="response", + agent_id="agent", + created_at="2026-09-09T00:00:00Z", + finish_reason="stop", + usage_details={"input_token_count": 3, "output_token_count": 2, "cache_read_input_token_count": 1}, + continuation_token=cast(ContinuationToken, deepcopy(application)), + additional_properties=deepcopy(application), + raw_representation=object(), + value=AliasCount(aliasCount=7), + ) + expected = response.to_dict() + payload = _wire(response) + loaded = load_agent_response(payload) + + assert loaded.to_dict() == expected + assert loaded.messages[0].contents[0].annotations == [citation] + assert loaded.messages[1].contents[0].items == response.messages[1].contents[0].items + ensure_response_format(AliasCount, CORRELATION_ID, loaded) + assert loaded.value == AliasCount(aliasCount=7) + assert loaded.to_dict() == expected + payload["additional_properties"]["provider_field"]["unknown"].append(9) + assert response.additional_properties == application + assert loaded.additional_properties == application + assert "raw_representation" not in payload + assert "raw_representation" not in payload["messages"][0] + + +def test_canonical_response_projection_tracks_public_constructor_fields() -> None: + response = AgentResponse(response_id="response", agent_id="agent", value=AliasCount(aliasCount=7)) + payload = _wire(response) + loaded = load_agent_response(payload) + # Derive the category from core's public signature rather than a copied response-field list. + for name, parameter in signature(AgentResponse).parameters.items(): + if parameter.kind not in (Parameter.POSITIONAL_OR_KEYWORD, Parameter.KEYWORD_ONLY): + continue + if name in ("value", "response_format", "raw_representation"): + continue + assert getattr(loaded, name) == getattr(response, name), name + + +@pytest.mark.parametrize("status", ["error", "already_completed"]) +def test_explicit_terminal_status_skips_typed_parsing_even_without_error_content(status: str) -> None: + response = AgentResponse( + messages=[Message("assistant", ["not JSON"])], + response_format=AliasCount, + additional_properties={"durable_status": status}, + ) + assert is_terminal_agent_response(response) + loaded = load_agent_response(_wire(response)) + + ensure_response_format(AliasCount, CORRELATION_ID, loaded) + + assert loaded.value is None + assert loaded.additional_properties == response.additional_properties + + +def test_accepted_acknowledgement_skips_validation_but_is_not_a_terminal_failure() -> None: + response = AgentResponse( + messages=[Message("assistant", ["Request accepted"])], + response_format=AliasCount, + additional_properties={"durable_status": "accepted"}, + ) + loaded = load_agent_response(_wire(response)) + + ensure_response_format(AliasCount, CORRELATION_ID, loaded) + + assert not is_terminal_agent_response(loaded) + assert loaded.value is None + + +@pytest.mark.parametrize("role", ["assistant", "system", "user", "developer", "tool"]) +def test_direct_legacy_errors_are_terminal_only_outside_tool_messages(role: str) -> None: + response: AgentResponse[Any] = AgentResponse( + messages=[Message(role, [Content.from_error(message="failure")])], + value={"aliasCount": 7}, + ) + loaded = load_agent_response(_wire(response)) + + assert is_terminal_agent_response(loaded) is (role != "tool") + ensure_response_format(AliasCount, CORRELATION_ID, loaded) + if role == "tool": + assert loaded.value == AliasCount(aliasCount=7) + else: + assert loaded.value == {"aliasCount": 7} + + +@pytest.mark.parametrize("valid", [False, True]) +@pytest.mark.parametrize("nested", [False, True]) +def test_recoverable_tool_errors_do_not_bypass_success_validation(valid: bool, nested: bool) -> None: + error = Content.from_error(message="retryable lookup failure") + content = Content.from_function_result("call", result=[error]) if nested else error + response = AgentResponse( + messages=[ + Message("tool", [content]), + Message("assistant", ['{"aliasCount":7}' if valid else "invalid structured result"]), + ] + ) + loaded = load_agent_response(_wire(response)) + + assert not is_terminal_agent_response(loaded) + if valid: + ensure_response_format(AliasCount, CORRELATION_ID, loaded) + assert loaded.value == AliasCount(aliasCount=7) + else: + with pytest.raises(ValueError): + ensure_response_format(AliasCount, CORRELATION_ID, loaded) + + +@pytest.mark.parametrize("version", [None, True, 0, 2, "1", {}, []]) +def test_unknown_or_malformed_codec_versions_fail_without_mutating_input(version: Any) -> None: + payload = {"type": "agent_response", "_durable_response_version": version, "future": [1]} + before = deepcopy(payload) + + with pytest.raises(ValueError, match="Unsupported durable response version"): + load_agent_response(payload) + + assert payload == before + + +def test_optional_supported_delivery_version_is_still_readable() -> None: + response = AgentResponse(messages=[Message("assistant", ["marked snapshot"])]) + payload = _wire(response) + payload["_durable_response_version"] = 1 + before = deepcopy(payload) + + assert load_agent_response(payload).to_dict() == response.to_dict() + assert payload == before + + +@pytest.mark.parametrize("payload", [[], "response", 1]) +def test_loader_rejects_unsupported_input_types(payload: Any) -> None: + with pytest.raises(TypeError, match="Unsupported type"): + load_agent_response(payload) + + +def test_loader_preserves_existing_instances_and_rejects_absent_input() -> None: + response = AgentResponse() + assert load_agent_response(response) is response + with pytest.raises(ValueError, match="cannot be None"): + load_agent_response(None) + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"invalid": "format"}, + {"value": 42}, + {"response_id": "not-an-envelope"}, + {"type": "text", "text": "not a response"}, + {"type": "custom_response"}, + {"messages": None}, + {"type": "custom_response", "messages": None}, + ], +) +def test_loader_rejects_nonresponse_mappings_without_mutating_input(payload: dict[str, Any]) -> None: + before = deepcopy(payload) + + with pytest.raises(ValueError, match="requires a response type or messages"): + load_agent_response(payload) + + assert payload == before + + +@pytest.mark.parametrize("response_type", [None, "", False, 1, [], {}]) +def test_loader_rejects_corrupt_response_types_even_with_valid_messages(response_type: Any) -> None: + with pytest.raises(ValueError, match="type must be a non-empty string"): + load_agent_response({"type": response_type, "messages": []}) + + +@pytest.mark.parametrize("response_type", [None, "agent_response", "custom_response"]) +@pytest.mark.parametrize("empty", [False, True]) +def test_loader_accepts_response_like_messages_with_or_without_type(response_type: str | None, empty: bool) -> None: + response = AgentResponse(messages=[] if empty else [Message("assistant", ["internal helper"])]) + payload = response.to_dict() + payload.pop("type", None) + if response_type is not None: + payload["type"] = response_type + + loaded = load_agent_response(payload) + + assert type(loaded) is AgentResponse + assert loaded.to_dict() == response.to_dict() + + +def test_loader_accepts_canonical_empty_response_without_messages() -> None: + assert load_agent_response({"type": "agent_response"}).to_dict() == AgentResponse().to_dict() + + +@pytest.mark.parametrize( + "messages", [{}, {"role": "assistant"}, "", "not messages", b"", 0, False, ["not a message"], [{"contents": []}]] +) +def test_loader_rejects_malformed_message_envelopes(messages: Any) -> None: + with pytest.raises(TypeError): + load_agent_response({"messages": messages}) + + +@pytest.mark.parametrize("content", [{"text": "missing type"}, {"type": None}, {"type": ""}, {"type": 1}]) +def test_loader_rejects_corrupt_content_type_instead_of_guessing_a_framework_shape(content: dict[str, Any]) -> None: + with pytest.raises(ValueError, match="requires 'type'"): + load_agent_response({"messages": [{"role": "assistant", "contents": [content]}]}) diff --git a/python/packages/durabletask/tests/test_retention.py b/python/packages/durabletask/tests/test_retention.py new file mode 100644 index 0000000..7c2a8ce --- /dev/null +++ b/python/packages/durabletask/tests/test_retention.py @@ -0,0 +1,775 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Tests for retention (ADR-0032, "Retention"). + +An explicit pressure budget evicts eligible transcript history independently of eager compaction +pruning. An exclusion made for token cost is not consent to delete the record, and an unreachable +protected floor reports capacity failure without deleting state. +""" + +import json +from collections.abc import AsyncIterator +from copy import deepcopy +from datetime import datetime, timedelta, timezone +from typing import Any, cast, get_args + +import pytest +from agent_framework import ( + Agent, + AgentResponse, + BaseChatClient, + ChatResponse, + ChatResponseUpdate, + Content, + Message, + ResponseStream, +) + +from agent_framework_durabletask import ( + AgentEntity, + AgentEntityStateProviderMixin, + DurableAgentState, + DurableAgentStateErrorResponse, + DurableAgentStateMessage, + DurableAgentStateRequest, + DurableAgentStateResponse, +) +from agent_framework_durabletask._retention import ( + DELIVERY_WINDOW_SECONDS, + HIGH_WATERMARK, + LOW_WATERMARK, + RetentionMode, + StateCapacityError, + _token_budget, + enforce_budget, + prunes_excluded, +) + +BUDGET = 40_000 +"""Small enough to keep these tests fast, large enough to hold a realistic conversation.""" + + +def _state(turns: int, *, chars: int = 400, excluded_before: int = 0, excluded_recent: int = 0) -> DurableAgentState: + """Build legacy transcript-delivered state with the given number of user/assistant turns. + + Args: + turns: How many exchanges to record. + chars: Size of each message's text. + + Keyword Args: + excluded_before: Mark this many leading messages as compaction-excluded, as a user's own + sliding window would. + excluded_recent: Mark this many of the most recent messages as compaction-excluded, as a + tool-result strategy can do without touching the oldest turns. + + Returns: + The populated state. + """ + # These manually appended responses use legacy history lookup. Version 2 fixtures must + # record independent mailbox results instead of treating transcript entries as delivery. + state = DurableAgentState(schema_version="1.2.0") + now = datetime.now(tz=timezone.utc) + # Space legacy turns a minute apart so their delivery windows have elapsed. Tests of live + # delivery explicitly refresh timestamps rather than depending on the test's running time. + marked = 0 + for index in range(turns): + occurred_at = now - timedelta(minutes=turns - index) + request = DurableAgentStateRequest( + correlation_id=f"c{index}", + created_at=occurred_at, + messages=[ + DurableAgentStateMessage.from_chat_message( + Message(role="user", contents=["u" * chars], message_id=f"u{index}") + ) + ], + ) + response = DurableAgentStateResponse( + correlation_id=f"c{index}", + created_at=occurred_at, + messages=[ + DurableAgentStateMessage.from_chat_message( + Message(role="assistant", contents=["a" * chars], message_id=f"a{index}") + ) + ], + ) + for entry in (request, response): + for stored in entry.messages: + if marked < excluded_before: + stored.extension_data = {"_excluded": True, "_excluded_reason": "sliding_window"} + marked += 1 + state.data.conversation_history.extend([request, response]) + + if excluded_recent: + stored_messages = [m for entry in state.data.conversation_history for m in entry.messages] + for stored in stored_messages[-excluded_recent:]: + stored.extension_data = {"_excluded": True, "_excluded_reason": "tool_result_compaction"} + return state + + +def _size(state: DurableAgentState) -> int: + return len(json.dumps(state.to_dict())) + + +def _message_ids(state: DurableAgentState) -> list[str]: + return [m.message_id or "" for entry in state.data.conversation_history for m in entry.messages] + + +class TestRetentionModes: + """The mode decides whether an exclusion may become a deletion.""" + + def test_only_follow_compaction_prunes_on_write(self) -> None: + modes = get_args(RetentionMode) + assert set(modes) == {"keep_all", "follow_compaction"} + for mode in modes: + assert prunes_excluded(mode) is (mode == "follow_compaction") + + def test_auto_is_rejected(self) -> None: + with pytest.raises(ValueError, match="retention"): + prunes_excluded(cast(Any, "auto")) + + def test_the_defaults_do_not_enable_deletion(self) -> None: + from agent_framework_durabletask import DurableAIAgentWorker + + worker = DurableAIAgentWorker(cast(Any, object())) + assert worker._retention == "keep_all" + assert worker._max_state_bytes is None + + +class TestBudgetEnforcement: + """Nothing happens until state is genuinely close to the limit.""" + + async def test_below_the_watermark_nothing_is_touched(self) -> None: + state = _state(turns=4) + before = state.to_json() + + removed = await enforce_budget(state, max_state_bytes=BUDGET) + + assert removed == 0 + assert state.to_json() == before + + async def test_over_the_watermark_evicts_to_the_low_watermark(self) -> None: + state = _state(turns=60) + assert _size(state) > BUDGET * HIGH_WATERMARK, "the fixture must start over the trigger" + + removed = await enforce_budget(state, max_state_bytes=BUDGET) + + assert removed > 0 + assert _size(state) <= BUDGET * LOW_WATERMARK, "eviction did not reach the low watermark" + + async def test_the_newest_turn_survives(self) -> None: + """Evicting the turn that just happened would defeat the point of running it.""" + state = _state(turns=60) + + await enforce_budget(state, max_state_bytes=BUDGET) + + assert _message_ids(state)[-2:] == ["u59", "a59"] + + async def test_eviction_is_hysteretic(self) -> None: + """Evicting to just under the trigger would evict again on every following turn.""" + state = _state(turns=60) + await enforce_budget(state, max_state_bytes=BUDGET) + + second = await enforce_budget(state, max_state_bytes=BUDGET) + + assert second == 0, "a second pass evicted again immediately, so there is no headroom" + + async def test_pressure_eviction_does_not_require_compaction_exclusions(self) -> None: + """An explicit budget can evict old groups without opting into eager pruning.""" + state = _state(turns=60) + + assert await enforce_budget(state, max_state_bytes=BUDGET) > 0 + + @pytest.mark.parametrize("turns", [0, 10]) + async def test_metadata_floor_fails_without_mutating_state(self, turns: int) -> None: + state = _state(turns=turns) + state.data.session = {"state": {"pending_approvals": ["p" * (BUDGET * 2)]}} + state.data.ingested_positions = {"source": 7} + before = state.to_json() + + with pytest.raises(StateCapacityError) as error: + await enforce_budget(state, max_state_bytes=BUDGET) + + assert error.value.floor_bytes > BUDGET + assert state.to_json() == before + + +class TestExclusionsAreNotConsentToDelete: + """A context decision must not silently become a storage decision.""" + + async def test_a_user_s_exclusions_survive_eviction(self) -> None: + """The budget is measured over a detached copy, so stored annotations are untouched. + + Exclusions are placed on recent messages here, which a tool-result strategy does, so they + sit inside the window eviction keeps. Had the annotation itself been the criterion they + would have gone regardless of where they were. + """ + state = _state(turns=60, excluded_recent=6) + + await enforce_budget(state, max_state_bytes=BUDGET) + + surviving = [ + stored + for entry in state.data.conversation_history + for stored in entry.messages + if (stored.extension_data or {}).get("_excluded") + ] + assert surviving, "every excluded message was evicted, so exclusion was treated as consent" + assert all((s.extension_data or {}).get("_excluded_reason") == "tool_result_compaction" for s in surviving) + + async def test_eviction_is_not_limited_to_what_compaction_excluded(self) -> None: + """The budget is computed over everything stored, not just the included messages. + + A user's own window can mark almost everything excluded. If those exclusions were left in + place the strategy would see a tiny included set, conclude it was already under budget, and + evict nothing while state kept growing. + """ + state = _state(turns=60, excluded_before=110) + + removed = await enforce_budget(state, max_state_bytes=BUDGET) + + assert removed > 0, "prior exclusions hid the real size and nothing was evicted" + + async def test_every_surviving_exclusion_keeps_its_annotation(self) -> None: + """Measuring the budget must not strip annotations off the messages it measured. + + To size the conversation, eviction clears ``_excluded`` on the message copies it hands to + the strategy. That is only safe while those really are copies. If the copy ever shared its + annotations with stored state, the clear would erase compaction's work from storage. + + Asserting merely that *some* exclusion survives is too weak to catch that: the newest + exchange is never a candidate, so its annotations would survive either way. This checks + every message that outlived eviction, which includes ones that were candidates. + """ + state = _state(turns=60, excluded_recent=40) + excluded_before_run = { + stored.message_id + for entry in state.data.conversation_history + for stored in entry.messages + if (stored.extension_data or {}).get("_excluded") + } + + removed = await enforce_budget(state, max_state_bytes=BUDGET) + assert removed > 0, "nothing was evicted, so the measuring path never ran" + + still_stored = { + stored.message_id: stored for entry in state.data.conversation_history for stored in entry.messages + } + survivors = excluded_before_run & still_stored.keys() + assert survivors, "every excluded message was evicted, so this proves nothing" + + stripped = [ + message_id + for message_id in survivors + if not (still_stored[message_id].extension_data or {}).get("_excluded") + ] + assert not stripped, f"eviction erased stored compaction annotations from {len(stripped)} message(s)" + + +class TestSingleOversizedTurn: + """Retention cannot save a conversation whose newest turn alone exceeds the budget.""" + + async def test_the_current_turn_is_never_evicted(self) -> None: + """An unretainable current exchange reports capacity failure without deleting state.""" + state = _state(turns=1, chars=BUDGET * 2) + before = state.to_json() + + with pytest.raises(StateCapacityError) as error: + await enforce_budget(state, max_state_bytes=BUDGET) + + assert error.value.floor_bytes == len(before) + assert state.to_json() == before + assert _message_ids(state) == ["u0", "a0"], "the turn that just ran was evicted" + + async def test_an_oversized_newest_turn_does_not_take_the_history_with_it(self) -> None: + state = _state(turns=10) + state.data.conversation_history[-1].messages = [ + DurableAgentStateMessage.from_chat_message( + Message(role="assistant", contents=["a" * (BUDGET * 2)], message_id="a9") + ) + ] + before = state.to_json() + + with pytest.raises(StateCapacityError): + await enforce_budget(state, max_state_bytes=BUDGET) + + assert state.to_json() == before, "capacity failure must preserve the entire original history" + + +class TestAResponseIsNotEvictedBeforeItsCallerReadsIt: + """A legacy caller reads its response by correlation id from transcript entries. + + Nothing tells the entity that a response was collected, so a turn completing is not permission + to delete the previous one. Evicting a response somebody is still polling for turns a run that + succeeded into a client timeout. + """ + + async def test_a_recent_response_is_not_evicted(self) -> None: + """The turn is early in the conversation, so oldest-first eviction reaches it. + + That is the whole point. Picking a recent turn would prove nothing, because eviction would + never have got that far and the test would pass with no protection at all. + """ + state = _state(turns=60) + # Second oldest turn, so it is squarely inside what eviction removes, but it completed + # seconds ago, so its caller may still be polling for it. + early = state.data.conversation_history[2:4] + for entry in early: + entry.created_at = datetime.now(tz=timezone.utc) + correlation = early[0].correlation_id + assert correlation is not None + original = state.try_get_agent_response(correlation) + assert original is not None + original_payload = deepcopy(original.to_dict()) + + removed = await enforce_budget(state, max_state_bytes=BUDGET) + + assert removed > 0, "nothing was evicted, so this proves nothing" + retained = state.try_get_agent_response(correlation) + assert retained is not None, "a recent response was evicted before its caller could read it" + assert retained.to_dict() == original_payload + + async def test_an_old_response_is_still_evictable(self) -> None: + """Protection has to expire, or a long conversation could never be trimmed at all.""" + state = _state(turns=60) + assert state.try_get_agent_response("c0") is not None + + removed = await enforce_budget(state, max_state_bytes=BUDGET) + + assert removed > 0 + assert state.try_get_agent_response("c0") is None, "an ancient response was kept forever" + + async def test_capacity_failure_preserves_every_recent_response(self) -> None: + """A full delivery window reports capacity failure instead of sacrificing responses.""" + state = _state(turns=60) + # Every turn happened just now, which is what a busy session looks like. + for entry in state.data.conversation_history: + entry.created_at = datetime.now(tz=timezone.utc) + before = state.to_json() + + with pytest.raises(StateCapacityError) as error: + await enforce_budget(state, max_state_bytes=BUDGET) + + assert error.value.floor_bytes == len(before) + assert state.to_json() == before + + async def test_a_failed_turn_is_protected_too(self) -> None: + """The caller waiting on a failed turn still needs to be told it failed.""" + state = _state(turns=60) + failure = DurableAgentStateErrorResponse( + correlation_id="boom", + created_at=datetime.now(tz=timezone.utc), + messages=[ + DurableAgentStateMessage.from_chat_message( + Message(role="assistant", contents=["it broke"], message_id="err0") + ) + ], + ) + # Early in the conversation, where eviction would otherwise reach it. + state.data.conversation_history.insert(2, failure) + + removed = await enforce_budget(state, max_state_bytes=BUDGET) + + assert removed > 0 + assert state.try_get_agent_response("boom") is not None + + +class TestMailboxDeliverySurvivesTranscriptEviction: + async def test_original_result_is_retained_until_expiry_and_completion_outlives_it(self) -> None: + state = DurableAgentState() + state.data.conversation_history = _state(turns=60).data.conversation_history + now = datetime.now(tz=timezone.utc) + for entry in state.data.conversation_history[:2]: + entry.created_at = now + response = AgentResponse( + messages=[Message("assistant", ["a" * 400], message_id="a0")], + additional_properties={"delivery": {"original": True}}, + ) + state.record_response("c0", response, delivery_window_seconds=DELIVERY_WINDOW_SECONDS, now=now) + mailbox = deepcopy(state.data.response_mailbox) + completed = deepcopy(state.data.completed_correlations) + assert _size(state) > BUDGET * HIGH_WATERMARK + + removed = await enforce_budget(state, max_state_bytes=BUDGET) + + assert removed > 0 + assert not {"u0", "a0"} & set(_message_ids(state)), "the recent transcript copy was not evicted" + restored = DurableAgentState.from_json(state.to_json()) + expiry = now + timedelta(seconds=DELIVERY_WINDOW_SECONDS) + restored.expire_responses(now=expiry - timedelta(microseconds=1)) + assert restored.data.response_mailbox == mailbox + assert restored.data.completed_correlations == completed + retained = restored.try_get_agent_response("c0") + assert retained is not None + assert retained.to_dict() == response.to_dict() + + restored.expire_responses(now=expiry) + + assert restored.data.response_mailbox == {} + assert restored.data.completed_correlations == completed + expired = DurableAgentState.from_json(restored.to_json()).try_get_agent_response("c0") + assert expired is not None + assert expired.additional_properties["durable_status"] == "already_completed" + assert expired.additional_properties["correlation_id"] == "c0" + assert expired.messages[0].contents[0].error_code == "response_expired" + + +def _tool_state(turns: int, *, chars: int = 400) -> DurableAgentState: + """Build a history of tool calls, which carry real bytes but no ``message.text``. + + This is the shape that broke the budget. A function call serializes to as much storage as + prose of the same length, but reading ``.text`` off it returns an empty string. + """ + state = DurableAgentState(schema_version="1.2.0") + now = datetime.now(tz=timezone.utc) + for index in range(turns): + occurred_at = now - timedelta(minutes=turns - index) + call: dict[str, Any] = { + "type": "function_call", + "call_id": f"call{index}", + "name": "lookup", + "arguments": json.dumps({"query": "q" * chars}), + } + result: dict[str, Any] = { + "type": "function_call", + "call_id": f"call{index}", + "name": "lookup", + "arguments": json.dumps({"result": "r" * chars}), + } + state.data.conversation_history.extend([ + DurableAgentStateRequest( + correlation_id=f"c{index}", + created_at=occurred_at, + messages=[ + DurableAgentStateMessage.from_chat_message( + Message(role="user", contents=[call], message_id=f"u{index}") + ) + ], + ), + DurableAgentStateResponse( + correlation_id=f"c{index}", + created_at=occurred_at, + messages=[ + DurableAgentStateMessage.from_chat_message( + Message(role="assistant", contents=[result], message_id=f"a{index}") + ) + ], + ), + ]) + return state + + +class TestTheBudgetDoesNotAssumeProse: + """A conversation of tool calls must be budgeted like any other. + + The budget converts bytes into tokens. Deriving that conversion from ``message.text`` made it + depend on the *kind* of content rather than its size, and a function call has no text at all. + A tool-only history therefore produced a budget of one token and evicted everything it was + permitted to touch, rather than evicting down to the watermark like any other conversation. + """ + + async def test_a_tool_only_history_keeps_roughly_what_prose_keeps(self) -> None: + prose = _state(turns=40) + tools = _tool_state(turns=40) + + await enforce_budget(prose, max_state_bytes=BUDGET) + await enforce_budget(tools, max_state_bytes=BUDGET) + + prose_left = len(_message_ids(prose)) + tools_left = len(_message_ids(tools)) + # Not identical, since the two shapes do not serialize to the same size per message, but + # the same order of magnitude. Before the fix this was 8 against 1. + assert tools_left > 2, "the budget retained only the protected newest exchange" + assert abs(prose_left - tools_left) <= max(2, prose_left // 2) + + async def test_a_tool_only_history_is_evicted_down_to_the_watermark(self) -> None: + state = _tool_state(turns=40) + + removed = await enforce_budget(state, max_state_bytes=BUDGET) + + assert removed > 0 + assert _size(state) <= BUDGET * LOW_WATERMARK + + async def test_the_budget_scales_with_bytes_not_text(self) -> None: + """Two histories of similar serialized size get similar budgets.""" + prose = _state(turns=10) + tools = _tool_state(turns=10) + + def budget_for(state: DurableAgentState) -> int: + origins = [(entry, m) for entry in state.data.conversation_history for m in entry.messages] + size = _size(state) + evictable = sum(len(json.dumps(m.to_dict())) for _, m in origins) + return _token_budget(origins, serialized_size=size, evictable_bytes=evictable, target_bytes=size // 2) + + prose_budget = budget_for(prose) + tools_budget = budget_for(tools) + + assert prose_budget > 1 + # The old formula gave exactly 1 here, whatever the tool payload weighed. + assert tools_budget > 1 + assert 0.4 < (tools_budget / prose_budget) < 2.5 + + +class TestTheAgentsInstructionsSurviveTheBudget: + """A system message is never evicted, however tight the budget gets. + + Core protects system groups in its first fallback but then has a *strict* fallback whose whole + job is to evict them when anchors alone exceed the budget. Relying on core's protection + therefore holds only until the budget is small enough to matter. Keeping system messages out + of the candidate set entirely makes them unevictable, and their bytes count as a floor. + """ + + def _with_system(self, turns: int, *, chars: int = 400) -> DurableAgentState: + state = _state(turns=turns, chars=chars) + anchor = DurableAgentStateRequest( + correlation_id="system-anchor", + created_at=datetime.now(tz=timezone.utc) - timedelta(minutes=turns + 5), + messages=[ + DurableAgentStateMessage.from_chat_message( + Message(role="system", contents=["S" * chars], message_id="system-0") + ) + ], + ) + state.data.conversation_history.insert(0, anchor) + return state + + def _system_count(self, state: DurableAgentState) -> int: + return sum(1 for entry in state.data.conversation_history for m in entry.messages if m.role == "system") + + async def test_the_system_message_survives_a_comfortable_budget(self) -> None: + state = self._with_system(turns=30) + + await enforce_budget(state, max_state_bytes=40_000) + + assert self._system_count(state) == 1 + + async def test_the_system_message_survives_a_tight_budget(self) -> None: + state = self._with_system(turns=30) + + removed = await enforce_budget(state, max_state_bytes=6_000) + + assert removed > 0 + assert self._system_count(state) == 1 + + async def test_the_system_message_survives_a_budget_it_cannot_fit(self) -> None: + """An unreachable protected floor leaves instructions and the entire history intact.""" + state = self._with_system(turns=30) + before = state.to_json() + + with pytest.raises(StateCapacityError): + await enforce_budget(state, max_state_bytes=1_500) + + assert state.to_json() == before + assert self._system_count(state) == 1 + + async def test_ordinary_messages_are_still_evicted_around_it(self) -> None: + state = self._with_system(turns=30) + + removed = await enforce_budget(state, max_state_bytes=6_000) + + surviving = _message_ids(state) + assert removed > 0 + assert "system-0" in surviving + + +class TestEvictionLeavesEvidence: + """A conversation that has lost content must say so in the state, not only in a log. + + Eviction is lossy and performed by the runtime rather than by the user. A warning is only + evidence to whoever happened to be watching at the time, which is nobody by the point someone + asks why an answer lost context. + """ + + async def test_nothing_is_recorded_when_nothing_is_evicted(self) -> None: + state = _state(turns=2) + + await enforce_budget(state, max_state_bytes=BUDGET) + + assert state.data.truncation is None + + async def test_eviction_is_recorded(self) -> None: + state = _state(turns=60) + + removed = await enforce_budget(state, max_state_bytes=BUDGET) + + assert removed > 0 + assert state.data.truncation is not None + assert state.data.truncation["evictedMessageCount"] == removed + assert state.data.truncation["firstEvictedAt"] + assert state.data.truncation["lastEvictedAt"] + + async def test_the_count_accumulates_across_evictions(self) -> None: + state = _state(turns=60) + + first = await enforce_budget(state, max_state_bytes=BUDGET) + for index in range(60, 120): + occurred_at = datetime.now(tz=timezone.utc) - timedelta(minutes=200 - index) + state.data.conversation_history.append( + DurableAgentStateRequest( + correlation_id=f"c{index}", + created_at=occurred_at, + messages=[ + DurableAgentStateMessage.from_chat_message( + Message(role="user", contents=["u" * 400], message_id=f"u{index}") + ) + ], + ) + ) + second = await enforce_budget(state, max_state_bytes=BUDGET) + + assert second > 0 + assert state.data.truncation is not None + assert state.data.truncation["evictedMessageCount"] == first + second + + async def test_the_record_survives_a_round_trip(self) -> None: + state = _state(turns=60) + + await enforce_budget(state, max_state_bytes=BUDGET) + restored = DurableAgentState.from_dict(json.loads(json.dumps(state.to_dict()))) + + assert restored.data.truncation == state.data.truncation + + +class TestStateShape: + """Eviction must leave durable state usable.""" + + async def test_bare_transcript_entries_emptied_by_eviction_are_removed(self) -> None: + state = _state(turns=60) + + removed = await enforce_budget(state, max_state_bytes=BUDGET) + + assert removed > 0 + assert all(entry.messages for entry in state.data.conversation_history) + + async def test_state_still_round_trips(self) -> None: + state = _state(turns=60) + + await enforce_budget(state, max_state_bytes=BUDGET) + + restored: Any = DurableAgentState.from_dict(state.to_dict()) + assert _message_ids(restored) == _message_ids(state) + + async def test_nothing_is_evicted_from_an_empty_conversation(self) -> None: + assert await enforce_budget(DurableAgentState(), max_state_bytes=BUDGET) == 0 + + +def test_watermarks_leave_room_to_work() -> None: + """The gap between them is what stops eviction running on every turn.""" + assert 0 < LOW_WATERMARK < HIGH_WATERMARK < 1 + + +class _VerboseClient(BaseChatClient): + """A client whose answers are long enough to reach the budget in a handful of turns.""" + + def __init__(self, *, reply_chars: int = 4_000) -> None: + super().__init__() + self._reply_chars = reply_chars + + def _inner_get_response(self, *, messages: Any, stream: bool, options: Any, **kwargs: Any) -> Any: + del options, kwargs + # Keyed off the question rather than a counter, so a retried call answers the same thing. + asked = next( + (m.text for m in reversed(list(messages)) if str(getattr(m.role, "value", m.role)) == "user"), + "?", + ) + body = f"answering:{asked} " + ("x" * self._reply_chars) + if stream: + + async def _updates() -> AsyncIterator[ChatResponseUpdate]: + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text(text=body)]) + + return ResponseStream(_updates(), finalizer=ChatResponse.from_updates) + + async def _response() -> ChatResponse: + return ChatResponse(messages=[Message(role="assistant", contents=[body])]) + + return _response() + + +class _EntityState(AgentEntityStateProviderMixin): + def __init__(self) -> None: + self._state_dict: dict[str, Any] = {} + + def _get_state_dict(self) -> dict[str, Any]: + return self._state_dict + + def _set_state_dict(self, state: dict[str, Any]) -> None: + # The real provider hands state to the SDK, which serializes it eagerly. + self._state_dict = json.loads(json.dumps(state)) + + def _get_session_id_from_entity(self) -> str: + return "retention-e2e" + + +class TestTheWholeLoopStaysUnderBudget: + """Drives the real entity, not just enforce_budget, because the value is in the wiring.""" + + LIMIT = 60_000 + TURNS = 20 + + async def _drive(self, **entity_kwargs: Any) -> tuple[_EntityState, list[str]]: + client = _VerboseClient() + agent = Agent(client=cast(Any, client), name="verbose") + provider = _EntityState() + entity = AgentEntity(agent, state_provider=provider, **entity_kwargs) + budget = entity_kwargs.get("max_state_bytes") + + replies: list[str] = [] + for turn in range(self.TURNS): + correlation_id = f"corr-{turn}" + result = await entity.run({"message": f"question {turn}", "correlationId": correlation_id}) + persisted = DurableAgentState.from_dict(provider._get_state_dict()) + polled = persisted.try_get_agent_response(correlation_id) + assert polled is not None + assert polled.to_dict() == result.to_dict() + replies.append(polled.text) + assert set(persisted.data.response_mailbox) == {correlation_id} + assert set(persisted.data.completed_correlations) == {f"corr-{index}" for index in range(turn + 1)} + if budget is not None: + assert _size(persisted) < int(budget * HIGH_WATERMARK) + + if turn < self.TURNS - 1: + # Simulate the next operation arriving after delivery expires, but only after + # polling this result. Let the entity remove the payload on its next operation; + # neither transcript timestamps nor completion receipts are changed here. + persisted.data.response_mailbox[correlation_id]["expiresAt"] = ( + datetime.now(tz=timezone.utc) - timedelta(seconds=1) + ).isoformat() + provider.replace_cached_state(persisted) + provider.persist_state() + return provider, replies + + @pytest.mark.parametrize("budget", [BUDGET, LIMIT]) + async def test_keep_all_with_a_budget_stays_bounded_across_many_turns(self, budget: int) -> None: + provider, _ = await self._drive(retention="keep_all", max_state_bytes=budget) + assert len(json.dumps(provider._get_state_dict())) <= budget + + async def test_follow_compaction_falls_back_to_pressure_eviction(self) -> None: + """With nothing to prune, only the shared pressure fallback can bound this run.""" + provider, _ = await self._drive(retention="follow_compaction", max_state_bytes=self.LIMIT) + state = DurableAgentState.from_dict(provider._get_state_dict()) + + assert len(json.dumps(provider._get_state_dict())) <= self.LIMIT + assert 2 <= len(_message_ids(state)) < self.TURNS * 2 + + async def test_every_turn_still_gets_its_own_answer(self) -> None: + """Eviction must not disturb the response the caller is waiting on.""" + _, replies = await self._drive(max_state_bytes=self.LIMIT) + assert [r.split(" x")[0] for r in replies] == [f"answering:question {i}" for i in range(self.TURNS)] + + async def test_history_is_actually_trimmed_not_just_small(self) -> None: + """Without this the bounded assertion above could pass for the wrong reason.""" + provider, _ = await self._drive(max_state_bytes=self.LIMIT) + state = DurableAgentState.from_dict(provider._get_state_dict()) + # Metadata-only envelopes are protected state, not retained transcript messages. + kept = len(_message_ids(state)) + assert 2 <= kept < self.TURNS * 2 + assert state.data.truncation is not None + assert state.data.truncation["evictedMessageCount"] == self.TURNS * 2 - kept + + async def test_keep_all_without_a_budget_lets_it_grow_past_the_limit(self) -> None: + """Proves the run is genuinely over budget, so the bounded case is a real result.""" + provider, _ = await self._drive(retention="keep_all", max_state_bytes=None) + state = DurableAgentState.from_dict(provider._get_state_dict()) + assert len(json.dumps(provider._get_state_dict())) > self.LIMIT + assert len(_message_ids(state)) == self.TURNS * 2 + assert state.data.truncation is None diff --git a/python/packages/durabletask/tests/test_retention_registration_dt.py b/python/packages/durabletask/tests/test_retention_registration_dt.py new file mode 100644 index 0000000..ef49ca0 --- /dev/null +++ b/python/packages/durabletask/tests/test_retention_registration_dt.py @@ -0,0 +1,313 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Registration-time validation and forwarding through the worker's real entity factory.""" + +from enum import Enum +from inspect import signature +from typing import Any, get_args +from unittest.mock import Mock, patch + +import pytest +from agent_framework import Agent, AgentExecutor, Executor, InMemoryHistoryProvider, WorkflowExecutor +from durabletask.azuremanaged.worker import DurableTaskSchedulerWorker + +import agent_framework_durabletask as durabletask +from agent_framework_durabletask import ( + DEFAULT_MAX_STATE_BYTES, + DEFAULT_RETENTION, + DELIVERY_WINDOW_SECONDS, + DTS_MAX_STATE_BYTES, + HIGH_WATERMARK, + INHERIT, + LOW_WATERMARK, + DurableAIAgentWorker, + Inherit, + RetentionMode, + StateBudgetOverride, + _configuration, + resolve_state_budget_override, + validate_history_providers, +) + + +def _agent(name: str = "assistant", *, ambiguous_history: bool = False) -> Agent: + client: Any = Mock(additional_properties={}, STORES_BY_DEFAULT=False) + providers = ( + [InMemoryHistoryProvider(source_id="first"), InMemoryHistoryProvider(source_id="second")] + if ambiguous_history + else [InMemoryHistoryProvider(source_id="primary")] + ) + return Agent(client=client, name=name, context_providers=providers) + + +def _workflow(name: str, *agents: Agent, child: Mock | None = None) -> Mock: + executors: dict[str, Mock] = {} + for index, agent in enumerate(agents): + node = Mock(spec=AgentExecutor) + node.id = f"node{index}" + node.agent = agent + executors[node.id] = node + if child is not None: + nested = Mock(spec=WorkflowExecutor) + nested.id = "child" + nested.workflow = child + executors[nested.id] = nested + activity = Mock(spec=Executor) + activity.id = "activity" + executors[activity.id] = activity + workflow = Mock() + workflow.name = name + workflow.executors = executors + return workflow + + +def _consumer_settings(grpc_worker: Mock, index: int = 0) -> dict[str, Any]: + entity_class = grpc_worker.add_entity.call_args_list[index].args[0] + with patch("agent_framework_durabletask._worker.AgentEntity") as consumer: + entity = entity_class() + consumer.assert_called_once() + kwargs = consumer.call_args.kwargs + assert kwargs["state_provider"] is entity + return dict(kwargs) + + +def _assert_settings(actual: dict[str, Any], **expected: Any) -> None: + assert {key: actual[key] for key in expected} == expected + + +def test_public_inheritance_contract_is_typed_and_exported() -> None: + assert isinstance(INHERIT, Enum) + assert INHERIT is Inherit.INHERIT + assert Inherit in get_args(StateBudgetOverride) + assert signature(DurableAIAgentWorker.add_agent).parameters["max_state_bytes"].default is INHERIT + assert signature(DurableAIAgentWorker.configure_workflow).parameters["max_state_bytes"].default is INHERIT + for name in _configuration.__all__: + assert name in durabletask.__all__ + assert getattr(durabletask, name) is getattr(_configuration, name) + assert callable(validate_history_providers) + + +@pytest.mark.parametrize("inherited", list(Inherit)) +def test_only_the_enum_inherits_a_budget(inherited: Inherit) -> None: + assert resolve_state_budget_override(inherited, 8192) == 8192 + assert resolve_state_budget_override(inherited, None) is None + assert resolve_state_budget_override(None, 8192) is None + with pytest.raises(ValueError, match="max_state_bytes"): + resolve_state_budget_override("inherit", 8192) # type: ignore[arg-type] + + +def test_worker_defaults_reach_the_entity_consumer() -> None: + grpc_worker = Mock() + worker = DurableAIAgentWorker(grpc_worker) + worker.add_agent(_agent()) + + assert DEFAULT_RETENTION == "keep_all" + assert DEFAULT_MAX_STATE_BYTES is None + _assert_settings( + _consumer_settings(grpc_worker), + retention=DEFAULT_RETENTION, + max_state_bytes=None, + high_watermark=HIGH_WATERMARK, + low_watermark=LOW_WATERMARK, + response_delivery_window_seconds=DELIVERY_WINDOW_SECONDS, + ) + + +@pytest.mark.parametrize("retention", get_args(RetentionMode)) +@pytest.mark.parametrize("budget,expected", [(None, None), (8192, 8192), ("backend_limit", DTS_MAX_STATE_BYTES)]) +def test_worker_pressure_budget_is_independent_of_retention( + retention: RetentionMode, budget: Any, expected: int | None +) -> None: + grpc_worker = Mock(spec=DurableTaskSchedulerWorker) if budget == "backend_limit" else Mock() + worker = DurableAIAgentWorker( + grpc_worker, + retention=retention, + max_state_bytes=budget, + high_watermark=0.95, + low_watermark=0.8, + response_delivery_window_seconds=120, + ) + worker.add_agent(_agent()) + + _assert_settings( + _consumer_settings(grpc_worker), + retention=retention, + max_state_bytes=expected, + high_watermark=0.95, + low_watermark=0.8, + response_delivery_window_seconds=120, + ) + + +@pytest.mark.parametrize("surface", ["agent", "workflow"]) +@pytest.mark.parametrize( + "overrides,expected", + [ + ({}, 8192), + ({"max_state_bytes": INHERIT}, 8192), + ({"max_state_bytes": None}, None), + ({"max_state_bytes": 4096}, 4096), + ({"max_state_bytes": "backend_limit"}, DTS_MAX_STATE_BYTES), + ], +) +def test_budget_override_distinguishes_omitted_and_disabled( + surface: str, overrides: dict[str, Any], expected: int | None +) -> None: + grpc_worker = ( + Mock(spec=DurableTaskSchedulerWorker) if overrides.get("max_state_bytes") == "backend_limit" else Mock() + ) + worker = DurableAIAgentWorker(grpc_worker, max_state_bytes=8192) + if surface == "agent": + worker.add_agent(_agent(), **overrides) + else: + worker.configure_workflow(_workflow("flow", _agent()), **overrides) + + assert _consumer_settings(grpc_worker)["max_state_bytes"] == expected + + +def test_per_agent_overrides_do_not_change_the_host_defaults_or_callbacks() -> None: + grpc_worker = Mock() + default_callback, specific_callback = Mock(), Mock() + worker = DurableAIAgentWorker( + grpc_worker, + callback=default_callback, + retention="follow_compaction", + max_state_bytes=8192, + high_watermark=0.95, + low_watermark=0.8, + response_delivery_window_seconds=120, + ) + worker.add_agent( + _agent("override"), + callback=specific_callback, + retention="keep_all", + max_state_bytes=None, + high_watermark=0.9, + low_watermark=0.6, + response_delivery_window_seconds=15, + ) + worker.add_agent(_agent("inherited")) + + _assert_settings( + _consumer_settings(grpc_worker), + callback=specific_callback, + retention="keep_all", + max_state_bytes=None, + high_watermark=0.9, + low_watermark=0.6, + response_delivery_window_seconds=15, + ) + _assert_settings( + _consumer_settings(grpc_worker, 1), + callback=default_callback, + retention="follow_compaction", + max_state_bytes=8192, + high_watermark=0.95, + low_watermark=0.8, + response_delivery_window_seconds=120, + ) + + +@pytest.mark.parametrize("retention", get_args(RetentionMode)) +def test_workflow_overrides_reach_every_new_nested_entity(retention: RetentionMode) -> None: + grpc_worker = Mock() + worker = DurableAIAgentWorker(grpc_worker, max_state_bytes=8192) + inner = _workflow("inner", _agent("inneragent")) + outer = _workflow("outer", _agent("outeragent"), child=inner) + worker.configure_workflow( + outer, + retention=retention, + max_state_bytes=None, + high_watermark=0.9, + low_watermark=0.6, + response_delivery_window_seconds=20, + ) + + assert worker.registered_agent_names == ["outer-node0", "inner-node0"] + for index in range(grpc_worker.add_entity.call_count): + _assert_settings( + _consumer_settings(grpc_worker, index), + retention=retention, + max_state_bytes=None, + high_watermark=0.9, + low_watermark=0.6, + response_delivery_window_seconds=20, + ) + + +_INVALID_SETTINGS: list[dict[str, Any]] = [ + {"retention": "auto"}, + {"retention": "invalid"}, + *({"max_state_bytes": value} for value in [0, -1, True, False, 1.5, "8192", "inherit"]), + *({"high_watermark": value} for value in [0, 1.1, True, float("nan"), float("inf")]), + *({"low_watermark": value} for value in [0, -0.1, True, float("nan"), float("inf")]), + {"high_watermark": 0.7, "low_watermark": 0.7}, + {"high_watermark": 0.6, "low_watermark": 0.7}, + *( + {"response_delivery_window_seconds": value} + for value in [0, -1, True, False, 1.5, "60", float("nan"), float("inf")] + ), +] + + +@pytest.mark.parametrize("settings", _INVALID_SETTINGS) +@pytest.mark.parametrize("surface", ["host", "agent", "workflow"]) +def test_invalid_settings_fail_before_registration(surface: str, settings: dict[str, Any]) -> None: + grpc_worker = Mock() + if surface == "host": + with pytest.raises(ValueError): + DurableAIAgentWorker(grpc_worker, **settings) + else: + worker = DurableAIAgentWorker(grpc_worker) + with pytest.raises(ValueError): + if surface == "agent": + worker.add_agent(_agent(), **settings) + else: + worker.configure_workflow(_workflow("flow", _agent()), **settings) + assert worker.registered_agent_names == [] + assert worker.registered_workflow_names == [] + assert worker._registered_orchestrations == {} + assert grpc_worker.mock_calls == [] + + +@pytest.mark.parametrize("surface", ["agent", "workflow", "nested_workflow"]) +def test_ambiguous_history_fails_before_any_registration(surface: str) -> None: + grpc_worker = Mock() + worker = DurableAIAgentWorker(grpc_worker) + agent = _agent(ambiguous_history=True) + original_providers = agent.context_providers + with pytest.raises(ValueError, match="primary"): + if surface == "agent": + worker.add_agent(agent) + elif surface == "workflow": + worker.configure_workflow(_workflow("flow", _agent("good"), agent)) + else: + worker.configure_workflow(_workflow("outer", _agent("good"), child=_workflow("inner", agent))) + + assert agent.context_providers is original_providers + assert all(isinstance(provider, InMemoryHistoryProvider) for provider in original_providers) + assert worker.registered_agent_names == [] + assert worker.registered_workflow_names == [] + assert worker._registered_orchestrations == {} + assert grpc_worker.mock_calls == [] + + +def test_registration_validates_without_replacing_the_users_history_provider() -> None: + grpc_worker = Mock() + worker = DurableAIAgentWorker(grpc_worker, retention="follow_compaction") + agent = _agent() + original_providers = agent.context_providers + with patch("agent_framework_durabletask._worker.AgentEntity") as consumer: + worker.add_agent(agent) + consumer.assert_not_called() + assert agent.context_providers is original_providers + assert isinstance(agent.context_providers[0], InMemoryHistoryProvider) + + +def test_backend_registration_failure_does_not_record_an_agent() -> None: + grpc_worker = Mock() + grpc_worker.add_entity.side_effect = RuntimeError("registration failed") + worker = DurableAIAgentWorker(grpc_worker) + with pytest.raises(RuntimeError, match="registration failed"): + worker.add_agent(_agent()) + assert worker.registered_agent_names == [] diff --git a/python/packages/durabletask/tests/test_retention_revision.py b/python/packages/durabletask/tests/test_retention_revision.py new file mode 100644 index 0000000..de4e0f9 --- /dev/null +++ b/python/packages/durabletask/tests/test_retention_revision.py @@ -0,0 +1,652 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Pressure-retention regressions for the independent ADR-0032 controls.""" + +import json +from copy import deepcopy +from datetime import datetime, timedelta, timezone +from typing import Any, get_args +from unittest.mock import AsyncMock, Mock + +import pytest +from agent_framework import CharacterEstimatorTokenizer, Content, Message, annotate_message_groups, included_token_count + +from agent_framework_durabletask import _retention as retention +from agent_framework_durabletask._durable_agent_state import ( + DurableAgentState, + DurableAgentStateCompaction, + DurableAgentStateEntry, + DurableAgentStateEntryJsonType, + DurableAgentStateErrorResponse, + DurableAgentStateMessage, + DurableAgentStateRequest, + DurableAgentStateResponse, + DurableAgentStateUnknownContent, + DurableAgentStateUsage, +) + +NOW = datetime(2026, 9, 8, 12, 0, 0, 123456, tzinfo=timezone.utc) +OLD = NOW - timedelta(hours=1) + + +@pytest.fixture(autouse=True) +def fixed_retention_clock(monkeypatch: pytest.MonkeyPatch) -> None: + clock = Mock(wraps=datetime) + clock.now.return_value = NOW + monkeypatch.setattr(retention, "datetime", clock) + + +def _message(message_id: str | None, role: str = "user", text: str = "x" * 400) -> DurableAgentStateMessage: + return DurableAgentStateMessage.from_chat_message(Message(role, [text], message_id=message_id)) + + +def _state(turns: int = 40, *, chars: int = 400) -> DurableAgentState: + state = DurableAgentState() + for index in range(turns): + state.data.conversation_history.extend([ + DurableAgentStateRequest(f"c{index}", OLD, [_message(f"u{index}", text="u" * chars)]), + DurableAgentStateResponse(f"c{index}", OLD, [_message(f"a{index}", "assistant", "a" * chars)]), + ]) + return state + + +def _ids(state: DurableAgentState) -> list[str | None]: + return [message.message_id for entry in state.data.conversation_history for message in entry.messages] + + +def _project_plain( + state: DurableAgentState, removed_ids: list[str | None], *, record: bool = True +) -> DurableAgentState: + """Independent byte oracle for fixtures containing only bare transcript envelopes.""" + projected = deepcopy(state) + removed = set(removed_ids) + for entry in projected.data.conversation_history: + entry.messages = [message for message in entry.messages if message.message_id not in removed] + projected.data.conversation_history = [entry for entry in projected.data.conversation_history if entry.messages] + if removed and record: + previous = projected.data.truncation or {} + projected.data.truncation = { + **previous, + "evictedMessageCount": previous.get("evictedMessageCount", 0) + len(removed_ids), + "firstEvictedAt": previous.get("firstEvictedAt", NOW.isoformat()), + "lastEvictedAt": NOW.isoformat(), + } + return projected + + +def _smallest_plain_prefix(state: DurableAgentState, target: int) -> tuple[int, DurableAgentState]: + eligible = _ids(state)[:-2] + for count in range(1, len(eligible) + 1): + projected = _project_plain(state, eligible[:count]) + if retention._serialized_size(projected) <= target: + return count, projected + raise AssertionError("fixture has no reachable prefix at the requested target") + + +def _delivery(state: DurableAgentState, correlations: list[str], *, payload_chars: int = 20) -> None: + # Use the real state data serializer, not a mock that could hide mailbox bytes from the floor. + data: Any = state.data + data.response_mailbox = { + correlation: { + "response": {"messages": [Message("assistant", ["r" * payload_chars]).to_dict()], "metadata": {"v": [1]}}, + "createdAt": NOW.isoformat(), + "expiresAt": (NOW + timedelta(seconds=retention.DELIVERY_WINDOW_SECONDS)).isoformat(), + "futureMailboxField": {"keep": True}, + } + for correlation in correlations + } + data.completed_correlations = { + correlation: {"completedAt": NOW.isoformat(), "futureReceiptField": [1, 3]} for correlation in correlations + } + assert state.to_dict()["data"]["responseMailbox"] == data.response_mailbox + assert state.to_dict()["data"]["completedCorrelations"] == data.completed_correlations + + +class TestConfiguration: + def test_public_aliases_and_non_deleting_defaults(self) -> None: + assert get_args(retention.RetentionMode) == ("keep_all", "follow_compaction") + budget_members = get_args(retention.StateBudget) + assert int in budget_members and type(None) in budget_members + assert any(get_args(member) == ("backend_limit",) for member in budget_members) + assert retention.DEFAULT_RETENTION == "keep_all" + assert retention.DEFAULT_MAX_STATE_BYTES is None + assert retention.DTS_MAX_STATE_BYTES == 1_048_576 + assert retention.HIGH_WATERMARK == 0.85 + assert retention.LOW_WATERMARK == 0.70 + assert retention.DELIVERY_WINDOW_SECONDS == 60 + assert { + "RetentionMode", + "StateBudget", + "StateCapacityError", + "resolve_state_budget", + "validate_retention", + } <= set(retention.__all__) + + @pytest.mark.parametrize("mode", ["keep_all", "follow_compaction"]) + @pytest.mark.parametrize("budget", [None, 1, 24_000, "backend_limit"]) + def test_pruning_and_pressure_are_independent(self, mode: Any, budget: Any) -> None: + retention.validate_retention(mode) + assert retention.prunes_excluded(mode) is (mode == "follow_compaction") + expected = retention.DTS_MAX_STATE_BYTES if budget == "backend_limit" else budget + assert retention.resolve_state_budget(budget, backend_limit=retention.DTS_MAX_STATE_BYTES) == expected + + @pytest.mark.parametrize( + "value", + [ + True, + False, + 0, + -1, + 1.5, + 1.0, + float("nan"), + float("inf"), + "auto", + "1", + "", + [], + {}, + (), + b"backend_limit", + object(), + ], + ) + def test_invalid_budgets_raise_value_error(self, value: Any) -> None: + with pytest.raises(ValueError, match="max_state_bytes"): + retention.resolve_state_budget(value) + + @pytest.mark.parametrize("limit", [None, True, False, 0, -1, 1.0, float("nan"), float("inf"), "1000", [], {}]) + def test_backend_limit_must_be_resolved_and_positive(self, limit: Any) -> None: + with pytest.raises(ValueError, match="backend_limit"): + retention.resolve_state_budget("backend_limit", backend_limit=limit) + + @pytest.mark.parametrize( + "mode", + ["auto", "", "KEEP_ALL", "follow-compaction", None, True, False, 1, 1.0, [], {}, (), b"keep_all", object()], + ) + def test_invalid_modes_raise_value_error(self, mode: Any) -> None: + with pytest.raises(ValueError, match="retention"): + retention.validate_retention(mode) + with pytest.raises(ValueError, match="retention"): + retention.prunes_excluded(mode) + + @pytest.mark.parametrize("name", ["high_watermark", "low_watermark"]) + @pytest.mark.parametrize( + "value", + [None, True, False, 0, -0.1, 1.1, float("nan"), float("inf"), float("-inf"), "0.8", [], {}, 1j, 10**400], + ) + def test_invalid_watermark_types_and_ranges(self, name: str, value: Any) -> None: + kwargs = {"high_watermark": 0.85, "low_watermark": 0.70, name: value} + with pytest.raises(ValueError, match="watermark"): + retention.validate_retention("keep_all", **kwargs) + + @pytest.mark.parametrize(("high", "low"), [(0.7, 0.7), (0.6, 0.7), (1, 1)]) + def test_watermarks_must_be_strictly_ordered(self, high: float, low: float) -> None: + with pytest.raises(ValueError, match="watermark"): + retention.validate_retention("keep_all", high, low) + + def test_high_watermark_may_equal_one(self) -> None: + retention.validate_retention("keep_all", 1, 0.5) + + @pytest.mark.parametrize("budget", [None, True, False, 0, -1, 1.0, "backend_limit", [], {}]) + async def test_enforcement_requires_a_resolved_integer(self, budget: Any) -> None: + state = _state(1) + before = state.to_json() + with pytest.raises(ValueError, match="max_state_bytes"): + await retention.enforce_budget(state, max_state_bytes=budget) + assert state.to_json() == before + + async def test_enforcement_validates_watermarks_even_below_pressure(self) -> None: + with pytest.raises(ValueError, match="watermark"): + await retention.enforce_budget(_state(1), max_state_bytes=1_000_000, high_watermark=float("nan")) + + +class TestProtectedFloor: + @pytest.mark.parametrize("empty_entries", [0, 30]) + async def test_metadata_only_floor_fails_without_any_mutation( + self, monkeypatch: pytest.MonkeyPatch, empty_entries: int + ) -> None: + state = _state(0) + state.data.session = {"approvals": "p" * 20_000} + state.data.conversation_history = [DurableAgentStateRequest(f"c{i}", OLD, []) for i in range(empty_entries)] + before = state.to_json() + history = state.data.conversation_history + strategy = Mock(side_effect=AssertionError("no eviction pass is permitted")) + monkeypatch.setattr(retention, "TokenBudgetComposedStrategy", strategy) + with pytest.raises(retention.StateCapacityError) as error: + await retention.enforce_budget(state, max_state_bytes=12_000) + assert error.value.size_bytes == error.value.floor_bytes == len(before) + assert error.value.max_state_bytes == 12_000 + assert "floor" in str(error.value) and "budget" in str(error.value) + assert state.to_json() == before + assert state.data.conversation_history is history + strategy.assert_not_called() + + async def test_floor_between_high_and_hard_limit_prevents_futile_deletion( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + state = _state() + state.data.session = {"protected": "p" * 9_000} + floor = retention._serialized_size(_project_plain(state, _ids(state)[:-2])) + assert 12_000 * 0.8 <= floor < 12_000 + before = state.to_json() + strategy = Mock(side_effect=AssertionError("floor is unreachable")) + monkeypatch.setattr(retention, "TokenBudgetComposedStrategy", strategy) + with pytest.raises(retention.StateCapacityError) as error: + await retention.enforce_budget(state, max_state_bytes=12_000, high_watermark=0.8, low_watermark=0.6) + assert error.value.floor_bytes == floor + assert state.to_json() == before + strategy.assert_not_called() + + async def test_truncation_cost_is_in_the_floor_before_any_eviction(self, monkeypatch: pytest.MonkeyPatch) -> None: + state = _state(12) + old_ids = _ids(state)[:-2] + without_record = retention._serialized_size(_project_plain(state, old_ids, record=False)) + floor = retention._serialized_size(_project_plain(state, old_ids)) + budget = (without_record + floor) // 2 + assert without_record < budget < floor + before = state.to_json() + strategy = Mock(side_effect=AssertionError("truncation cannot fit")) + monkeypatch.setattr(retention, "TokenBudgetComposedStrategy", strategy) + with pytest.raises(retention.StateCapacityError) as error: + await retention.enforce_budget(state, max_state_bytes=budget, high_watermark=1, low_watermark=0.5) + assert error.value.floor_bytes == floor + assert state.to_json() == before + strategy.assert_not_called() + + async def test_unreachable_low_uses_the_reachable_floor_below_high(self) -> None: + state = _state(25) + state.data.session = {"protected": "p" * 9_000} + expected = _project_plain(state, _ids(state)[:-2]) + assert 13_000 * 0.5 < retention._serialized_size(expected) < 13_000 * 0.9 + removed = await retention.enforce_budget(state, max_state_bytes=13_000, high_watermark=0.9, low_watermark=0.5) + assert removed == 48 + assert state.to_dict() == expected.to_dict() + + async def test_all_recent_legacy_results_are_absolute_protections(self, monkeypatch: pytest.MonkeyPatch) -> None: + state = _state() + for entry in state.data.conversation_history: + entry.created_at = NOW + before = state.to_json() + strategy = Mock(side_effect=AssertionError("recent results cannot be sacrificed")) + monkeypatch.setattr(retention, "TokenBudgetComposedStrategy", strategy) + with pytest.raises(retention.StateCapacityError): + await retention.enforce_budget(state, max_state_bytes=12_000) + assert state.to_json() == before + strategy.assert_not_called() + + async def test_unknown_entry_payloads_contribute_to_the_floor(self, monkeypatch: pytest.MonkeyPatch) -> None: + state = _state() + unknown_kind: Any = "futureKind" + state.data.conversation_history.insert( + 0, DurableAgentStateEntry(unknown_kind, "opaque", OLD, [_message("opaque", text="p" * 30_000)]) + ) + before = state.to_json() + strategy = Mock(side_effect=AssertionError("unknown state cannot be evicted")) + monkeypatch.setattr(retention, "TokenBudgetComposedStrategy", strategy) + with pytest.raises(retention.StateCapacityError) as error: + await retention.enforce_budget(state, max_state_bytes=12_000) + assert error.value.floor_bytes > 30_000 + assert state.to_json() == before + strategy.assert_not_called() + + async def test_mailbox_receipts_and_control_fields_survive_transcript_eviction(self) -> None: + state = _state() + _delivery(state, ["c0", "c1"]) + state.data.session = {"service_session_id": "branch", "state": {"approvals": ["keep"]}} + state.data.ingested_positions = {"source": 99} + state.data.extension_data = {"customIds": ["opaque-id"], "futureControl": {"keep": [1, 3]}} + state.data.conversation_history[1].created_at = NOW + before = deepcopy(state.to_dict()["data"]) + removed = await retention.enforce_budget(state, max_state_bytes=16_000) + assert removed > 0 and "a0" not in _ids(state) + after = state.to_dict()["data"] + for field in ("responseMailbox", "completedCorrelations", "session", "ingestedPositions", "extensionData"): + assert after[field] == before[field] + + @pytest.mark.parametrize("has_mailbox", [False, True]) + async def test_only_independently_completed_recent_results_are_evictable(self, has_mailbox: bool) -> None: + state = _state() + _delivery(state, ["c0"]) + data: Any = state.data + if not has_mailbox: + data.response_mailbox.clear() + for entry in state.data.conversation_history[:4]: + entry.created_at = NOW + before = deepcopy(data.completed_correlations) + assert await retention.enforce_budget(state, max_state_bytes=12_000) > 0 + assert "a0" not in _ids(state) + assert {"u1", "a1"} <= set(_ids(state)) + assert data.completed_correlations == before + + @pytest.mark.parametrize("field", ["response_mailbox", "completed_correlations"]) + async def test_delivery_records_alone_can_fill_the_floor(self, field: str) -> None: + state = _state() + _delivery(state, ["c0"], payload_chars=30_000 if field == "response_mailbox" else 1) + if field == "completed_correlations": + data: Any = state.data + data.completed_correlations["c0"]["futureReceiptField"] = "p" * 30_000 + before = state.to_json() + with pytest.raises(retention.StateCapacityError) as error: + await retention.enforce_budget(state, max_state_bytes=12_000) + assert error.value.floor_bytes > 30_000 + assert state.to_json() == before + + async def test_entry_schema_and_usage_metadata_are_not_transcript_capacity(self) -> None: + state = _state() + request = state.data.conversation_history[0] + assert isinstance(request, DurableAgentStateRequest) + request.response_schema = {"largeControl": "s" * 2_000} + request.orchestration_id = "workflow" + response = state.data.conversation_history[1] + assert isinstance(response, DurableAgentStateResponse) + response.usage = DurableAgentStateUsage(input_token_count=10, extensionData={"opaque": "keep"}) + before = [deepcopy(entry.to_dict()) for entry in (request, response)] + assert await retention.enforce_budget(state, max_state_bytes=10_000) > 0 + retained = [entry for entry in state.data.conversation_history if entry.correlation_id == "c0"] + assert len(retained) == 2 + for entry, original in zip(retained, before): + assert entry.messages == [] + original["messages"] = [] + assert entry.to_dict() == original + + +class TestSelectionAndMeasurements: + @pytest.mark.parametrize("kind", list(DurableAgentStateEntryJsonType)) + async def test_each_known_transcript_kind_is_storage_eligible(self, kind: DurableAgentStateEntryJsonType) -> None: + state = _state(0) + for index in range(30): + state.data.conversation_history.append( + DurableAgentStateEntry(kind, f"old-{index}", OLD, [_message(f"old-{index}")]) + ) + state.data.conversation_history.extend(_state(1).data.conversation_history) + assert await retention.enforce_budget(state, max_state_bytes=8_000) > 0 + assert "old-0" not in _ids(state) + assert _ids(state)[-2:] == ["u0", "a0"] + + async def test_expired_runtime_error_content_is_evictable(self) -> None: + state = _state(0) + for index in range(35): + error = Message("assistant", [Content.from_error(message="e" * 500)], message_id=f"error-{index}") + state.data.conversation_history.append( + DurableAgentStateErrorResponse( + f"failed-{index}", OLD, [DurableAgentStateMessage.from_chat_message(error)] + ) + ) + assert await retention.enforce_budget(state, max_state_bytes=10_000) > 0 + assert "error-0" not in _ids(state) + assert "error-34" in _ids(state) + assert retention._serialized_size(state) <= 7_000 + + @pytest.mark.parametrize("recent", [False, True]) + async def test_error_delivery_protection_applies_only_inside_legacy_window(self, recent: bool) -> None: + state = _state() + occurred_at = NOW - timedelta(seconds=30 if recent else retention.DELIVERY_WINDOW_SECONDS) + failure = DurableAgentStateErrorResponse("failed", occurred_at.replace(tzinfo=None), [_message("failure")]) + state.data.conversation_history.insert(0, failure) + assert await retention.enforce_budget(state, max_state_bytes=12_000) > 0 + assert ("failure" in _ids(state)) is recent + + async def test_custom_high_watermark_controls_trigger_and_low_controls_target(self) -> None: + state = _state(16) + before = state.to_json() + budget = len(before) * 2 + assert await retention.enforce_budget(state, max_state_bytes=budget, high_watermark=0.75) == 0 + assert state.to_json() == before + count, expected = _smallest_plain_prefix(state, int(budget * 0.25)) + removed = await retention.enforce_budget(state, max_state_bytes=budget, high_watermark=0.4, low_watermark=0.25) + assert removed == count + assert state.to_dict() == expected.to_dict() + + @pytest.mark.parametrize("previous_count", [0, 9, 99, 999]) + async def test_low_target_includes_truncation_and_does_not_over_evict(self, previous_count: int) -> None: + state = _state() + if previous_count: + state.data.truncation = { + "evictedMessageCount": previous_count, + "firstEvictedAt": OLD.isoformat(), + "lastEvictedAt": OLD.isoformat(), + "futureEvidence": {"keep": [1, 3]}, + } + count, expected = _smallest_plain_prefix(state, 18_000) + removed = await retention.enforce_budget(state, max_state_bytes=20_000, high_watermark=0.95, low_watermark=0.9) + assert removed == count + assert state.to_dict() == expected.to_dict() + assert ( + await retention.enforce_budget(state, max_state_bytes=20_000, high_watermark=0.95, low_watermark=0.9) == 0 + ) + + async def test_actual_bytes_correct_an_optimistic_plan_without_halving_the_target( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + state = _state() + count, expected = _smallest_plain_prefix(state, 8_400) + prefix_sizes = retention._prefix_sizes + + def optimistic_sizes(*args: Any, **kwargs: Any) -> list[int]: + return [size - 1_500 for size in prefix_sizes(*args, **kwargs)] + + factory = Mock(wraps=retention.TokenBudgetComposedStrategy) + monkeypatch.setattr(retention, "_prefix_sizes", optimistic_sizes) + monkeypatch.setattr(retention, "TokenBudgetComposedStrategy", factory) + assert await retention.enforce_budget(state, max_state_bytes=12_000) == count + assert state.to_dict() == expected.to_dict() + assert 2 <= factory.call_count <= 3 + + async def test_mixed_unicode_and_large_tool_payloads_choose_the_smallest_atomic_prefix(self) -> None: + state = _state(0) + for index in range(30): + body = "\u754c\U0001f680" * 200 if index % 2 else "plain" * 100 + call = Content.from_function_call(call_id=f"call-{index}", name="lookup", arguments=json.dumps({"q": body})) + result = Content.from_function_result(call_id=f"call-{index}", result={"records": [body]}) + state.data.conversation_history.append( + DurableAgentStateResponse( + f"tools-{index}", + OLD, + [ + DurableAgentStateMessage.from_chat_message( + Message("assistant", [call], message_id=f"call-{index}") + ), + DurableAgentStateMessage.from_chat_message( + Message("tool", [result], message_id=f"result-{index}") + ), + ], + ) + ) + state.data.conversation_history.extend(_state(1).data.conversation_history) + candidates = _ids(state)[:-2] + expected_count = 0 + for count in range(2, len(candidates) + 1, 2): + projected = _project_plain(state, candidates[:count]) + if retention._serialized_size(projected) <= 28_000: + expected_count = count + break + assert 0 < expected_count < len(candidates) + expected = _project_plain(state, candidates[:expected_count]) + assert await retention.enforce_budget(state, max_state_bytes=40_000) == expected_count + assert state.to_dict() == expected.to_dict() + assert retention._serialized_size(state) == len(state.to_json().encode("utf-8")) + + def test_token_budget_uses_core_tokens_not_escaped_json_or_text_length(self) -> None: + message = Message( + "assistant", + [Content.from_function_call(call_id="call", name="lookup", arguments=json.dumps({"q": "\u754c" * 100}))], + message_id="tool", + ) + state = _state(0) + entry = DurableAgentStateResponse("tools", OLD, [DurableAgentStateMessage.from_chat_message(message)]) + entry.messages.append(_message("unicode", text="\U0001f680\u754c" * 100)) + state.data.conversation_history.append(entry) + origins: list[tuple[DurableAgentStateEntry, DurableAgentStateMessage]] = [ + (entry, stored) for stored in entry.messages + ] + messages = [deepcopy(stored).to_chat_message() for stored in entry.messages] + annotate_message_groups(messages, tokenizer=CharacterEstimatorTokenizer()) + tokens = included_token_count(messages) + persisted_bytes = sum(len(json.dumps(stored.to_dict())) for stored in entry.messages) + assert retention._token_budget( + origins, + serialized_size=persisted_bytes + 1_000, + evictable_bytes=persisted_bytes, + target_bytes=1_000 + persisted_bytes // 2, + ) == max((persisted_bytes // 2) * tokens // persisted_bytes, 1) + + async def test_summaries_do_not_replace_the_newest_exchange(self) -> None: + state = _state() + newest = state.data.conversation_history[-2:] + state.data.conversation_history.append(DurableAgentStateCompaction(NOW, [_message("summary")], "summary-cid")) + assert retention._newest_exchange(state.data.conversation_history) == newest + assert await retention.enforce_budget(state, max_state_bytes=12_000) > 0 + assert {"u39", "a39"} <= set(_ids(state)) + + async def test_unknown_entries_and_already_empty_envelopes_remain_opaque(self) -> None: + state = _state() + future_kind: Any = "futureKind" + opaque = DurableAgentStateEntry( + future_kind, + "unknown", + OLD, + [DurableAgentStateMessage("assistant", [DurableAgentStateUnknownContent({})], message_id="opaque")], + ) + empty = DurableAgentStateRequest("metadata-only", OLD, [], response_schema={"keep": True}) + state.data.conversation_history[:0] = [opaque, empty] + before = [deepcopy(entry.to_dict()) for entry in (opaque, empty)] + assert await retention.enforce_budget(state, max_state_bytes=12_000) > 0 + assert [entry.to_dict() for entry in state.data.conversation_history[:2]] == before + + +class TestAtomicityAndIsolation: + @pytest.mark.parametrize("non_contiguous", [False, True]) + async def test_reasoning_call_and_result_are_one_oldest_group(self, non_contiguous: bool) -> None: + state = _state(20) + group = [ + Message("assistant", [Content.from_text_reasoning(text="reason" * 100)], message_id="reason"), + Message( + "assistant", [Content.from_function_call(call_id="t", name="tool", arguments="{}")], message_id="call" + ), + Message("tool", [Content.from_function_result(call_id="t", result="r" * 500)], message_id="result"), + ] + if non_contiguous: + group.insert(2, Message("user", ["gap"], message_id="gap")) + state.data.conversation_history.insert( + 0, DurableAgentStateResponse("tools", OLD, [DurableAgentStateMessage.from_chat_message(m) for m in group]) + ) + budget = retention._serialized_size(state) - 1 + assert await retention.enforce_budget(state, max_state_bytes=budget, high_watermark=1, low_watermark=0.99) == 3 + assert not {"reason", "call", "result"} & set(_ids(state)) + assert ("gap" in _ids(state)) is non_contiguous + + async def test_system_intersection_protects_the_entire_persisted_group(self) -> None: + state = _state() + messages = [_message("policy", "system"), _message("linked", "assistant")] + for message in messages: + message.extension_data = {"_group": {"id": "atomic-policy"}} + state.data.conversation_history.insert(0, DurableAgentStateRequest("policy", OLD, messages)) + before = [deepcopy(message.to_dict()) for message in messages] + assert await retention.enforce_budget(state, max_state_bytes=12_000) > 0 + assert [message.to_dict() for message in state.data.conversation_history[0].messages] == before + + async def test_current_exchange_protects_its_non_contiguous_tool_declaration(self) -> None: + state = _state() + call = Message( + "assistant", [Content.from_function_call(call_id="t", name="tool", arguments="{}")], message_id="call" + ) + state.data.conversation_history.insert( + 0, DurableAgentStateResponse("earlier", OLD, [DurableAgentStateMessage.from_chat_message(call)]) + ) + result = Message("tool", [Content.from_function_result(call_id="t", result="result")], message_id="result") + state.data.conversation_history[-1].messages.append(DurableAgentStateMessage.from_chat_message(result)) + assert await retention.enforce_budget(state, max_state_bytes=12_000) > 0 + assert {"call", "result", "u39", "a39"} <= set(_ids(state)) + + async def test_nested_annotations_and_payloads_never_alias_planning_copies( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + state = _state() + original_messages = [message for entry in state.data.conversation_history for message in entry.messages] + for index, message in enumerate(original_messages): + message.extension_data = { + "_excluded": True, + "_exclude_reason": "user_compaction", + "_group": {"id": f"saved-{index}", "token_count": 999_999, "future": {"values": [1, 2]}}, + "future": {"values": [1, 3]}, + } + original = [deepcopy(message.to_dict()) for message in original_messages] + converter = DurableAgentStateMessage.to_chat_message + strategy_class = retention.TokenBudgetComposedStrategy + + def aliasing_converter(stored: DurableAgentStateMessage) -> Message: + converted: Message = converter(stored) + if stored.extension_data is not None: + converted.additional_properties = stored.extension_data + return converted + + def mutating_strategy_factory(**kwargs: Any) -> Any: + strategy = strategy_class(**kwargs) + + async def mutate_and_evict(messages: list[Message]) -> bool: + for message in messages[:-1]: + message.additional_properties["future"]["values"].append("planning-only") + return await strategy(messages) + + return mutate_and_evict + + monkeypatch.setattr(DurableAgentStateMessage, "to_chat_message", aliasing_converter) + monkeypatch.setattr(retention, "TokenBudgetComposedStrategy", mutating_strategy_factory) + assert await retention.enforce_budget(state, max_state_bytes=24_000) > 0 + assert [message.to_dict() for message in original_messages] == original + by_id = {message["messageId"]: message for message in original} + survivors = [message for entry in state.data.conversation_history for message in entry.messages] + assert len(survivors) > 2 + assert all(message.to_dict() == by_id[message.message_id] for message in survivors) + + @pytest.mark.parametrize("ids", [[None, None], ["duplicate", "duplicate"]]) + async def test_missing_or_duplicate_message_ids_do_not_alias_eviction_origins(self, ids: list[str | None]) -> None: + state = _state() + for index, entry in enumerate(state.data.conversation_history): + entry.messages[0].message_id = ids[index % 2] + before = len(_ids(state)) + removed = await retention.enforce_budget(state, max_state_bytes=12_000) + assert 0 < removed < before - 2 + assert len(_ids(state)) == before - removed + assert _ids(state)[-2:] == ids + + async def test_strategy_is_deterministic_and_has_no_user_strategies(self, monkeypatch: pytest.MonkeyPatch) -> None: + strategy = Mock(wraps=retention.TokenBudgetComposedStrategy) + monkeypatch.setattr(retention, "TokenBudgetComposedStrategy", strategy) + first = _state() + second = deepcopy(first) + assert await retention.enforce_budget(first, max_state_bytes=12_000) > 0 + assert await retention.enforce_budget(second, max_state_bytes=12_000) > 0 + assert first.to_dict() == second.to_dict() + assert 2 <= strategy.call_count <= 6 + for call in strategy.call_args_list: + assert call.kwargs["strategies"] == [] + assert isinstance(call.kwargs["tokenizer"], CharacterEstimatorTokenizer) + + async def test_unsatisfied_strategy_stops_after_three_passes_and_rolls_back( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + state = _state() + before = state.to_json() + strategy = AsyncMock(return_value=False) + factory = Mock(return_value=strategy) + monkeypatch.setattr(retention, "TokenBudgetComposedStrategy", factory) + with pytest.raises(retention.StateCapacityError): + await retention.enforce_budget(state, max_state_bytes=12_000) + assert strategy.await_count == 3 + assert state.to_json() == before + + async def test_strategy_failure_cannot_leak_annotation_changes(self, monkeypatch: pytest.MonkeyPatch) -> None: + state = _state() + before = state.to_json() + + async def failing_strategy(messages: list[Message]) -> bool: + messages[0].additional_properties["poison"] = {"mutated": True} + messages[0].contents.clear() + raise RuntimeError("injected strategy failure") + + monkeypatch.setattr(retention, "TokenBudgetComposedStrategy", Mock(return_value=failing_strategy)) + with pytest.raises(RuntimeError, match="injected strategy failure"): + await retention.enforce_budget(state, max_state_bytes=12_000) + assert state.to_json() == before diff --git a/python/packages/durabletask/tests/test_retention_telemetry.py b/python/packages/durabletask/tests/test_retention_telemetry.py new file mode 100644 index 0000000..912832b --- /dev/null +++ b/python/packages/durabletask/tests/test_retention_telemetry.py @@ -0,0 +1,534 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Real SDK measurements of staged retention, isolated from the global meter provider.""" + +import asyncio +import json +from collections.abc import Iterator +from copy import deepcopy +from datetime import datetime, timedelta, timezone +from typing import Any +from unittest.mock import AsyncMock, Mock + +import pytest +from agent_framework import Agent, Message +from opentelemetry.metrics import NoOpMeterProvider +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import Histogram, InMemoryMetricReader, Metric, Sum +from test_durable_history_provider import RecordingChatClient + +from agent_framework_durabletask import ( + AgentEntity, + AgentEntityStateProviderMixin, + DurableAgentState, + DurableAgentStateMessage, + DurableAgentStateRequest, + DurableAgentStateResponse, + DurableHistoryProvider, + _history_provider, +) +from agent_framework_durabletask import _retention as retention +from agent_framework_durabletask import _retention_telemetry as telemetry +from agent_framework_durabletask._history_provider import ( + DurableHistoryBinding, + bind_durable_history, + unbind_durable_history, +) + +NOW = datetime(2026, 9, 11, 12, 0, 0, 123456, tzinfo=timezone.utc) +BUDGET = 12_000 +PREFIX = "durable.retention." + + +@pytest.fixture +def reader(monkeypatch: pytest.MonkeyPatch) -> Iterator[InMemoryMetricReader]: + metric_reader = InMemoryMetricReader() + provider = MeterProvider(metric_readers=[metric_reader], shutdown_on_exit=False) + monkeypatch.setattr(telemetry, "get_meter", provider.get_meter) + telemetry._instruments.cache_clear() + clock = Mock(wraps=datetime) + clock.now.return_value = NOW + monkeypatch.setattr(retention, "datetime", clock) + try: + yield metric_reader + finally: + telemetry._instruments.cache_clear() + provider.shutdown() + + +def _metrics(reader: InMemoryMetricReader) -> dict[str, Metric]: + data = reader.get_metrics_data() + if data is None: + return {} + result = {} + for resource in data.resource_metrics: + for scope in resource.scope_metrics: + assert scope.scope.name == "agent_framework.durabletask" + for metric in scope.metrics: + assert metric.name.startswith(PREFIX) + result[metric.name.removeprefix(PREFIX)] = metric + return result + + +def _counter(metric: Metric, attributes: dict[str, Any], value: int) -> None: + assert isinstance(metric.data, Sum) + assert metric.data.is_monotonic + matches = [point for point in metric.data.data_points if point.attributes == attributes] + assert len(matches) == 1 + assert matches[0].value == value + + +def _histogram(metric: Metric, attributes: dict[str, Any], total: int, count: int = 1) -> None: + assert metric.unit == "By" + assert isinstance(metric.data, Histogram) + matches = [point for point in metric.data.data_points if point.attributes == attributes] + assert len(matches) == 1 + assert matches[0].count == count + assert matches[0].sum == total + + +def _attributes(mechanism: str = "pressure", outcome: str = "staged") -> dict[str, Any]: + return {"mechanism": mechanism, "outcome": outcome, "commit_status": "not_attempted"} + + +def _state(turns: int = 40) -> DurableAgentState: + state = DurableAgentState() + for index in range(turns): + for role, kind in (("user", DurableAgentStateRequest), ("assistant", DurableAgentStateResponse)): + state.data.conversation_history.append( + kind( + correlation_id=f"private-correlation-{index}", + created_at=NOW - timedelta(days=1), + messages=[ + DurableAgentStateMessage.from_chat_message( + Message(role, ["private payload " * 30], message_id=f"private-{role}-{index}") + ) + ], + ) + ) + return state + + +def _size(state: DurableAgentState) -> int: + return len(json.dumps(state.to_dict())) + + +def _counts(state: DurableAgentState) -> tuple[int, int]: + history = state.data.conversation_history + return sum(len(entry.messages) for entry in history), len(history) + + +class _Storage(AgentEntityStateProviderMixin): + def __init__(self, state: DurableAgentState, failure: BaseException | None = None) -> None: + self.raw = state.to_dict() + self.failure = failure + self.attempts: list[dict[str, Any]] = [] + + def _get_state_dict(self) -> dict[str, Any]: + return deepcopy(self.raw) + + def _set_state_dict(self, state: dict[str, Any]) -> None: + self.attempts.append(deepcopy(state)) + if self.failure is not None: + raise self.failure + self.raw = json.loads(json.dumps(state)) + + def _get_session_id_from_entity(self) -> str: + return "private-session" + + +def _entity( + storage: _Storage, *, budget: int | None = BUDGET, history: DurableHistoryProvider | None = None +) -> AgentEntity: + client: Any = RecordingChatClient() + return AgentEntity( + Agent(client=client, context_providers=[history] if history is not None else None), + state_provider=storage, + max_state_bytes=budget, + ) + + +async def test_under_budget_records_one_unchanged_size_pair(reader: InMemoryMetricReader) -> None: + state = _state(1) + before = state.to_dict() + assert await retention.enforce_budget(state, max_state_bytes=BUDGET) == 0 + assert state.to_dict() == before + metrics = _metrics(reader) + assert set(metrics) == {"evaluations", "budget", "state.size"} + attrs = _attributes(outcome="below_threshold") + _counter(metrics["evaluations"], attrs, 1) + _histogram(metrics["budget"], attrs, BUDGET) + _histogram(metrics["state.size"], {**attrs, "phase": "before"}, _size(state)) + _histogram(metrics["state.size"], {**attrs, "phase": "after"}, _size(state)) + + +async def test_pressure_reports_only_applied_plan_and_exact_bytes(reader: InMemoryMetricReader) -> None: + state = _state() + before_bytes = _size(state) + before_messages, before_entries = _counts(state) + removed = await retention.enforce_budget(state, max_state_bytes=BUDGET) + assert removed > 0 + after_messages, after_entries = _counts(state) + assert removed == before_messages - after_messages + assert state.data.truncation is not None + assert state.data.truncation["evictedMessageCount"] == removed + metrics = _metrics(reader) + assert set(metrics) == { + "evaluations", + "budget", + "state.size", + "removed_messages", + "removed_entries", + "reclaimed_bytes", + } + attrs = _attributes() + _counter(metrics["evaluations"], attrs, 1) + _counter(metrics["removed_messages"], attrs, removed) + _counter(metrics["removed_entries"], attrs, before_entries - after_entries) + _counter(metrics["reclaimed_bytes"], attrs, before_bytes - _size(state)) + _histogram(metrics["budget"], attrs, BUDGET) + _histogram(metrics["state.size"], {**attrs, "phase": "before"}, before_bytes) + _histogram(metrics["state.size"], {**attrs, "phase": "after"}, _size(state)) + + +@pytest.mark.parametrize("floor", [True, False]) +async def test_capacity_failure_never_reports_detached_deletion( + reader: InMemoryMetricReader, monkeypatch: pytest.MonkeyPatch, floor: bool +) -> None: + state = _state() + if floor: + state.data.session = {"private_control": "p" * BUDGET} + strategy = Mock(side_effect=AssertionError("floor must be checked before planning")) + else: + + async def insufficient_plan(messages: list[Message]) -> bool: + messages[0].additional_properties["_excluded"] = True + return True + + # A detached plan deletes something but cannot reach the byte target. + strategy = Mock(return_value=AsyncMock(side_effect=insufficient_plan)) + monkeypatch.setattr(retention, "TokenBudgetComposedStrategy", strategy) + before = state.to_dict() + with pytest.raises(retention.StateCapacityError): + await retention.enforce_budget(state, max_state_bytes=BUDGET) + assert state.to_dict() == before + assert strategy.call_count == (0 if floor else 3) + metrics = _metrics(reader) + assert set(metrics) == {"evaluations", "budget", "state.size", "capacity_failures"} + attrs = _attributes(outcome="protected_floor" if floor else "unreachable_target") + _counter(metrics["evaluations"], attrs, 1) + _counter(metrics["capacity_failures"], attrs, 1) + _histogram(metrics["budget"], attrs, BUDGET) + for phase in ("before", "after"): + _histogram(metrics["state.size"], {**attrs, "phase": phase}, _size(state)) + + +@pytest.mark.parametrize("failure", [None, OSError("private write failure"), asyncio.CancelledError()]) +async def test_entity_write_outcomes_remain_unconfirmed_and_failure_rolls_back( + reader: InMemoryMetricReader, failure: BaseException | None +) -> None: + storage = _Storage(_state(), failure) + entity = _entity(storage) + original = entity.state + before = deepcopy(storage.raw) + if failure is None: + await entity.run({"message": "new", "correlationId": "private-current"}) + assert storage.raw != before + else: + with pytest.raises(type(failure)) as caught: + await entity.run({"message": "new", "correlationId": "private-current"}) + assert caught.value is failure + assert entity.state is original + assert entity.state.to_dict() == before + assert storage.raw == before + assert len(storage.attempts) == 1 + attempted = DurableAgentState.from_dict(storage.attempts[0]) + assert attempted.data.truncation is not None + removed = attempted.data.truncation["evictedMessageCount"] + assert removed > 0 + metrics = _metrics(reader) + _counter(metrics["removed_messages"], _attributes(), removed) + attrs = { + "outcome": "returned" if failure is None else "failed", + "commit_status": "unknown", + "deletion_staged": True, + } + _counter(metrics["operations"], attrs, 1) + _counter(metrics["write_attempts"], {**attrs, "stage": "set_state"}, 1) + assert telemetry._current(attempted) is None + + +async def test_floor_failure_has_no_host_write_attempt(reader: InMemoryMetricReader) -> None: + state = _state() + state.data.session = {"session_id": "private-session", "state": {"private_control": "p" * BUDGET}} + storage = _Storage(state) + entity = _entity(storage) + original = entity.state + with pytest.raises(retention.StateCapacityError): + await entity.run({"message": "new", "correlationId": "private-current"}) + assert entity.state is original + assert storage.attempts == [] + metrics = _metrics(reader) + assert "write_attempts" not in metrics + assert "removed_messages" not in metrics + _counter( + metrics["operations"], + {"outcome": "failed", "commit_status": "not_attempted", "deletion_staged": False}, + 1, + ) + + +async def test_eager_only_flush_measures_after_truncation_and_never_claims_confirmation( + reader: InMemoryMetricReader, +) -> None: + state = _state(5) + for entry in state.data.conversation_history[:2]: + entry.messages[0].extension_data = {"_excluded": True} + storage = _Storage(state) + history = DurableHistoryProvider(prune_excluded=True) + # Real entity -> provider flush -> set_state path, with pressure budgeting disabled. + await _entity(storage, budget=None, history=history).run({"message": "new", "correlationId": "private-current"}) + metrics = _metrics(reader) + assert "budget" not in metrics + assert "capacity_failures" not in metrics + attrs = _attributes("eager") + _counter(metrics["evaluations"], attrs, 1) + _counter(metrics["removed_messages"], attrs, 2) + _counter(metrics["removed_entries"], attrs, 2) + # This flush precedes the new append/mailbox, so derive its independent boundary oracle. + after = deepcopy(state) + after.data.conversation_history = after.data.conversation_history[2:] + after.data.truncation = storage.raw["data"]["truncation"] + _histogram(metrics["state.size"], {**attrs, "phase": "before"}, _size(state)) + _histogram(metrics["state.size"], {**attrs, "phase": "after"}, _size(after)) + _counter(metrics["reclaimed_bytes"], attrs, _size(state) - _size(after)) + _counter(metrics["operations"], {"outcome": "returned", "commit_status": "unknown", "deletion_staged": True}, 1) + + +async def test_eager_protected_exclusions_do_not_serialize_or_count_deletion( + reader: InMemoryMetricReader, monkeypatch: pytest.MonkeyPatch +) -> None: + state = _state(1) + for entry in state.data.conversation_history: + entry.messages[0].extension_data = {"_excluded": True} + storage = _Storage(state) + history = DurableHistoryProvider(prune_excluded=True) + token = bind_durable_history(DurableHistoryBinding(storage)) + try: + bag: dict[str, Any] = {} + await history.get_messages(None, state=bag) + serialized = Mock(side_effect=AssertionError("protected exclusions need no telemetry serialization")) + monkeypatch.setattr(_history_provider, "eager_state_size", serialized) + before = storage.state.to_dict() + history.flush(bag) + assert storage.state.to_dict() == before + serialized.assert_not_called() + finally: + unbind_durable_history(token) + metrics = _metrics(reader) + assert set(metrics) == {"evaluations"} + _counter(metrics["evaluations"], _attributes("eager", "protected"), 1) + + +async def test_concurrent_scopes_do_not_share_write_status_or_deletion(reader: InMemoryMetricReader) -> None: + ready = asyncio.Event() + release = asyncio.Event() + deleting = _Storage(_state()) + small = _Storage(_state(1)) + + async def evict() -> None: + with telemetry.retention_operation(deleting.state): + await retention.enforce_budget(deleting.state, max_state_bytes=BUDGET) + ready.set() + await release.wait() + deleting.persist_state() + + async def check() -> None: + await ready.wait() + with telemetry.retention_operation(small.state): + await retention.enforce_budget(small.state, max_state_bytes=BUDGET) + release.set() + + await asyncio.gather(evict(), check()) + metrics = _metrics(reader) + assert len(metrics["operations"].data.data_points) == 2 + _counter(metrics["operations"], {"outcome": "returned", "commit_status": "unknown", "deletion_staged": True}, 1) + _counter( + metrics["operations"], + {"outcome": "returned", "commit_status": "not_attempted", "deletion_staged": False}, + 1, + ) + assert len(metrics["write_attempts"].data.data_points) == 1 + + +async def test_nested_scope_state_identity_and_closed_inherited_context(reader: InMemoryMetricReader) -> None: + outer = _Storage(_state(1)) + inner = _Storage(_state(1)) + release = asyncio.Event() + + async def inherited() -> None: + await release.wait() + # A task copied the ContextVar, but its originating operation has ended. + assert telemetry._current(outer.state) is None + await retention.enforce_budget(outer.state, max_state_bytes=BUDGET) + outer.persist_state() + + with telemetry.retention_operation(outer.state): + await retention.enforce_budget(outer.state, max_state_bytes=BUDGET) + assert telemetry._current(inner.state) is None + inner.persist_state() # A different provider must not mark outer as written. + with telemetry.retention_operation(inner.state): + await retention.enforce_budget(inner.state, max_state_bytes=BUDGET) + inner.persist_state() + assert telemetry._current(outer.state) is not None + task = asyncio.create_task(inherited()) + release.set() + await task + metrics = _metrics(reader) + _counter(metrics["evaluations"], _attributes(outcome="below_threshold"), 3) + _counter( + metrics["operations"], + {"outcome": "returned", "commit_status": "not_attempted", "deletion_staged": False}, + 1, + ) + _counter(metrics["operations"], {"outcome": "returned", "commit_status": "unknown", "deletion_staged": False}, 1) + _counter( + metrics["write_attempts"], + {"stage": "set_state", "outcome": "returned", "commit_status": "unknown", "deletion_staged": False}, + 1, + ) + + +async def test_dimensions_are_exact_bounded_values_and_never_state_data(reader: InMemoryMetricReader) -> None: + await _entity(_Storage(_state())).run({"message": "private input", "correlationId": "private-request"}) + dimensions: dict[str, set[Any]] = { + "mechanism": {"pressure", "eager"}, + "outcome": {"staged", "returned"}, + "commit_status": {"not_attempted", "unknown"}, + "phase": {"before", "after"}, + "stage": {"set_state"}, + "deletion_staged": {True, False}, + } + for metric in _metrics(reader).values(): + for point in metric.data.data_points: + assert point.attributes is not None + for name, value in point.attributes.items(): + assert name in dimensions + assert value in dimensions[name] + + +@pytest.mark.parametrize("mechanism", ["eager", "pressure"]) +async def test_telemetry_on_off_and_broken_meter_do_not_change_state_or_return( + reader: InMemoryMetricReader, monkeypatch: pytest.MonkeyPatch, mechanism: str +) -> None: + initial = _state() + for entry in initial.data.conversation_history[:2]: + entry.messages[0].extension_data = {"_excluded": True} + + async def snapshot() -> tuple[int, dict[str, Any]]: + state = deepcopy(initial) + if mechanism == "pressure": + removed = await retention.enforce_budget(state, max_state_bytes=BUDGET) + else: + storage = _Storage(state) + state = storage.state + DurableHistoryProvider._prune( + DurableHistoryBinding(storage), + [(entry, entry.messages[0]) for entry in state.data.conversation_history[:2]], + ) + assert state.data.truncation is not None + removed = state.data.truncation["evictedMessageCount"] + return removed, state.to_dict() + + expected = await snapshot() + assert expected[0] > 0 + assert _metrics(reader) + for get_meter in (NoOpMeterProvider().get_meter, Mock(side_effect=RuntimeError("broken instrumentation"))): + monkeypatch.setattr(telemetry, "get_meter", get_meter) + telemetry._instruments.cache_clear() + assert await snapshot() == expected + + +async def test_no_budget_and_no_eager_deletion_does_not_initialize_metrics( + monkeypatch: pytest.MonkeyPatch, +) -> None: + meter = Mock(side_effect=AssertionError("ordinary writes must not initialize retention instruments")) + monkeypatch.setattr(telemetry, "get_meter", meter) + telemetry._instruments.cache_clear() + await _entity(_Storage(_state(1)), budget=None).run({"message": "new", "correlationId": "private-current"}) + meter.assert_not_called() + + +def test_commit_serialization_failure_keeps_not_attempted_status( + reader: InMemoryMetricReader, monkeypatch: pytest.MonkeyPatch +) -> None: + storage = _Storage(_state(1)) + state = storage.state + failure = ValueError("private serialization failure") + with pytest.raises(ValueError) as caught, telemetry.retention_operation(state): + telemetry.record_retention(state, mechanism="eager", outcome="staged", removed_messages=1) + monkeypatch.setattr(DurableAgentState, "to_dict", Mock(side_effect=failure)) + storage.persist_state() + assert caught.value is failure + assert storage.attempts == [] + attrs = {"outcome": "failed", "commit_status": "not_attempted", "deletion_staged": True} + metrics = _metrics(reader) + _counter(metrics["operations"], attrs, 1) + _counter(metrics["write_attempts"], {**attrs, "stage": "serialization"}, 1) + + +async def test_eager_then_pressure_in_one_operation_accumulates_without_double_counting( + reader: InMemoryMetricReader, +) -> None: + storage = _Storage(_state()) + state = storage.state + before_messages, before_entries = _counts(state) + initial_bytes = _size(state) + with telemetry.retention_operation(state): + DurableHistoryProvider._prune( + DurableHistoryBinding(storage), + [(entry, entry.messages[0]) for entry in state.data.conversation_history[:2]], + ) + eager_bytes = _size(state) + pressure_removed = await retention.enforce_budget(state, max_state_bytes=BUDGET) + assert pressure_removed > 0 + storage.persist_state() + assert state.data.truncation is not None + assert state.data.truncation["evictedMessageCount"] == pressure_removed + 2 + after_messages, after_entries = _counts(state) + assert before_messages - after_messages == before_entries - after_entries == pressure_removed + 2 + metrics = _metrics(reader) + for metric in ("removed_messages", "removed_entries"): + _counter(metrics[metric], _attributes("eager"), 2) + _counter(metrics[metric], _attributes(), pressure_removed) + _counter(metrics["reclaimed_bytes"], _attributes("eager"), initial_bytes - eager_bytes) + _counter(metrics["reclaimed_bytes"], _attributes(), eager_bytes - _size(state)) + _counter(metrics["operations"], {"outcome": "returned", "commit_status": "unknown", "deletion_staged": True}, 1) + + +def test_explicit_noop_skips_eager_serialization(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(telemetry, "get_meter", NoOpMeterProvider().get_meter) + telemetry._instruments.cache_clear() + serialized = Mock(side_effect=AssertionError("no-op instrumentation must not serialize")) + monkeypatch.setattr(DurableAgentState, "to_dict", serialized) + try: + assert telemetry.eager_state_size(DurableAgentState()) is None + serialized.assert_not_called() + finally: + telemetry._instruments.cache_clear() + + +async def test_recording_failure_does_not_replace_capacity_error( + reader: InMemoryMetricReader, monkeypatch: pytest.MonkeyPatch +) -> None: + instruments = telemetry._instruments() + broken = Mock(side_effect=RuntimeError("reader failure")) + monkeypatch.setattr(instruments.evaluations, "add", broken) + state = _state(1) + state.data.session = {"private_control": "p" * BUDGET} + before = state.to_dict() + with pytest.raises(retention.StateCapacityError) as caught: + await retention.enforce_budget(state, max_state_bytes=BUDGET) + assert caught.value.size_bytes == _size(state) + assert state.to_dict() == before + broken.assert_called_once() diff --git a/python/packages/durabletask/tests/test_revision_contract.py b/python/packages/durabletask/tests/test_revision_contract.py new file mode 100644 index 0000000..a893da3 --- /dev/null +++ b/python/packages/durabletask/tests/test_revision_contract.py @@ -0,0 +1,230 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Regression tests for the revised durable execution and history contract.""" + +import json +from copy import deepcopy +from datetime import datetime, timedelta, timezone +from typing import Any + +import pytest +from agent_framework import Agent, AgentSession, HistoryProvider, InMemoryHistoryProvider, Message +from test_durable_history_provider import RecordingChatClient + +from agent_framework_durabletask import ( + AgentEntity, + AgentEntityStateProviderMixin, + DurableAgentState, + DurableAgentStateErrorResponse, + DurableAgentStateMessage, + DurableAgentStateRequest, + RunRequest, +) +from agent_framework_durabletask._history_provider import ensure_durable_history +from agent_framework_durabletask._retention import DEFAULT_MAX_STATE_BYTES, DEFAULT_RETENTION, enforce_budget + + +class JsonStateProvider(AgentEntityStateProviderMixin): + """Storage boundary that never aliases staged state and supports cold reloads.""" + + def __init__(self, raw: dict[str, Any] | None = None) -> None: + self.raw = deepcopy(raw or {}) + self.writes = 0 + self.fail_writes = False + + def _get_state_dict(self) -> dict[str, Any]: + return deepcopy(self.raw) + + def _set_state_dict(self, state: dict[str, Any]) -> None: + if self.fail_writes: + raise OSError("injected commit failure") + self.raw = json.loads(json.dumps(state)) + self.writes += 1 + + def _get_session_id_from_entity(self) -> str: + return "revision-session" + + +class ExternalHistory(HistoryProvider): + def __init__(self) -> None: + super().__init__("external") + self.messages: list[Message] = [] + + async def get_messages(self, session_id: str | None, **kwargs: Any) -> list[Message]: + return list(self.messages) + + async def save_messages(self, session_id: str | None, messages: Any, **kwargs: Any) -> None: + self.messages.extend(messages) + + +def make_agent(client: Any, providers: list[Any] | None = None) -> Agent: + return Agent(client=client, name="revision", context_providers=providers) + + +def test_retention_defaults_do_not_enable_deletion() -> None: + assert DEFAULT_RETENTION == "keep_all" + assert DEFAULT_MAX_STATE_BYTES is None + + +def test_multiple_primary_providers_fail_without_mutating_agent() -> None: + providers = [InMemoryHistoryProvider("first"), InMemoryHistoryProvider("second")] + agent = make_agent(RecordingChatClient(), providers) + with pytest.raises(ValueError, match="primary"): + ensure_durable_history(agent) + assert agent.context_providers == providers + + +@pytest.mark.parametrize("context", [["bad"], [1], [{}, None], "not-a-list", {}]) +def test_malformed_context_is_rejected_at_the_request_boundary(context: Any) -> None: + with pytest.raises(ValueError, match="contextMessages"): + RunRequest.from_dict({"message": "input", "correlationId": "c0", "contextMessages": context}) + + +def test_empty_projection_does_not_become_the_unfiltered_input() -> None: + request = RunRequest(message="must not leak", correlation_id="empty", context_messages=[]) + restored = RunRequest.from_dict(request.to_dict()) + assert restored.context_messages == [] + assert DurableAgentStateRequest.from_run_request(restored).messages == [] + + +async def test_original_response_survives_transcript_mutation_and_cold_reload() -> None: + provider = JsonStateProvider() + client = RecordingChatClient() + entity = AgentEntity(make_agent(client), state_provider=provider) + response = await entity.run({"message": "first", "correlationId": "c0"}) + original = response.to_dict() + entity.state.data.conversation_history.clear() + entity.persist_state() + restored = AgentEntity(make_agent(client), state_provider=JsonStateProvider(provider.raw)) + duplicate = await restored.run({"message": "first", "correlationId": "c0"}) + assert duplicate.to_dict() == original + assert len(client.received_messages) == 1 + + +async def test_external_history_needs_no_contentless_request_mirror() -> None: + external = ExternalHistory() + provider = JsonStateProvider() + entity = AgentEntity(make_agent(RecordingChatClient(), [external]), state_provider=provider) + await entity.run({"message": "first", "correlationId": "external-0"}) + assert [m.text for m in external.messages] == ["first", "reply-1"] + assert entity.state.data.conversation_history == [] + assert entity.state.try_get_agent_response("external-0") is not None + + +async def test_external_reset_does_not_claim_to_clear_an_untouched_store() -> None: + external = ExternalHistory() + provider = JsonStateProvider() + entity = AgentEntity(make_agent(RecordingChatClient(), [external]), state_provider=provider) + await entity.run({"message": "first", "correlationId": "external-0"}) + before = deepcopy(provider.raw) + with pytest.raises(NotImplementedError, match="external"): + entity.reset() + assert provider.raw == before + + +def test_sparse_context_is_not_lost_after_an_older_position_was_skipped() -> None: + provider = JsonStateProvider() + entity = AgentEntity(make_agent(RecordingChatClient()), state_provider=provider) + + def deliver(positions: list[int]) -> list[int]: + messages = [ + DurableAgentStateMessage.from_chat_message( + Message("user", [str(position)], message_id=f"wf_source_{position}") + ) + for position in positions + ] + return [int(m.text) for m in entity._drop_already_stored(messages)] + + assert deliver([1, 3]) == [1, 3] + entity.state.data.conversation_history.clear() + entity.persist_state() + entity = AgentEntity(make_agent(RecordingChatClient()), state_provider=JsonStateProvider(provider.raw)) + assert deliver([2, 4]) == [2, 4] + + +async def test_unreachable_protected_floor_does_not_destroy_old_history() -> None: + state = DurableAgentState() + old = datetime.now(timezone.utc) - timedelta(hours=1) + for index in range(10): + state.data.conversation_history.append( + DurableAgentStateRequest( + correlation_id=f"old-{index}", + created_at=old, + messages=[DurableAgentStateMessage.from_chat_message(Message("user", ["x" * 1000]))], + ) + ) + state.data.session = {"protected": "s" * 30_000} + before = state.to_json() + with pytest.raises(ValueError, match="[Cc]apacity|budget|floor"): + await enforce_budget(state, max_state_bytes=12_000) + assert state.to_json() == before + + +async def test_old_failed_turns_are_storage_candidates_not_model_history() -> None: + state = DurableAgentState() + old = datetime.now(timezone.utc) - timedelta(hours=1) + for index in range(80): + state.data.conversation_history.append( + DurableAgentStateErrorResponse( + correlation_id=f"failed-{index}", + created_at=old, + messages=[DurableAgentStateMessage.from_chat_message(Message("assistant", ["e" * 400]))], + ) + ) + assert await enforce_budget(state, max_state_bytes=12_000) > 0 + assert len(state.to_json()) < 12_000 + + +async def test_commit_failure_does_not_leave_an_in_memory_completed_request() -> None: + provider = JsonStateProvider() + client = RecordingChatClient() + entity = AgentEntity(make_agent(client), state_provider=provider) + provider.fail_writes = True + with pytest.raises(OSError, match="commit failure"): + await entity.run({"message": "first", "correlationId": "c0"}) + assert provider.raw == {} + assert entity.state.try_get_agent_response("c0") is None + provider.fail_writes = False + await entity.run({"message": "first", "correlationId": "c0"}) + assert len(client.received_messages) == 2 + assert provider.writes == 1 + + +async def test_inactive_service_id_is_not_sent_to_a_client_owned_run() -> None: + from agent_framework import AgentResponse + + seen: list[Any] = [] + + class ServiceAgent: + name = "service" + client = type("Client", (), {"STORES_BY_DEFAULT": True})() + context_providers: list[Any] = [] + + def create_session(self, **kwargs: Any) -> AgentSession: + return AgentSession(**kwargs) + + async def run(self, *, session: AgentSession, stream: bool = False, **kwargs: Any) -> AgentResponse: + if stream: + raise TypeError("stream is not supported") + seen.append(session.service_session_id) + if kwargs["options"].get("store", True): + session.service_session_id = "service-branch" + return AgentResponse(messages=[Message("assistant", ["ok"])]) + + provider = JsonStateProvider() + for index, store in enumerate([True, False, True]): + agent: Any = ServiceAgent() + entity = AgentEntity(agent, state_provider=provider) + await entity.run({"message": f"m{index}", "correlationId": f"c{index}", "options": {"store": store}}) + provider = JsonStateProvider(provider.raw) + assert seen == [None, None, "service-branch"] + + +def test_unknown_state_data_and_entry_fields_survive_round_trip() -> None: + state = DurableAgentState("1.2.0").to_dict() + state["futureRoot"] = {"opaque": [1, 2]} + state["data"]["futureData"] = {"opaque": [3, 4]} + state["data"]["conversationHistory"] = [ + {"$type": "futureKind", "correlationId": "x", "createdAt": "2026-01-01T00:00:00+00:00", "extra": 9} + ] + assert DurableAgentState.from_dict(state).to_dict() == state diff --git a/python/packages/durabletask/tests/test_state_fidelity_review.py b/python/packages/durabletask/tests/test_state_fidelity_review.py new file mode 100644 index 0000000..0dee7d7 --- /dev/null +++ b/python/packages/durabletask/tests/test_state_fidelity_review.py @@ -0,0 +1,432 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Transcript and mailbox fidelity across a real JSON storage boundary.""" + +import json +from copy import deepcopy +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, cast, get_args, get_type_hints + +import jsonschema +import pytest +from agent_framework import AgentResponse, Content, Message +from pydantic import BaseModel, Field + +from agent_framework_durabletask._constants import ContentTypes +from agent_framework_durabletask._durable_agent_state import ( + DurableAgentState, + DurableAgentStateContent, + DurableAgentStateEntryJsonType, + DurableAgentStateMessage, + DurableAgentStateRequest, + DurableAgentStateResponse, + DurableAgentStateTextContent, + DurableAgentStateUnknownContent, + DurableAgentStateUsage, +) +from agent_framework_durabletask._message_identity import message_identity +from agent_framework_durabletask._models import RunRequest +from agent_framework_durabletask._response_utils import ensure_response_format + +NOW = datetime(2026, 9, 9, tzinfo=timezone.utc) + + +@pytest.fixture(scope="module") +def schema() -> dict[str, Any]: + path = Path(__file__).resolve().parents[4] / "schemas" / "durable-agent-entity-state.json" + return json.loads(path.read_text(encoding="utf-8")) + + +def _state(message: Message) -> DurableAgentState: + state = DurableAgentState() + state.data.conversation_history = [ + DurableAgentStateResponse.from_run_response( + "response", AgentResponse(messages=[message], created_at=NOW.isoformat()) + ) + ] + return state + + +def _stored_message(state: DurableAgentState) -> DurableAgentStateMessage: + return state.data.conversation_history[0].messages[0] + + +def _cold(state: DurableAgentState, schema: dict[str, Any]) -> DurableAgentState: + payload = json.loads(state.to_json()) + jsonschema.Draft202012Validator(schema, format_checker=jsonschema.FormatChecker()).validate(payload) + return DurableAgentState.from_dict(payload) + + +@pytest.mark.parametrize("kind", get_args(get_type_hints(Content.__init__)["type"])) +def test_every_core_kind_preserves_metadata_in_transcript(kind: Any, schema: dict[str, Any]) -> None: + citation = { + "type": "citation", + "title": "Source", + "url": "https://example.test/source", + "annotated_regions": [{"type": "text_span", "start_index": 0, "end_index": 4, "future": [1]}], + "additional_properties": {"type": "opaque", "nested": [2]}, + } + # Exercise cross-subtype fields too. Core's constructor, not a copied list of kinds, + # defines the category; typed durable fields must coexist with the remaining fields. + content = Content( + kind, + text="body", + uri="data:image/png;base64,AA==", + call_id="call", + name="lookup", + file_id="file", + vector_store_id="vector", + usage_details={"input_token_count": 0}, + protected_data="protected", + informational_only=True, + id="content", + exception="retry", + annotations=[cast(Any, citation)], + additional_properties={"type": "text", "future": {"keep": [3]}}, + raw_representation=object(), + ) + original = Message( + "developer", + [content], + author_name="author", + message_id="id", + additional_properties={"nested": {"keep": [4]}}, + raw_representation=object(), + ) + restored = _stored_message(_cold(_state(original), schema)).to_chat_message() + + assert restored.to_dict() == original.to_dict() + restored.additional_properties["nested"]["keep"].append(5) + assert original.additional_properties["nested"]["keep"] == [4] + + +def test_typed_content_mapping_covers_shared_schema_kinds(schema: dict[str, Any]) -> None: + known = {value for key, value in vars(ContentTypes).items() if key.isupper()} + branches = schema["$defs"]["chatContentItem"]["oneOf"] + declared = { + schema["$defs"][branch["$ref"].split("/")[-1]]["properties"]["$type"]["const"] + for branch in branches + if "$ref" in branch + } + assert declared == known + opaque = next(branch for branch in branches if "$ref" not in branch) + assert set(opaque["properties"]["$type"]["not"]["enum"]) == known + assert {cls.type for cls in DurableAgentStateContent.__subclasses__() if cls.type} == known + for kind in known - {"unknown"}: + core_kind = { + "functionCall": "function_call", + "functionResult": "function_result", + "hostedFile": "hosted_file", + "hostedVectorStore": "hosted_vector_store", + "reasoning": "text_reasoning", + }.get(kind, kind) + content = Content( + core_kind, + text="text", + uri="https://example.test", + call_id="c", + name="f", + file_id="file", + vector_store_id="vector", + usage_details={}, + ) + stored = DurableAgentStateContent.from_ai_content(content) + assert stored.type == kind + assert not isinstance(stored, DurableAgentStateUnknownContent) + + +def test_function_result_retains_binary_and_text_items_without_a_transcript_mirror(schema: dict[str, Any]) -> None: + content = Content.from_function_result( + "call", + result=[Content.from_text("answer"), Content.from_data(b"\x00\xff", "image/png")], + exception="recoverable", + additional_properties={"future": [1]}, + ) + original = Message("tool", [content]) + state = _state(original) + raw = _stored_message(state).to_dict()["contents"][0] + + assert raw["$type"] == "functionResult" + assert raw["result"] == content.result + overlay = raw["extensionData"]["coreContent"] + assert overlay["items"] == content.to_dict()["items"] + assert not {"type", "call_id", "result"} & overlay.keys() + assert not {"coreMessage", "core_message"} & _stored_message(state).to_dict().keys() + restored = _stored_message(_cold(state, schema)).to_chat_message() + assert restored.to_dict() == original.to_dict() + assert restored.contents[0].items is not None + assert isinstance(restored.contents[0].items[1], Content) + + +def test_current_known_content_changes_win_over_metadata(schema: dict[str, Any]) -> None: + original = Message("assistant", [Content.from_text("original", additional_properties={"nested": [1]})]) + cold = _cold(_state(original), schema) + stored = _stored_message(cold) + assert isinstance(stored.contents[0], DurableAgentStateTextContent) + raw = stored.to_dict()["contents"][0] + assert "text" not in raw["extensionData"]["coreContent"] + stored.contents[0].text = "edited" + restored = _stored_message(_cold(cold, schema)).to_chat_message() + assert restored.text == "edited" + assert restored.contents[0].additional_properties == {"nested": [1]} + restored.contents[0].additional_properties["nested"].append(2) + assert stored.to_dict()["contents"][0]["extensionData"]["coreContent"]["additional_properties"] == {"nested": [1]} + + +@pytest.mark.parametrize("arguments", [None, {}, {"type": "text", "opaque": [1]}, '{"x":1}', "{unfinished"]) +def test_arguments_remain_exact_not_reparsed_or_reformatted(arguments: Any, schema: dict[str, Any]) -> None: + original = Message("assistant", [Content.from_function_call("call", "f", arguments=arguments)]) + restored = _stored_message(_cold(_state(original), schema)).to_chat_message() + assert restored.to_dict() == original.to_dict() + + +def test_uri_without_media_type_and_partial_usage_validate(schema: dict[str, Any]) -> None: + original = Message( + "user", [Content.from_uri("https://example.test/file"), Content.from_usage({"input_token_count": 0})] + ) + restored = _cold(_state(original), schema) + assert _stored_message(restored).to_chat_message().to_dict() == original.to_dict() + usage = DurableAgentStateUsage.from_dict({"inputTokenCount": 0, "future": {"nested": [1]}}) + assert usage.to_dict() == {"inputTokenCount": 0, "future": {"nested": [1]}} + assert usage.to_usage_details() == {"input_token_count": 0} + + +@pytest.mark.parametrize("kind", list(DurableAgentStateEntryJsonType)) +def test_unknown_fields_are_owned_by_actual_entry_subtype(kind: str, schema: dict[str, Any]) -> None: + entry: dict[str, Any] = { + "$type": kind, + "createdAt": NOW.isoformat(), + "messages": [ + { + "role": "assistant", + "futureMessage": {"nested": [1]}, + "contents": [ + {"$type": "text", "text": "hello", "callId": {"future": [2]}, "futureContent": {"nested": [3]}} + ], + } + ], + "futureEntry": {"nested": [4]}, + } + if kind != "request": + entry.update(responseType="future-format", responseSchema={"future": [5]}, orchestrationId="future-id") + if kind not in ("response", "errorResponse"): + entry["usage"] = {"future": {"nested": [6]}} + else: + entry["usage"] = {"inputTokenCount": 1, "future": {"nested": [6]}} + payload = {"schemaVersion": "2.1.0", "data": {"conversationHistory": [entry]}} + before = deepcopy(payload) + loaded = DurableAgentState.from_dict(payload) + assert _cold(loaded, schema).to_dict() == before + serialized = loaded.to_dict() + serialized["data"]["conversationHistory"][0]["messages"][0]["futureMessage"]["nested"].append(9) + assert payload == before + assert loaded.to_dict() == before + + +def test_core_context_preserves_nested_future_items_before_consumer_filtering(schema: dict[str, Any]) -> None: + raw: dict[str, Any] = { + "role": "tool", + "future_message": {"nested": [1]}, + "contents": [ + { + "type": "function_result", + "call_id": "call", + "result": "text", + "items": [ + {"type": "text", "text": "text", "future_content": {"nested": [2]}}, + Content.from_data(b"data", "image/png").to_dict(), + ], + "future_outer": {"nested": [3]}, + } + ], + "additional_properties": {"nested": [4]}, + } + request = RunRequest("", "c", context_messages=[raw]) + entry = DurableAgentStateRequest.from_run_request(request) + assert entry.messages[0].message_id is None + assert entry.messages[0].ingestion_identity == message_identity(entry.messages[0].to_chat_message()) + state = DurableAgentState() + state.data.conversation_history = [entry] + loaded = _cold(state, schema) + stored = _stored_message(loaded).to_dict() + assert stored["future_message"] == raw["future_message"] + overlay = stored["contents"][0]["extensionData"]["coreContent"] + assert overlay["items"] == raw["contents"][0]["items"] + assert overlay["future_outer"] == {"nested": [3]} + items = _stored_message(loaded).to_chat_message().contents[0].items + assert items is not None + assert items[0].text == "text" + assert loaded.to_dict() == state.to_dict() + + +def test_future_entry_and_content_are_opaque_even_with_unfamiliar_shapes(schema: dict[str, Any]) -> None: + state = _state(Message("assistant", ["hello"])) + raw = state.to_dict() + future_entry = {"$type": "futureEntry", "messages": {"futureShape": [None]}, "usage": [1]} + future_content = {"$type": "futureContent", "payload": None, "items": {"futureShape": [2]}} + raw["data"]["conversationHistory"].append(future_entry) + raw["data"]["conversationHistory"][0]["messages"][0]["contents"].append(future_content) + loaded = _cold(DurableAgentState.from_dict(raw), schema) + assert loaded.to_dict() == raw + assert loaded.data.conversation_history[-1].messages == [] + + +@pytest.mark.parametrize("level", ["history", "messages", "contents"]) +@pytest.mark.parametrize("malformed", [None, False, 7, "text", {}, [None], [42], ["text"]]) +def test_malformed_transcript_containers_fail_instead_of_dropping_data(level: str, malformed: Any) -> None: + raw = _state(Message("assistant", ["hello"])).to_dict() + data = raw["data"] + if level == "history": + data["conversationHistory"] = malformed + elif level == "messages": + data["conversationHistory"][0]["messages"] = malformed + else: + data["conversationHistory"][0]["messages"][0]["contents"] = malformed + with pytest.raises(ValueError): + DurableAgentState.from_dict(raw) + + +@pytest.mark.parametrize("number", [float("nan"), float("inf"), float("-inf")]) +@pytest.mark.parametrize("location", ["root", "session", "content"]) +def test_nonfinite_json_is_rejected_including_unknown_fields(number: float, location: str) -> None: + raw = _state(Message("assistant", ["hello"])).to_dict() + target = raw + if location == "session": + raw["data"]["session"] = target = {} + elif location == "content": + target = raw["data"]["conversationHistory"][0]["messages"][0]["contents"][0] + target["future"] = {"nested": [number]} + with pytest.raises(ValueError): + DurableAgentState.from_dict(raw) + with pytest.raises(ValueError): + DurableAgentState.from_json(json.dumps(raw)) + + +def _mailbox() -> dict[str, Any]: + state = DurableAgentState() + state.record_response("c", AgentResponse(messages=[Message("assistant", ["42"])]), delivery_window_seconds=60) + raw = state.to_dict() + raw["data"]["responseMailbox"]["c"].update(createdAt="2026-09-09T00:00:00Z", expiresAt="2099-01-01T00:00:00Z") + raw["data"]["completedCorrelations"]["c"]["completedAt"] = "2026-09-09T00:00:00Z" + return raw + + +def test_poll_uses_versioned_loader_preserving_null_and_future_envelope_fields(schema: dict[str, Any]) -> None: + raw = _mailbox() + response = raw["data"]["responseMailbox"]["c"]["response"] + response.update(value=None, future_response={"nested": [1]}) + response["messages"][0]["future_message"] = {"nested": [2]} + response["messages"][0]["contents"][0]["future_content"] = {"nested": [3]} + loaded = _cold(DurableAgentState.from_dict(raw), schema) + result = loaded.try_get_agent_response("c") + assert type(result) is AgentResponse + assert result.value is None + result.messages[0].contents[0].text = "changed" + assert loaded.to_dict() == raw + loaded.expire_responses(now=datetime(2100, 1, 1, tzinfo=timezone.utc)) + expired = loaded.try_get_agent_response("c") + assert expired is not None + assert expired.additional_properties["durable_status"] == "already_completed" + + +def test_poll_preserves_value_by_name_marker(schema: dict[str, Any]) -> None: + class Aliased(BaseModel): + count: int = Field(validation_alias="inputCount", serialization_alias="outputCount") + + raw = _mailbox() + response = raw["data"]["responseMailbox"]["c"]["response"] + response.update(value={"count": 7}, _durable_value_by_name=True) + result = _cold(DurableAgentState.from_dict(raw), schema).try_get_agent_response("c") + assert result is not None + ensure_response_format(Aliased, "c", result) + assert result.value == Aliased(inputCount=7) + + +@pytest.mark.parametrize("field", ["createdAt", "expiresAt", "completedAt"]) +@pytest.mark.parametrize( + "timestamp", + [ + None, + "2026-09-09", + "2026-09-09T00:00:00", + "2026-09-09T00:00:00+00:00\n", + "2026-02-30T00:00:00Z", + "2026-09-09T00:00:00+00:60", + ], +) +def test_new_delivery_timestamps_require_valid_rfc3339(field: str, timestamp: Any) -> None: + raw = _mailbox() + collection = "completedCorrelations" if field == "completedAt" else "responseMailbox" + raw["data"][collection]["c"][field] = timestamp + with pytest.raises(ValueError): + DurableAgentState.from_dict(raw) + + +def test_legacy_timestamp_tolerance_and_scalar_migration_remain_unchanged() -> None: + raw = { + "schemaVersion": "1.2.0", + "data": { + "ingestedPositions": {"source": 7}, + "conversationHistory": [ + {"$type": "request", "createdAt": "2026-09-09", "messages": []}, + ], + }, + } + state = DurableAgentState.from_dict(raw) + before = state.to_dict() + with pytest.raises(ValueError, match="delivery evidence"): + state.prepare_for_write(delivery_window_seconds=60) + assert state.to_dict() == before + + +@pytest.mark.parametrize( + ("kind", "field"), + [ + ("search_tool_result", "items"), + ("code_interpreter_tool_call", "inputs"), + ("code_interpreter_tool_result", "outputs"), + ("shell_tool_result", "outputs"), + ("function_approval_request", "function_call"), + ("function_approval_response", "function_call"), + ], +) +def test_nested_core_edges_are_reconstructed_in_transcript(kind: str, field: str, schema: dict[str, Any]) -> None: + nested = Content.from_function_result( + "call", result=[Content.from_text("nested"), Content.from_data(b"data", "image/png")] + ) + content_class: Any = Content + content: Content = content_class(kind, **{field: nested if field == "function_call" else [nested]}) + original = Message("tool", [content]) + restored = _stored_message(_cold(_state(original), schema)).to_chat_message() + assert restored.to_dict() == original.to_dict() + inner = getattr(restored.contents[0], field) + inner = inner if field == "function_call" else inner[0] + assert isinstance(inner, Content) + assert inner.items is not None + assert isinstance(inner.items[1], Content) + + +@pytest.mark.parametrize("invalid", [None, "text", {}, [None], ["text"], [42]]) +def test_mailbox_core_contents_reject_malformed_containers(invalid: Any) -> None: + raw = _mailbox() + raw["data"]["responseMailbox"]["c"]["response"]["messages"][0]["contents"] = invalid + with pytest.raises(ValueError): + DurableAgentState.from_dict(raw) + + +def test_z_delivery_timestamp_is_normalized_before_fromisoformat(monkeypatch: pytest.MonkeyPatch) -> None: + import agent_framework_durabletask._durable_agent_state as state_module + + class Python310Datetime(datetime): + @classmethod + def fromisoformat(cls, value: str) -> "Python310Datetime": + assert not value.endswith(("Z", "z")), "Python 3.10 does not accept Z directly" + return super().fromisoformat(value) + + raw = _mailbox() + monkeypatch.setattr(state_module, "datetime", Python310Datetime) + loaded = DurableAgentState.from_dict(raw) + assert loaded.try_get_agent_response("c") is not None + loaded.expire_responses(now=datetime(2100, 1, 1, tzinfo=timezone.utc)) + assert not loaded.data.response_mailbox diff --git a/python/packages/durabletask/tests/test_state_followup_review.py b/python/packages/durabletask/tests/test_state_followup_review.py new file mode 100644 index 0000000..708c7f4 --- /dev/null +++ b/python/packages/durabletask/tests/test_state_followup_review.py @@ -0,0 +1,412 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Focused state and detached migration regressions, without entity ownership changes.""" + +import json +from copy import deepcopy +from datetime import datetime, timezone +from pathlib import Path +from typing import Any +from unittest.mock import Mock + +import jsonschema +import pytest +from agent_framework import AgentResponse, Message + +from agent_framework_durabletask import _durable_agent_state as state_module +from agent_framework_durabletask import migrate_legacy_state, state_snapshot_digest +from agent_framework_durabletask._durable_agent_state import ( + DurableAgentState, + DurableAgentStateContent, + DurableAgentStateTextContent, +) +from agent_framework_durabletask._message_identity import message_identity +from agent_framework_durabletask._response_utils import is_terminal_agent_response, load_agent_response +from agent_framework_durabletask._workflows.naming import workflow_message_id + +NOW = datetime(2026, 9, 9, 12, tzinfo=timezone.utc) +OLD = datetime(2024, 1, 1, tzinfo=timezone.utc) +SESSION_ID = "dafx-agent:original-session" + + +@pytest.fixture(scope="module") +def schema() -> dict[str, Any]: + path = Path(__file__).resolve().parents[4] / "schemas" / "durable-agent-entity-state.json" + return json.loads(path.read_text(encoding="utf-8")) + + +def _validate(payload: dict[str, Any], schema: dict[str, Any]) -> None: + jsonschema.Draft202012Validator(schema, format_checker=jsonschema.FormatChecker()).validate(payload) + + +def _source(*, version: str = "1.1.0", contents: list[dict[str, Any]] | None = None) -> dict[str, Any]: + return { + "schemaVersion": version, + "data": { + "conversationHistory": [ + { + "$type": "request", + "correlationId": "turn", + "createdAt": OLD.isoformat(), + "messages": [ + { + "role": "user", + "messageId": "custom-id", + "contents": contents if contents is not None else [{"$type": "text", "text": "retained"}], + } + ], + } + ] + }, + } + + +def _migrate(source: dict[str, Any], *, evidence: dict[str, Any] | None = None) -> DurableAgentState: + return migrate_legacy_state( + source, + source_digest=state_snapshot_digest(source), + source_session_id=SESSION_ID, + migration_id="followup-migration", + ownership_transfer_id="authorized-transfer", + delivery_window_seconds=60, + delivery_evidence=evidence, + now=NOW, + ) + + +def _evidence(source: dict[str, Any], messages: list[Message]) -> dict[str, Any]: + return { + "sourceDigest": state_snapshot_digest(source), + "evidenceId": "complete-journal", + "complete": True, + "messages": [message.to_dict() for message in messages], + } + + +@pytest.mark.parametrize( + "invalid", + [ + pytest.param({1: "numeric", "1": "string"}, id="colliding-keys"), + pytest.param({False: "boolean-key"}, id="boolean-key"), + pytest.param({None: "null-key"}, id="null-key"), + pytest.param((1, "tuple"), id="tuple"), + pytest.param(object(), id="object"), + pytest.param(float("nan"), id="nan"), + pytest.param(float("inf"), id="infinity"), + pytest.param(float("-inf"), id="negative-infinity"), + ], +) +@pytest.mark.parametrize("location", ["root", "session", "content"]) +def test_ordinary_state_rejects_non_json_before_encoding(invalid: Any, location: str) -> None: + raw = _source(version="2.0.0") + target = raw + if location == "session": + target = {"session_id": SESSION_ID, "state": {}} + raw["data"]["session"] = target + elif location == "content": + target = raw["data"]["conversationHistory"][0]["messages"][0]["contents"][0] + target["future"] = {"nested": [invalid]} + + with pytest.raises(ValueError, match="strict JSON"): + DurableAgentState.from_dict(raw) + + # The write boundary uses the same validation, not only migration's digest. + state = DurableAgentState() + state.unknown_fields["future"] = {"nested": [invalid]} + with pytest.raises(ValueError, match="strict JSON"): + state.to_dict() + with pytest.raises(ValueError, match="strict JSON"): + state_snapshot_digest(raw) + + +def test_strict_json_snapshot_preserves_valid_values_and_detaches_them() -> None: + raw = _source(version="2.0.0") + raw["future"] = {"1": [None, False, 0, 0.0, "", [], {}, "雪"]} + before = deepcopy(raw) + state = DurableAgentState.from_dict(raw) + assert state.to_dict() == before + assert json.dumps(state.to_dict(), sort_keys=True) == json.dumps(before, sort_keys=True) + raw["future"]["1"].append("caller edit") + detached = state.to_dict() + detached["future"]["1"].append("consumer edit") + assert state.to_dict() == before + + +def test_ordinary_state_rejects_cycles_as_invalid_json() -> None: + raw = _source(version="2.0.0") + raw["cycle"] = raw + with pytest.raises(ValueError, match="strict JSON"): + DurableAgentState.from_dict(raw) + + +@pytest.mark.parametrize("invalid", [{1: "numeric", "1": "string"}, (1, 2), float("nan")]) +def test_mailbox_snapshot_rejects_non_json_without_staging_a_completion( + invalid: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + response = AgentResponse(messages=[]) + payload = {"type": "agent_response", "messages": [], "future": {"nested": invalid}} + monkeypatch.setattr(state_module, "serialize_agent_response", lambda _: payload) + state = DurableAgentState() + before = state.to_dict() + with pytest.raises(ValueError, match="strict JSON"): + state.record_response("done", response, delivery_window_seconds=60, now=NOW) + assert state.to_dict() == before + + +def test_subtype_null_exceptions_match_the_shared_schema(schema: dict[str, Any]) -> None: + definitions = schema["$defs"] + known = { + definitions[branch["$ref"].split("/")[-1]]["properties"]["$type"]["const"]: definitions[ + branch["$ref"].split("/")[-1] + ] + for branch in definitions["chatContentItem"]["oneOf"] + if "$ref" in branch + } + subclasses = {cls.type: cls for cls in DurableAgentStateContent.__subclasses__() if cls.type} + assert subclasses.keys() == known.keys() + nullable_fields = {kind: cls._NULLABLE_FIELDS for kind, cls in subclasses.items() if cls._NULLABLE_FIELDS} + assert nullable_fields == {"unknown": {"content"}, "functionResult": {"result"}} + for kind, fields in nullable_fields.items(): + for field in fields: + jsonschema.Draft202012Validator(known[kind]["properties"][field]).validate(None) + assert "content" in known["unknown"]["required"] + assert "text" in known["text"]["required"] + + +@pytest.mark.parametrize("opaque", [None, False, 0, 0.0, "", [], {}]) +def test_unknown_falsey_payloads_preserve_required_content_and_opaque_extensions( + opaque: Any, schema: dict[str, Any], monkeypatch: pytest.MonkeyPatch +) -> None: + raw = _source( + version="2.0.0", + contents=[ + { + "$type": "unknown", + "content": opaque, + "future": {"null": None, "flag": False, "count": 0}, + "extensionData": {"coreContent": {"type": "future.module.Content", "items": {"opaque": None}}}, + } + ], + ) + _validate(raw, schema) + state = DurableAgentState.from_json(json.dumps(raw)) + persisted = json.loads(state.to_json()) + _validate(persisted, schema) + assert persisted == raw + stored = state.data.conversation_history[0].messages[0].contents[0] + loader = Mock(side_effect=AssertionError("Opaque content must not interpret extensionData.coreContent")) + monkeypatch.setattr(state_module, "load_agent_response", loader) + for restored in (stored.to_ai_content(), stored.to_core_content()): + assert restored.type == "unknown" + assert restored.additional_properties == {"content": opaque} + assert type(restored.additional_properties["content"]) is type(opaque) + loader.assert_not_called() + assert state.to_dict() == raw + + +@pytest.mark.parametrize("result", [None, False, 0, "", [], {}]) +def test_function_result_nullable_payload_survives_schema_and_cold_roundtrip( + result: Any, schema: dict[str, Any] +) -> None: + raw = _source(version="2.0.0", contents=[{"$type": "functionResult", "callId": "call", "result": result}]) + _validate(raw, schema) + state = DurableAgentState.from_json(json.dumps(raw)) + persisted = json.loads(state.to_json()) + _validate(persisted, schema) + assert persisted == raw + assert "result" in persisted["data"]["conversationHistory"][0]["messages"][0]["contents"][0] + + +def test_optional_nonnullable_fields_are_omitted_without_dropping_empty_text_or_zero_flags( + schema: dict[str, Any], +) -> None: + contents: list[dict[str, Any]] = [ + {"$type": "reasoning", "text": None}, + {"$type": "uri", "uri": "https://example.test", "mediaType": None}, + {"$type": "data", "uri": "data:text/plain,", "mediaType": None}, + {"$type": "functionCall", "callId": "call", "name": "f", "arguments": None}, + {"$type": "text", "text": "", "extensionData": {"flag": False, "count": 0}}, + {"$type": "usage", "usage": {"inputTokenCount": 0}}, + ] + state = DurableAgentState.from_dict(_source(version="2.0.0", contents=contents)) + persisted = state.to_dict() + _validate(persisted, schema) + actual = persisted["data"]["conversationHistory"][0]["messages"][0]["contents"] + expected = [{key: value for key, value in item.items() if value is not None} for item in contents] + assert actual == expected + + +def test_required_text_is_not_silently_omitted_on_write() -> None: + with pytest.raises(ValueError, match="requires a text string"): + DurableAgentStateTextContent(text=None).to_persisted_dict() + + +def test_future_raw_content_preserves_nulls_and_never_interprets_extensions(schema: dict[str, Any]) -> None: + content = { + "$type": "futureContent", + "content": None, + "payload": [False, 0, {}], + "extensionData": {"coreContent": {"type": "text", "text": "not authoritative"}}, + } + raw = _source(version="2.0.0", contents=[content]) + state = DurableAgentState.from_json(json.dumps(raw)) + _validate(state.to_dict(), schema) + restored = state.data.conversation_history[0].messages[0].contents[0].to_core_content() + assert restored.type == "unknown" + assert restored.additional_properties == {"content": content} + restored.additional_properties["content"]["payload"].append("consumer edit") + assert state.to_dict() == raw + + +@pytest.mark.parametrize("version", ["2.0.1", "2.1.0", "2.999.0"]) +def test_future_revision_reads_but_cannot_prepare_a_write_preserving_all_control_state( + version: str, schema: dict[str, Any] +) -> None: + raw = _source(version=version) + raw["futureRoot"] = {"opaque": [None, False, 0]} + raw["data"].update( + futureData={"opaque": [None]}, + session={"session_id": SESSION_ID, "state": {"provider": {"thread": "original", "value": None}}}, + extensionData={"opaque": [None]}, + ingestedPositions={"producer": 3}, + ingestedMessages={"custom-id": None, "exact": ["a" * 64]}, + completedCorrelations={"done": {"completedAt": NOW.isoformat(), "future": None}}, + responseMailbox={ + "done": { + "createdAt": NOW.isoformat(), + "expiresAt": "2099-01-01T00:00:00+00:00", + "future": None, + "response": {"type": "agent_response", "messages": [], "future": {"opaque": None}}, + } + }, + truncation={"evictedMessageCount": 1, "firstEvictedAt": OLD.isoformat(), "lastEvictedAt": NOW.isoformat()}, + ) + raw["data"]["conversationHistory"].append({"$type": "futureEntry", "messages": {"opaque": None}}) + _validate(raw, schema) + before = deepcopy(raw) + state = DurableAgentState.from_json(json.dumps(raw)) + assert state.try_get_agent_response("done") is not None + with pytest.raises(ValueError, match="Only 2.0.0 is writable"): + state.prepare_for_write(delivery_window_seconds=60) + assert state.to_dict() == raw == before + + +def test_exact_current_revision_is_writable_without_mutation() -> None: + state = DurableAgentState.from_dict(_source(version=DurableAgentState.SCHEMA_VERSION)) + before = state.to_dict() + state.prepare_for_write(delivery_window_seconds=60) + assert state.to_dict() == before + + +@pytest.mark.parametrize("version", ["1.0.0", "1.999.0"]) +def test_legacy_write_rejection_message_remains_unchanged(version: str) -> None: + state = DurableAgentState.from_dict(_source(version=version)) + before = state.to_dict() + with pytest.raises(ValueError) as error: + state.prepare_for_write(delivery_window_seconds=60) + assert str(error.value) == ( + "Legacy state is read-only in this runtime. Keep it on its original deployment or use explicit " + "migration into a separate isolated-v2 entity. Legacy ingestedPositions require recorded delivery evidence." + ) + assert state.to_dict() == before + + +@pytest.mark.parametrize("version", ["2", "2.", "2.1.0-preview", "2.1.0\n", "2.\u0661.0", "3.0.0"]) +def test_version_admission_requires_a_complete_supported_ascii_version(version: str) -> None: + with pytest.raises(ValueError, match="Unsupported.*schemaVersion"): + DurableAgentState.from_dict(_source(version=version)) + + +@pytest.mark.parametrize("kind", ["errorResponse", "response"]) +def test_migrated_text_only_failure_retains_http_terminal_classification(kind: str) -> None: + source = _source() + source["data"]["conversationHistory"] = [ + { + "$type": kind, + "correlationId": "done", + "createdAt": OLD.isoformat(), + "messages": [{"role": "assistant", "contents": [{"$type": "text", "text": "legacy text only"}]}], + } + ] + before = deepcopy(source) + legacy_response = DurableAgentState.from_dict(source).try_get_agent_response("done") + assert legacy_response is not None + state = DurableAgentState.from_json(_migrate(source).to_json()) + payload = state.data.response_mailbox["done"]["response"] + delivered = load_agent_response(payload) + for response in (legacy_response, delivered): + assert response.text == "legacy text only" + assert all(content.type == "text" for message in response.messages for content in message.contents) + # HTTP polling branches on this predicate, even when there is no error Content. + assert is_terminal_agent_response(response) is (kind == "errorResponse") + assert response.additional_properties == ({"durable_status": "error"} if kind == "errorResponse" else {}) + assert state.data.completed_correlations["done"] == { + "completedAt": NOW.isoformat(), + "legacy": True, + **({"outcome": "failed"} if kind == "errorResponse" else {}), + } + assert state.to_dict()["data"]["conversationHistory"] == source["data"]["conversationHistory"] + assert source == before + + +@pytest.mark.parametrize("retained_contents", [[], [{"$type": "text", "text": "pruned portion"}]]) +@pytest.mark.parametrize("journal_kind", ["empty", "other-custom", "workflow"]) +def test_complete_journal_cannot_omit_retained_custom_request_identity( + retained_contents: list[dict[str, Any]], journal_kind: str +) -> None: + source = _source(contents=retained_contents) + messages: list[Message] = [] + if journal_kind == "other-custom": + messages = [Message("user", ["accepted"], message_id="other-custom")] + elif journal_kind == "workflow": + source["data"]["ingestedPositions"] = {"upstream": 3} + messages = [Message("user", ["accepted"], message_id=workflow_message_id("upstream", 3))] + evidence = _evidence(source, messages) + before_source, before_evidence = deepcopy(source), deepcopy(evidence) + with pytest.raises(ValueError, match="must include every retained legacy custom request message ID"): + _migrate(source, evidence=evidence) + assert source == before_source + assert evidence == before_evidence + + +@pytest.mark.parametrize("retained_contents", [[], [{"$type": "text", "text": "pruned portion"}]]) +def test_custom_journal_compares_identity_not_the_pruned_body(retained_contents: list[dict[str, Any]]) -> None: + source = _source(contents=retained_contents) + original = Message("user", ["complete original accepted input"], message_id="custom-id") + evidence = _evidence(source, [original]) + before_source, before_evidence = deepcopy(source), deepcopy(evidence) + state = DurableAgentState.from_json(_migrate(source, evidence=evidence).to_json()) + assert state.data.ingested_messages == {"custom-id": [message_identity(original)]} + assert state.to_dict()["data"]["conversationHistory"] == source["data"]["conversationHistory"] + assert source == before_source + assert evidence == before_evidence + + +def test_no_journal_keeps_custom_identity_markers_without_fabricating_workflow_receipts() -> None: + source = _source(contents=[]) + source["data"]["conversationHistory"][0]["messages"].append({ + "role": "user", + "messageId": workflow_message_id("upstream", 3), + "contents": [], + }) + source["data"]["ingestedMessages"] = {"existing-exact": ["a" * 64], "existing-marker": None} + before = deepcopy(source) + state = DurableAgentState.from_json(_migrate(source).to_json()) + assert state.data.ingested_messages == { + "custom-id": None, + "existing-exact": ["a" * 64], + "existing-marker": None, + } + assert state.data.completed_correlations == {} + assert state.data.response_mailbox == {} + assert source == before + + +def test_empty_complete_journal_is_valid_when_no_retained_custom_requests_contradict_it() -> None: + source = _source(contents=[]) + history = source["data"]["conversationHistory"] + history[0]["messages"][0]["messageId"] = workflow_message_id("upstream", 3) + history.append({"$type": "futureEntry", "messages": {"messageId": "opaque-not-a-request"}}) + state = DurableAgentState.from_json(_migrate(source, evidence=_evidence(source, [])).to_json()) + assert state.data.ingested_messages == {} + assert state.to_dict()["data"]["conversationHistory"] == history diff --git a/python/packages/durabletask/tests/test_state_layout_admission.py b/python/packages/durabletask/tests/test_state_layout_admission.py new file mode 100644 index 0000000..8b73a74 --- /dev/null +++ b/python/packages/durabletask/tests/test_state_layout_admission.py @@ -0,0 +1,135 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Reject known incompatible delivery layouts instead of reopening completed work.""" + +import json +from copy import deepcopy +from typing import Any +from unittest.mock import Mock + +import pytest +from agent_framework import Agent, AgentResponse, Message +from test_durable_history_provider import RecordingChatClient +from test_revision_contract import JsonStateProvider + +from agent_framework_durabletask import AgentEntity, DurableAgentState + + +def _foreign_state(*, expired: bool = False) -> dict[str, Any]: + receipt: dict[str, Any] = { + "correlationId": "completed", + "outcome": "succeeded", + "completedAt": "2026-01-01T00:00:00Z", + "resultState": "unavailable" if expired else "available", + } + results = { + "completed": { + "correlationId": "completed", + "outcome": "succeeded", + "completedAt": "2026-01-01T00:00:00Z", + "response": {"messages": [{"role": "assistant", "contents": [{"$type": "text", "text": "answer"}]}]}, + } + } + if expired: + receipt["resultUnavailableAt"] = "2026-01-01T00:01:00Z" + results = {} + return { + "schemaVersion": "2.0.0", + "data": { + "conversationHistory": [], + "terminalResults": results, + "completionReceipts": {"completed": receipt}, + "historyBinding": {"version": 1, "ownerKind": "durableState", "providerKey": "example.history"}, + }, + } + + +@pytest.mark.parametrize("expired", [False, True], ids=["available", "expired"]) +@pytest.mark.parametrize("json_boundary", [False, True], ids=["dict", "json"]) +def test_incompatible_layout_is_rejected_even_with_the_same_schema_version(expired: bool, json_boundary: bool) -> None: + payload = _foreign_state(expired=expired) + before = deepcopy(payload) + with pytest.raises(ValueError, match="incompatible.*delivery|delivery.*incompatible"): + if json_boundary: + DurableAgentState.from_json(json.dumps(payload)) + else: + DurableAgentState.from_dict(payload) + assert payload == before + + +@pytest.mark.parametrize("field", ["terminalResults", "completionReceipts"]) +@pytest.mark.parametrize("value", [{}, None, []], ids=["empty", "null", "malformed"]) +@pytest.mark.parametrize("mixed", [False, True]) +def test_reserved_alternate_containers_never_hide_as_optional_metadata(field: str, value: Any, mixed: bool) -> None: + state = DurableAgentState() + if mixed: + state.record_response("native", AgentResponse(messages=[]), delivery_window_seconds=3600) + raw = state.to_dict() + raw["data"][field] = value + before = deepcopy(raw) + with pytest.raises(ValueError, match="incompatible.*delivery|delivery.*incompatible"): + DurableAgentState.from_dict(raw) + assert raw == before + + +@pytest.mark.parametrize("operation", ["run", "reset", "expire_responses"]) +@pytest.mark.parametrize("expired", [False, True]) +async def test_entity_refuses_incompatible_state_before_model_calls_or_writes(operation: str, expired: bool) -> None: + provider = JsonStateProvider(_foreign_state(expired=expired)) + before = deepcopy(provider.raw) + client: Any = RecordingChatClient() + entity = AgentEntity(Agent(client=client), state_provider=provider) + with pytest.raises(ValueError, match="incompatible.*delivery|delivery.*incompatible"): + if operation == "run": + await entity.run({"message": "do not repeat", "correlationId": "completed"}) + else: + getattr(entity, operation)() + assert client.received_messages == [] + assert provider.writes == 0 and provider.raw == before + + +@pytest.mark.parametrize("field", ["terminalResults", "completionReceipts"]) +@pytest.mark.parametrize("operation", ["read", "serialize", "write", "run"]) +async def test_cached_state_cannot_bypass_delivery_layout_admission(field: str, operation: str) -> None: + provider = JsonStateProvider() + provider.state.data.unknown_fields[field] = {"completed": {"outcome": "succeeded"}} + state = provider.state + before = deepcopy(state.data.unknown_fields) + client: Any = RecordingChatClient() + entity = AgentEntity(Agent(client=client), state_provider=provider) + with pytest.raises(ValueError, match="incompatible.*delivery|delivery.*incompatible"): + if operation == "read": + state.try_get_agent_response("completed") + elif operation == "serialize": + state.to_dict() + elif operation == "write": + state.prepare_for_write(delivery_window_seconds=60) + else: + await entity.run({"message": "do not repeat", "correlationId": "completed"}) + assert state.data.unknown_fields == before + assert provider.writes == 0 and client.received_messages == [] + + +def test_native_empty_delivery_and_unrelated_nested_metadata_remain_supported() -> None: + state = DurableAgentState() + raw = state.to_dict() + raw["data"]["futureMetadata"] = {"terminalResults": {}, "completionReceipts": {"opaque": False}} + raw["data"]["session"] = {"session_id": "test", "state": {"application": {"terminalResults": [1]}}} + raw["application"] = {"completionReceipts": None} + restored = DurableAgentState.from_dict(raw) + assert restored.to_dict() == raw + assert restored.try_get_agent_response("absent") is None + restored.record_response( + "native", AgentResponse(messages=[Message("assistant", ["native answer"])]), delivery_window_seconds=3600 + ) + cold = DurableAgentState.from_json(restored.to_json()) + response = cold.try_get_agent_response("native") + assert response is not None and response.text == "native answer" + + +def test_alternate_completion_is_rejected_before_any_response_deserialization(monkeypatch: pytest.MonkeyPatch) -> None: + loader = Mock(side_effect=AssertionError("Unsupported state must not be interpreted as a response")) + monkeypatch.setattr("agent_framework_durabletask._durable_agent_state.load_agent_response", loader) + with pytest.raises(ValueError, match="incompatible.*delivery|delivery.*incompatible"): + DurableAgentState.from_dict(_foreign_state()) + loader.assert_not_called() diff --git a/python/packages/durabletask/tests/test_state_migration_review.py b/python/packages/durabletask/tests/test_state_migration_review.py new file mode 100644 index 0000000..42caa2c --- /dev/null +++ b/python/packages/durabletask/tests/test_state_migration_review.py @@ -0,0 +1,653 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Detached explicit migration contracts. No hosts, providers or backends are needed.""" + +import hashlib +import json +from copy import deepcopy +from datetime import datetime, timedelta, timezone +from typing import Any + +import pytest +from agent_framework import AgentResponse, Content, Message +from typing_extensions import Self + +from agent_framework_durabletask import migrate_legacy_state, state_snapshot_digest +from agent_framework_durabletask._durable_agent_state import DurableAgentState, DurableAgentStateEntryJsonType +from agent_framework_durabletask._message_identity import message_identity +from agent_framework_durabletask._retention import StateCapacityError +from agent_framework_durabletask._workflows.naming import workflow_message_id + +NOW = datetime(2026, 9, 9, 12, tzinfo=timezone.utc) +OLD = datetime(2024, 1, 1, tzinfo=timezone.utc) +SESSION_ID = "dafx-agent:original-session" +WINDOW = 60 + + +def _entry(kind: str, correlation: str, *, message_id: str = "custom-id") -> dict[str, Any]: + return { + "$type": kind, + "correlationId": correlation, + "createdAt": OLD.isoformat(), + "messages": [ + { + "role": "user" if kind == "request" else "assistant", + "messageId": message_id, + "contents": [{"$type": "text", "text": "retained portion"}], + } + ], + } + + +def _source() -> dict[str, Any]: + return { + "schemaVersion": "1.1.0", + "data": { + "conversationHistory": [ + _entry("request", "done"), + _entry("response", "done", message_id="answer-id"), + ] + }, + } + + +def _migrate(source: dict[str, Any], **overrides: Any) -> DurableAgentState: + options: dict[str, Any] = { + "source_digest": state_snapshot_digest(source), + "source_session_id": SESSION_ID, + "migration_id": "migration-1", + "ownership_transfer_id": "transfer-1", + "delivery_window_seconds": WINDOW, + "now": NOW, + } + options.update(overrides) + return migrate_legacy_state(source, **options) + + +def _message(position: int, *, producer: str = "upstream", text: str = "accepted") -> Message: + return Message( + "user", + [Content.from_text(text, additional_properties={"nested": {"labels": ["original"]}})], + message_id=workflow_message_id(producer, position), + author_name="author", + additional_properties={"nested": {"labels": ["message"]}}, + ) + + +def _evidence(source: dict[str, Any], messages: list[Message]) -> dict[str, Any]: + return { + "sourceDigest": state_snapshot_digest(source), + "evidenceId": "operator-journal-1", + "complete": True, + "messages": [message.to_dict() for message in messages], + } + + +def _cold(state: DurableAgentState) -> DurableAgentState: + return DurableAgentState.from_json(json.dumps(state.to_dict(), allow_nan=False)) + + +def _existing_delivery() -> dict[str, Any]: + state = DurableAgentState() + state.record_response( + "done", + AgentResponse(messages=[Message("assistant", ["original mailbox"])]), + delivery_window_seconds=WINDOW, + now=OLD, + ) + return state.to_dict()["data"] + + +def test_source_digest_uses_complete_strict_canonical_utf8_json() -> None: + source: dict[str, Any] = { + "schemaVersion": "1.1.0", + "data": {"conversationHistory": [], "future": ["é", "雪", 0, False]}, + } + expected = hashlib.sha256( + json.dumps(source, sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False).encode("utf-8") + ).hexdigest() + reordered: dict[str, Any] = { + "data": {"future": source["data"]["future"], "conversationHistory": []}, + "schemaVersion": "1.1.0", + } + assert state_snapshot_digest(source) == state_snapshot_digest(reordered) == expected + assert expected != hashlib.sha256(json.dumps(source, sort_keys=True).encode()).hexdigest() + reordered["data"]["future"] = list(reversed(source["data"]["future"])) + assert state_snapshot_digest(reordered) != expected + + +@pytest.mark.parametrize("invalid", [float("nan"), float("inf"), float("-inf"), (1, 2), {1: "key"}, {1, 2}]) +def test_digest_rejects_non_json_and_nonfinite_nested_values(invalid: Any) -> None: + source = _source() + source["data"]["future"] = {"nested": invalid} + with pytest.raises(ValueError, match="strict JSON"): + state_snapshot_digest(source) + + +def test_digest_rejects_cycles_and_nonobject_source() -> None: + source: dict[str, Any] = {} + source["cycle"] = source + with pytest.raises(ValueError, match="strict JSON"): + state_snapshot_digest(source) + with pytest.raises(ValueError, match="JSON object"): + state_snapshot_digest([]) # type: ignore[arg-type] + + +@pytest.mark.parametrize("kind", list(DurableAgentStateEntryJsonType)) +def test_only_recorded_response_kinds_backfill_completion(kind: str) -> None: + source = _source() + source["data"]["conversationHistory"] = [_entry(kind, "done")] + result = _cold(_migrate(source)) + is_response = kind in (DurableAgentStateEntryJsonType.RESPONSE, DurableAgentStateEntryJsonType.ERROR_RESPONSE) + assert ("done" in result.data.completed_correlations) is is_response + assert ("done" in result.data.response_mailbox) is is_response + assert result.to_dict()["data"]["conversationHistory"] == source["data"]["conversationHistory"] + if is_response: + assert result.data.completed_correlations["done"] == { + "completedAt": NOW.isoformat(), + "legacy": True, + **({"outcome": "failed"} if kind == DurableAgentStateEntryJsonType.ERROR_RESPONSE else {}), + } + mailbox = result.data.response_mailbox["done"] + assert mailbox["createdAt"] == NOW.isoformat() + assert mailbox["expiresAt"] == (NOW + timedelta(seconds=WINDOW)).isoformat() + assert mailbox["response"]["messages"][0]["contents"][0]["text"] == "retained portion" + assert mailbox["response"]["created_at"] == OLD.isoformat() + + +def test_partial_and_contentless_recorded_responses_are_not_claimed_as_originals() -> None: + source = _source() + history = source["data"]["conversationHistory"] + history[1]["usage"] = {"inputTokenCount": 3, "futureUsage": {"keep": [1]}} + history.append(_entry("response", "empty")) + history[-1]["messages"] = [] + history.append(_entry("request", "old-pruned-without-response")) + history[-1]["messages"][0]["contents"] = [] + history.append(_entry("request", "unfinished")) + source["data"]["truncation"] = {"evictedMessageCount": 20, "future": [1]} + result = _cold(_migrate(source)) + + assert set(result.data.completed_correlations) == {"done", "empty"} + assert set(result.data.response_mailbox) == {"done", "empty"} + assert result.data.response_mailbox["empty"]["response"]["messages"] == [] + assert result.data.response_mailbox["done"]["response"]["usage_details"] == {"input_token_count": 3} + assert all(record["legacy"] is True for record in result.data.completed_correlations.values()) + assert result.try_get_agent_response("old-pruned-without-response") is None + assert result.try_get_agent_response("unfinished") is None + assert result.to_dict()["data"]["conversationHistory"] == history + + +@pytest.mark.parametrize("keep_mailbox", [False, True]) +def test_existing_completion_and_mailbox_are_not_overwritten_or_reopened(keep_mailbox: bool) -> None: + source = _source() + delivery = _existing_delivery() + delivery["completedCorrelations"]["done"]["future"] = {"keep": [1]} + source["data"]["completedCorrelations"] = delivery["completedCorrelations"] + if keep_mailbox: + delivery["responseMailbox"]["done"]["response"]["futureResponse"] = {"keep": [2]} + source["data"]["responseMailbox"] = delivery["responseMailbox"] + result = _cold(_migrate(source)) + assert result.data.completed_correlations == delivery["completedCorrelations"] + assert result.data.response_mailbox == (delivery["responseMailbox"] if keep_mailbox else {}) + + +def test_existing_mailbox_without_receipt_is_preserved_with_completion_backfill() -> None: + source = _source() + source["data"]["responseMailbox"] = _existing_delivery()["responseMailbox"] + result = _cold(_migrate(source)) + assert result.data.response_mailbox == source["data"]["responseMailbox"] + assert result.data.completed_correlations["done"] == { + "completedAt": OLD.isoformat(), + "legacy": True, + "outcome": "succeeded", + } + + +def test_sparse_journal_preserves_exact_revisions_not_an_inferred_prefix_after_cold_reload() -> None: + source = _source() + source["data"]["ingestedPositions"] = {"upstream": 3} + # Neither the original accepted input nor the missing position can be recovered + # from this compacted/pruned transcript. It must not contribute fingerprints. + source["data"]["conversationHistory"][0]["messages"] = [ + {"role": "user", "messageId": workflow_message_id("upstream", 3), "contents": []} + ] + first, third, revision = _message(1), _message(3), _message(3, text="accepted revision") + assert first.message_id is not None + assert third.message_id is not None + source["data"]["ingestedMessages"] = {third.message_id: [message_identity(third)]} + evidence = _evidence(source, [revision, first, third]) + before_source, before_evidence = deepcopy(source), deepcopy(evidence) + result = _cold(_migrate(source, delivery_evidence=evidence)) + + assert result.data.ingested_messages == { + first.message_id: [message_identity(first)], + third.message_id: [message_identity(third), message_identity(revision)], + } + assert workflow_message_id("upstream", 0) not in result.data.ingested_messages + assert workflow_message_id("upstream", 2) not in result.data.ingested_messages + fingerprints = result.data.ingested_messages[third.message_id] + assert fingerprints is not None + assert message_identity(_message(3, text="new revision")) not in fingerprints + assert result.data.ingested_positions == {"upstream": 3} + assert result.data.unknown_fields["migration"]["evidenceId"] == "operator-journal-1" + assert source == before_source and evidence == before_evidence + evidence["messages"][0]["contents"][0]["additional_properties"]["nested"]["labels"].append("caller edit") + assert result.data.ingested_messages[third.message_id] == [message_identity(third), message_identity(revision)] + + +def test_multiple_producers_allow_sparse_zero_based_and_out_of_order_journal() -> None: + source = _source() + source["data"]["ingestedPositions"] = {"first_with_underscores": 9, "other": 0} + # Supply the accepted custom input explicitly, not a fingerprint inferred from its retained portion. + custom = Message("user", ["complete original accepted input"], message_id="custom-id") + messages = [_message(9, producer="first_with_underscores"), _message(0, producer="other"), custom] + result = _cold(_migrate(source, delivery_evidence=_evidence(source, messages))) + assert result.data.ingested_messages == {message.message_id: [message_identity(message)] for message in messages} + + +@pytest.mark.parametrize("position", [0, 3]) +def test_scalar_positions_without_complete_journal_fail_even_with_retained_messages(position: int) -> None: + source = _source() + source["data"]["ingestedPositions"] = {"upstream": position} + source["data"]["conversationHistory"][0]["messages"][0]["messageId"] = workflow_message_id("upstream", position) + source["data"]["ingestedMessages"] = {workflow_message_id("upstream", position): ["a" * 64]} + before = deepcopy(source) + with pytest.raises(ValueError, match="recorded delivery evidence.*old engine"): + _migrate(source) + assert source == before + + +def test_no_scalar_no_journal_uses_only_custom_id_markers_and_preserves_exact_receipts() -> None: + source = _source() + messages = source["data"]["conversationHistory"][0]["messages"] + messages.extend([ + {"role": "user", "messageId": "cleared-custom", "contents": []}, + {"role": "user", "messageId": workflow_message_id("upstream", 3), "contents": []}, + {"role": "user", "messageId": "already-exact", "contents": []}, + ]) + source["data"]["ingestedMessages"] = {"already-exact": ["b" * 64, "a" * 64], "old-marker": None} + result = _cold(_migrate(source)) + assert result.data.ingested_messages == { + "already-exact": ["b" * 64, "a" * 64], + "old-marker": None, + "custom-id": None, + "cleared-custom": None, + } + assert "answer-id" not in result.data.ingested_messages + + +@pytest.mark.parametrize("identity", [None, "", " "]) +def test_anonymous_legacy_request_ids_are_preserved_without_receipts(identity: Any) -> None: + source = _source() + message = source["data"]["conversationHistory"][0]["messages"][0] + if identity is None: + message.pop("messageId") + else: + message["messageId"] = identity + result = _cold(_migrate(source)) + assert result.data.ingested_messages == {} + assert result.to_dict()["data"]["conversationHistory"][0] == source["data"]["conversationHistory"][0] + + +def test_complete_custom_journal_replaces_identity_marker_with_content_sensitive_revisions() -> None: + source = _source() + source["data"]["ingestedMessages"] = {"custom-id": None} + old = Message("user", ["original"], message_id="custom-id") + revised = Message("user", ["revision"], message_id="custom-id") + result = _cold(_migrate(source, delivery_evidence=_evidence(source, [old, revised]))) + assert result.data.ingested_messages == {"custom-id": [message_identity(old), message_identity(revised)]} + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("sourceDigest", "a" * 64), + ("evidenceId", " "), + ("evidenceId", None), + ("complete", False), + ("complete", 1), + ("complete", "true"), + ("messages", {}), + ("extra", []), + ], +) +def test_evidence_binding_envelope_completeness_and_extra_fields_are_validated(field: str, value: Any) -> None: + source = _source() + source["data"]["ingestedPositions"] = {"upstream": 3} + evidence = _evidence(source, [_message(3)]) + evidence[field] = value + before = deepcopy(evidence) + with pytest.raises(ValueError, match="[Rr]ecorded delivery evidence"): + _migrate(source, delivery_evidence=evidence) + assert evidence == before + + +@pytest.mark.parametrize("field", ["sourceDigest", "evidenceId", "complete", "messages"]) +def test_all_evidence_fields_are_required(field: str) -> None: + source = _source() + evidence = _evidence(source, []) + del evidence[field] + with pytest.raises(ValueError, match="requires exactly"): + _migrate(source, delivery_evidence=evidence) + + +@pytest.mark.parametrize( + "messages", + [[], [_message(1)], [_message(4)], [_message(3, producer="other")], [_message(3), _message(0, producer="extra")]], +) +def test_evidence_workflow_producer_set_and_maxima_must_match(messages: list[Message]) -> None: + source = _source() + source["data"]["ingestedPositions"] = {"upstream": 3} + with pytest.raises(ValueError, match="producers and maximum positions"): + _migrate(source, delivery_evidence=_evidence(source, messages)) + + +@pytest.mark.parametrize( + "positions", + [None, [], False, {"upstream": True}, {"upstream": -1}, {"upstream": 1.5}, {"upstream": "3"}, {"": 0}, {" ": 0}], +) +def test_all_legacy_cursor_entries_require_named_producers_and_nonbool_nonnegative_ints(positions: Any) -> None: + source = _source() + source["data"]["ingestedPositions"] = positions + with pytest.raises(ValueError, match="ingestedPositions"): + _migrate(source, delivery_evidence=_evidence(source, [])) + + +@pytest.mark.parametrize( + "identity", [None, "", " ", True, 7, "wf_upstream_-1", "wf_upstream_true", "wf__3", "wf_upstream_3\n"] +) +def test_evidence_rejects_missing_blank_nonstring_and_malformed_workflow_ids(identity: Any) -> None: + source = _source() + source["data"]["ingestedPositions"] = {"upstream": 3} + evidence = _evidence(source, [_message(3)]) + evidence["messages"][0]["message_id"] = identity + with pytest.raises(ValueError): + _migrate(source, delivery_evidence=evidence) + + +def test_exact_duplicate_evidence_is_rejected_but_revisions_are_not() -> None: + source = _source() + source["data"]["ingestedPositions"] = {"upstream": 3} + message = _message(3) + with pytest.raises(ValueError, match="duplicate message ID/fingerprint"): + _migrate(source, delivery_evidence=_evidence(source, [message, message])) + + +@pytest.mark.parametrize("change", ["unknown-field", "raw-representation", "wrong-contents", "bad-content", "bad-role"]) +def test_journal_never_hashes_a_lossy_or_malformed_message_projection(change: str) -> None: + source = _source() + source["data"]["ingestedPositions"] = {"upstream": 3} + evidence = _evidence(source, [_message(3)]) + message = evidence["messages"][0] + if change == "unknown-field": + message["future_unrecognized_message_field"] = {"must-not-disappear": [1]} + elif change == "raw-representation": + message["raw_representation"] = {"must-not-disappear": [1]} + elif change == "wrong-contents": + message["contents"] = {} + elif change == "bad-content": + message["contents"] = [{"type": ""}] + else: + message["role"] = "" + with pytest.raises(ValueError): + _migrate(source, delivery_evidence=evidence) + + +@pytest.mark.parametrize("mutation", ["author", "role", "text", "message-metadata", "content-metadata"]) +def test_journal_fingerprints_cover_complete_canonical_inputs(mutation: str) -> None: + source = _source() + source["data"]["ingestedPositions"] = {"upstream": 3} + original = _message(3) + changed = deepcopy(original) + if mutation == "author": + changed.author_name = "different" + elif mutation == "role": + changed.role = "assistant" + elif mutation == "text": + changed.contents[0].text = "different" + elif mutation == "message-metadata": + changed.additional_properties["nested"]["labels"].append("different") + else: + changed.contents[0].additional_properties["nested"]["labels"].append("different") + custom = Message("user", ["complete original accepted input"], message_id="custom-id") + result = _cold(_migrate(source, delivery_evidence=_evidence(source, [original, changed, custom]))) + assert message_identity(original) != message_identity(changed) + assert original.message_id is not None + assert result.data.ingested_messages == { + original.message_id: [message_identity(original), message_identity(changed)], + "custom-id": [message_identity(custom)], + } + + +@pytest.mark.parametrize( + "receipts", + [ + {"custom": []}, + {"custom": ["short"]}, + {"custom": ["A" * 64]}, + {"custom": ["g" * 64]}, + {"custom": [True]}, + {"custom": ["a" * 64, "a" * 64]}, + {"custom": "a" * 64}, + {" ": ["a" * 64]}, + {"wf_upstream_3": None}, + ], +) +def test_existing_receipts_reject_invalid_fingerprint_shapes_and_workflow_markers(receipts: Any) -> None: + source = _source() + source["data"]["ingestedMessages"] = receipts + with pytest.raises(ValueError): + _migrate(source) + + +@pytest.mark.parametrize("receipt", [{"other-custom": None}, {"wf_upstream_3": ["a" * 64]}]) +def test_complete_journal_must_include_existing_identity_and_exact_receipts(receipt: Any) -> None: + source = _source() + source["data"]["ingestedPositions"] = {"upstream": 3} + source["data"]["ingestedMessages"] = receipt + with pytest.raises(ValueError, match="inconsistent with existing"): + _migrate(source, delivery_evidence=_evidence(source, [_message(3)])) + + +def test_raw_unknown_nested_fields_order_session_and_source_are_preserved() -> None: + source = _source() + source["futureRoot"] = {"nested": [None, {"keep": "雪"}]} + data = source["data"] + data["futureData"] = {"nested": [1]} + data["session"] = { + "session_id": SESSION_ID, + "service_session_id": "original-provider-thread", + "state": {"external-store": {"provider-key": "original", "nested": [2]}}, + "futureSession": [3], + } + data["extensionData"] = {"keep": [4]} + history = data["conversationHistory"] + history[0]["futureEntry"] = {"nested": [5]} + history[0]["messages"][0]["futureMessage"] = {"nested": [6]} + history[0]["messages"][0]["contents"][0]["futureContent"] = {"nested": [7]} + history[1]["usage"] = {"inputTokenCount": 0, "futureUsage": {"nested": [8]}} + history.append({"$type": "futureEntryKind", "opaque": [None], "messages": {"unknown": [9]}}) + history[0]["messages"][0]["contents"].append({"$type": "futureContentKind", "payload": None}) + before = deepcopy(source) + migrated = _migrate(source) + result = _cold(migrated) + serialized = result.to_dict() + assert serialized["futureRoot"] == source["futureRoot"] + for key in ("futureData", "session", "extensionData", "conversationHistory"): + assert serialized["data"][key] == data[key] + assert serialized["schemaVersion"] == "2.0.0" + assert source == before + + assert migrated.data.session is not None + migrated.data.session["state"]["external-store"]["nested"].append("result edit") + migrated.data.conversation_history[0].messages[0].unknown_fields["futureMessage"]["nested"].append("result edit") + source["futureRoot"]["nested"].append("source edit") + source["data"]["session"]["state"]["external-store"]["nested"].append("source edit") + assert result.to_dict() == serialized + assert before["data"]["session"]["state"]["external-store"]["nested"] == [2] + assert migrated.unknown_fields["futureRoot"] == before["futureRoot"] + + +@pytest.mark.parametrize( + "session", + [ + None, + {}, + {"state": {"keep": [1]}}, + {"session_id": "", "state": {"keep": [1]}}, + {"session_id": " ", "state": {"keep": [1]}}, + ], +) +def test_missing_logical_session_identity_is_filled_without_random_or_provider_state_reset(session: Any) -> None: + source = _source() + source["data"]["session"] = session + result = _cold(_migrate(source)) + expected = deepcopy(session) if session is not None else {} + expected["session_id"] = SESSION_ID + expected.setdefault("state", {}) + assert result.data.session == expected + + +@pytest.mark.parametrize("session", [{"session_id": "destination-id"}, {"session_id": True}, []]) +def test_conflicting_or_malformed_session_identity_fails(session: Any) -> None: + source = _source() + source["data"]["session"] = session + with pytest.raises(ValueError, match="session"): + _migrate(source) + + +def test_metadata_exact_contract_fixed_now_repeatability_and_parent_owned_idempotency() -> None: + source = _source() + before = deepcopy(source) + first, second = _migrate(source), _migrate(source) + assert first.to_dict() == second.to_dict() + assert first is not second + assert first.data.unknown_fields["migration"] == { + "id": "migration-1", + "sourceDigest": state_snapshot_digest(source), + "sourceSessionId": SESSION_ID, + "ownershipTransferId": "transfer-1", + "createdAt": NOW.isoformat(), + } + assert first.data.session == {"session_id": SESSION_ID, "state": {}} + later = _migrate(source, now=NOW + timedelta(days=1)) + assert later.data.response_mailbox["done"]["expiresAt"] != first.data.response_mailbox["done"]["expiresAt"] + with pytest.raises(ValueError, match="never a v2 source"): + _migrate(first.to_dict()) + assert source == before + + +def test_one_utc_clock_capture_for_all_backfills(monkeypatch: pytest.MonkeyPatch) -> None: + from agent_framework_durabletask import _state_migration as migration_module + + calls: list[Any] = [] + + class Clock(datetime): + @classmethod + def now(cls, tz: Any = None) -> Self: + calls.append(tz) + return cls(2026, 9, 9, 12, tzinfo=timezone.utc) + + source = _source() + source["data"]["conversationHistory"].append(_entry("response", "another")) + monkeypatch.setattr(migration_module, "datetime", Clock) + result = _migrate(source, now=None) + assert calls == [timezone.utc] + assert {record["createdAt"] for record in result.data.response_mailbox.values()} == {NOW.isoformat()} + assert {record["completedAt"] for record in result.data.completed_correlations.values()} == {NOW.isoformat()} + + +def test_rfc3339_z_existing_delivery_reloads_without_python311_fromisoformat(monkeypatch: pytest.MonkeyPatch) -> None: + from agent_framework_durabletask import _durable_agent_state as state_module + + class Python310Datetime(datetime): + @classmethod + def fromisoformat(cls, value: str) -> Self: + assert not value.endswith(("Z", "z")) + return super().fromisoformat(value) + + source = _source() + delivery = _existing_delivery() + delivery["responseMailbox"]["done"].update(createdAt="2024-01-01T00:00:00Z", expiresAt="2024-01-01T00:01:00z") + delivery["completedCorrelations"]["done"]["completedAt"] = "2024-01-01T00:00:00Z" + source["data"].update( + responseMailbox=delivery["responseMailbox"], completedCorrelations=delivery["completedCorrelations"] + ) + monkeypatch.setattr(state_module, "datetime", Python310Datetime) + result = _cold(_migrate(source)) + assert result.data.response_mailbox == delivery["responseMailbox"] + assert result.data.completed_correlations == delivery["completedCorrelations"] + + +@pytest.mark.parametrize("version", ["2.0.0", "2.3.0", "3.0.0", "1", None, True]) +def test_migration_is_legacy_only(version: Any) -> None: + source = _source() + source["schemaVersion"] = version + with pytest.raises(ValueError, match="only legacy"): + _migrate(source) + + +@pytest.mark.parametrize("digest", ["a" * 64, "A" * 64, "short", None]) +def test_source_digest_must_match_the_unmodified_snapshot(digest: Any) -> None: + with pytest.raises(ValueError, match="source_digest"): + _migrate(_source(), source_digest=digest) + + +@pytest.mark.parametrize("field", ["source_session_id", "migration_id", "ownership_transfer_id"]) +@pytest.mark.parametrize("value", [None, "", " ", True, 1]) +def test_parent_identifiers_must_be_nonblank_strings(field: str, value: Any) -> None: + with pytest.raises(ValueError, match=field): + _migrate(_source(), **{field: value}) + + +@pytest.mark.parametrize("field", ["delivery_window_seconds", "max_state_bytes"]) +@pytest.mark.parametrize("value", [True, False, 0, -1, 1.5, "60"]) +def test_grace_and_resolved_budget_must_be_positive_nonbool_integers(field: str, value: Any) -> None: + with pytest.raises(ValueError, match=field): + _migrate(_source(), **{field: value}) + + +@pytest.mark.parametrize("now", [datetime(2026, 9, 9), "2026-09-09T00:00:00Z", False]) +def test_now_requires_an_aware_datetime(now: Any) -> None: + with pytest.raises(ValueError, match="offset-aware"): + _migrate(_source(), now=now) + + +def test_grace_overflow_is_rejected_even_without_a_recorded_response() -> None: + source: dict[str, Any] = {"schemaVersion": "1.0.0", "data": {"conversationHistory": []}} + with pytest.raises(ValueError, match="bounded grace"): + _migrate(source, delivery_window_seconds=10**100) + + +def test_reserved_migration_metadata_is_not_silently_overwritten() -> None: + source = _source() + source["data"]["migration"] = {"unknown-owner": [1]} + before = deepcopy(source) + with pytest.raises(ValueError, match="reserved migration metadata"): + _migrate(source) + assert source == before + + +def test_budget_includes_metadata_receipts_mailbox_session_and_ascii_escaped_unknowns_without_pruning() -> None: + source = _source() + source["data"]["ingestedPositions"] = {"upstream": 3} + source["futureRoot"] = {"unicode": "雪😀" * 20} + custom = Message("user", ["complete original accepted input"], message_id="custom-id") + messages = [_message(1), _message(3), custom] + evidence = _evidence(source, messages) + before_source, before_evidence = deepcopy(source), deepcopy(evidence) + result = _migrate(source, delivery_evidence=evidence) + assert result.data.ingested_messages == {message.message_id: [message_identity(message)] for message in messages} + size = len(json.dumps(result.to_dict(), allow_nan=False)) + assert size > len(json.dumps(result.to_dict(), ensure_ascii=False, allow_nan=False).encode("utf-8")) + assert _migrate(source, delivery_evidence=evidence, max_state_bytes=size).to_dict() == result.to_dict() + without_metadata = result.to_dict() + del without_metadata["data"]["migration"] + for budget in (size - 1, len(json.dumps(without_metadata, allow_nan=False))): + with pytest.raises(StateCapacityError) as error: + _migrate(source, delivery_evidence=evidence, max_state_bytes=budget) + assert error.value.size_bytes == error.value.floor_bytes == size + assert error.value.max_state_bytes == error.value.target_bytes == budget + assert source == before_source and evidence == before_evidence + assert result.to_dict()["data"]["conversationHistory"] == source["data"]["conversationHistory"] + assert "truncation" not in result.to_dict()["data"] diff --git a/python/packages/durabletask/tests/test_state_schema.py b/python/packages/durabletask/tests/test_state_schema.py new file mode 100644 index 0000000..e87393b --- /dev/null +++ b/python/packages/durabletask/tests/test_state_schema.py @@ -0,0 +1,368 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Validate real versioned state against the shared transcript and delivery contract.""" + +import json +from copy import deepcopy +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import jsonschema +import pytest +from agent_framework import AgentResponse, Message + +from agent_framework_durabletask import ( + DurableAgentState, + DurableAgentStateCompaction, + DurableAgentStateErrorResponse, + DurableAgentStateMessage, + DurableAgentStateRequest, + DurableAgentStateResponse, +) +from agent_framework_durabletask._durable_agent_state import DurableAgentStateEntryJsonType +from agent_framework_durabletask._message_identity import message_identity + +SCHEMA_PATH = Path(__file__).resolve().parents[4] / "schemas" / "durable-agent-entity-state.json" + + +@pytest.fixture(scope="module") +def schema() -> dict[str, Any]: + return json.loads(SCHEMA_PATH.read_text(encoding="utf-8")) + + +def _validate(payload: dict[str, Any], schema: dict[str, Any]) -> None: + jsonschema.Draft202012Validator(schema, format_checker=jsonschema.FormatChecker()).validate(payload) + + +def _populated_state() -> DurableAgentState: + """Build real version 2 transcript, delivery, ingestion and opaque session state.""" + now = datetime.now(tz=timezone.utc) + request_message = Message(role="user", contents=["hello"], message_id="wf_input_0") + request = DurableAgentStateRequest( + correlation_id="c0", + created_at=now, + messages=[DurableAgentStateMessage.from_chat_message(request_message)], + ) + core_response = AgentResponse( + messages=[Message(role="assistant", contents=["hi"], author_name="writer", message_id="wf_writer_1")], + response_id="response-0", + agent_id="writer", + created_at=now.isoformat(), + finish_reason="stop", + usage_details={"input_token_count": 2, "output_token_count": 1, "total_token_count": 3}, + additional_properties={"provider": {"metadata": [1, 2]}}, + ) + response = DurableAgentStateResponse.from_run_response("c0", core_response) + # Annotations are what carry compaction state across a round-trip. + response.messages[0].extension_data = {"_excluded": True, "_excluded_reason": "sliding_window"} + + state = DurableAgentState() + state.data.conversation_history.extend([request, response]) + state.data.session = {"type": "session", "session_id": "@dafx-writer@run-1", "state": {"compaction": {}}} + state.data.ingested_messages = {"wf_input_0": [message_identity(request_message)], "legacy-known-id": None} + state.record_response("c0", core_response, delivery_window_seconds=60, now=now) + return state + + +def test_the_schema_itself_is_valid(schema: dict[str, Any]) -> None: + jsonschema.Draft202012Validator.check_schema(schema) + + +def test_empty_state_validates(schema: dict[str, Any]) -> None: + _validate(DurableAgentState().to_dict(), schema) + + +def test_populated_state_validates(schema: dict[str, Any]) -> None: + _validate(_populated_state().to_dict(), schema) + + +def test_message_identity_and_annotations_are_declared(schema: dict[str, Any]) -> None: + """Both are load-bearing for compaction, so an implementer must be told to round-trip them.""" + properties = schema["$defs"]["chatMessage"]["properties"] + + assert "messageId" in properties + assert "extensionData" in properties + + +def test_delivery_and_exact_ingestion_fields_are_declared(schema: dict[str, Any]) -> None: + properties = schema["$defs"]["data"]["properties"] + assert {"responseMailbox", "completedCorrelations", "ingestedMessages"} <= properties.keys() + assert properties["ingestedPositions"]["deprecated"] is True + + +def test_session_is_left_opaque(schema: dict[str, Any]) -> None: + """The two runtimes serialize sessions differently, so the shared schema must not fix a shape. + + .NET produces ``conversationId`` plus ``stateBag``. Python produces ``session_id``, + ``service_session_id`` and ``state``. Declaring either one would make the other invalid. + """ + session = schema["$defs"]["data"]["properties"]["session"] + + assert "properties" not in session, "the schema pins one runtime's session shape" + + dotnet_shaped = { + "schemaVersion": DurableAgentState.SCHEMA_VERSION, + "data": {"conversationHistory": [], "session": {"conversationId": "abc", "stateBag": {}}}, + } + _validate(dotnet_shaped, schema) + + +def test_state_survives_a_round_trip_through_the_schema(schema: dict[str, Any]) -> None: + """Serialize, validate, restore, and confirm the compaction-critical fields came back.""" + payload = _populated_state().to_dict() + _validate(payload, schema) + + restored = DurableAgentState.from_dict(payload) + stored = restored.data.conversation_history[1].messages[0] + + assert stored.message_id == "wf_writer_1" + assert (stored.extension_data or {}).get("_excluded") is True + assert restored.data.ingested_messages == payload["data"]["ingestedMessages"] + delivered = restored.try_get_agent_response("c0") + assert isinstance(delivered, AgentResponse) + assert delivered.to_dict() == payload["data"]["responseMailbox"]["c0"]["response"] + + +def _entry_of_each_kind() -> DurableAgentState: + """State containing each known entry kind, including compaction without a correlation.""" + now = datetime.now(tz=timezone.utc) + state = _populated_state() + state.data.conversation_history.append( + DurableAgentStateErrorResponse( + correlation_id="c1", + created_at=now, + messages=[ + DurableAgentStateMessage.from_chat_message( + Message(role="assistant", contents=["it broke"], message_id="err0") + ) + ], + ) + ) + state.data.conversation_history.append( + DurableAgentStateCompaction( + created_at=now, + messages=[ + DurableAgentStateMessage.from_chat_message( + Message(role="assistant", contents=["summary"], message_id="sum0") + ) + ], + ) + ) + return state + + +def test_every_entry_kind_validates(schema: dict[str, Any]) -> None: + payload = _entry_of_each_kind().to_dict() + + _validate(payload, schema) + + kinds = {entry["$type"] for entry in payload["data"]["conversationHistory"]} + assert kinds == {kind.value for kind in DurableAgentStateEntryJsonType} + + +def test_an_entry_without_a_correlation_omits_the_field(schema: dict[str, Any]) -> None: + """A compaction entry answers no request, so it has no correlation to record. + + Written as an absent field rather than an explicit null. `null` would type the field as + something other than a string wherever a reader looks at it, which the schema rejects and + which a stricter cross-language reader would too. + """ + payload = _entry_of_each_kind().to_dict() + + compaction = next(e for e in payload["data"]["conversationHistory"] if e["$type"] == "compaction") + + assert "correlationId" not in compaction + _validate(payload, schema) + + +def test_the_discriminator_is_required(schema: dict[str, Any]) -> None: + """An entry that does not say what it is must not validate. + + The four entry schemas existed before but nothing referenced them, so `conversationHistory` + accepted any loosely entry-shaped object and `$type` was documentation rather than contract. + """ + payload = { + "schemaVersion": DurableAgentState.SCHEMA_VERSION, + "data": {"conversationHistory": [{"createdAt": datetime.now(tz=timezone.utc).isoformat(), "messages": []}]}, + } + + with pytest.raises(jsonschema.ValidationError): + _validate(payload, schema) + + +def test_future_entries_and_unknown_properties_validate_and_round_trip(schema: dict[str, Any]) -> None: + payload = _populated_state().to_dict() + payload["schemaVersion"] = "2.7.3" + payload["futureRoot"] = {"nested": [1, {"opaque": True}]} + payload["data"]["futureData"] = {"nested": [None, "keep"]} + payload["data"]["conversationHistory"][0]["futureEntry"] = {"nested": [2, 3]} + payload["data"]["responseMailbox"]["c0"]["futureDelivery"] = {"nested": [4, 5]} + payload["data"]["completedCorrelations"]["c0"]["futureReceipt"] = {"nested": [6, 7]} + payload["data"]["conversationHistory"].append({ + "$type": "futureKind", + "payload": {"owned": [1, 2]}, + "messages": {"futureShape": True}, + }) + _validate(payload, schema) + restored = DurableAgentState.from_json(json.dumps(payload)) + assert restored.to_dict() == payload + + +def test_opaque_entry_branch_excludes_exactly_the_known_discriminators(schema: dict[str, Any]) -> None: + kinds = {kind.value for kind in DurableAgentStateEntryJsonType} + opaque = schema["$defs"]["opaqueConversationEntry"] + assert set(opaque["properties"]["$type"]["not"]["enum"]) == kinds + entries = schema["$defs"]["data"]["properties"]["conversationHistory"]["items"]["oneOf"] + typed_kinds = { + definition["properties"]["$type"]["const"] + for entry in entries + if "const" in (definition := schema["$defs"][entry["$ref"].split("/")[-1]])["properties"]["$type"] + } + assert typed_kinds == kinds + + +@pytest.mark.parametrize("kind", list(DurableAgentStateEntryJsonType)) +@pytest.mark.parametrize( + "invalid_fields", + [ + {"messages": "not-an-array"}, + {"messages": [{"contents": []}]}, + {"correlationId": 17}, + {"createdAt": False}, + ], +) +def test_known_entries_cannot_bypass_their_contract_as_opaque_entries( + schema: dict[str, Any], kind: DurableAgentStateEntryJsonType, invalid_fields: dict[str, Any] +) -> None: + payload = { + "schemaVersion": DurableAgentState.SCHEMA_VERSION, + "data": {"conversationHistory": [{"$type": kind.value, **invalid_fields}]}, + } + with pytest.raises(jsonschema.ValidationError): + _validate(payload, schema) + + +@pytest.mark.parametrize("kind", [None, False, 0, "", [], {}]) +def test_invalid_discriminators_are_not_future_entry_kinds(schema: dict[str, Any], kind: Any) -> None: + payload = { + "schemaVersion": DurableAgentState.SCHEMA_VERSION, + "data": {"conversationHistory": [{"$type": kind}]}, + } + with pytest.raises(jsonschema.ValidationError): + _validate(payload, schema) + + +def test_default_schema_version_matches_the_distinct_version_two_writer(schema: dict[str, Any]) -> None: + assert schema["properties"]["schemaVersion"]["default"] == DurableAgentState.SCHEMA_VERSION == "2.0.0" + + +@pytest.mark.parametrize("version", ["1.0.0", "1.1.0", "1.2.0", "2.0.0", "2.7.3"]) +def test_reader_versions_remain_valid_without_implicit_upgrade(schema: dict[str, Any], version: str) -> None: + payload = {"schemaVersion": version, "data": {"conversationHistory": []}} + _validate(payload, schema) + assert DurableAgentState.from_json(json.dumps(payload)).to_dict() == payload + + +@pytest.mark.parametrize("version", [None, False, 2, "", "0.1.0", "3.0.0", "2.0", "2.0.0-preview", "2.0.0\n"]) +def test_schema_rejects_unsupported_or_malformed_versions(schema: dict[str, Any], version: Any) -> None: + with pytest.raises(jsonschema.ValidationError): + _validate({"schemaVersion": version, "data": {}}, schema) + + +@pytest.mark.parametrize("field", ["schemaVersion", "data"]) +def test_root_fields_are_required(schema: dict[str, Any], field: str) -> None: + payload = DurableAgentState().to_dict() + del payload[field] + with pytest.raises(jsonschema.ValidationError): + _validate(payload, schema) + + +@pytest.mark.parametrize("version", ["1.0.0", "1.1.0"]) +def test_legacy_scalar_positions_remain_readable(schema: dict[str, Any], version: str) -> None: + payload = {"schemaVersion": version, "data": {"conversationHistory": [], "ingestedPositions": {"executor": 3}}} + _validate(payload, schema) + assert DurableAgentState.from_json(json.dumps(payload)).to_dict() == payload + + +@pytest.mark.parametrize("field", ["responseMailbox", "completedCorrelations", "ingestedMessages"]) +@pytest.mark.parametrize("value", [None, False, 0, "", [], "not-an-object"]) +def test_delivery_containers_are_typed(schema: dict[str, Any], field: str, value: Any) -> None: + payload = {"schemaVersion": DurableAgentState.SCHEMA_VERSION, "data": {field: value}} + with pytest.raises(jsonschema.ValidationError): + _validate(payload, schema) + + +@pytest.mark.parametrize( + ("record_name", "required_field"), + [ + ("responseMailbox", "response"), + ("responseMailbox", "createdAt"), + ("responseMailbox", "expiresAt"), + ("completedCorrelations", "completedAt"), + ], +) +def test_delivery_record_fields_are_required(schema: dict[str, Any], record_name: str, required_field: str) -> None: + payload = _populated_state().to_dict() + del payload["data"][record_name]["c0"][required_field] + with pytest.raises(jsonschema.ValidationError): + _validate(payload, schema) + + +@pytest.mark.parametrize( + ("record_name", "field"), + [("responseMailbox", "createdAt"), ("responseMailbox", "expiresAt"), ("completedCorrelations", "completedAt")], +) +@pytest.mark.parametrize("value", [None, False, 0, "not-a-timestamp"]) +def test_delivery_timestamps_are_validated(schema: dict[str, Any], record_name: str, field: str, value: Any) -> None: + payload = _populated_state().to_dict() + payload["data"][record_name]["c0"][field] = value + with pytest.raises(jsonschema.ValidationError): + _validate(payload, schema) + + +def test_delivery_timestamp_shape_is_checked_without_optional_format_extras(schema: dict[str, Any]) -> None: + payload = _populated_state().to_dict() + payload["data"]["responseMailbox"]["c0"]["expiresAt"] = "not-a-timestamp" + with pytest.raises(jsonschema.ValidationError): + jsonschema.Draft202012Validator(schema).validate(payload) + + +@pytest.mark.parametrize("response", [None, [], "{}", {}, {"$type": "response", "messages": []}]) +def test_mailbox_requires_inline_core_response_json(schema: dict[str, Any], response: Any) -> None: + payload = _populated_state().to_dict() + payload["data"]["responseMailbox"]["c0"]["response"] = response + with pytest.raises(jsonschema.ValidationError): + _validate(payload, schema) + + +@pytest.mark.parametrize("legacy", [None, 0, 1, "true", [], {}]) +def test_legacy_receipt_marker_is_boolean(schema: dict[str, Any], legacy: Any) -> None: + payload = _populated_state().to_dict() + payload["data"]["completedCorrelations"]["c0"]["legacy"] = legacy + with pytest.raises(jsonschema.ValidationError): + _validate(payload, schema) + + +@pytest.mark.parametrize("fingerprints", [None, [], ["a" * 64, "b" * 64]]) +def test_ingestion_accepts_hash_lists_or_legacy_known_id_markers(schema: dict[str, Any], fingerprints: Any) -> None: + payload = {"schemaVersion": DurableAgentState.SCHEMA_VERSION, "data": {"ingestedMessages": {"id": fingerprints}}} + _validate(payload, schema) + + +@pytest.mark.parametrize("fingerprints", [False, 0, "a" * 64, {}, [None], [1], ["a" * 64, False]]) +def test_ingestion_rejects_invalid_receipts(schema: dict[str, Any], fingerprints: Any) -> None: + payload = {"schemaVersion": DurableAgentState.SCHEMA_VERSION, "data": {"ingestedMessages": {"id": fingerprints}}} + with pytest.raises(jsonschema.ValidationError): + _validate(payload, schema) + + +def test_opaque_session_preserves_owner_message_shaped_data(schema: dict[str, Any]) -> None: + session = {"owner": "external", "state": {"provider": {"messages": [{"custom": "keep"}], "cursor": [1, 2]}}} + payload = { + "schemaVersion": DurableAgentState.SCHEMA_VERSION, + "data": {"conversationHistory": [], "session": session}, + } + original = deepcopy(payload) + _validate(payload, schema) + assert DurableAgentState.from_json(json.dumps(payload)).to_dict() == original diff --git a/python/packages/durabletask/tests/test_subworkflow_orchestration.py b/python/packages/durabletask/tests/test_subworkflow_orchestration.py index 3743396..0607f1e 100644 --- a/python/packages/durabletask/tests/test_subworkflow_orchestration.py +++ b/python/packages/durabletask/tests/test_subworkflow_orchestration.py @@ -33,6 +33,7 @@ _try_unwrap_subworkflow_input, _unpack_subworkflow_result, ) +from agent_framework_durabletask._workflows.protocol import unwrap_workflow_input, wrap_workflow_input from agent_framework_durabletask._workflows.serialization import ( SUBWORKFLOW_RESULT_KEY, deserialize_value, @@ -94,7 +95,11 @@ def test_wraps_message_in_marker(self) -> None: _prepare_subworkflow_task(ctx, executor, "payload", "child-id", _CHILD_ADDRESS) args, _ = ctx.call_sub_orchestrator.call_args - child_input = args[1] + assert args[1] == wrap_workflow_input({ + SUBWORKFLOW_INPUT_KEY: serialize_value("payload"), + SUBWORKFLOW_ADDRESS_KEY: _CHILD_ADDRESS, + }) + child_input = unwrap_workflow_input(args[1]) # The wrapped payload round-trips back to the original message. assert deserialize_value(child_input[SUBWORKFLOW_INPUT_KEY]) == "payload" # The address marker rides alongside so the child can build respond URLs. @@ -330,7 +335,11 @@ def _dispatch( captured: list[dict[str, str]] = [] def _call_sub(name: str, input_: dict[str, object], *, instance_id: str) -> str: # noqa: ARG001 - captured.append(cast("dict[str, str]", input_[SUBWORKFLOW_ADDRESS_KEY])) + child_input = unwrap_workflow_input(input_) + assert input_ == wrap_workflow_input(child_input) + assert set(child_input) == {SUBWORKFLOW_INPUT_KEY, SUBWORKFLOW_ADDRESS_KEY} + assert deserialize_value(child_input[SUBWORKFLOW_INPUT_KEY]) == f"msg-{len(captured)}" + captured.append(cast("dict[str, str]", child_input[SUBWORKFLOW_ADDRESS_KEY])) return f"task::{instance_id}" ctx = Mock() diff --git a/python/packages/durabletask/tests/test_terminal_history_review.py b/python/packages/durabletask/tests/test_terminal_history_review.py new file mode 100644 index 0000000..1f7bf41 --- /dev/null +++ b/python/packages/durabletask/tests/test_terminal_history_review.py @@ -0,0 +1,608 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Terminal delivery versus local model history, using real core hooks and JSON storage.""" + +from __future__ import annotations + +import json +from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence +from copy import deepcopy +from datetime import datetime, timezone +from inspect import signature +from typing import Any, cast + +import pytest +from agent_framework import ( + GROUP_ANNOTATION_KEY, + GROUP_ID_KEY, + Agent, + AgentResponse, + AgentResponseUpdate, + AgentSession, + BaseChatClient, + ChatMiddlewareLayer, + ChatResponse, + ChatResponseUpdate, + Content, + ContextProvider, + FunctionInvocationLayer, + InMemoryHistoryProvider, + Message, + ResponseStream, + SessionContext, + SupportsAgentRun, + annotate_message_groups, + tool, +) +from pydantic import BaseModel, ValidationError + +from agent_framework_durabletask import AgentEntity, AgentEntityStateProviderMixin, DurableHistoryProvider, RunRequest +from agent_framework_durabletask._callbacks import AgentCallbackContext +from agent_framework_durabletask._durable_agent_state import ( + DurableAgentState, + DurableAgentStateErrorResponse, + DurableAgentStateMessage, + DurableAgentStateResponse, +) +from agent_framework_durabletask._history_provider import ( + POSITIONS_KEY, + WORKING_BUFFER_KEY, + DurableHistoryBinding, + bind_durable_history, + unbind_durable_history, +) +from agent_framework_durabletask._message_identity import message_identity +from agent_framework_durabletask._response_utils import is_terminal_agent_response, serialize_agent_response + + +def _json(value: Any) -> Any: + return json.loads(json.dumps(value, allow_nan=False)) + + +class _JsonState(AgentEntityStateProviderMixin): + def __init__(self, raw: dict[str, Any] | None = None) -> None: + self.raw = _json(raw or {}) + self.writes = 0 + + def _get_state_dict(self) -> dict[str, Any]: + return _json(self.raw) + + def _set_state_dict(self, state: dict[str, Any]) -> None: + self.raw = _json(state) + self.writes += 1 + + def _get_session_id_from_entity(self) -> str: + return "terminal-review" + + +def _seed() -> dict[str, Any]: + state = DurableAgentState() + message = Message( + "assistant", + ["previous valid answer"], + message_id="previous-answer", + additional_properties={"provider_metadata": {"keep": [None, False, 7]}}, + ) + state.data.conversation_history.append( + DurableAgentStateResponse( + "previous", datetime(2026, 1, 1, tzinfo=timezone.utc), [DurableAgentStateMessage.from_chat_message(message)] + ) + ) + state.data.session = AgentSession(session_id="terminal-review").to_dict() + state.data.session["state"]["foreign"] = {"opaque": [None, False, {"keep": "original"}]} + return _json(state.to_dict()) + + +def _mailbox(provider: _JsonState, correlation: str) -> dict[str, Any]: + return provider.raw["data"]["responseMailbox"][correlation]["response"] + + +def _request() -> dict[str, Any]: + message = Message("user", ["first input"], message_id="input-id") + return { + "message": "first input", + "correlationId": "first", + "contextMessages": [message.to_dict()], + "contextMessageIds": ["input-occurrence"], + } + + +def _error_message() -> Message: + return Message( + "assistant", + [Content.from_error(message="original model failure", error_code="model_error"), "original terminal text"], + message_id="terminal-output", + author_name="review", + additional_properties={"provider_metadata": {"keep": [1, None, False]}}, + ) + + +class _Count(BaseModel): + count: int + + +class _LegacyAgent: + name = "legacy-review" + id = "legacy-review" + description = None + + def __init__(self, response: AgentResponse[Any]) -> None: + self.response = response + self.inputs: list[list[Message]] = [] + + # Deliberately no stream or **kwargs: exercise the existing signature fallback. + async def run(self, messages: list[Message], *, options: Mapping[str, Any]) -> AgentResponse[Any]: + self.inputs.append(deepcopy(messages)) + return self.response + + +@pytest.mark.parametrize("text", ['{"count":"not an integer"}', "not JSON", '{"count":7}']) +async def test_lazy_value_failure_is_mailbox_only_and_never_legacy_history(text: str) -> None: + original = AgentResponse(messages=[Message("assistant", [text])], response_format=_Count) + valid = text == '{"count":7}' + assert not is_terminal_agent_response(original) + if not valid: + with pytest.raises(ValidationError): + _ = deepcopy(original).value + agent = _LegacyAgent(original) + seed = _seed() + provider = _JsonState(seed) + request = RunRequest("first input", "first", response_format=_Count) + + response = await AgentEntity(cast(SupportsAgentRun, agent), state_provider=provider).run(request) + + assert provider.writes == 1 and len(agent.inputs) == 1 + if valid: + assert response is original and response.value == _Count(count=7) + else: + assert response is not original + assert response.additional_properties["durable_status"] == "error" + assert response.messages[0].contents[0].error_code == "ValidationError" + assert response.text.startswith("ValidationError:") + assert original.text == text + assert _mailbox(provider, "first") == _json(serialize_agent_response(response)) + entries = provider.raw["data"]["conversationHistory"] + assert entries[0] == seed["data"]["conversationHistory"][0] + assert [entry["$type"] for entry in entries[1:]] == (["request", "response"] if valid else ["request"]) + + cold_provider = _JsonState(provider.raw) + next_agent = _LegacyAgent(AgentResponse(messages=[Message("assistant", ["next answer"])])) + cold = AgentEntity(cast(SupportsAgentRun, next_agent), state_provider=cold_provider) + duplicate = await cold.run(request) + assert serialize_agent_response(duplicate) == _mailbox(provider, "first") + assert next_agent.inputs == [] and cold_provider.writes == 0 + await cold.run({"message": "second input", "correlationId": "second"}) + assert [message.text for message in next_agent.inputs[0]] == [ + "previous valid answer", + "first input", + *([text] if valid else []), + "second input", + ] + assert not any(content.type == "error" for message in next_agent.inputs[0] for content in message.contents) + assert _mailbox(cold_provider, "first") == _mailbox(provider, "first") + + +class _ScriptedClient(FunctionInvocationLayer, ChatMiddlewareLayer, BaseChatClient): + def __init__(self, replies: Sequence[Message]) -> None: + super().__init__(middleware=[]) + self.replies = deepcopy(list(replies)) + self.inputs: list[list[Message]] = [] + + def _inner_get_response( + self, *, messages: Sequence[Message], stream: bool, options: Mapping[str, Any], **kwargs: Any + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + self.inputs.append(deepcopy(list(messages))) + message = deepcopy(self.replies[len(self.inputs) - 1]) + response = ChatResponse( + messages=[message], + response_id=f"model-{len(self.inputs)}", + usage_details={"input_token_count": 3, "output_token_count": 2}, + conversation_id="service-history" if options.get("store") else None, + finish_reason="tool_calls" if any(c.type == "function_call" for c in message.contents) else "stop", + ) + if stream: + + async def updates() -> AsyncIterable[ChatResponseUpdate]: + yield ChatResponseUpdate( + role=cast(Any, message.role), + contents=message.contents, + message_id=message.message_id, + author_name=message.author_name, + additional_properties=deepcopy(message.additional_properties), + response_id=response.response_id, + conversation_id=response.conversation_id, + finish_reason=response.finish_reason, + ) + assert response.usage_details is not None + yield ChatResponseUpdate(contents=[Content.from_usage(response.usage_details)]) + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) + + async def get() -> ChatResponse: + return response + + return get() + + +class _NonStreamingAgent(Agent): + def run(self, *args: Any, **kwargs: Any) -> Any: + if kwargs.get("stream"): + raise TypeError("stream is not supported") + return super().run(*args, **kwargs) + + +class _ObservedHistory(DurableHistoryProvider): + def __init__(self, **kwargs: Any) -> None: + super().__init__(prune_excluded=False, **kwargs) + self.responses: list[AgentResponse[Any]] = [] + self.buffers: list[list[Message]] = [] + + async def after_run(self, *, context: SessionContext, state: dict[str, Any], **kwargs: Any) -> None: + assert isinstance(context.response, AgentResponse), "core adapts per-call ChatResponse before the hook" + self.responses.append(deepcopy(context.response)) + await super().after_run(context=context, state=state, **kwargs) + self.buffers.append(deepcopy(state.get(WORKING_BUFFER_KEY, []))) + + +@pytest.mark.parametrize("per_call", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("terminal", [False, True], ids=["successful-control", "terminal-error"]) +async def test_real_core_terminal_outputs_are_not_replayed_after_json_reload( + per_call: bool, stream: bool, terminal: bool +) -> None: + output = _error_message() if terminal else Message("assistant", ["valid answer"], message_id="valid-output") + before_output = deepcopy(output.to_dict()) + client = _ScriptedClient([output]) + history = _ObservedHistory() + agent_type = Agent if stream else _NonStreamingAgent + agent = agent_type( + client=client, context_providers=[history], require_per_service_call_history_persistence=per_call + ) + seed = _seed() + provider = _JsonState(seed) + request = _request() + original_request = deepcopy(request) + + response = await AgentEntity(agent, state_provider=provider).run(request) + + assert is_terminal_agent_response(response) is terminal + assert response.text == output.text and len(client.inputs) == 1 and provider.writes == 1 + assert [content.to_dict() for content in response.messages[0].contents] == [c.to_dict() for c in output.contents] + assert output.to_dict() == before_output and request == original_request + assert _mailbox(provider, "first") == _json(serialize_agent_response(response)) + entries = provider.raw["data"]["conversationHistory"] + assert entries[0] == seed["data"]["conversationHistory"][0] + assert [entry["$type"] for entry in entries[1:]] == ["request", "errorResponse" if terminal else "response"] + saved_response = provider.state.data.conversation_history[-1] + assert isinstance(saved_response, DurableAgentStateResponse) and saved_response.usage is not None + assert saved_response.usage.to_usage_details()["input_token_count"] == 3 + assert len(history.responses) == 1 + assert [m.text for m in history.buffers[0]] == [ + "previous valid answer", + "first input", + *([] if terminal else [output.text]), + ] + assert provider.raw["data"]["session"]["state"] == seed["data"]["session"]["state"] + message = Message.from_dict(request["contextMessages"][0]) + assert provider.raw["data"]["ingestedMessages"] == {"input-occurrence": [message_identity(message)]} + + cold_provider = _JsonState(provider.raw) + cold_client = _ScriptedClient([Message("assistant", ["next answer"])]) + cold = AgentEntity( + agent_type(client=cold_client, require_per_service_call_history_persistence=per_call), + state_provider=cold_provider, + ) + assert (await cold.run(request)).to_dict() == _mailbox(provider, "first") + assert cold_client.inputs == [] and cold_provider.writes == 0 + await cold.run({"message": "second input", "correlationId": "second"}) + assert [m.text for m in cold_client.inputs[0]] == [ + "previous valid answer", + "first input", + *([] if terminal else [output.text]), + "second input", + ] + assert not any(c.type == "error" for m in cold_client.inputs[0] for c in m.contents) + assert _mailbox(cold_provider, "first") == _mailbox(provider, "first") + + +class _GroupAfter(ContextProvider): + def __init__(self, source_id: str) -> None: + super().__init__("last-after") + self.history_source = source_id + self.groups: dict[str, dict[str, Any]] = {} + + async def after_run(self, *, session: AgentSession, **kwargs: Any) -> None: + buffer = session.state[self.history_source][WORKING_BUFFER_KEY] + annotate_message_groups(buffer, force_reannotate=True) + for message in buffer: + if any(c.type in ("function_call", "function_result") for c in message.contents): + message.additional_properties["last_after"] = {"keep": [None, False, 1]} + self.groups[message.message_id] = deepcopy(message.additional_properties[GROUP_ANNOTATION_KEY]) + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_later_terminal_call_keeps_prior_tool_pair_and_final_hook_group_metadata(stream: bool) -> None: + tool_invocations: list[str] = [] + + @tool(name="lookup", approval_mode="never_require") + def lookup(key: str) -> str: + """Return a deterministic lookup result.""" + tool_invocations.append(key) + return f"value:{key}" + + call = Message( + "assistant", [Content.from_function_call("call-1", "lookup", arguments='{"key":"kept"}')], message_id="call" + ) + client = _ScriptedClient([call, _error_message()]) + history = _ObservedHistory() + last_after = _GroupAfter(history.source_id) + agent_type = Agent if stream else _NonStreamingAgent + provider = _JsonState(_seed()) + response = await AgentEntity( + agent_type( + client=client, + tools=[lookup], + context_providers=[last_after, history], + require_per_service_call_history_persistence=True, + ), + state_provider=provider, + ).run(_request()) + + assert is_terminal_agent_response(response) and response.text == "original terminal text" + assert tool_invocations == ["kept"] and len(client.inputs) == 2 + assert len(history.responses) == 2 + assert not is_terminal_agent_response(history.responses[0]) and is_terminal_agent_response(history.responses[1]) + assert [entry.json_type for entry in provider.state.data.conversation_history[1:]] == [ + "request", + "response", + "request", + "errorResponse", + ] + assert len(last_after.groups) == 2 + assert len({group[GROUP_ID_KEY] for group in last_after.groups.values()}) == 1 + assert not any(c.type == "error" for batch in history.buffers for m in batch for c in m.contents) + assert _mailbox(provider, "first") == _json(serialize_agent_response(response)) + + cold_client = _ScriptedClient([Message("assistant", ["next answer"])]) + cold_provider = _JsonState(provider.raw) + await AgentEntity( + agent_type(client=cold_client, require_per_service_call_history_persistence=True), state_provider=cold_provider + ).run({"message": "second input", "correlationId": "second"}) + replayed = cold_client.inputs[0] + assert [m.text for m in replayed] == ["previous valid answer", "first input", "", "", "second input"] + calls = [c for m in replayed for c in m.contents if c.type == "function_call"] + results = [c for m in replayed for c in m.contents if c.type == "function_result"] + assert len(calls) == len(results) == 1 and calls[0].call_id == results[0].call_id == "call-1" + assert results[0].result == "value:kept" and tool_invocations == ["kept"] + for message in replayed: + if message.message_id in last_after.groups: + assert message.additional_properties[GROUP_ANNOTATION_KEY] == last_after.groups[message.message_id] + assert message.additional_properties["last_after"] == {"keep": [None, False, 1]} + assert _mailbox(cold_provider, "first") == _mailbox(provider, "first") + + +@pytest.mark.parametrize("per_call", [False, True]) +@pytest.mark.parametrize("store_inputs", [False, True]) +@pytest.mark.parametrize("store_outputs", [False, True]) +@pytest.mark.parametrize("service_owned", [False, True]) +async def test_terminal_core_hooks_respect_storage_flags_and_service_ownership( + per_call: bool, store_inputs: bool, store_outputs: bool, service_owned: bool +) -> None: + seed = _seed() + provider = _JsonState(seed) + history = _ObservedHistory(store_inputs=store_inputs, store_outputs=store_outputs) + client = _ScriptedClient([_error_message()]) + agent = Agent(client=client, context_providers=[history], require_per_service_call_history_persistence=per_call) + request = {**_request(), "options": {"store": service_owned}} + + response = await AgentEntity(agent, state_provider=provider).run(request) + + assert is_terminal_agent_response(response) and response.text == "original terminal text" + entries = provider.raw["data"]["conversationHistory"] + assert entries[0] == seed["data"]["conversationHistory"][0] + expected = ( + [] if service_owned else (["request"] if store_inputs else []) + (["errorResponse"] if store_outputs else []) + ) + assert [entry["$type"] for entry in entries[1:]] == expected + assert bool(provider.raw["data"].get("ingestedMessages")) is (store_inputs and not service_owned) + assert [m.text for m in client.inputs[0]] == ([] if service_owned else ["previous valid answer"]) + ["first input"] + assert _mailbox(provider, "first") == _json(serialize_agent_response(response)) + assert provider.writes == 1 + + +@pytest.mark.parametrize("kind", ["error", "already_completed", "tool-error", "approval", "success"]) +@pytest.mark.parametrize("store_context", [False, True]) +@pytest.mark.parametrize("context_sources", [None, set(), {"selected"}]) +async def test_shared_classifier_and_context_masks_exclude_only_terminal_output_batches( + kind: str, store_context: bool, context_sources: set[str] | None +) -> None: + metadata = {"durable_status": kind} if kind in ("error", "already_completed") else {} + messages = [Message("assistant", ["not structured JSON"])] + if kind == "tool-error": + messages = [ + Message("assistant", [Content.from_function_call("call", "lookup", arguments="{}")]), + Message( + "tool", [Content.from_function_result("call", result="failed"), Content.from_error(message="tool")] + ), + Message("assistant", ["recovered"]), + ] + elif kind == "approval": + messages[0].contents.append( + Content.from_function_approval_request( + "approval", Content.from_function_call("call", "lookup", arguments="{}") + ) + ) + response = AgentResponse(messages=messages, additional_properties=metadata) + terminal = is_terminal_agent_response(response) + assert terminal is (kind in ("error", "already_completed")) + before = deepcopy(response.to_dict()) + provider = _JsonState(_seed()) + history = DurableHistoryProvider( + store_context_messages=store_context, store_context_from=context_sources, prune_excluded=False + ) + context = SessionContext(input_messages=[Message("user", ["accepted input"])]) + for source in ("selected", "other"): + context.extend_messages(source, [Message("user", [f"context-{source}"])]) + context.extend_messages(history, [Message("assistant", ["must not duplicate own history"])]) + context._response = response + state: dict[str, Any] = {} + token = bind_durable_history(DurableHistoryBinding(provider, "current")) + try: + await history.after_run(agent=None, session=None, context=context, state=state) + expected_inputs = [ + f"context-{source}" + for source in ("selected", "other") + if store_context and (context_sources is None or source in context_sources) + ] + ["accepted input"] + expected = ["previous valid answer", *expected_inputs, *([] if terminal else [m.text for m in messages])] + assert [m.text for m in state[WORKING_BUFFER_KEY]] == expected + entry = provider.state.data.conversation_history[-1] + assert isinstance(entry, DurableAgentStateErrorResponse) is terminal + assert len(state[POSITIONS_KEY]) == len(expected) + snapshot = _json(provider.state.to_dict()) + history.flush(state) + history.flush(state) + assert provider.state.to_dict() == snapshot, "terminal outputs must not be resurrected as compaction entries" + assert [m.text for m in await history.get_messages("terminal-review", state={})] == expected + assert provider.writes == 0 and response.to_dict() == before + finally: + unbind_durable_history(token) + + +@pytest.mark.parametrize("working_state", [False, True]) +async def test_aggregated_terminal_batch_does_not_erase_or_duplicate_prior_tool_pair(working_state: bool) -> None: + provider = _JsonState(_seed()) + history = DurableHistoryProvider(prune_excluded=False) + pair = [ + Message("assistant", [Content.from_function_call("call", "lookup", arguments="{}")], message_id="call"), + Message("tool", [Content.from_function_result("call", result="kept")], message_id="result"), + ] + annotate_message_groups(pair, force_reannotate=True) + original_pair = [deepcopy(m.to_dict()) for m in pair] + state: dict[str, Any] | None = {} if working_state else None + binding = DurableHistoryBinding(provider, "current") + token = bind_durable_history(binding) + try: + history._append_messages(binding, pair, state=state, response=AgentResponse(messages=pair)) + prior_entry = deepcopy(provider.state.data.conversation_history[-1].to_dict()) + terminal = AgentResponse(messages=[*deepcopy(pair), _error_message()]) + history._append_messages(binding, terminal.messages, state=state, response=terminal) + if state is not None: + history.flush(state) + assert isinstance(provider.state.data.conversation_history[-1], DurableAgentStateErrorResponse) + assert provider.state.data.conversation_history[-2].to_dict() == prior_entry + provider.persist_state() + finally: + unbind_durable_history(token) + + cold = _JsonState(provider.raw) + token = bind_durable_history(DurableHistoryBinding(cold, "next")) + try: + replayed = await history.get_messages("terminal-review", state={}) + finally: + unbind_durable_history(token) + assert [m.message_id for m in replayed] == ["previous-answer", "call", "result"] + assert [m.to_dict() for m in replayed[1:]] == original_pair + assert [m.to_dict() for m in pair] == original_pair + + +@pytest.mark.parametrize("per_call", [False, True]) +async def test_real_core_pending_approval_still_skips_lazy_typed_validation(per_call: bool) -> None: + output = Message( + "assistant", + [ + "approval needed, not JSON", + Content.from_function_approval_request( + "approval", Content.from_function_call("call", "lookup", arguments="{}") + ), + ], + ) + client = _ScriptedClient([output]) + provider = _JsonState() + response = await AgentEntity( + Agent(client=client, require_per_service_call_history_persistence=per_call), state_provider=provider + ).run(RunRequest("first input", "first", response_format=_Count)) + + assert len(client.inputs) == 1 and not is_terminal_agent_response(response) + assert response.text == output.text and len(response.user_input_requests) == 1 + with pytest.raises(ValidationError): + _ = deepcopy(response).value + assert "value" not in _mailbox(provider, "first") + assert [entry["$type"] for entry in provider.raw["data"]["conversationHistory"]] == ["request", "response"] + + +class _ExternalHistory(InMemoryHistoryProvider): + """A custom primary keeps its own transcript semantics, including terminal messages.""" + + +async def test_external_primary_is_not_rewritten_or_given_a_second_durable_transcript() -> None: + seed = _seed() + external = _ExternalHistory("external") + history_message = Message("assistant", ["external prior"], message_id="external-prior") + seed["data"]["session"]["state"]["external"] = { + "messages": [history_message.to_dict()], + "opaque": {"keep": [None, False, 1]}, + } + provider = _JsonState(seed) + client = _ScriptedClient([_error_message()]) + agent = Agent(client=client, context_providers=[external]) + entity = AgentEntity(agent, state_provider=provider) + + response = await entity.run(_request()) + + assert entity.agent is agent and agent.context_providers == [external] + assert response.text == "original terminal text" + assert provider.raw["data"]["conversationHistory"] == seed["data"]["conversationHistory"] + saved = AgentSession.from_dict(provider.raw["data"]["session"]).state["external"] + assert saved["opaque"] == {"keep": [None, False, 1]} + assert saved["messages"][0].to_dict() == history_message.to_dict() + assert [m.text for m in saved["messages"]] == ["external prior", "first input", "original terminal text"] + + cold_client = _ScriptedClient([Message("assistant", ["next answer"])]) + cold_provider = _JsonState(provider.raw) + await AgentEntity( + Agent(client=cold_client, context_providers=[_ExternalHistory("external")]), state_provider=cold_provider + ).run({"message": "second input", "correlationId": "second"}) + # No promise to filter an opaque primary: only DurableHistoryProvider owns the new policy. + assert [m.text for m in cold_client.inputs[0]] == [ + "external prior", + "first input", + "original terminal text", + "second input", + ] + assert provider.raw["data"]["conversationHistory"] == cold_provider.raw["data"]["conversationHistory"] + assert _mailbox(cold_provider, "first") == _mailbox(provider, "first") + + +class _Callback: + def __init__(self) -> None: + self.responses: list[AgentResponse[Any]] = [] + self.contexts: list[AgentCallbackContext] = [] + self.updates: list[AgentResponseUpdate] = [] + + async def on_streaming_response_update(self, update: AgentResponseUpdate, context: AgentCallbackContext) -> None: + self.updates.append(update) + + async def on_agent_response(self, response: AgentResponse[Any], context: AgentCallbackContext) -> None: + self.responses.append(response) + self.contexts.append(context) + response.messages[0].contents[0].text = "callback copy only" + + +async def test_nonstreaming_signature_fallback_and_final_callback_contract_are_unchanged() -> None: + assert list(signature(_LegacyAgent.run).parameters) == ["self", "messages", "options"] + original = AgentResponse(messages=[Message("assistant", ['{"count":7}'])], response_format=_Count) + agent = _LegacyAgent(original) + callback = _Callback() + provider = _JsonState() + + response = await AgentEntity(cast(SupportsAgentRun, agent), callback=callback, state_provider=provider).run( + RunRequest("first input", "first", response_format=_Count) + ) + + assert len(agent.inputs) == len(callback.responses) == 1 and callback.updates == [] + assert response is original and response.value == _Count(count=7) + assert response.text == '{"count":7}' and callback.responses[0].text == "callback copy only" + assert callback.responses[0] is not response and callback.responses[0]._response_format is _Count + assert callback.contexts == [AgentCallbackContext("legacy-review", "first", "terminal-review", "first input")] + assert _mailbox(provider, "first") == _json(serialize_agent_response(response)) diff --git a/python/packages/durabletask/tests/test_workflow_agent_contract_review.py b/python/packages/durabletask/tests/test_workflow_agent_contract_review.py new file mode 100644 index 0000000..36939cb --- /dev/null +++ b/python/packages/durabletask/tests/test_workflow_agent_contract_review.py @@ -0,0 +1,715 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Core 1.16 agent yields, approval barriers and locally declared structured output.""" + +from __future__ import annotations + +import asyncio +import json +from copy import deepcopy +from datetime import datetime, timezone +from typing import Any, cast +from unittest.mock import Mock, patch +from uuid import UUID + +import pytest +from agent_framework import ( + AgentExecutor, + AgentExecutorResponse, + AgentResponse, + AgentSession, + Content, + Executor, + Message, + Workflow, + WorkflowBuilder, + WorkflowContext, + WorkflowExecutor, + handler, +) +from durabletask.task import CompletableTask, OrchestrationContext +from pydantic import BaseModel, Field + +from agent_framework_durabletask import AgentEntity, AgentEntityStateProviderMixin, serialize_agent_response +from agent_framework_durabletask._workflows.activity import execute_workflow_activity +from agent_framework_durabletask._workflows.dt_context import DurableTaskWorkflowContext +from agent_framework_durabletask._workflows.orchestrator import ( + SOURCE_HITL_RESPONSE, + ExecutorResult, + TaskMetadata, + TaskType, + _collect_hitl_requests, + _prepare_agent_task, + _process_agent_response, + _WorkflowDeliveryLedger, + run_workflow_orchestrator, +) +from agent_framework_durabletask._workflows.protocol import unwrap_workflow_input +from agent_framework_durabletask._workflows.serialization import deserialize_value + + +class Answer(BaseModel): + answer: int = Field(validation_alias="inputAnswer", serialization_alias="outputAnswer") + + +class _Agent: + description = None + + def __init__(self, name: str, responses: list[AgentResponse], response_format: Any = None) -> None: + self.id = self.name = name + self.responses = iter(responses) + self.default_options = {"response_format": response_format} + self.inputs: list[list[Message]] = [] + + def create_session(self, **kwargs: Any) -> AgentSession: + return AgentSession(**kwargs) + + async def run( + self, messages: list[Message], *, session: AgentSession | None = None, **kwargs: Any + ) -> AgentResponse: + self.inputs.append(deepcopy(messages)) + return next(self.responses) + + +def _agent(name: str, responses: list[AgentResponse] | None = None, response_format: Any = None) -> AgentExecutor: + agent: Any = _Agent(name, responses or [], response_format) + return AgentExecutor(agent, id=name) + + +def _response(text: str = "done", **kwargs: Any) -> AgentResponse: + return AgentResponse(messages=[Message("assistant", [text])], **kwargs) + + +def _approval(request_id: str) -> Content: + return Content.from_function_approval_request( + request_id, Content.from_function_call(f"call-{request_id}", "lookup", arguments={"flag": False}) + ) + + +def _pending(requests: list[Content]) -> AgentResponse: + # Non-request content is intentionally suppressed, just as in core non-streaming mode. + return AgentResponse(messages=[Message("assistant", [Content.from_text("not final"), *requests])]) + + +def _wire(response: AgentResponse) -> dict[str, Any]: + return json.loads(json.dumps(serialize_agent_response(response), allow_nan=False)) + + +class _Adapter: + """Use actual host adapters and task wrappers, mocking only native scheduling.""" + + def __init__(self, kind: str, *, replay: bool = False) -> None: + self.kind = kind + self.pending: list[tuple[str, Any, tuple[Any, ...], dict[str, Any]]] = [] + self.calls: list[tuple[str, Any, tuple[Any, ...], dict[str, Any]]] = [] + self.statuses: list[dict[str, Any]] = [] + self.ordinal = 0 + if kind == "dt": + self.native = Mock(spec=OrchestrationContext) + self.context: Any = DurableTaskWorkflowContext(self.native) + else: + df = pytest.importorskip("azure.durable_functions") + module = pytest.importorskip("agent_framework_azurefunctions._workflow_af_context") + self.native = Mock(spec=df.DurableOrchestrationContext) + self.native.task_all.side_effect = lambda tasks: tasks + self.context = module.AzureFunctionsWorkflowContext(self.native) + self.native.instance_id = "contract-run" + self.native.is_replaying = replay + self.native.current_utc_datetime = datetime(2026, 9, 9, tzinfo=timezone.utc) + self.native.new_uuid.side_effect = [str(UUID(int=i)) for i in range(1, 100)] + self.native.call_entity.side_effect = lambda *a, **kw: self.schedule("entity", *a, **kw) + self.native.call_activity.side_effect = lambda *a, **kw: self.schedule("activity", *a, **kw) + self.native.call_sub_orchestrator.side_effect = lambda *a, **kw: self.schedule("child", *a, **kw) + self.native.wait_for_external_event.side_effect = lambda *a, **kw: self.schedule("event", *a, **kw) + self.native.set_custom_status.side_effect = lambda status: self.statuses.append(deepcopy(status)) + + def schedule(self, kind: str, *args: Any, **kwargs: Any) -> Any: + if self.kind == "dt": + task: Any = CompletableTask() + else: + from azure.durable_functions.models.actions.NoOpAction import NoOpAction + from azure.durable_functions.models.Task import AtomicTask + + task = AtomicTask(self.ordinal, NoOpAction()) + self.ordinal += 1 + call = (kind, task, args, kwargs) + self.pending.append(call) + self.calls.append(call) + return task + + def complete(self, yielded: Any, *values: Any) -> Any: + assert len(values) == len(self.pending) + pending, self.pending = self.pending, [] + for (_, task, _, _), value in zip(pending, values, strict=True): + value = json.loads(json.dumps(value, allow_nan=False)) + if self.kind == "dt": + task.complete(value) + else: + task.set_value(is_error=False, value=value) + if isinstance(yielded, list): + return [self.context.get_task_result(task) for task in yielded] + return self.context.get_task_result(yielded) + + def payload(self) -> dict[str, Any]: + assert len(self.pending) == 1 + kind, _, args, _ = self.pending[0] + assert kind == "entity" + assert args[1] == "run" + return json.loads(json.dumps(args[2], allow_nan=False)) + + def activity_input(self) -> str: + kind, _, args, kwargs = self.pending[0] + assert kind == "activity" + return args[1] if len(args) > 1 else kwargs["input"] + + def finish(self, generator: Any, yielded: Any, *values: Any) -> Any: + with pytest.raises(StopIteration) as completed: + generator.send(self.complete(yielded, *values)) + return deserialize_value(completed.value.value) + + +@pytest.mark.parametrize("adapter", ["dt", "af"]) +def test_one_agent_output_matches_core_type_count_and_designation(adapter: str) -> None: + response = _response("answer", response_id="actual-response") + core = WorkflowBuilder(name="core", start_executor=_agent("A", [response])).build() + + async def core_run() -> list[Any]: + return (await core.run("question")).get_outputs() + + expected: list[Any] = asyncio.run(core_run()) + assert len(expected) == 1 + assert isinstance(expected[0], AgentResponse) + + a = _agent("A") + workflow = WorkflowBuilder(name="review", start_executor=a).build() + host = _Adapter(adapter) + generator = run_workflow_orchestrator(host.context, workflow, "question") + yielded = next(generator) + output = host.finish(generator, yielded, _wire(response)) + assert len(output) == len(expected) + assert type(output[0]) is AgentResponse + assert output[0].to_dict() == expected[0].to_dict() + + +@pytest.mark.parametrize("adapter", ["dt", "af"]) +@pytest.mark.parametrize("designation", ["output", "intermediate", "hidden"]) +def test_agent_yields_follow_real_workflow_designation_and_streaming_gate(adapter: str, designation: str) -> None: + a, b = _agent("A"), _agent("B") + workflow = ( + WorkflowBuilder( + name="review", + start_executor=a, + output_from=[a, b] if designation == "output" else [b], + intermediate_output_from=[a] if designation == "intermediate" else [], + ) + .add_edge(a, b) + .build() + ) + host = _Adapter(adapter) + generator = run_workflow_orchestrator(host.context, workflow, "question") + yielded = next(generator) + yielded = generator.send(host.complete(yielded, _wire(_response("A")))) + assert [Message.from_dict(m).text for m in host.payload()["contextMessages"]] == ["question", "A"] + output = host.finish(generator, yielded, _wire(_response("B"))) + assert [value.text for value in output] == (["A", "B"] if designation == "output" else ["B"]) + if adapter == "af": + assert all("events" not in status for status in host.statuses) + else: + events = host.statuses[-1]["events"] + yields = [event for event in events if event["type"] in ("output", "intermediate")] + assert [(event["executor_id"], event["type"]) for event in yields] == ( + [("A", designation), ("B", "output")] if designation != "hidden" else [("B", "output")] + ) + + +class _InspectEnvelope(Executor): + def __init__(self) -> None: + super().__init__(id="inspect") + + @handler + async def inspect_response( + self, response: AgentExecutorResponse, ctx: WorkflowContext[None, AgentResponse] + ) -> None: + assert isinstance(response.agent_response.value, Answer) + assert response.agent_response.value.answer == 42 + ctx.set_state("verified", True) + ctx.state.delete("remove") + await ctx.yield_output(response.agent_response) + + +@pytest.mark.parametrize("adapter", ["dt", "af"]) +@pytest.mark.parametrize("retained", [False, True]) +def test_declared_model_survives_entity_wire_condition_activity_and_output(adapter: str, retained: bool) -> None: + response = ( + _response("not structured text", value=Answer(inputAnswer=42)) if retained else _response('{"inputAnswer":42}') + ) + a, inspect_response = _agent("A", response_format=Answer), _InspectEnvelope() + observed: list[AgentExecutorResponse] = [] + + def condition(value: AgentExecutorResponse) -> bool: + observed.append(value) + assert isinstance(value.agent_response.value, Answer) + return value.agent_response.value.answer == 42 + + workflow = ( + WorkflowBuilder(name="review", start_executor=a, output_from=[a, inspect_response]) + .add_edge(a, inspect_response, condition=condition) + .build() + ) + state = {"remove": True, "keep": False} + host = _Adapter(adapter) + generator = run_workflow_orchestrator(host.context, workflow, "question", state) + yielded = next(generator) + assert "response_format" not in host.payload() + payload = _wire(response) + # Persisted Python class names are not a source of declared response types. + payload["response_format"] = "untrusted.module:Model" + with patch("importlib.import_module", side_effect=AssertionError("must not resolve wire types")): + yielded = generator.send(host.complete(yielded, payload)) + assert isinstance(observed[0].agent_response.value, Answer) + encoded_input = host.activity_input() + restored = deserialize_value(json.loads(encoded_input)["message"]) + assert isinstance(restored, AgentExecutorResponse) + assert isinstance(restored.agent_response.value, Answer) + result = execute_workflow_activity(inspect_response, encoded_input, workflow) + output = host.finish(generator, yielded, result) + assert len(output) == 2 + assert all(isinstance(value, AgentResponse) for value in output) + # Generated agent output is portable JSON; an explicit activity yield retains + # the existing arbitrary-object checkpoint contract. + assert output[0].value == {"answer": 42} + assert isinstance(output[1].value, Answer) and output[1].value.answer == 42 + assert state == {"keep": False, "verified": True} + + +class _InspectChild(Executor): + def __init__(self) -> None: + super().__init__(id="inspect-child") + + @handler + async def inspect_response(self, response: AgentResponse, ctx: WorkflowContext[None, str]) -> None: + assert type(response) is AgentResponse + await ctx.yield_output(response.text) + + +@pytest.mark.parametrize("adapter", ["dt", "af"]) +@pytest.mark.parametrize("direct", [False, True]) +def test_child_ending_in_agent_forwards_agent_response_not_executor_envelope(adapter: str, direct: bool) -> None: + inner = WorkflowBuilder(name="inner", start_executor=_agent("A")).build() + child = WorkflowExecutor(inner, id="child", allow_direct_output=direct) + inspector = _InspectChild() + builder = WorkflowBuilder(name="outer", start_executor=child, output_from=[child] if direct else [inspector]) + if not direct: + builder.add_edge(child, inspector) + outer = builder.build() + host = _Adapter(adapter) + generator = run_workflow_orchestrator(host.context, outer, "question") + yielded = next(generator) + _, _, _, kwargs = host.pending[0] + child_input = unwrap_workflow_input(kwargs["input"] if adapter == "dt" else kwargs["input_"]) + child_host = _Adapter(adapter) + child_host.native.instance_id = kwargs["instance_id"] + child_generator = run_workflow_orchestrator(child_host.context, inner, child_input) + child_yielded = next(child_generator) + with pytest.raises(StopIteration) as completed: + child_generator.send(child_host.complete(child_yielded, _wire(_response("child answer")))) + child_result = completed.value.value + assert type(deserialize_value(child_result["outputs"])[0]) is AgentResponse + if direct: + output = host.finish(generator, yielded, child_result) + assert len(output) == 1 and type(output[0]) is AgentResponse + else: + yielded = generator.send(host.complete(yielded, child_result)) + encoded_input = host.activity_input() + assert type(deserialize_value(json.loads(encoded_input)["message"])) is AgentResponse + result = execute_workflow_activity(inspector, encoded_input, outer) + assert host.finish(generator, yielded, result) == ["child answer"] + + +@pytest.mark.parametrize("adapter", ["dt", "af"]) +@pytest.mark.parametrize("count", [1, 2]) +@pytest.mark.parametrize("reply_kind", ["approval", "results", "mixed"]) +@pytest.mark.parametrize("batch", [False, True]) +def test_approval_barrier_resumes_once_after_all_answers_with_core_content_and_role( + adapter: str, count: int, reply_kind: str, batch: bool +) -> None: + request_ids = [f"request-{i}" for i in range(count)] + requests = [_approval(request_id) for request_id in request_ids] + replies = [ + Content.from_function_result(f"call-request-{i}", result=False if i == 0 else 0) + if reply_kind == "results" or (reply_kind == "mixed" and i == 0) + else request.to_function_approval_response(False) + for i, request in enumerate(requests) + ] + expected_role = "tool" if all(reply.type == "function_result" for reply in replies) else "user" + core_a, core_b = _agent("A", [_pending(requests), _response("final")]), _agent("B", [_response("B")]) + core_agent_a = cast(_Agent, core_a.agent) + core_agent_b = cast(_Agent, core_b.agent) + core = WorkflowBuilder(name="core", start_executor=core_a, output_from=[core_b]).add_edge(core_a, core_b).build() + + async def core_run() -> None: + pending = await core.run("question") + assert pending.get_outputs() == [] + assert [event.request_id for event in pending.get_request_info_events()] == [r.id for r in requests] + if batch: + await core.run( + responses={request_id: reply for request_id, reply in zip(request_ids, replies, strict=True)} + ) + assert len(core_agent_b.inputs) == 1 + else: + for index, (request_id, reply) in enumerate(zip(request_ids, replies, strict=True)): + await core.run(responses={request_id: reply}) + assert len(core_agent_b.inputs) == int(index == count - 1) + + asyncio.run(core_run()) + expected_input = core_agent_a.inputs[-1] + assert len(expected_input) == 1 and expected_input[0].role == expected_role + + a, b = _agent("A", response_format=Answer), _agent("B") + workflow = WorkflowBuilder(name="review", start_executor=a, output_from=[b]).add_edge(a, b).build() + before = deepcopy(a._session.to_dict()) + host = _Adapter(adapter) + generator = run_workflow_orchestrator(host.context, workflow, "question") + yielded = next(generator) + first_payload = host.payload() + first_entity = host.native.call_entity.call_args.args[0] + yielded = generator.send(host.complete(yielded, _wire(_pending(requests)))) + for index, (request, reply) in enumerate(zip(requests, replies, strict=True)): + assert host.pending[0][0] == "event" + assert host.pending[0][2] == (request.id,) + assert host.native.call_entity.call_count == 1 + waiting = host.statuses[-1] + assert waiting["state"] == "waiting_for_human_input" + assert list(waiting["pending_requests"]) == [r.id for r in requests[index:]] + event = waiting["pending_requests"][request.id] + assert deserialize_value(event["data"]) == request + assert event["response_type"] == f"{Content.__module__}:{Content.__name__}" + yielded = generator.send(host.complete(yielded, reply.to_dict())) + assert host.native.call_entity.call_count == 2 + assert host.native.call_activity.call_count == 0 + resumed = host.payload() + assert resumed["correlationId"] != first_payload["correlationId"] + assert host.native.call_entity.call_args.args[0] == first_entity + assert resumed["contextMessages"] == [message.to_dict() for message in expected_input] + assert len(resumed["contextMessageIds"]) == 1 + assert resumed["message"] == "" + final = _response("final", value=Answer(inputAnswer=42)) + yielded = generator.send(host.complete(yielded, _wire(final))) + downstream = host.payload() + assert [Message.from_dict(m).to_dict() for m in downstream["contextMessages"]] == [ + *[message.to_dict() for message in expected_input], + *[message.to_dict() for message in final.messages], + ] + assert downstream["contextMessageIds"][0] == resumed["contextMessageIds"][0] + output = host.finish(generator, yielded, _wire(_response("B"))) + assert len(output) == 1 and output[0].text == "B" + assert a._pending_agent_requests == {} and a._pending_responses_to_agent == [] and a._cache == [] + assert a._session.to_dict() == before + if adapter == "dt": + events = host.statuses[-1]["events"] + assert [event["request_id"] for event in events if event["type"] == "request_info"] == [r.id for r in requests] + assert [event["executor_id"] for event in events if event["type"] == "output"] == ["B"] + + +@pytest.mark.parametrize("adapter", ["dt", "af"]) +@pytest.mark.parametrize("bad", [None, False, {"type": ""}, {"__pickled__": "bad", "__type__": "bad:Type"}]) +def test_invalid_agent_reply_does_not_consume_request_and_can_be_corrected(adapter: str, bad: Any) -> None: + request = _approval("request") + a = _agent("A", response_format=Answer) + workflow = WorkflowBuilder(name="review", start_executor=a).build() + host = _Adapter(adapter) + generator = run_workflow_orchestrator(host.context, workflow, "question") + yielded = generator.send(host.complete(next(generator), _wire(_pending([request])))) + yielded = generator.send(host.complete(yielded, bad)) + assert host.pending[0][2] == ("request",) + assert host.native.call_entity.call_count == 1 + assert "request" in host.statuses[-1]["pending_requests"] + yielded = generator.send(host.complete(yielded, request.to_function_approval_response(False).to_dict())) + assert host.native.call_entity.call_count == 2 + assert host.finish(generator, yielded, _wire(_response(value=Answer(inputAnswer=42))))[0].value == {"answer": 42} + + +def test_duplicate_unknown_and_out_of_order_replies_keep_accumulated_content_and_prepare_is_atomic() -> None: + a, host, ledger = _agent("A"), _Adapter("dt"), _WorkflowDeliveryLedger(instance_id="contract-run") + metadata = TaskMetadata("A", "question", "start", TaskType.AGENT) + _prepare_agent_task(host.context, a, "A", "question", "review", ledger, metadata) + requests = [_approval("first"), _approval("second")] + result = _process_agent_response(_wire(_pending(requests)), "A", "question", ledger, metadata) + assert result.output_message is None + + def respond(request_id: str, response: Content) -> Any: + message = {"request_id": request_id, "response": response.to_dict(), "response_type": "unsafe.module:Type"} + meta = TaskMetadata("A", message, f"{SOURCE_HITL_RESPONSE}_{request_id}", TaskType.AGENT) + return _prepare_agent_task(host.context, a, "A", message, "review", ledger, meta) + + second = requests[1].to_function_approval_response(False) + assert respond("second", second) is None + snapshot = ledger.fork() + assert respond("second", second) is None + assert respond("unknown", second) is None + assert ledger == snapshot + first = Content.from_function_result("call-first", result=0) + with ( + patch.object(host.context, "prepare_agent_task", side_effect=OSError("prepare failed")), + pytest.raises(OSError, match="prepare failed"), + ): + respond("first", first) + assert ledger == snapshot + assert respond("first", first) is not None + wire = host.native.call_entity.call_args.args[2] + assert wire["contextMessages"] == [Message("user", [second, first]).to_dict()] + assert ledger.pending_agent_requests == ledger.pending_agent_responses == {} + assert host.native.call_entity.call_count == 2 + + +class _SessionAgent(_Agent): + """Exercise the entity's session branch without a network model dependency.""" + + def __init__(self, responses: list[AgentResponse]) -> None: + super().__init__("A", responses, Answer) + self.context_providers: list[Any] = [] + self.default_options["store"] = True + self.sessions: list[tuple[str, Any, dict[str, Any]]] = [] + + async def run( + self, messages: list[Message], *, session: AgentSession | None = None, **kwargs: Any + ) -> AgentResponse: + assert session is not None + self.sessions.append((session.session_id, session.service_session_id, deepcopy(session.state))) + session.service_session_id = "service-conversation" + session.state["application"] = {"pending": False, "turn": len(self.sessions)} + return await super().run(messages, session=session, **kwargs) + + +class _StateProvider(AgentEntityStateProviderMixin): + def __init__(self, raw: dict[str, Any] | None = None) -> None: + self.raw = raw or {} + + def _get_state_dict(self) -> dict[str, Any]: + return json.loads(json.dumps(self.raw)) + + def _set_state_dict(self, state: dict[str, Any]) -> None: + self.raw = json.loads(json.dumps(state, allow_nan=False)) + + def _get_session_id_from_entity(self) -> str: + return "contract-run" + + def _get_entity_name_from_entity(self) -> str: + return "dafx-review-a" + + +@pytest.mark.parametrize("adapter", ["dt", "af"]) +async def test_actual_entity_cold_resume_preserves_service_session_and_fresh_correlation(adapter: str) -> None: + request = _approval("real-request") + agent: Any = _SessionAgent([_pending([request]), _response(value=Answer(inputAnswer=42))]) + a = AgentExecutor(agent, id="A") + workflow = WorkflowBuilder(name="review", start_executor=a).build() + provider = _StateProvider() + host = _Adapter(adapter) + generator = run_workflow_orchestrator(host.context, workflow, "question") + yielded = next(generator) + first_payload = host.payload() + response = await AgentEntity(agent, state_provider=provider).run(first_payload) + yielded = generator.send(host.complete(yielded, _wire(response))) + yielded = generator.send(host.complete(yielded, request.to_function_approval_response(False).to_dict())) + resumed_payload = host.payload() + assert resumed_payload["correlationId"] != first_payload["correlationId"] + cold_provider = _StateProvider(provider.raw) + response = await AgentEntity(agent, state_provider=cold_provider).run(resumed_payload) + output = host.finish(generator, yielded, _wire(response)) + assert len(agent.sessions) == 2 + assert agent.sessions[0][0] == agent.sessions[1][0] + assert agent.sessions[1][1] == "service-conversation" + assert agent.sessions[1][2]["application"] == {"pending": False, "turn": 1} + assert len(agent.inputs[1]) == 1 + assert agent.inputs[1][0].contents == [request.to_function_approval_response(False)] + assert output[0].value == {"answer": 42} + + +@pytest.mark.parametrize("status", ["error", "already_completed"]) +def test_terminal_failure_precedes_structured_parse_and_does_not_emit_output(status: str) -> None: + a = _agent("A", response_format=Answer) + host = _Adapter("dt") + workflow = WorkflowBuilder(name="review", start_executor=a).build() + generator = run_workflow_orchestrator(host.context, workflow, "question") + yielded = next(generator) + payload = _wire(_response("invalid json", additional_properties={"durable_status": status})) + with ( + patch( + "agent_framework_durabletask._workflows.orchestrator.ensure_response_format", + side_effect=AssertionError("must not parse terminal response"), + ), + pytest.raises(RuntimeError, match="expired durable response|terminal runtime error"), + ): + generator.send(host.complete(yielded, payload)) + assert host.native.call_entity.call_count == 1 + + +def test_unconfigured_mock_workflow_does_not_accidentally_designate_every_agent() -> None: + a = _agent("A") + workflow = Mock(spec=Workflow) + workflow.name = "review" + workflow.executors = {"A": a} + workflow.start_executor_id = "A" + workflow.edge_groups = [] + workflow.max_iterations = 5 + host = _Adapter("dt") + generator = run_workflow_orchestrator(host.context, workflow, "question") + assert host.finish(generator, next(generator), _wire(_response())) == [] + + +@pytest.mark.parametrize("adapter", ["dt", "af"]) +def test_two_approval_rounds_rebuild_on_replay_without_duplicate_requests_or_correlations(adapter: str) -> None: + a = _agent("A", response_format=Answer) + workflow = WorkflowBuilder(name="review", start_executor=a).build() + + def run(replay: bool) -> tuple[list[dict[str, Any]], list[AgentResponse]]: + host = _Adapter(adapter, replay=replay) + generator = run_workflow_orchestrator(host.context, workflow, "question") + yielded = next(generator) + wires = [host.payload()] + for request_id in ["first-round", "second-round"]: + request = _approval(request_id) + yielded = generator.send(host.complete(yielded, _wire(_pending([request])))) + assert host.pending[0][2] == (request_id,) + yielded = generator.send(host.complete(yielded, request.to_function_approval_response(False).to_dict())) + wires.append(host.payload()) + outputs = host.finish(generator, yielded, _wire(_response(value=Answer(inputAnswer=42)))) + assert len({wire["correlationId"] for wire in wires}) == 3 + assert len({wire["contextMessageIds"][0] for wire in wires[1:]}) == 2 + if replay: + assert host.statuses == [] + elif adapter == "dt": + events = host.statuses[-1]["events"] + assert [event["request_id"] for event in events if event["type"] == "request_info"] == [ + "first-round", + "second-round", + ] + assert len([event for event in events if event["type"] == "output"]) == 1 + return wires, outputs + + live_wires, live = run(False) + replay_wires, replay = run(True) + # RunRequest's existing created_at default is wall-clock metadata, not an + # orchestration-generated correlation or occurrence identity. + assert [{key: value for key, value in wire.items() if key != "created_at"} for wire in live_wires] == [ + {key: value for key, value in wire.items() if key != "created_at"} for wire in replay_wires + ] + assert len(live) == len(replay) == 1 + assert live[0].value == replay[0].value == {"answer": 42} + assert a._pending_agent_requests == {} and a._pending_responses_to_agent == [] + + +@pytest.mark.parametrize("adapter", ["dt", "af"]) +@pytest.mark.parametrize("response_kind", ["function_approval_response", "function_result"]) +def test_wrong_response_identity_is_rejected_without_losing_real_request(adapter: str, response_kind: str) -> None: + request = _approval("actual-request") + wrong = ( + _approval("unknown-request").to_function_approval_response(True) + if response_kind == "function_approval_response" + else Content.from_function_result("unknown-call", result=False) + ) + a = _agent("A") + workflow = WorkflowBuilder(name="review", start_executor=a).build() + host = _Adapter(adapter) + generator = run_workflow_orchestrator(host.context, workflow, "question") + yielded = generator.send(host.complete(next(generator), _wire(_pending([request])))) + yielded = generator.send(host.complete(yielded, wrong.to_dict())) + assert host.pending[0][2] == ("actual-request",) + assert host.native.call_entity.call_count == 1 + response = Content.from_function_result( + "call-actual-request", result={"type": "untrusted.module:Value", "flag": False} + ) + with patch("importlib.import_module", side_effect=AssertionError("must not import external content types")): + yielded = generator.send(host.complete(yielded, response.to_dict())) + assert host.payload()["contextMessages"] == [Message("tool", [response]).to_dict()] + assert len(host.finish(generator, yielded, _wire(_response()))) == 1 + + +class _TwoInputs(Executor): + def __init__(self) -> None: + super().__init__(id="source") + + @handler + async def send_inputs(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message("first") + await ctx.send_message("second") + + +@pytest.mark.parametrize("adapter", ["dt", "af"]) +def test_sequential_agent_completion_uses_same_output_and_structured_contract(adapter: str) -> None: + source, a = _TwoInputs(), _agent("A", response_format=Answer) + workflow = WorkflowBuilder(name="review", start_executor=source, output_from=[a]).add_edge(source, a).build() + host = _Adapter(adapter) + generator = run_workflow_orchestrator(host.context, workflow, "question") + yielded = next(generator) + result = execute_workflow_activity(source, host.activity_input(), workflow) + yielded = generator.send(host.complete(yielded, result)) + assert host.payload()["message"] == "first" + yielded = generator.send(host.complete(yielded, _wire(_response(value=Answer(inputAnswer=1))))) + assert host.payload()["message"] == "second" + outputs = host.finish(generator, yielded, _wire(_response(value=Answer(inputAnswer=2)))) + assert [output.value for output in outputs] == [{"answer": 1}, {"answer": 2}] + + +@pytest.mark.parametrize("declared", [None, "untrusted.module:Model", {"type": "json_object"}, int]) +def test_only_locally_declared_pydantic_classes_trigger_reconstruction(declared: Any) -> None: + a = _agent("A", response_format=declared) + workflow = WorkflowBuilder(name="review", start_executor=a).build() + host = _Adapter("dt") + generator = run_workflow_orchestrator(host.context, workflow, "question") + yielded = next(generator) + with patch( + "agent_framework_durabletask._workflows.orchestrator.ensure_response_format", + side_effect=AssertionError("must only parse locally declared Pydantic classes"), + ): + outputs = host.finish(generator, yielded, _wire(_response(value={"answer": 42}))) + assert outputs[0].value == {"answer": 42} + + +@pytest.mark.parametrize("agent_first", [False, True]) +def test_agent_activity_request_id_collision_cannot_silently_overwrite_either_request(agent_first: bool) -> None: + def result(task_type: TaskType) -> ExecutorResult: + return ExecutorResult( + executor_id=task_type.value, + output_message=None, + activity_result={"pending_request_info_events": [{"request_id": "collision", "data": task_type.value}]}, + task_type=task_type, + ) + + first, second = (TaskType.AGENT, TaskType.ACTIVITY) if agent_first else (TaskType.ACTIVITY, TaskType.AGENT) + pending: dict[str, Any] = {} + _collect_hitl_requests(result(first), pending) + with pytest.raises(ValueError, match="collides"): + _collect_hitl_requests(result(second), pending) + assert pending["collision"].source_executor_id == first.value + + +@pytest.mark.parametrize("request_id", [None, "", "duplicate"]) +def test_malformed_agent_request_ids_fail_before_registering_partial_batch(request_id: str | None) -> None: + host, ledger, a = _Adapter("dt"), _WorkflowDeliveryLedger(), _agent("A") + metadata = TaskMetadata("A", "question", "start", TaskType.AGENT) + _prepare_agent_task(host.context, a, "A", "question", "review", ledger, metadata) + malformed = Content("function_call", id=request_id, user_input_request=True, call_id="call") + requests = [_approval("duplicate"), malformed] + with pytest.raises(ValueError, match="without an id|duplicate user input request"): + _process_agent_response(_wire(_pending(requests)), "A", "question", ledger, metadata) + assert ledger.pending_agent_requests == ledger.pending_agent_responses == {} + + +@pytest.mark.parametrize("reply", ["clarification", {"type": "function_result", "call_id": "external", "result": None}]) +def test_general_core_content_requests_use_text_coercion_and_preserve_null_results(reply: Any) -> None: + request = Content("function_call", id="input", call_id="external", user_input_request=True) + host, a = _Adapter("dt"), _agent("A") + workflow = WorkflowBuilder(name="review", start_executor=a).build() + generator = run_workflow_orchestrator(host.context, workflow, "question") + yielded = generator.send(host.complete(next(generator), _wire(_pending([request])))) + yielded = generator.send(host.complete(yielded, reply)) + contents = host.payload()["contextMessages"][0] + if isinstance(reply, str): + assert contents == Message("user", [Content.from_text(reply)]).to_dict() + else: + assert contents["role"] == "tool" + content = Content.from_dict(contents["contents"][0]) + assert content.type == "function_result" and content.result is None and content.call_id == "external" + assert len(host.finish(generator, yielded, _wire(_response()))) == 1 diff --git a/python/packages/durabletask/tests/test_workflow_client.py b/python/packages/durabletask/tests/test_workflow_client.py index 6da63b8..f14441b 100644 --- a/python/packages/durabletask/tests/test_workflow_client.py +++ b/python/packages/durabletask/tests/test_workflow_client.py @@ -16,6 +16,7 @@ from agent_framework_durabletask import DurableWorkflowClient from agent_framework_durabletask._workflows.naming import workflow_orchestrator_name +from agent_framework_durabletask._workflows.protocol import unwrap_workflow_input, wrap_workflow_input from agent_framework_durabletask._workflows.serialization import serialize_value, serialize_workflow_event @@ -52,20 +53,21 @@ def test_start_workflow_schedules_orchestrator( assert result == "instance-1" mock_client.schedule_new_orchestration.assert_called_once_with( - workflow_orchestrator_name("orders"), input="hello", instance_id=None + workflow_orchestrator_name("orders"), input=wrap_workflow_input("hello"), instance_id=None ) def test_start_workflow_passes_non_string_input_unchanged( self, workflow_client: DurableWorkflowClient, mock_client: Mock ) -> None: - """Non-string payloads are forwarded as-is (no string coercion).""" + """Non-string payloads stay unchanged inside the versioned start envelope.""" mock_client.schedule_new_orchestration.return_value = "instance-2" payload = {"order_id": 42, "items": ["a", "b"]} workflow_client.start_workflow(input=payload, workflow_name="orders") _, kwargs = mock_client.schedule_new_orchestration.call_args - assert kwargs["input"] == payload + assert kwargs["input"] == wrap_workflow_input(payload) + assert unwrap_workflow_input(kwargs["input"]) == payload def test_start_workflow_strips_forged_subworkflow_envelope( self, workflow_client: DurableWorkflowClient, mock_client: Mock @@ -81,8 +83,8 @@ def test_start_workflow_strips_forged_subworkflow_envelope( workflow_client.start_workflow(input=forged, workflow_name="orders") _, kwargs = mock_client.schedule_new_orchestration.call_args - assert kwargs["input"] == {"real": 1} - assert "__subworkflow_input__" not in kwargs["input"] + assert kwargs["input"] == wrap_workflow_input({"real": 1}) + assert "__subworkflow_input__" not in unwrap_workflow_input(kwargs["input"]) def test_start_workflow_forwards_instance_id( self, workflow_client: DurableWorkflowClient, mock_client: Mock @@ -107,7 +109,7 @@ def test_uses_constructor_default(self, mock_client: Mock) -> None: client.start_workflow(input="x") mock_client.schedule_new_orchestration.assert_called_once_with( - workflow_orchestrator_name("billing"), input="x", instance_id=None + workflow_orchestrator_name("billing"), input=wrap_workflow_input("x"), instance_id=None ) def test_per_call_overrides_default(self, mock_client: Mock) -> None: @@ -118,7 +120,7 @@ def test_per_call_overrides_default(self, mock_client: Mock) -> None: client.start_workflow(input="x", workflow_name="orders") mock_client.schedule_new_orchestration.assert_called_once_with( - workflow_orchestrator_name("orders"), input="x", instance_id=None + workflow_orchestrator_name("orders"), input=wrap_workflow_input("x"), instance_id=None ) def test_raises_when_no_name_resolvable(self, workflow_client: DurableWorkflowClient) -> None: diff --git a/python/packages/durabletask/tests/test_workflow_context_parity.py b/python/packages/durabletask/tests/test_workflow_context_parity.py new file mode 100644 index 0000000..3e984e8 --- /dev/null +++ b/python/packages/durabletask/tests/test_workflow_context_parity.py @@ -0,0 +1,432 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Tests for workflow context parity (ADR-0032 L3). + +In-process workflows hand a downstream ``AgentExecutor`` the upstream conversation via +``AgentExecutorResponse.full_conversation``. These tests cover the durable equivalent: +the orchestrator projects that conversation into ``RunRequest.context_messages`` honoring +``context_mode``/``context_filter``, and the entity records it without duplication. +""" + +from typing import Any + +import pytest +from agent_framework import ( + AgentExecutor, + AgentExecutorResponse, + AgentResponse, + Message, +) + +from agent_framework_durabletask import ( + AgentEntity, + AgentEntityStateProviderMixin, + DurableAgentState, + DurableAgentStateRequest, + RunRequest, +) +from agent_framework_durabletask._message_identity import message_identity +from agent_framework_durabletask._workflows.orchestrator import ( + _build_context_messages, + build_agent_executor_response, +) + + +class _StubAgent: + """Minimal agent stand-in for constructing an AgentExecutor.""" + + def __init__(self, name: str = "stub") -> None: + self.name = name + self.id = name + self.description = None + + async def run(self, messages: Any = None, **kwargs: Any) -> AgentResponse: + return AgentResponse(messages=[Message(role="assistant", contents=["ok"])]) + + def create_session(self, **kwargs: Any) -> Any: + from agent_framework import AgentSession + + return AgentSession() + + +class _InMemoryStateProvider(AgentEntityStateProviderMixin): + def __init__(self, *, session_id: str = "wf-session", entity_name: str = "") -> None: + self._session_id = session_id + self._entity_name = entity_name + self._state_dict: dict[str, Any] = {} + + def _get_state_dict(self) -> dict[str, Any]: + return self._state_dict + + def _set_state_dict(self, state: dict[str, Any]) -> None: + self._state_dict = state + + def _get_session_id_from_entity(self) -> str: + return self._session_id + + def _get_entity_name_from_entity(self) -> str: + return self._entity_name + + +def _upstream_response(*, texts: list[str], agent_text: str) -> AgentExecutorResponse: + conversation = [Message(role="user", contents=[t], message_id=f"m{i}") for i, t in enumerate(texts)] + agent_message = Message(role="assistant", contents=[agent_text], message_id="agent-msg") + conversation.append(agent_message) + return AgentExecutorResponse( + executor_id="upstream", + agent_response=AgentResponse(messages=[agent_message]), + full_conversation=conversation, + ) + + +def _stub_agent() -> Any: + """Return the stub agent typed loosely. + + It implements the parts of the agent protocol these tests exercise but not its full signature, + so the type is relaxed here rather than at every call site. + """ + return _StubAgent() + + +class TestContextProjection: + """The orchestrator projects upstream conversation per context_mode.""" + + def test_full_mode_forwards_entire_conversation(self) -> None: + executor = AgentExecutor(_stub_agent(), id="downstream") + upstream = _upstream_response(texts=["first", "second"], agent_text="reply") + + projected = _build_context_messages(executor, upstream) + + assert projected is not None + assert len(projected) == 3 + + def test_last_agent_mode_forwards_only_agent_messages(self) -> None: + executor = AgentExecutor(_stub_agent(), id="downstream", context_mode="last_agent") + upstream = _upstream_response(texts=["first", "second"], agent_text="reply") + + projected = _build_context_messages(executor, upstream) + + assert projected is not None + assert len(projected) == 1 + + def test_custom_mode_uses_context_filter(self) -> None: + executor = AgentExecutor( + _stub_agent(), + id="downstream", + context_mode="custom", + context_filter=lambda messages: messages[-2:], + ) + upstream = _upstream_response(texts=["first", "second"], agent_text="reply") + + projected = _build_context_messages(executor, upstream) + + assert projected is not None + assert len(projected) == 2 + + def test_non_agent_input_has_no_upstream_context(self) -> None: + """The first node receives raw input, so there is no conversation to forward.""" + executor = AgentExecutor(_stub_agent(), id="downstream") + + assert _build_context_messages(executor, "plain input") is None + + +class TestEntityContextIngestion: + """The entity records forwarded context and does not duplicate it.""" + + def _request(self, messages: list[Message], correlation_id: str) -> RunRequest: + return RunRequest( + message=messages[-1].text or "", + correlation_id=correlation_id, + context_messages=[m.to_dict() for m in messages], + ) + + def test_context_messages_become_request_messages(self) -> None: + messages = [ + Message(role="user", contents=["hello"], message_id="m0"), + Message(role="assistant", contents=["hi"], message_id="m1"), + ] + + entry = DurableAgentStateRequest.from_run_request(self._request(messages, "corr-0")) + + assert [m.message_id for m in entry.messages] == ["m0", "m1"] + + def test_repeated_context_is_not_duplicated(self) -> None: + """A node that runs twice in a cycle must not re-record the same conversation.""" + provider = _InMemoryStateProvider() + entity = AgentEntity(_stub_agent(), state_provider=provider) + + first = [Message(role="user", contents=["hello"], message_id="m0")] + initial = DurableAgentStateRequest.from_run_request(self._request(first, "corr-0")) + initial.messages = entity._drop_already_stored(initial.messages) + entity.state.data.conversation_history.append(initial) + assert entity.state.data.ingested_messages == {"m0": [message_identity(first[0])]} + + repeated = [ + Message(role="user", contents=["hello"], message_id="m0"), + Message(role="assistant", contents=["new"], message_id="m1"), + ] + entry = DurableAgentStateRequest.from_run_request(self._request(repeated, "corr-1")) + entry.messages = entity._drop_already_stored(entry.messages) + + assert [m.message_id for m in entry.messages] == ["m1"] + + def test_fully_duplicate_context_stays_empty(self) -> None: + """A repeated projection must not re-ingest its final message as a new input.""" + provider = _InMemoryStateProvider() + entity = AgentEntity(_stub_agent(), state_provider=provider) + + messages = [Message(role="user", contents=["hello"], message_id="m0")] + initial = DurableAgentStateRequest.from_run_request(self._request(messages, "corr-0")) + initial.messages = entity._drop_already_stored(initial.messages) + entity.state.data.conversation_history.append(initial) + assert entity.state.data.ingested_messages == {"m0": [message_identity(messages[0])]} + + entry = DurableAgentStateRequest.from_run_request(self._request(messages, "corr-1")) + entry.messages = entity._drop_already_stored(entry.messages) + + assert entry.messages == [] + + @pytest.mark.parametrize("message_id", ["m0", "wf_upstream_3"]) + @pytest.mark.parametrize("transcript_contents", [["hello"], []], ids=["retained", "pruned"]) + def test_transcript_without_receipt_does_not_suppress_first_delivery( + self, message_id: str, transcript_contents: list[str] + ) -> None: + """A retained transcript alone is not evidence that an input was delivered.""" + entity = AgentEntity(_stub_agent(), state_provider=_InMemoryStateProvider()) + transcript = [Message(role="user", contents=transcript_contents, message_id=message_id)] + entity.state.data.conversation_history.append( + DurableAgentStateRequest.from_run_request(self._request(transcript, "legacy")) + ) + assert entity.state.data.ingested_messages == {} + + incoming = [Message(role="user", contents=["hello"], message_id=message_id)] + entry = DurableAgentStateRequest.from_run_request(self._request(incoming, "first-delivery")) + entry.messages = entity._drop_already_stored(entry.messages) + + assert [message.to_chat_message().to_dict() for message in entry.messages] == [incoming[0].to_dict()] + assert entity.state.data.ingested_messages == {message_id: [message_identity(incoming[0])]} + repeated = DurableAgentStateRequest.from_run_request(self._request(incoming, "repeat-delivery")) + assert entity._drop_already_stored(repeated.messages) == [] + + def test_repeated_context_does_not_duplicate_message_ids(self) -> None: + """A cycle that re-delivers the whole upstream conversation must not collide ids.""" + provider = _InMemoryStateProvider() + entity = AgentEntity(_stub_agent(), state_provider=provider) + + messages = [Message(role="user", contents=["hello"], message_id="m0")] + for index in range(3): + entry = DurableAgentStateRequest.from_run_request(self._request(messages, f"corr-{index}")) + entry.messages = entity._drop_already_stored(entry.messages) + entity.state.data.conversation_history.append(entry) + + stored_ids = [ + m.message_id for entry in entity.state.data.conversation_history for m in entry.messages if m.message_id + ] + assert len(stored_ids) == len(set(stored_ids)), f"duplicate message ids persisted: {stored_ids}" + + @pytest.mark.parametrize("application_id", [None, "opaque", "wf_source_0"]) + def test_occurrence_ids_keep_equal_new_events_and_drop_only_repeat_delivery( + self, application_id: str | None + ) -> None: + entity = AgentEntity(_stub_agent(), state_provider=_InMemoryStateProvider()) + original = Message("assistant", ["same"], message_id=application_id) + before = original.to_dict() + occurrence_ids = ["occurrence-first", "occurrence-second"] + request = RunRequest( + message="same", + correlation_id="first", + context_messages=[before, before], + context_message_ids=occurrence_ids, + ) + restored = RunRequest.from_dict(request.to_dict()) + entry = DurableAgentStateRequest.from_run_request(restored) + entry.messages = entity._drop_already_stored(entry.messages, occurrence_ids=restored.context_message_ids) + entity.state.data.conversation_history.append(entry) + assert [message.message_id for message in entry.messages] == [application_id, application_id] + assert [message.ingestion_occurrence for message in entry.messages] == occurrence_ids + assert [message.to_chat_message().to_dict() for message in entry.messages] == [before, before] + + repeated = DurableAgentStateRequest.from_run_request(restored) + assert entity._drop_already_stored(repeated.messages, occurrence_ids=restored.context_message_ids) == [] + new_request = RunRequest( + message="same", correlation_id="new", context_messages=[before], context_message_ids=["occurrence-third"] + ) + new_entry = DurableAgentStateRequest.from_run_request(new_request) + kept = entity._drop_already_stored(new_entry.messages, occurrence_ids=new_request.context_message_ids) + assert [message.message_id for message in kept] == [application_id] + assert entity.state.data.ingested_messages == { + identity: [message_identity(original)] for identity in [*occurrence_ids, "occurrence-third"] + } + assert original.to_dict() == before + + +class TestWorkflowConversationIdentity: + """The legacy text helper still assigns IDs to messages it creates. + + Production completions preserve application messages and use separate occurrence IDs. + These compatibility checks exercise only the helper and the legacy receiver fallback. + """ + + def _cycle_ids(self) -> list[str]: + conversation: Any = "start" + for node in ["A", "B", "A", "B"]: + conversation = build_agent_executor_response(node, f"{node} says", None, conversation) + return [m.message_id or "" for m in conversation.full_conversation] + + def test_every_built_message_carries_an_id(self) -> None: + response = build_agent_executor_response("writer", "drafted", None, "start") + + ids = [m.message_id for m in response.full_conversation] + assert all(ids), f"a message went out without an id: {ids}" + + def test_ids_stay_unique_around_a_cycle(self) -> None: + ids = self._cycle_ids() + + assert all(ids), f"a message went out without an id: {ids}" + assert len(ids) == len(set(ids)), f"ids collided around the cycle: {ids}" + + def test_ids_are_replay_stable(self) -> None: + """The orchestrator rebuilds this conversation on replay, so the ids must not move.""" + assert self._cycle_ids() == self._cycle_ids() + + def test_a_revisited_node_records_only_what_is_new(self) -> None: + provider = _InMemoryStateProvider() + entity = AgentEntity(_stub_agent(), state_provider=provider) + + def _deliver_to_a(context: list[Message], correlation_id: str) -> int: + request = RunRequest( + message=context[-1].text or "", + correlation_id=correlation_id, + context_messages=[m.to_dict() for m in context], + ) + entry = DurableAgentStateRequest.from_run_request(request) + entry.messages = entity._drop_already_stored(entry.messages) + entity.state.data.conversation_history.append(entry) + return len(entry.messages) + + conversation: Any = "start" + conversation = build_agent_executor_response("A", "a1", None, conversation) + conversation = build_agent_executor_response("B", "b1", None, conversation) + first = _deliver_to_a(list(conversation.full_conversation), "corr-1") + + conversation = build_agent_executor_response("A", "a2", None, conversation) + conversation = build_agent_executor_response("B", "b2", None, conversation) + second = _deliver_to_a(list(conversation.full_conversation), "corr-2") + + assert first == 3, f"expected the first delivery to be recorded whole, got {first}" + assert second == 2, f"expected only the two new messages, got {second} of 5 delivered" + + +class TestDedupSurvivesRetention: + """Duplicate detection must not depend on the messages still being there. + + Retention deletes oldest-first, which removes exactly the ids an identity check relies on. The + orchestrator's own conversation is never evicted, so it re-sends them, and an entity comparing + against stored ids would re-record precisely what was just deleted. + """ + + def _deliver(self, entity: AgentEntity, context: list[Message], correlation_id: str) -> int: + request = RunRequest( + message=context[-1].text or "", + correlation_id=correlation_id, + context_messages=[m.to_dict() for m in context], + ) + entry = DurableAgentStateRequest.from_run_request(request) + entry.messages = entity._drop_already_stored(entry.messages) + entity.state.data.conversation_history.append(entry) + return len(entry.messages) + + def test_evicted_context_is_not_re_ingested(self) -> None: + provider = _InMemoryStateProvider() + entity = AgentEntity(_stub_agent(), state_provider=provider) + + conversation: Any = "start" + conversation = build_agent_executor_response("A", "a1", None, conversation) + conversation = build_agent_executor_response("B", "b1", None, conversation) + self._deliver(entity, list(conversation.full_conversation), "corr-1") + + # Retention deletes the oldest messages, taking their ids with them. + entity.state.data.conversation_history.clear() + + conversation = build_agent_executor_response("A", "a2", None, conversation) + conversation = build_agent_executor_response("B", "b2", None, conversation) + recorded = self._deliver(entity, list(conversation.full_conversation), "corr-2") + + assert recorded == 2, f"expected only the two new messages after eviction, got {recorded} of 5" + + def test_the_mark_is_kept_per_executor(self) -> None: + """A fan-out gives two branches the same position, so one global mark would conflate them.""" + provider = _InMemoryStateProvider() + entity = AgentEntity(_stub_agent(), state_provider=provider) + + conversation: Any = "start" + conversation = build_agent_executor_response("A", "a1", None, conversation) + conversation = build_agent_executor_response("B", "b1", None, conversation) + self._deliver(entity, list(conversation.full_conversation), "corr-1") + + receipts = entity.state.data.ingested_messages + assert set(receipts) == {"wf_input_0", "wf_A_1", "wf_B_2"} + assert all(values and len(values) == 1 for values in receipts.values()) + + def test_the_mark_round_trips_through_durable_state(self) -> None: + provider = _InMemoryStateProvider() + entity = AgentEntity(_stub_agent(), state_provider=provider) + + conversation: Any = build_agent_executor_response("A", "a1", None, "start") + self._deliver(entity, list(conversation.full_conversation), "corr-1") + entity.persist_state() + + restored = DurableAgentState.from_dict(provider._get_state_dict()) + assert restored.data.ingested_messages == entity.state.data.ingested_messages + assert restored.data.ingested_messages + + +class TestCoreSessionIdentity: + """The id handed to core must identify one entity, not one workflow run.""" + + def test_workflow_nodes_do_not_share_a_core_session_id(self) -> None: + """Nodes of one workflow share the entity key and differ only by entity name. + + An external history provider keys its storage on the core session id, so taking the key + alone would file every node's conversation under one entry. + """ + writer = _InMemoryStateProvider(session_id="run-1", entity_name="dafx-writer") + reviewer = _InMemoryStateProvider(session_id="run-1", entity_name="dafx-reviewer") + + assert writer.session_id == reviewer.session_id + assert writer.core_session_id != reviewer.core_session_id, f"both nodes resolved to {writer.core_session_id}" + + def test_core_session_id_falls_back_to_the_key(self) -> None: + """State providers predating the entity-name hook keep working.""" + assert _InMemoryStateProvider(session_id="solo").core_session_id == "solo" + + +class TestRunRequestRoundTrip: + """Context messages and their separate occurrence IDs survive the entity wire format.""" + + def test_context_messages_round_trip(self) -> None: + messages = [Message(role="user", contents=["hello"], message_id="m0")] + request = RunRequest( + message="hello", + correlation_id="corr-0", + context_messages=[m.to_dict() for m in messages], + context_message_ids=["occurrence-0"], + ) + + restored = RunRequest.from_dict(request.to_dict()) + + assert restored.context_messages is not None + assert len(restored.context_messages) == 1 + assert restored.context_messages[0]["message_id"] == "m0" + assert restored.context_message_ids == ["occurrence-0"] + assert request.to_dict()["contextMessageIds"] == ["occurrence-0"] + + def test_absent_context_messages_stay_none(self) -> None: + request = RunRequest(message="hello", correlation_id="corr-0") + + restored = RunRequest.from_dict(request.to_dict()) + + assert restored.context_messages is None + assert restored.context_message_ids is None + assert "contextMessages" not in request.to_dict() + assert "contextMessageIds" not in request.to_dict() diff --git a/python/packages/durabletask/tests/test_workflow_deltas.py b/python/packages/durabletask/tests/test_workflow_deltas.py new file mode 100644 index 0000000..8095cb7 --- /dev/null +++ b/python/packages/durabletask/tests/test_workflow_deltas.py @@ -0,0 +1,1241 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Source-side workflow deltas, driven through projection and generator dispatch.""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Generator +from copy import deepcopy +from datetime import datetime, timezone +from typing import Any +from unittest.mock import Mock + +import pytest +from agent_framework import ( + AgentExecutor, + AgentExecutorResponse, + AgentResponse, + AgentSession, + Content, + Executor, + Message, + Workflow, +) +from agent_framework._workflows._edge import EdgeGroup, FanInEdgeGroup, FanOutEdgeGroup, SingleEdgeGroup + +from agent_framework_durabletask._message_identity import message_identity +from agent_framework_durabletask._workflows.orchestrator import ( + _AGENT_TASK_MESSAGE_PREVIEW_LIMIT, + _build_context_messages, + _prepare_agent_task, + _WorkflowDeliveryLedger, + build_agent_executor_response, + run_workflow_orchestrator, +) +from agent_framework_durabletask._workflows.serialization import deserialize_value, serialize_value + + +class _StubAgent: + name = "stub" + id = "stub" + description = None + + def create_session(self, **kwargs: Any) -> AgentSession: + return AgentSession(**kwargs) + + async def run(self, messages: Any = None, **kwargs: Any) -> AgentResponse: + raise AssertionError("The recording host must not invoke a model") + + +def _agent(executor_id: str = "target", **kwargs: Any) -> AgentExecutor: + agent: Any = _StubAgent() + return AgentExecutor(agent, id=executor_id, **kwargs) + + +def _message(position: int, producer: str = "source", text: str | None = None) -> Message: + return Message( + "assistant", [text if text is not None else f"{producer}-{position}"], message_id=f"wf_{producer}_{position}" + ) + + +def _response( + messages: list[Message], producer: str = "source", *, latest: list[Message] | None = None +) -> AgentExecutorResponse: + return AgentExecutorResponse( + executor_id=producer, + agent_response=AgentResponse(messages=messages[-1:] if latest is None else latest), + full_conversation=list(messages), + ) + + +def _ids(call: dict[str, Any]) -> list[str | None]: + """Read application IDs, which are independent of delivery occurrences.""" + assert call["contextMessages"] is not None + return [message.get("message_id") for message in call["contextMessages"]] + + +def _occurrences(call: dict[str, Any]) -> list[str]: + ids = call["contextMessageIds"] + assert isinstance(ids, list) + assert len(ids) == len(call["contextMessages"]) + assert all(isinstance(value, str) and value.startswith("wf:occurrence:") for value in ids) + return ids + + +def _texts(call: dict[str, Any]) -> list[str]: + assert call["contextMessages"] is not None + return [Message.from_dict(message).text for message in call["contextMessages"]] + + +class _RecordingHost: + """Return recorded task outcomes while capturing the adapter-boundary payloads.""" + + supports_event_streaming = False + current_utc_datetime = datetime(2026, 1, 1, tzinfo=timezone.utc) + + def __init__( + self, + *, + instance_id: str = "run", + is_replaying: bool = False, + activities: dict[str, list[dict[str, Any]]] | None = None, + agent_reply: str | None = None, + ) -> None: + self.instance_id = instance_id + self.is_replaying = is_replaying + self.calls: list[dict[str, Any]] = [] + self.activity_inputs: list[dict[str, Any]] = [] + self.waited_for: list[str] = [] + self.batch_sizes: list[int] = [] + self.statuses: list[Any] = [] + self.fail_prepare = False + self._activities = {name: iter(results) for name, results in (activities or {}).items()} + self._agent_reply = agent_reply + + def prepare_agent_task( + self, + executor_id: str, + message: str, + orchestration_instance_id: str, + context_messages: list[dict[str, Any]] | None = None, + context_message_ids: list[str] | None = None, + ) -> AgentResponse: + assert (context_messages is None) == (context_message_ids is None) + if context_messages is not None: + assert context_message_ids is not None + assert len(context_messages) == len(context_message_ids) + # JSON round-trip the complete adapter arguments, not just a count of messages. + self.calls.append( + json.loads( + json.dumps( + { + "executorId": executor_id, + "message": message, + "instanceId": orchestration_instance_id, + "contextMessages": context_messages, + "contextMessageIds": context_message_ids, + }, + allow_nan=False, + ) + ) + ) + if self.fail_prepare: + raise OSError("injected preparation failure") + reply = self._agent_reply if self._agent_reply is not None else f"reply-{len(self.calls)}" + return AgentResponse(messages=[Message("assistant", [reply])]) + + def prepare_activity_task(self, activity_name: str, input_json: str) -> str: + payload = json.loads(input_json) + self.activity_inputs.append(payload) + return json.dumps(next(self._activities[payload["executor_id"]])) + + def call_sub_orchestrator(self, name: str, input: Any, instance_id: str | None = None) -> Any: + raise AssertionError("These workflows have no child orchestrations") + + def task_all(self, tasks: list[Any]) -> list[Any]: + self.batch_sizes.append(len(tasks)) + return tasks + + def task_any(self, tasks: list[Any]) -> Any: + raise AssertionError("These workflows do not race tasks") + + def wait_for_external_event(self, name: str) -> str: + self.waited_for.append(name) + return "approved" + + def create_timer(self, fire_at: datetime) -> Any: + raise AssertionError("These workflows have no timers") + + def set_custom_status(self, status: Any) -> None: + self.statuses.append(deepcopy(status)) + + def new_uuid(self) -> str: + raise AssertionError("Message identity must not require UUIDs") + + def cancel_task(self, task: Any) -> None: + raise AssertionError("These workflows do not cancel tasks") + + def get_task_result(self, task: Any) -> Any: + return task + + +def _dispatch( + host: _RecordingHost, executor: AgentExecutor, message: Any, ledger: _WorkflowDeliveryLedger +) -> dict[str, Any]: + _prepare_agent_task(host, executor, executor.id, message, "delta", ledger) + return host.calls[-1] + + +def _workflow(nodes: list[Any], edges: list[EdgeGroup], *, max_iterations: int = 20) -> Any: + # The graph container is passive here. Real executors and edge groups exercise + # the orchestrator's production classification, routing and task grouping. + workflow = Mock(spec=Workflow) + workflow.name = "delta" + workflow.start_executor_id = nodes[0].id + workflow.executors = {node.id: node for node in nodes} + workflow.edge_groups = edges + workflow.max_iterations = max_iterations + return workflow + + +def _activity(executor_id: str) -> Mock: + executor = Mock(spec=Executor) + executor.id = executor_id + executor.input_types = [str] + return executor + + +def _activity_result(messages: list[Any], target: str | None = "target", *, request: bool = False) -> dict[str, Any]: + result: dict[str, Any] = { + "sent_messages": [{"message": serialize_value(message), "target_id": target} for message in messages] + } + if request: + result["pending_request_info_events"] = [ + {"request_id": "approval", "source_executor_id": "gate", "data": "review"} + ] + return result + + +def _finish(orchestration: Generator[Any, Any, Any], yielded: Any) -> Any: + while True: + try: + yielded = orchestration.send(yielded) + except StopIteration as completed: + return completed.value + + +def _run(host: _RecordingHost, workflow: Any) -> Any: + orchestration = run_workflow_orchestrator(host, workflow, "start") + return _finish(orchestration, next(orchestration)) + + +def test_message_identity_uses_canonical_full_message_json() -> None: + original = Message( + "assistant", + [{"type": "function_call", "call_id": "call", "name": "lookup", "arguments": {"b": 2, "a": 1}}], + message_id="custom-id", + author_name="author", + additional_properties={"nested": {"z": "世界", "a": 1}}, + raw_representation=object(), + ) + reordered = Message( + "assistant", + [{"arguments": {"a": 1, "b": 2}, "name": "lookup", "call_id": "call", "type": "function_call"}], + message_id="custom-id", + author_name="author", + additional_properties={"nested": {"a": 1, "z": "世界"}}, + ) + canonical = json.dumps( + original.to_dict(), sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False + ) + expected = hashlib.sha256(canonical.encode("utf-8")).hexdigest() + + assert message_identity(original) == expected == message_identity(reordered) + assert message_identity(Message.from_dict(json.loads(original.to_json()))) == expected + assert original.message_id == "custom-id" + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("message_id", "other-id"), + ("role", "system"), + ("author_name", "other-author"), + ("contents", [{"type": "text", "text": "changed"}]), + ("contents", [{"type": "text", "text": "second"}, {"type": "text", "text": "first"}]), + ("additional_properties", {"_is_summary": True}), + ], +) +def test_message_identity_detects_meaningful_changes(field: str, value: Any) -> None: + original = Message("assistant", ["first", "second"], message_id="custom-id", author_name="author") + modified = original.to_dict() + modified[field] = value + + assert message_identity(original) != message_identity(Message.from_dict(modified)) + + +@pytest.mark.parametrize("mode", ["full", "last_agent", "custom"]) +def test_empty_projection_is_explicit_and_does_not_leak_raw_response(mode: str) -> None: + secret = Message("assistant", ["unselected secret " * 10_000], message_id="secret") + upstream = _response([] if mode == "full" else [_message(1)], latest=[] if mode == "last_agent" else [secret]) + executor = _agent(context_mode=mode, context_filter=(lambda messages: []) if mode == "custom" else None) + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + + assert _build_context_messages(executor, upstream) == [] + call = _dispatch(host, executor, upstream, ledger) + + assert call["contextMessages"] == [] + assert call["message"] == "" + assert "secret" not in json.dumps(call) + assert ledger.sent == {} + + +def test_projection_remains_stateless_and_does_not_stamp_filter_input() -> None: + original = Message("user", ["anonymous"]) + upstream = _response([original]) + executor = _agent(context_mode="custom", context_filter=lambda messages: [m for m in messages if not m.message_id]) + expected = [original.to_dict()] + + assert _build_context_messages(executor, upstream) == expected + call = _dispatch(_RecordingHost(), executor, upstream, _WorkflowDeliveryLedger()) + assert _ids(call) == [None] + assert len(_occurrences(call)) == 1 + assert _build_context_messages(executor, upstream) == expected + assert original.message_id is None + + +def test_no_context_raw_requests_are_not_deduplicated_or_truncated() -> None: + executor = _agent() + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + prompt = "a new request " * 1000 + + assert _build_context_messages(executor, prompt) is None + for _ in range(2): + call = _dispatch(host, executor, prompt, ledger) + assert call["contextMessages"] is None + assert call["message"] == prompt + assert ledger.sent == {} + + +def test_last_agent_delta_preserves_all_selected_assistant_and_tool_messages() -> None: + executor = _agent(context_mode="last_agent") + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + first = Message("assistant", ["answer"], message_id="wf_source_5") + second = Message( + "tool", [{"type": "function_result", "call_id": "call", "result": "result"}], message_id="wf_source_6" + ) + upstream = _response([_message(0), first, second], latest=[first, second]) + + call = _dispatch(host, executor, upstream, ledger) + assert call["contextMessages"] == [first.to_dict(), second.to_dict()] + assert call["message"] == "" + assert _ids(_dispatch(host, executor, upstream, ledger)) == [] + assert ledger.sent == { + "target": set(zip(_occurrences(call), [message_identity(first), message_identity(second)], strict=True)) + } + + +def test_missing_custom_filter_fails_instead_of_forwarding_unfiltered_input() -> None: + executor = _agent(context_mode="custom", context_filter=lambda messages: messages) + executor._context_filter = None + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + + with pytest.raises(ValueError, match="context_filter"): + _dispatch(host, executor, _response([_message(1)]), ledger) + assert host.calls == [] + assert ledger == _WorkflowDeliveryLedger() + + +def test_empty_selection_does_not_mark_unselected_positions_delivered() -> None: + executor = _agent(context_mode="custom", context_filter=lambda messages: [] if len(messages) == 1 else messages[:1]) + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + + assert _ids(_dispatch(host, executor, _response([_message(1)]), ledger)) == [] + assert _ids(_dispatch(host, executor, _response([_message(1), _message(2)]), ledger)) == ["wf_source_1"] + + +def test_sparse_custom_selection_delivers_previously_skipped_lower_positions() -> None: + positions = [0, 2] + executor = _agent(context_mode="custom", context_filter=lambda messages: [messages[i] for i in positions]) + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + upstream = _response([_message(i) for i in [1, 2, 3, 4]]) + + first = _dispatch(host, executor, upstream, ledger) + positions[:] = [1, 3] + second = _dispatch(host, executor, upstream, ledger) + repeated = _dispatch(host, executor, upstream, ledger) + + assert _ids(first) == ["wf_source_1", "wf_source_3"] + assert _ids(second) == ["wf_source_2", "wf_source_4"] + assert _ids(repeated) == [] + assert repeated["message"] == "" + assert len(set(_occurrences(first) + _occurrences(second))) == 4 + assert ledger.sent == { + "target": { + (occurrence, message_identity(Message.from_dict(message))) + for call in [first, second] + for occurrence, message in zip(_occurrences(call), call["contextMessages"], strict=True) + } + } + + +def test_reordered_projection_preserves_new_message_order_without_a_cursor() -> None: + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + positions = [2, 3] + executor = _agent(context_mode="custom", context_filter=lambda messages: [messages[i] for i in positions]) + upstream = _response([_message(i) for i in [4, 2, 3, 1]]) + + first = _dispatch(host, executor, upstream, ledger) + positions[:] = [0, 1, 2, 3] + call = _dispatch(host, executor, upstream, ledger) + + assert _ids(call) == ["wf_source_4", "wf_source_2"] + assert set(_occurrences(first)).isdisjoint(_occurrences(call)) + assert call["message"] == "source-2" + + +def test_fanout_delivery_is_independent_for_each_target() -> None: + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + left, right = _agent("left"), _agent("right") + first = _response([_message(1), _message(3)]) + next_projection = _response([_message(2), _message(4), first.full_conversation[0]], latest=[]) + + left_first = _dispatch(host, left, first, ledger) + right_next = _dispatch(host, right, next_projection, ledger) + left_next = _dispatch(host, left, next_projection, ledger) + right_first = _dispatch(host, right, first, ledger) + assert _ids(left_first) == ["wf_source_1", "wf_source_3"] + assert _ids(right_next) == ["wf_source_2", "wf_source_4", "wf_source_1"] + assert _ids(left_next) == ["wf_source_2", "wf_source_4"] + assert _ids(right_first) == ["wf_source_3"] + assert _occurrences(left_first) == [_occurrences(right_next)[-1], *_occurrences(right_first)] + assert _occurrences(left_next) == _occurrences(right_next)[:2] + + +def test_fanin_tracks_each_messages_producer_not_the_immediate_sender() -> None: + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + executor = _agent() + common = _message(0, "input") + first = [ + _response([common, _message(100, "A")], "relay"), + _response([common, _message(1, "B")], "relay"), + ] + second = [ + _response([common, _message(99, "A"), first[0].full_conversation[-1]], "other-relay", latest=[]), + _response([common, _message(0, "B"), first[1].full_conversation[-1]], "other-relay", latest=[]), + ] + + projected = [m.to_dict() for response in first for m in response.full_conversation] + assert _build_context_messages(executor, first) == projected + assert _ids(_dispatch(host, executor, first, ledger)) == ["wf_input_0", "wf_A_100", "wf_B_1"] + assert _ids(_dispatch(host, executor, second, ledger)) == ["wf_A_99", "wf_B_0"] + + +@pytest.mark.parametrize( + "message_id", + [ + "wf_source_7", + "wf:external:" + "a" * 64, + "wf:projection:" + "b" * 64, + "custom-id", + "wf_not_a_position", + "wf:external:not-a-hash", + "wf:projection:not-a-hash", + ], +) +def test_same_id_content_changes_are_delivered_and_exact_repeats_are_not(message_id: str) -> None: + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + executor = _agent() + original = Message("assistant", ["old"], message_id=message_id) + changed = Message("assistant", ["new"], message_id=message_id) + + source = _response([original]) + first = _dispatch(host, executor, source, ledger) + assert _texts(first) == ["old"] + copied = _agent(context_mode="custom", context_filter=lambda messages: deepcopy(messages)) + assert _occurrences(_dispatch(host, copied, source, ledger)) == [] + redacted = _agent(context_mode="custom", context_filter=lambda messages: [deepcopy(changed)]) + call = _dispatch(host, redacted, source, ledger) + assert _texts(call) == ["new"] + assert _ids(call) == [message_id] + assert _occurrences(call) == _occurrences(first) + assert _occurrences(_dispatch(host, redacted, source, ledger)) == [] + assert _occurrences(_dispatch(host, executor, source, ledger)) == [] + independent = _dispatch(host, executor, _response([deepcopy(original)]), ledger) + assert _ids(independent) == [message_id] + assert set(_occurrences(first)).isdisjoint(_occurrences(independent)) + assert original.message_id == changed.message_id == message_id + + +def test_nontext_updates_and_repeated_ids_with_different_contents_are_not_lost() -> None: + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + executor = _agent() + original = Message( + "tool", [{"type": "function_result", "call_id": "call", "result": {"answer": 1}}], message_id="m" + ) + changed = Message("tool", [{"type": "function_result", "call_id": "call", "result": {"answer": 2}}], message_id="m") + source = _response([original, deepcopy(original), changed]) + call = _dispatch(host, executor, source, ledger) + + assert call["contextMessages"] == [message.to_dict() for message in source.full_conversation] + assert _ids(call) == ["m"] * 3 + assert len(set(_occurrences(call))) == 3 + assert call["message"] == "" + assert _occurrences(_dispatch(host, executor, source, ledger)) == [] + assert original.message_id == changed.message_id == "m" + + +def test_distinct_custom_ids_do_not_globally_deduplicate_equal_text() -> None: + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + executor = _agent() + + occurrences: list[str] = [] + for message_id in ["first-request", "second-request"]: + call = _dispatch(host, executor, _response([Message("user", ["again"], message_id=message_id)]), ledger) + assert _ids(call) == [message_id] + occurrences.extend(_occurrences(call)) + assert len(set(occurrences)) == 2 + + +@pytest.mark.parametrize("mode", ["full", "last_agent", "custom"]) +@pytest.mark.parametrize("batch", [False, True]) +def test_equal_custom_ids_from_different_producers_have_distinct_transport_identities(mode: str, batch: bool) -> None: + original = Message( + "assistant", + ["approved"], + message_id="custom-id", + author_name="reviewer", + additional_properties={"nested": {"decision": "approved"}}, + ) + before = original.to_dict() + sources = [_response([deepcopy(original)], producer) for producer in ["left", "right"]] + executor = _agent( + context_mode=mode, + context_filter=(lambda messages: [m for m in messages if m.message_id == "custom-id"]) + if mode == "custom" + else None, + ) + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + deliveries: list[Any] = [sources] if batch else sources + calls = [_dispatch(host, executor, message, ledger) for message in deliveries] + transported = [message for call in calls for message in call["contextMessages"]] + + assert transported == [before, before] + # Equal application payloads still represent two independently produced events. + assert len({identity for call in calls for identity in _occurrences(call)}) == 2 + assert len({message_identity(Message.from_dict(message)) for message in transported}) == 1 + assert _ids(_dispatch(host, executor, list(reversed(sources)), ledger)) == [] + assert [source.full_conversation[0].to_dict() for source in sources] == [before, before] + assert _build_context_messages(executor, sources) == [before, before] + + +def test_custom_id_scopes_use_unambiguous_producer_and_id_addresses() -> None: + sources = [ + _response([Message("assistant", ["approved"], message_id=message_id)], producer) + for producer, message_id in [("left_part", "id"), ("left", "part_id")] + ] + call = _dispatch(_RecordingHost(), _agent(), sources, _WorkflowDeliveryLedger()) + + assert _ids(call) == ["id", "part_id"] + assert len(set(_occurrences(call))) == 2 + + +def test_mixed_custom_and_anonymous_messages_keep_each_producers_identity() -> None: + custom = Message("assistant", ["approved"], message_id="custom-id") + anonymous = Message("user", ["same"]) + sources = [_response(deepcopy([custom, anonymous]), producer) for producer in ["left", "right"]] + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + executor = _agent() + + first = _dispatch(host, executor, sources, ledger) + assert _ids(first) == ["custom-id", None, "custom-id", None] + assert len(set(_occurrences(first))) == 4 + assert _ids(_dispatch(host, executor, sources, ledger)) == [] + assert custom.message_id == "custom-id" + assert anonymous.message_id is None + + +def test_custom_source_identity_survives_chained_copies_and_serialization() -> None: + original = Message("assistant", ["approved"], message_id="custom-id", additional_properties={"label": "original"}) + before = original.to_dict() + upstream = _response([original], "origin") + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + executor = _agent() + + first = _dispatch(host, executor, upstream, ledger) + assert _ids(first) == ["custom-id"] + forwarded = build_agent_executor_response("relay", "reply", None, upstream) + forwarded = deserialize_value(json.loads(json.dumps(serialize_value(forwarded)))) + ledger.identify(forwarded, upstream) + assert _ids(_dispatch(host, executor, forwarded, ledger)) == ["wf_relay_1"] + assert ledger.identify(forwarded)[0][0] == _occurrences(first)[0] + next_hop = build_agent_executor_response("next", "reply", None, forwarded) + assert _ids(_dispatch(host, executor, next_hop, ledger)) == ["wf_next_2"] + assert forwarded.full_conversation[0].to_dict() == next_hop.full_conversation[0].to_dict() == before + assert original.to_dict() == upstream.full_conversation[0].to_dict() == before + + +@pytest.mark.parametrize("message_id", ["wf_origin_7", "wf:external:" + "a" * 64, "wf:projection:" + "b" * 64]) +def test_workflow_shaped_application_ids_do_not_conflate_independent_relay_outputs(message_id: str) -> None: + original = Message("assistant", ["approved"], message_id=message_id) + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + executor = _agent() + + first = _dispatch(host, executor, _response([original], "left"), ledger) + assert _ids(first) == [message_id] + copied = Message.from_dict(host.calls[-1]["contextMessages"][0]) + source = _response([copied], "right") + second = _dispatch(host, executor, source, ledger) + assert _ids(second) == [message_id] + assert set(_occurrences(first)).isdisjoint(_occurrences(second)) + assert _occurrences(_dispatch(host, executor, source, ledger)) == [] + forwarded = build_agent_executor_response("relay", "reply", None, _response([copied], "right")) + assert forwarded.full_conversation[0].message_id == message_id + assert original.message_id == message_id + + +def test_anonymous_equal_text_is_identified_by_source_position_without_mutating_callers() -> None: + executor = _agent(context_mode="custom", context_filter=lambda messages: messages) + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + originals = [Message("user", ["again"]) for _ in range(3)] + before = [m.to_dict() for m in originals] + + first = _dispatch(host, executor, _response(originals[:2], latest=[]), ledger) + source = _response(originals, latest=[]) + second = _dispatch(host, executor, source, ledger) + assert _ids(first) == [None, None] + assert _ids(second) == [None] + assert len(set(_occurrences(first) + _occurrences(second))) == 3 + copied = _agent(context_mode="custom", context_filter=lambda messages: deepcopy(messages)) + assert _occurrences(_dispatch(host, copied, source, ledger)) == [] + assert [m.to_dict() for m in originals] == before + assert all(m.message_id is None for m in originals) + + +def test_detached_anonymous_copies_do_not_guess_positions_from_equal_text() -> None: + executor = _agent( + context_mode="custom", context_filter=lambda messages: [Message.from_dict(messages[-1].to_dict())] + ) + originals = [Message("user", ["again"]), Message("user", ["again"])] + + def replay() -> list[dict[str, Any]]: + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + _dispatch(host, executor, _response(originals[:1]), ledger) + _dispatch(host, executor, _response(originals), ledger) + return host.calls + + calls = replay() + assert [_texts(call) for call in calls] == [["again"], ["again"]] + assert _ids(calls[0]) == _ids(calls[1]) == [None] + assert _occurrences(calls[0]) != _occurrences(calls[1]) + assert calls == replay() + assert all(m.message_id is None for m in originals) + + +def test_anonymous_projection_reordering_uses_original_positions() -> None: + def project(messages: list[Message]) -> list[Message]: + return [messages[2], messages[0]] if len(messages) == 3 else [messages[3], messages[1]] + + executor = _agent( + context_mode="custom", + context_filter=project, + ) + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + originals = [Message("user", [str(i)]) for i in range(4)] + + first = _dispatch(host, executor, _response(originals[:3], latest=[]), ledger) + source = _response(originals, latest=[]) + second = _dispatch(host, executor, source, ledger) + assert _texts(first) == ["2", "0"] + assert _texts(second) == ["3", "1"] + assert _ids(first) == _ids(second) == [None, None] + assert len(set(_occurrences(first) + _occurrences(second))) == 4 + assert _occurrences(_dispatch(host, executor, source, ledger)) == [] + assert all(m.message_id is None for m in originals) + + +def test_reused_anonymous_object_at_two_source_positions_keeps_both_occurrences() -> None: + original = Message("user", ["again"]) + call = _dispatch(_RecordingHost(), _agent(), _response([original, original]), _WorkflowDeliveryLedger()) + + assert _ids(call) == [None, None] + assert len(set(_occurrences(call))) == 2 + assert original.message_id is None + + +def test_anonymous_same_position_in_different_producers_does_not_collide() -> None: + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + executor = _agent() + sources = [_response([Message("user", ["same"])], producer) for producer in ["left", "right"]] + + first = _dispatch(host, executor, sources, ledger) + assert _ids(first) == [None, None] + assert len(set(_occurrences(first))) == 2 + assert _ids(_dispatch(host, executor, list(reversed(sources)), ledger)) == [] + assert all(source.full_conversation[0].message_id is None for source in sources) + + +def test_anonymous_ids_remain_stable_when_forwarded_around_a_cycle() -> None: + original = Message("user", ["source input"]) + upstream = _response([original], "origin") + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + executor = _agent() + + first = _dispatch(host, executor, upstream, ledger) + assert _ids(first) == [None] + forwarded = build_agent_executor_response("relay", "reply", None, upstream) + assert _ids(_dispatch(host, executor, forwarded, ledger)) == ["wf_relay_1"] + assert ledger.identify(forwarded)[0][0] == _occurrences(first)[0] + assert original.message_id is None + assert forwarded.full_conversation[0].message_id is None + + +def test_synthesized_anonymous_messages_are_distinct_per_handoff_and_replay_stable() -> None: + executor = _agent( + context_mode="custom", + context_filter=lambda messages: [Message("system", ["summary"]), Message("system", ["summary"])], + ) + upstream = _response([_message(1)]) + + def replay() -> list[dict[str, Any]]: + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + for _ in range(2): + _dispatch(host, executor, deepcopy(upstream), ledger) + return host.calls + + first, repeated = replay() + assert _ids(first) == _ids(repeated) == [None, None] + assert len(set(_occurrences(first) + _occurrences(repeated))) == 4 + assert [first, repeated] == replay() + assert _texts(first) == _texts(repeated) == ["summary", "summary"] + + +def test_synthesized_message_with_explicit_id_is_a_new_occurrence_each_handoff() -> None: + executor = _agent( + context_mode="custom", + context_filter=lambda messages: [Message("system", [f"summary-{len(messages)}"], message_id="summary")], + ) + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + + source = _response([_message(1)]) + first = _dispatch(host, executor, source, ledger) + repeated = _dispatch(host, executor, source, ledger) + changed = _dispatch(host, executor, _response([_message(1), _message(2)]), ledger) + assert _texts(first) == _texts(repeated) == ["summary-1"] + assert _texts(changed) == ["summary-2"] + assert all(_ids(call) == ["summary"] for call in [first, repeated, changed]) + assert len({identity for call in [first, repeated, changed] for identity in _occurrences(call)}) == 3 + + +def test_last_agent_projection_without_original_position_gets_a_stable_handoff_identity() -> None: + latest = Message("assistant", ["response absent from full_conversation"]) + upstream = _response([], latest=[latest]) + executor = _agent(context_mode="last_agent") + + first = _dispatch(_RecordingHost(), executor, upstream, _WorkflowDeliveryLedger()) + replay = _dispatch(_RecordingHost(), executor, deepcopy(upstream), _WorkflowDeliveryLedger()) + assert first == replay + assert len(_occurrences(first)) == 1 + assert _ids(first) == [None] + assert latest.message_id is None + + +def test_preparation_failure_does_not_mark_delivery_or_consume_synthetic_ordinal() -> None: + executor = _agent( + context_mode="custom", context_filter=lambda messages: [Message("system", ["summary"]), *messages] + ) + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + upstream = _response([_message(1)]) + host.fail_prepare = True + + with pytest.raises(OSError, match="preparation failure"): + _dispatch(host, executor, upstream, ledger) + assert ledger == _WorkflowDeliveryLedger() + + host.fail_prepare = False + retried = _dispatch(host, executor, upstream, ledger) + assert retried == host.calls[0] + assert len(ledger.sent["target"]) == 2 + assert ledger.handoffs == {"target": 1} + + +@pytest.mark.parametrize("bad_value", [float("inf"), float("nan")]) +def test_serialization_failure_does_not_partially_record_a_batch(bad_value: Any) -> None: + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + executor = _agent() + invalid = Message("user", ["bad"], message_id="invalid", additional_properties={"nested": {"value": bad_value}}) + + with pytest.raises((TypeError, ValueError)): + _dispatch(host, executor, _response([_message(1), invalid]), ledger) + assert host.calls == [] + assert ledger == _WorkflowDeliveryLedger() + assert _ids(_dispatch(host, executor, _response([_message(1)]), ledger)) == ["wf_source_1"] + + +def test_projection_can_exclude_non_json_source_values() -> None: + invalid = Message("user", ["bad"], additional_properties={"nested": {"value": float("nan")}}) + selected = Message("user", ["selected"]) + executor = _agent( + context_mode="custom", context_filter=lambda messages: [Message.from_dict(messages[-1].to_dict())] + ) + + call = _dispatch(_RecordingHost(), executor, _response([invalid, selected]), _WorkflowDeliveryLedger()) + assert len(_occurrences(call)) == 1 + assert _ids(call) == [None] + assert _texts(call) == ["selected"] + + +def test_preview_uses_only_new_selected_text_and_never_the_large_raw_response() -> None: + selected = _message(1, text="selected") + excluded = _message(2, text="unselected secret " * 10_000) + upstream = _response([selected, excluded], latest=[excluded]) + executor = _agent(context_mode="custom", context_filter=lambda messages: messages[:1]) + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + + first = _dispatch(host, executor, upstream, ledger) + assert first["message"] == "selected" + assert "secret" not in json.dumps(first) + repeated = _dispatch(host, executor, upstream, ledger) + assert repeated["contextMessages"] == [] + assert repeated["message"] == "" + assert len(json.dumps(repeated)) < 200 + + +def test_large_new_context_retains_its_contents_but_has_a_bounded_preview() -> None: + latest = _message(1, text="large selected input " * 10_000) + call = _dispatch(_RecordingHost(), _agent(), _response([latest]), _WorkflowDeliveryLedger()) + + assert len(call["message"]) == _AGENT_TASK_MESSAGE_PREVIEW_LIMIT + assert call["message"] == latest.text[:_AGENT_TASK_MESSAGE_PREVIEW_LIMIT] + assert call["contextMessages"] == [latest.to_dict()] + + +def test_eight_hundred_turn_payload_contains_only_new_context_and_a_bounded_envelope() -> None: + host = _RecordingHost() + ledger = _WorkflowDeliveryLedger() + executor = _agent() + upstream: Any = "initial prompt" + for turn in range(800): + upstream = build_agent_executor_response("source", f"turn-{turn}:" + "x" * 700, None, upstream) + _dispatch(host, executor, upstream, ledger) + + final_call = host.calls[-1] + latest = upstream.full_conversation[-1] + latest_bytes = len(json.dumps([latest.to_dict()]).encode("utf-8")) + payload_bytes = len(json.dumps(final_call).encode("utf-8")) + projected = _build_context_messages(executor, upstream) + full_bytes = len(json.dumps(projected).encode("utf-8")) + assert final_call["contextMessages"] == [latest.to_dict()] + assert len(_occurrences(final_call)) == 1 + assert payload_bytes <= latest_bytes + _AGENT_TASK_MESSAGE_PREVIEW_LIMIT + 200 + assert full_bytes > 100 * payload_bytes + assert "initial prompt" not in json.dumps(final_call) + + repeated = _dispatch(host, executor, upstream, ledger) + assert repeated["contextMessages"] == [] + assert repeated["message"] == "" + assert len(json.dumps(repeated).encode("utf-8")) < 200 + + +def test_generator_preserves_independent_events_between_parallel_and_sequential_agent_tasks() -> None: + projections = [_response([_message(i) for i in positions]) for positions in ([1, 3], [2, 4], [4, 1])] + workflow = _workflow([_activity("source"), _agent()], []) + host = _RecordingHost(activities={"source": [_activity_result(projections)]}) + + assert _run(host, workflow) == [] + assert host.batch_sizes == [1, 1] + assert [_ids(call) for call in host.calls] == [ + ["wf_source_1", "wf_source_3"], + ["wf_source_2", "wf_source_4"], + ["wf_source_4", "wf_source_1"], + ] + assert len({identity for call in host.calls for identity in _occurrences(call)}) == 6 + assert host.calls[-1]["message"] == "source-1" + + +@pytest.mark.parametrize("representation", ["typed", "serialized", "restored"]) +@pytest.mark.parametrize("sequential", [False, True]) +@pytest.mark.parametrize( + ("status", "error_code", "include_text", "reason"), + [ + ("error", "ValueError", True, "a terminal runtime error"), + (None, "ValueError", True, "a terminal runtime error"), + (None, "", False, "a terminal runtime error"), + ("error", None, True, "a terminal runtime error"), + ("already_completed", "response_expired", False, "an expired durable response"), + ("already_completed", None, False, "an expired durable response"), + (None, "response_expired", False, "an expired durable response"), + ], +) +def test_generator_terminal_agent_result_stops_pending_and_downstream_dispatch( + representation: str, sequential: bool, status: str | None, error_code: str | None, include_text: bool, reason: str +) -> None: + secret = "private request and exception details" + contents = ( + [ + Content.from_error( + message=secret, + error_code=error_code, + error_details=secret, + additional_properties={"future_error_metadata": {"opaque": [secret]}}, + ) + ] + if error_code is not None + else [] + ) + if include_text: + contents.append(Content.from_text(f"ValueError: {secret}")) + properties: dict[str, Any] = {"correlation_id": "retained-call", "future_metadata": {"opaque": [secret]}} + if status is not None: + properties["durable_status"] = status + response = AgentResponse( + messages=[Message("system" if reason == "an expired durable response" else "assistant", contents)], + additional_properties=properties, + ) + wire = json.loads(response.to_json()) + assert wire["type"] == "agent_response" + restored = AgentResponse.from_dict(deepcopy(wire)) + assert restored.to_dict() == wire + assert restored.additional_properties == properties + payload: Any = {"typed": response, "serialized": wire, "restored": restored}[representation] + before = deepcopy(wire) + + workflow = _workflow([_activity("source"), _agent("A"), _agent("B")], [SingleEdgeGroup("A", "B")]) + host = _RecordingHost(activities={"source": [_activity_result([secret, "second", "must not run"], "A")]}) + orchestration = run_workflow_orchestrator(host, workflow, "start") + yielded = orchestration.send(next(orchestration)) + if sequential: + orchestration.send(yielded) + + with pytest.raises(RuntimeError) as failure: + orchestration.send(payload if sequential else [payload]) + + assert str(failure.value) == f"Agent executor 'A' returned {reason}." + assert secret not in str(failure.value) + assert [call["executorId"] for call in host.calls] == ["delta-A"] * (2 if sequential else 1) + assert [call["message"] for call in host.calls] == ([secret, "second"] if sequential else [secret]) + assert wire == before == response.to_dict() == restored.to_dict() + with pytest.raises(StopIteration): + next(orchestration) + + +def test_generator_checks_raw_error_before_deserializing_unknown_wire_fields() -> None: + response = AgentResponse( + messages=[Message("assistant", [Content.from_error(error_code="ValueError"), "exception text"])] + ) + wire = json.loads(response.to_json()) + wire["future_response_field"] = {"opaque": True} + wire["messages"][0]["contents"][0]["future_content_field"] = {"opaque": True} + before = deepcopy(wire) + host = _RecordingHost() + workflow = _workflow([_agent("A"), _agent("B")], [SingleEdgeGroup("A", "B")]) + orchestration = run_workflow_orchestrator(host, workflow, "start") + next(orchestration) + + with pytest.raises(RuntimeError, match="Agent executor 'A' returned a terminal runtime error"): + orchestration.send([wire]) + + assert [call["executorId"] for call in host.calls] == ["delta-A"] + assert wire == before + + +@pytest.mark.parametrize("serialized", [False, True]) +@pytest.mark.parametrize("tool_error", [False, True]) +def test_generator_normal_response_and_recovered_tool_errors_still_flow(serialized: bool, tool_error: bool) -> None: + messages: list[Message] = [] + if tool_error: + error = Content.from_error(message="recoverable tool error", error_code="ValueError") + messages.append( + Message( + "tool", + [error, Content.from_function_result("call", result=[error], exception="recoverable tool error")], + ) + ) + # A tool result may also appear in an assistant message, still as tool data. + messages.append(Message("assistant", [Content.from_function_result("call", result=[error])])) + messages.append(Message("assistant", ["approved"])) + response = AgentResponse(messages=messages, additional_properties={"future_metadata": {"error": "not a status"}}) + before = response.to_dict() + payload = json.loads(response.to_json()) if serialized else response + host = _RecordingHost() + workflow = _workflow([_agent("A"), _agent("B")], [SingleEdgeGroup("A", "B")]) + orchestration = run_workflow_orchestrator(host, workflow, "start") + next(orchestration) + + assert _finish(orchestration, orchestration.send([payload])) == [] + assert [call["executorId"] for call in host.calls] == ["delta-A", "delta-B"] + assert host.calls[-1]["contextMessages"] == [Message("user", ["start"]).to_dict(), *before["messages"]] + assert _ids(host.calls[-1]) == [None] * (1 + len(messages)) + assert len(set(_occurrences(host.calls[-1]))) == 1 + len(messages) + assert response.to_dict() == before + + +@pytest.mark.parametrize("response_type", [None, "application_result"]) +@pytest.mark.parametrize("structured", [False, True]) +def test_generator_lightweight_dict_is_not_mistaken_for_a_durable_failure( + response_type: str | None, structured: bool +) -> None: + payload: dict[str, Any] = { + "text": "ValueError: ordinary application text", + "error": "application data", + "additional_properties": {"durable_status": "error"}, + "messages": [Message("assistant", [Content.from_error(error_code="ValueError")]).to_dict()], + } + if response_type is not None: + payload["type"] = response_type + if structured: + payload["value"] = {"error": "ordinary structured output"} + before = deepcopy(payload) + host = _RecordingHost() + workflow = _workflow([_agent("A"), _agent("B")], [SingleEdgeGroup("A", "B")]) + orchestration = run_workflow_orchestrator(host, workflow, "start") + next(orchestration) + + assert _finish(orchestration, orchestration.send([payload])) == [] + assert [call["executorId"] for call in host.calls] == ["delta-A", "delta-B"] + assert _texts(host.calls[-1]) == ["start", json.dumps(payload["value"]) if structured else payload["text"]] + assert payload == before + + +@pytest.mark.parametrize("repeat_input", [False, True]) +@pytest.mark.parametrize("pause_between", [False, True]) +def test_generator_independent_strings_deliver_equal_outputs_as_new_turns( + repeat_input: bool, pause_between: bool +) -> None: + inputs = ["first request", "first request" if repeat_input else "second request"] + results = ( + [_activity_result(inputs[:1], "A", request=True), _activity_result(inputs[1:], "A")] + if pause_between + else [_activity_result(inputs, "A")] + ) + workflow = _workflow( + [_activity("gate"), _agent("A"), _agent("B", context_mode="last_agent")], + [SingleEdgeGroup("A", "B")], + ) + live = _RecordingHost(activities={"gate": results}, agent_reply="approved") + replay = _RecordingHost(is_replaying=True, activities={"gate": results}, agent_reply="approved") + + assert _run(live, workflow) == _run(replay, workflow) == [] + assert live.calls == replay.calls + producer_calls = [call for call in live.calls if call["executorId"] == "delta-A"] + consumer_calls = [call for call in live.calls if call["executorId"] == "delta-B"] + assert [call["message"] for call in producer_calls] == inputs + assert all(call["contextMessages"] is None for call in producer_calls) + assert [_ids(call) for call in consumer_calls] == [[None], [None]] + assert len({identity for call in consumer_calls for identity in _occurrences(call)}) == 2 + assert [_texts(call) for call in consumer_calls] == [["approved"], ["approved"]] + assert live.waited_for == replay.waited_for == (["approval"] if pause_between else []) + + +def test_generator_output_positions_survive_shorter_and_empty_incoming_conversations() -> None: + inputs = [_response([_message(position) for position in range(length)]) for length in [0, 4, 1, 0, 8]] + workflow = _workflow( + [_activity("source"), _agent("A"), _agent("B", context_mode="last_agent")], + [SingleEdgeGroup("A", "B")], + ) + activities = {"source": [_activity_result(inputs, "A")]} + live = _RecordingHost(activities=activities, agent_reply="approved") + replay = _RecordingHost(is_replaying=True, activities=activities, agent_reply="approved") + + assert _run(live, workflow) == _run(replay, workflow) == [] + assert live.calls == replay.calls + consumer_calls = [call for call in live.calls if call["executorId"] == "delta-B"] + assert [_ids(call) for call in consumer_calls] == [[None]] * len(inputs) + assert len({identity for call in consumer_calls for identity in _occurrences(call)}) == len(inputs) + assert [_texts(call) for call in consumer_calls] == [["approved"]] * len(inputs) + + +def test_generator_same_producer_on_independent_branches_assigns_distinct_output_positions() -> None: + workflow = _workflow( + [_agent("source"), _agent("left"), _agent("right"), _agent("A"), _agent("B", context_mode="last_agent")], + [ + FanOutEdgeGroup("source", ["left", "right"]), + SingleEdgeGroup("left", "A"), + SingleEdgeGroup("right", "A"), + SingleEdgeGroup("A", "B"), + ], + ) + live = _RecordingHost(agent_reply="approved") + replay = _RecordingHost(is_replaying=True, agent_reply="approved") + + assert _run(live, workflow) == _run(replay, workflow) == [] + assert live.calls == replay.calls + producer_calls = [call for call in live.calls if call["executorId"] == "delta-A"] + assert [_ids(call) for call in producer_calls] == [[None, None, None], [None]] + assert [_texts(call) for call in producer_calls] == [["start", "approved", "approved"], ["approved"]] + assert len({identity for call in producer_calls for identity in _occurrences(call)}) == 4 + consumer_calls = [call for call in live.calls if call["executorId"] == "delta-B"] + assert [_ids(call) for call in consumer_calls] == [[None], [None]] + assert len({identity for call in consumer_calls for identity in _occurrences(call)}) == 2 + assert [_texts(call) for call in consumer_calls] == [["approved"], ["approved"]] + + +def test_generator_fanin_keeps_repeated_outputs_from_each_producer_and_replays_identically() -> None: + workflow = _workflow( + [_activity("source"), _agent("left"), _agent("right"), _agent("join", context_mode="last_agent")], + [FanOutEdgeGroup("source", ["left", "right"]), FanInEdgeGroup(["left", "right"], "join")], + ) + activities = {"source": [_activity_result(["first request", "second request"], None)]} + live = _RecordingHost(activities=activities, agent_reply="approved") + replay = _RecordingHost(is_replaying=True, activities=activities, agent_reply="approved") + + assert _run(live, workflow) == _run(replay, workflow) == [] + assert live.calls == replay.calls + assert live.batch_sizes == [1, 2, 1] + joined = [call for call in live.calls if call["executorId"] == "delta-join"] + assert len(joined) == 1 + assert _ids(joined[0]) == [None] * 4 + assert len(set(_occurrences(joined[0]))) == 4 + assert _texts(joined[0]) == ["approved"] * 4 + + +@pytest.mark.parametrize("batch", [False, True]) +def test_generator_custom_id_collisions_are_scoped_on_the_wire_and_replay_stable(batch: bool) -> None: + sources = [ + _response([Message("assistant", ["approved"], message_id="custom-id")], producer) + for producer in ["left", "right"] + ] + deliveries: list[Any] = [sources] if batch else sources + workflow = _workflow([_activity("source"), _agent(context_mode="last_agent")], []) + activities = {"source": [_activity_result(deliveries)]} + live = _RecordingHost(activities=activities) + replay = _RecordingHost(is_replaying=True, activities=activities) + + assert _run(live, workflow) == _run(replay, workflow) == [] + assert live.calls == replay.calls + assert [message_id for call in live.calls for message_id in _ids(call)] == [ + "custom-id", + "custom-id", + ] + assert len({identity for call in live.calls for identity in _occurrences(call)}) == 2 + assert [text for call in live.calls for text in _texts(call)] == ["approved", "approved"] + assert [source.full_conversation[0].message_id for source in sources] == ["custom-id", "custom-id"] + + +def _cycle_workflow() -> Any: + return _workflow( + [_agent("A"), _agent("B")], + [ + SingleEdgeGroup("A", "B", condition=lambda response: len(response.full_conversation) < 6), + SingleEdgeGroup("B", "A", condition=lambda response: len(response.full_conversation) < 6), + ], + ) + + +def test_generator_replay_rebuilds_the_same_cycle_delta_sequence() -> None: + workflow = _cycle_workflow() + live, replay = _RecordingHost(), _RecordingHost(is_replaying=True) + + assert _run(live, workflow) == _run(replay, workflow) == [] + assert live.calls == replay.calls + assert live.calls[0]["contextMessages"] is None + assert [_ids(call) for call in live.calls[1:]] == [[None] * count for count in [2, 3, 2, 2]] + assert [_texts(call) for call in live.calls[1:]] == [ + ["start", "reply-1"], + ["start", "reply-1", "reply-2"], + ["reply-2", "reply-3"], + ["reply-3", "reply-4"], + ] + first_b, first_a, next_b, next_a = [_occurrences(call) for call in live.calls[1:]] + assert first_a[:2] == first_b + assert next_b[0] == first_a[-1] + assert next_a[0] == next_b[-1] + assert set(first_b).isdisjoint(next_b) + assert set(first_a).isdisjoint(next_a) + assert live.statuses + assert replay.statuses == [] + + +def test_interleaved_live_runs_do_not_share_delivery_on_retained_executors() -> None: + workflow = _cycle_workflow() + first, second = _RecordingHost(instance_id="first"), _RecordingHost(instance_id="second") + first_run = run_workflow_orchestrator(first, workflow, "start") + second_run = run_workflow_orchestrator(second, workflow, "start") + first_yield = next(first_run) + second_yield = next(second_run) + first_yield = first_run.send(first_yield) + second_yield = second_run.send(second_yield) + + assert _finish(first_run, first_yield) == _finish(second_run, second_yield) == [] + assert [call["contextMessages"] for call in first.calls] == [call["contextMessages"] for call in second.calls] + assert all(call["instanceId"] == "first" for call in first.calls) + assert all(call["instanceId"] == "second" for call in second.calls) + assert len(_ids(first.calls[-1])) == len(_ids(second.calls[-1])) == 2 + assert {identity for call in first.calls[1:] for identity in _occurrences(call)}.isdisjoint( + identity for call in second.calls[1:] for identity in _occurrences(call) + ) + + +def test_generator_fanout_fanin_and_cycle_preserve_producer_identity() -> None: + workflow = _workflow( + [_agent("source"), _agent("left"), _agent("right"), _agent("join")], + [ + FanOutEdgeGroup("source", ["left", "right"]), + FanInEdgeGroup(["left", "right"], "join"), + SingleEdgeGroup("join", "join", condition=lambda response: len(response.full_conversation) < 8), + ], + ) + host = _RecordingHost() + + assert _run(host, workflow) == [] + assert host.batch_sizes == [1, 2, 1, 1] + assert [_ids(call) for call in host.calls[1:]] == [[None] * count for count in [2, 2, 4, 1]] + assert [_texts(call) for call in host.calls[1:]] == [ + ["start", "reply-1"], + ["start", "reply-1"], + ["start", "reply-1", "reply-2", "reply-3"], + ["reply-4"], + ] + left, right, joined, repeated = [_occurrences(call) for call in host.calls[1:]] + assert left == right == joined[:2] + assert len(set(joined + repeated)) == 5 + + +def test_generator_hitl_resume_keeps_independent_activity_events_distinct_on_replay() -> None: + workflow = _workflow([_activity("gate"), _agent()], []) + results = [ + _activity_result([_response([_message(1), _message(3)])], request=True), + _activity_result([_response([_message(3), _message(2), _message(4), _message(1)])]), + ] + live = _RecordingHost(activities={"gate": results}) + replay = _RecordingHost(is_replaying=True, activities={"gate": results}) + + assert _run(live, workflow) == _run(replay, workflow) == [] + assert live.calls == replay.calls + assert [_ids(call) for call in live.calls] == [ + ["wf_source_1", "wf_source_3"], + ["wf_source_3", "wf_source_2", "wf_source_4", "wf_source_1"], + ] + assert set(_occurrences(live.calls[0])).isdisjoint(_occurrences(live.calls[1])) + assert live.waited_for == replay.waited_for == ["approval"] + assert deserialize_value(live.activity_inputs[1]["message"])["response"] == "approved" + assert live.activity_inputs[1]["source_executor_ids"] == ["__hitl_response___approval"] + assert any(status["state"] == "waiting_for_human_input" for status in live.statuses) diff --git a/python/packages/durabletask/tests/test_workflow_dispatch_revision.py b/python/packages/durabletask/tests/test_workflow_dispatch_revision.py new file mode 100644 index 0000000..cd87a79 --- /dev/null +++ b/python/packages/durabletask/tests/test_workflow_dispatch_revision.py @@ -0,0 +1,327 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Workflow dispatch through the real shim, request serializer and DurableTask adapter.""" + +from __future__ import annotations + +import json +from copy import deepcopy +from datetime import datetime, timezone +from typing import Any +from unittest.mock import Mock +from uuid import UUID + +import pytest +from agent_framework import AgentExecutor, AgentExecutorResponse, AgentResponse, AgentSession, Content, Message +from durabletask.task import CompletableTask, OrchestrationContext + +from agent_framework_durabletask import DurableAgentStateRequest, RunRequest +from agent_framework_durabletask._executors import DurableAgentExecutor, DurableAgentTask +from agent_framework_durabletask._shim import DurableAIAgent +from agent_framework_durabletask._workflows.dt_context import DurableTaskWorkflowContext +from agent_framework_durabletask._workflows.orchestrator import ( + _AGENT_TASK_MESSAGE_PREVIEW_LIMIT, + _prepare_agent_task, + _WorkflowDeliveryLedger, + build_agent_executor_response, +) + + +class _StubAgent: + name = "stub" + id = "stub" + description = None + + def create_session(self, **kwargs: Any) -> AgentSession: + return AgentSession(**kwargs) + + async def run(self, messages: Any = None, **kwargs: Any) -> AgentResponse: + raise AssertionError("Dispatch must schedule an entity, not invoke a model") + + +class _CaptureExecutor(DurableAgentExecutor[RunRequest]): + """Capture dispatch without replacing the inherited get_run_request implementation.""" + + def __init__(self) -> None: + self.requests: list[RunRequest] = [] + + def generate_unique_id(self) -> str: + return str(UUID(int=len(self.requests) + 1)) + + def run_durable_agent( + self, agent_name: str, run_request: RunRequest, session: AgentSession | None = None + ) -> RunRequest: + self.requests.append(run_request) + return run_request + + +def _agent(**kwargs: Any) -> AgentExecutor: + stub: Any = _StubAgent() + return AgentExecutor(stub, id="target", **kwargs) + + +def _upstream(messages: list[Message]) -> AgentExecutorResponse: + return AgentExecutorResponse( + executor_id="source", + agent_response=AgentResponse(messages=messages[-1:]), + full_conversation=list(messages), + ) + + +def _context(calls: int = 2) -> tuple[DurableTaskWorkflowContext, Mock, list[CompletableTask[Any]]]: + host = Mock(spec=OrchestrationContext) + host.instance_id = "dispatch-revision-run" + host.is_replaying = False + host.current_utc_datetime = datetime(2026, 1, 1, tzinfo=timezone.utc) + host.new_uuid.side_effect = [str(UUID(int=index + 1)) for index in range(calls)] + children: list[CompletableTask[Any]] = [CompletableTask() for _ in range(calls)] + host.call_entity.side_effect = children + return DurableTaskWorkflowContext(host), host, children + + +def _dispatch( + context: DurableTaskWorkflowContext, + host: Mock, + executor: AgentExecutor, + message: Any, + ledger: _WorkflowDeliveryLedger, +) -> tuple[DurableAgentTask, dict[str, Any]]: + task = _prepare_agent_task(context, executor, executor.id, message, "dispatch-revision", ledger) + assert isinstance(task, DurableAgentTask) + assert not task.is_complete + _, operation, payload = host.call_entity.call_args.args + assert operation == "run" + # This is the actual executor's RunRequest.to_dict(), not a reconstruction of its arguments. + wire = json.loads(json.dumps(payload, allow_nan=False)) + assert wire["orchestrationId"] == context.instance_id + assert wire["correlationId"] == str(UUID(int=host.call_entity.call_count)) + assert host.new_uuid.call_count == host.call_entity.call_count + if "contextMessages" in wire: + assert len(wire["contextMessageIds"]) == len(wire["contextMessages"]) + assert all(isinstance(identity, str) and identity for identity in wire["contextMessageIds"]) + else: + assert "contextMessageIds" not in wire + host.signal_entity.assert_not_called() + return task, wire + + +@pytest.mark.parametrize("preview", ["", "unselected logging preview"]) +def test_shim_preserves_explicit_empty_context_in_the_real_run_request(preview: str) -> None: + executor = _CaptureExecutor() + agent = DurableAIAgent(executor, "target") + + request = agent.run(preview, context_messages=[], context_message_ids=[]) + wire = json.loads(json.dumps(request.to_dict())) + + assert executor.requests == [request] + assert wire["contextMessages"] == [] + assert wire["contextMessageIds"] == [] + restored = RunRequest.from_dict(wire) + assert restored.context_messages == [] + assert restored.context_message_ids == [] + assert DurableAgentStateRequest.from_run_request(restored).messages == [] + + +def test_shim_does_not_preprocess_or_drop_raw_context_type_fields() -> None: + context_messages = [ + { + "type": "message", + "role": "tool", + "message_id": "wf_source_0", + "contents": [ + { + "type": "function_result", + "call_id": "lookup-1", + "result": {"type": "application_payload", "items": [0, False, None, "世界"]}, + "future_content_field": {"type": "opaque", "items": []}, + }, + ], + "future_message_field": {"type": "opaque", "items": []}, + }, + ] + before = deepcopy(context_messages) + executor = _CaptureExecutor() + + request = DurableAIAgent(executor, "target").run( + "", context_messages=context_messages, context_message_ids=["occurrence-0"] + ) + wire = json.loads(json.dumps(request.to_dict(), allow_nan=False)) + + assert wire["message"] == "" + assert wire["contextMessages"] == before + assert wire["contextMessageIds"] == ["occurrence-0"] + assert RunRequest.from_dict(wire).context_messages == before + assert RunRequest.from_dict(wire).context_message_ids == ["occurrence-0"] + assert context_messages == before + + +@pytest.mark.parametrize( + ("messages", "expected"), + [ + pytest.param(None, "", id="none"), + pytest.param([], "", id="empty-list"), + pytest.param("standalone", "standalone", id="text"), + pytest.param(Message("user", ["standalone"]), "standalone", id="message"), + pytest.param(["first", "second"], "first\nsecond", id="text-list"), + ], +) +def test_shim_without_context_retains_standalone_text_normalization(messages: Any, expected: str) -> None: + executor = _CaptureExecutor() + + request = DurableAIAgent(executor, "target").run(messages, context_messages=None) + + assert request.message == expected + assert request.context_messages is None + assert "contextMessages" not in request.to_dict() + assert executor.requests == [request] + + +@pytest.mark.parametrize( + "messages", + [ + pytest.param("", id="empty-text"), + pytest.param(Message("user", []), id="contentless-message"), + pytest.param( + Message("tool", [Content.from_function_result("lookup-1", result={"answer": 42})]), + id="nontext-message", + ), + ], +) +def test_shim_without_context_still_rejects_nontext_inputs(messages: Any) -> None: + executor = _CaptureExecutor() + + with pytest.raises(ValueError, match="only supports text message inputs"): + DurableAIAgent(executor, "target").run(messages, context_messages=None) + + assert executor.requests == [] + + +def test_custom_empty_projection_reaches_the_dt_entity_as_an_empty_list() -> None: + context, host, _ = _context() + executor = _agent(context_mode="custom", context_filter=lambda messages: []) + excluded = Message("assistant", ["unselected secret" * 1000], message_id="wf_source_0") + ledger = _WorkflowDeliveryLedger() + + _, wire = _dispatch(context, host, executor, _upstream([excluded]), ledger) + + assert wire["message"] == "" + assert wire["contextMessages"] == [] + assert "unselected secret" not in json.dumps(wire) + assert DurableAgentStateRequest.from_run_request(RunRequest.from_dict(wire)).messages == [] + assert ledger.sent == {} + assert ledger.handoffs == {"target": 1} + host.call_entity.assert_called_once() + + +def test_fully_duplicate_projection_reaches_the_dt_entity_on_the_second_call() -> None: + context, host, _ = _context() + executor = _agent() + messages = [ + Message("user", ["question"], message_id="wf_source_0"), + Message("assistant", ["answer"], message_id="wf_source_1"), + ] + upstream = _upstream(messages) + expected = [message.to_dict() for message in messages] + ledger = _WorkflowDeliveryLedger() + + _, first = _dispatch(context, host, executor, upstream, ledger) + assert first["contextMessages"] == expected + assert len(set(first["contextMessageIds"])) == 2 + assert set(first["contextMessageIds"]).isdisjoint(message.message_id for message in messages) + _, repeated = _dispatch(context, host, executor, upstream, ledger) + + assert repeated["contextMessages"] == [] + assert repeated["contextMessageIds"] == [] + assert repeated["message"] == "" + assert first["correlationId"] != repeated["correlationId"] + assert DurableAgentStateRequest.from_run_request(RunRequest.from_dict(repeated)).messages == [] + assert len(ledger.sent["target"]) == 2 + assert ledger.handoffs == {"target": 2} + assert [message.to_dict() for message in messages] == expected + assert host.call_entity.call_count == 2 + + +def test_tool_only_projection_survives_dt_dispatch_and_request_parsing() -> None: + context, host, children = _context() + result = {"type": "lookup_result", "items": [{"answer": 0, "label": "世界"}], "flags": [False, None]} + message = Message( + "tool", + [Content.from_function_result("lookup-1", result=result)], + message_id="wf_source_0", + author_name="lookup", + additional_properties={"provider": {"type": "context", "labels": []}}, + ) + expected = message.to_dict() + ledger = _WorkflowDeliveryLedger() + + task, wire = _dispatch(context, host, _agent(), _upstream([message]), ledger) + + assert wire["message"] == "" + assert wire["contextMessages"] == [expected] + request = RunRequest.from_json(json.dumps(wire)) + assert request.context_message_ids == wire["contextMessageIds"] + assert request.context_message_ids != [message.message_id] + entry = DurableAgentStateRequest.from_run_request(request) + assert len(entry.messages) == 1 + forwarded = entry.messages[0].to_chat_message() + assert isinstance(forwarded, Message) + assert forwarded.role == "tool" + assert forwarded.message_id == message.message_id + assert forwarded.text == "" + assert len(forwarded.contents) == 1 + assert forwarded.contents[0].type == "function_result" + assert forwarded.contents[0].call_id == "lookup-1" + assert forwarded.contents[0].result == message.contents[0].result + assert json.loads(forwarded.contents[0].result) == result + assert message.to_dict() == expected + + assert not children[0].is_complete + children[0].complete(AgentResponse(messages=[Message("assistant", ["received"])]).to_dict()) + assert task.is_complete and not task.is_failed + assert context.get_task_result(task).text == "received" + + +def test_standalone_dt_input_is_not_truncated_or_deduplicated() -> None: + context, host, _ = _context() + executor = _agent() + ledger = _WorkflowDeliveryLedger() + prompt = "standalone input " * 1000 + + for _ in range(2): + _, wire = _dispatch(context, host, executor, prompt, ledger) + assert wire["message"] == prompt + assert "contextMessages" not in wire + assert RunRequest.from_dict(wire).context_messages is None + + assert ledger.sent == {} + assert host.call_entity.call_count == 2 + + +def test_eight_hundred_turns_have_a_bounded_real_dt_request_envelope() -> None: + context, host, _ = _context(calls=801) + executor = _agent() + ledger = _WorkflowDeliveryLedger() + upstream: Any = "initial prompt" + wire: dict[str, Any] = {} + for turn in range(800): + upstream = build_agent_executor_response("source", f"turn-{turn}:" + "x" * 1600, None, upstream) + _, wire = _dispatch(context, host, executor, upstream, ledger) + + latest = upstream.full_conversation[-1] + assert wire["contextMessages"] == [latest.to_dict()] + assert len(wire["contextMessageIds"]) == 1 + assert wire["message"] == latest.text[:_AGENT_TASK_MESSAGE_PREVIEW_LIMIT] + assert len(wire["message"]) == _AGENT_TASK_MESSAGE_PREVIEW_LIMIT + payload_bytes = len(json.dumps(wire).encode("utf-8")) + context_bytes = len(json.dumps([latest.to_dict()]).encode("utf-8")) + full_bytes = len(json.dumps([message.to_dict() for message in upstream.full_conversation]).encode("utf-8")) + assert payload_bytes <= context_bytes + _AGENT_TASK_MESSAGE_PREVIEW_LIMIT + 512 + assert full_bytes > 100 * payload_bytes + assert "initial prompt" not in json.dumps(wire) + + _, repeated = _dispatch(context, host, executor, upstream, ledger) + assert repeated["contextMessages"] == [] + assert repeated["contextMessageIds"] == [] + assert repeated["message"] == "" + assert len(json.dumps(repeated).encode("utf-8")) < 512 + assert host.call_entity.call_count == 801 diff --git a/python/packages/durabletask/tests/test_workflow_output_boundaries_review.py b/python/packages/durabletask/tests/test_workflow_output_boundaries_review.py new file mode 100644 index 0000000..e5c3956 --- /dev/null +++ b/python/packages/durabletask/tests/test_workflow_output_boundaries_review.py @@ -0,0 +1,304 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Parent output selection and portable generated responses at public boundaries.""" + +from __future__ import annotations + +import asyncio +import json +from datetime import date +from types import SimpleNamespace +from typing import Any +from unittest.mock import Mock, patch + +import pytest +from agent_framework import ( + AgentExecutorResponse, + AgentResponse, + Workflow, + WorkflowBuilder, + WorkflowEvent, + WorkflowExecutor, +) +from agent_framework._workflows import _checkpoint_encoding +from durabletask.client import TaskHubGrpcClient +from pydantic import BaseModel, Field +from test_workflow_agent_contract_review import _Adapter, _agent, _InspectChild, _response, _wire + +from agent_framework_durabletask import DurableWorkflowClient, deserialize_workflow_output, serialize_agent_response +from agent_framework_durabletask._response_utils import load_agent_response +from agent_framework_durabletask._workflows.activity import execute_workflow_activity +from agent_framework_durabletask._workflows.orchestrator import _FORWARDING_PROVENANCE, run_workflow_orchestrator +from agent_framework_durabletask._workflows.protocol import unwrap_workflow_input +from agent_framework_durabletask._workflows.serialization import ( + deserialize_value, + deserialize_workflow_event, + serialize_value, + serialize_workflow_agent_response, +) + + +def _finish_raw(host: _Adapter, generator: Any, yielded: Any, value: Any) -> Any: + with pytest.raises(StopIteration) as completed: + generator.send(host.complete(yielded, value)) + return json.loads(json.dumps(completed.value.value, allow_nan=False)) + + +def _nested(direct: bool, designation: str) -> tuple[Workflow, Workflow, _InspectChild]: + progress = _agent("progress", [_response("progress")]) + answer = _agent("answer", [_response("answer")]) + inner = ( + WorkflowBuilder( + name="inner", start_executor=progress, output_from=[answer], intermediate_output_from=[progress] + ) + .add_edge(progress, answer) + .build() + ) + child = WorkflowExecutor(inner, id="child", allow_direct_output=direct) + sink = _InspectChild() + options: dict[str, Any] = {} + if designation != "omitted": + options = { + "output_from": [child, sink] if designation == "output" else [sink], + "intermediate_output_from": [child] if designation == "intermediate" else [], + } + outer = WorkflowBuilder(name="outer", start_executor=child, **options).add_edge(child, sink).build() + return outer, inner, sink + + +@pytest.mark.parametrize("adapter", ["dt", "af"]) +@pytest.mark.parametrize("direct", [False, True]) +@pytest.mark.parametrize("designation", ["hidden", "intermediate", "output", "omitted"]) +async def test_child_outputs_follow_parent_yield_policy_and_core_events( + adapter: str, direct: bool, designation: str +) -> None: + core, _, _ = _nested(direct, designation) + expected = await core.run("question") + outer, inner, sink = _nested(direct, designation) + host = _Adapter(adapter) + generator = run_workflow_orchestrator(host.context, outer, "question") + yielded = next(generator) + kind, _, _, kwargs = host.pending[0] + assert kind == "child" + child_input = unwrap_workflow_input(kwargs["input"] if adapter == "dt" else kwargs["input_"]) + child_host = _Adapter(adapter) + child_host.native.instance_id = kwargs["instance_id"] + child_generator = run_workflow_orchestrator(child_host.context, inner, child_input) + child_yielded = next(child_generator) + child_yielded = child_generator.send(child_host.complete(child_yielded, _wire(_response("progress")))) + child_result = _finish_raw(child_host, child_generator, child_yielded, _wire(_response("answer"))) + assert child_result["outputs"][0]["_durable_agent_response"] == 1 + + if direct: + raw = _finish_raw(host, generator, yielded, child_result) + host.native.call_activity.assert_not_called() + else: + yielded = generator.send(host.complete(yielded, child_result)) + activity_input = host.activity_input() + assert type(deserialize_value(json.loads(activity_input)["message"])) is AgentResponse + activity_result = await asyncio.to_thread(execute_workflow_activity, sink, activity_input, outer) + raw = _finish_raw(host, generator, yielded, activity_result) + + def snapshot(value: Any) -> Any: + return serialize_agent_response(value) if isinstance(value, AgentResponse) else value + + assert [snapshot(value) for value in deserialize_workflow_output(raw)] == [ + snapshot(value) for value in expected.get_outputs() + ] + if adapter == "af": + assert all("events" not in status for status in [*host.statuses, *child_host.statuses]) + return + + events = [deserialize_workflow_event(event) for event in host.statuses[-1]["events"]] + actual_yields = [event for event in events if event.type in ("output", "intermediate")] + expected_yields = [event for event in expected if event.type in ("output", "intermediate")] + assert all(isinstance(event, WorkflowEvent) for event in actual_yields) + assert [(e.type, e.executor_id, snapshot(e.data)) for e in actual_yields] == [ + (e.type, e.executor_id, snapshot(e.data)) for e in expected_yields + ] + child_events = [event for event in events if event.executor_id == "child"] + assert child_events[0].type == "executor_invoked" + assert child_events[-1].type == "executor_completed" + # Core forwards inner intermediate events even when the child node's own + # direct yields are hidden, and outputs precede that forwarded progress. + assert child_events[-2].type == "intermediate" and child_events[-2].data.text == "progress" + + +@pytest.mark.parametrize("adapter", ["dt", "af"]) +@pytest.mark.parametrize("designation", ["output", "intermediate", "hidden"]) +def test_worker_local_model_is_typed_in_conditions_but_never_pickled_for_generated_yields( + adapter: str, designation: str, monkeypatch: pytest.MonkeyPatch +) -> None: + class WorkerAnswer(BaseModel): + answer: int = Field(validation_alias="inputAnswer", serialization_alias="outputAnswer") + day: date + + a = _agent("A", response_format=WorkerAnswer) + b = _agent("B") + observed: list[AgentExecutorResponse] = [] + + def condition(value: AgentExecutorResponse) -> bool: + assert type(value) is AgentExecutorResponse + assert isinstance(value.agent_response.value, WorkerAnswer) + assert value.agent_response.value.day == date(2026, 9, 9) + observed.append(value) + return True + + workflow = ( + WorkflowBuilder( + name="portable", + start_executor=a, + output_from=[a, b] if designation == "output" else [b], + intermediate_output_from=[a] if designation == "intermediate" else [], + ) + .add_edge(a, b, condition=condition) + .build() + ) + response = _response("not JSON", value=WorkerAnswer(inputAnswer=42, day=date(2026, 9, 9))) + setattr(response, _FORWARDING_PROVENANCE, ("private", response.messages)) + external = serialize_workflow_agent_response(response) + assert _FORWARDING_PROVENANCE not in json.dumps(external) + assert getattr(response, _FORWARDING_PROVENANCE) == ("private", response.messages) + host = _Adapter(adapter) + generator = run_workflow_orchestrator(host.context, workflow, "question") + yielded = next(generator) + no_pickle = Mock(side_effect=AssertionError("Generated responses must not require worker classes")) + monkeypatch.setattr(_checkpoint_encoding, "_pickle_to_base64", no_pickle) + monkeypatch.setattr(_checkpoint_encoding, "_base64_to_unpickle", no_pickle) + yielded = generator.send(host.complete(yielded, _wire(response))) + assert len(observed) == 1 + raw = _finish_raw(host, generator, yielded, _wire(_response("last", value=False))) + assert "__pickled__" not in json.dumps([raw, host.statuses]) + assert _FORWARDING_PROVENANCE not in json.dumps([raw, host.statuses]) + assert "WorkerAnswer" not in json.dumps([raw, host.statuses]) + with patch("importlib.import_module", side_effect=AssertionError("Client must not import stored response types")): + output = deserialize_workflow_output(raw) + events = [deserialize_workflow_event(event) for event in host.statuses[-1].get("events", [])] + assert all(type(value) is AgentResponse for value in output) + assert output[-1].value is False + if designation == "output": + assert output[0].value == {"answer": 42, "day": "2026-09-09"} + assert raw[0]["response"]["_durable_value_by_name"] is True + if adapter == "dt" and designation != "hidden": + emitted = next(event for event in events if event.executor_id == "A" and event.type == designation) + assert type(emitted.data) is AgentResponse + assert emitted.data.value == {"answer": 42, "day": "2026-09-09"} + no_pickle.assert_not_called() + + +@pytest.mark.parametrize("value", [False, None, {"wireAlias": 0, "nullable": None}]) +async def test_public_client_returns_response_values_and_streamed_events_without_type_resolution(value: Any) -> None: + response = load_agent_response({"type": "agent_response", "messages": [], "value": value}) + host = _Adapter("dt") + workflow = WorkflowBuilder(name="portable", start_executor=_agent("A")).build() + generator = run_workflow_orchestrator(host.context, workflow, "question") + raw = _finish_raw(host, generator, next(generator), _wire(response)) + state = SimpleNamespace( + name="dafx-portable", + runtime_status=SimpleNamespace(name="COMPLETED"), + serialized_output=json.dumps(raw), + serialized_custom_status=json.dumps(host.statuses[-1]), + ) + native = Mock(spec=TaskHubGrpcClient) + native.wait_for_orchestration_completion.return_value = state + native.get_orchestration_state.return_value = state + client = DurableWorkflowClient(native, workflow_name="portable") + with ( + patch("importlib.import_module", side_effect=AssertionError("No response type imports")), + patch.object(_checkpoint_encoding, "_base64_to_unpickle", side_effect=AssertionError("No response pickle")), + ): + output = client.await_workflow_output("contract-run") + events = [event async for event in client.stream_workflow("contract-run")] + assert len(output) == 1 and type(output[0]) is AgentResponse + emitted = [event.data for event in events if event.type == "output"] + assert len(emitted) == 1 and type(emitted[0]) is AgentResponse + for restored in [output[0], emitted[0]]: + assert restored.value == value and type(restored.value) is type(value) + assert "value" in serialize_agent_response(restored) + + +def test_known_envelopes_recurse_only_through_codec_containers_not_response_application_data() -> None: + application = { + "type": "worker.only:Model", + "__pickled__": "application data, not a pickle", + "__type__": "application:type", + "nested": {"_durable_agent_response": 99, "response": {"type": "business"}}, + } + response = load_agent_response({ + "type": "agent_response", + "messages": [], + "value": application, + "additional_properties": application, + }) + envelope = serialize_workflow_agent_response(response) + plain_response_dict = {"type": "agent_response", "messages": [], "value": False} + # Plain lists/dicts produced by the core encoder can carry a known envelope. + container = serialize_value({"outputs": [serialize_workflow_agent_response(_response(value=False))]}) + container["outputs"].extend([envelope, plain_response_dict]) + with ( + patch("importlib.import_module", side_effect=AssertionError("Application types are not imported")), + patch.object(_checkpoint_encoding, "_base64_to_unpickle", side_effect=AssertionError("Data is not pickle")), + ): + restored = deserialize_workflow_output(json.loads(json.dumps(container))) + first, second, plain = restored["outputs"] + assert type(first) is AgentResponse and first.value is False + assert type(second) is AgentResponse and second.value == application + assert second.additional_properties == application + assert type(plain) is dict and plain == plain_response_dict + + +@pytest.mark.parametrize( + "envelope", + [ + *[{"_durable_agent_response": version, "response": {}} for version in [None, False, True, 0, 2, 1.0, "1"]], + {"_durable_agent_response": 1}, + {"_durable_agent_response": 1, "response": None}, + {"_durable_agent_response": 1, "response": []}, + {"_durable_agent_response": 1, "response": {}, "extra": False}, + {"_durable_agent_response": 1, "response": {}, "__pickled__": "bad", "__type__": "worker:Type"}, + ], +) +def test_invalid_known_envelope_rejected_before_core_decoder(envelope: Any) -> None: + with ( + patch( + "agent_framework_durabletask._workflows.serialization.decode_checkpoint_value", + side_effect=AssertionError("Malformed response envelope must not reach the generic decoder"), + ), + pytest.raises(ValueError, match="workflow agent response envelope"), + ): + deserialize_workflow_output([{"nested": envelope}]) + + +@pytest.mark.parametrize("payload", [{}, {"type": ""}, {"type": "agent_response", "messages": False}]) +def test_known_envelope_still_validates_base_response_fields(payload: Any) -> None: + with pytest.raises((ValueError, TypeError)): + deserialize_value({"_durable_agent_response": 1, "response": payload}) + + +def test_stored_response_type_and_format_are_not_client_constructor_instructions() -> None: + envelope = { + "_durable_agent_response": 1, + "response": { + "type": "worker.only:Response", + "response_format": "worker.only:Model", + "messages": [], + "value": {"type": "business.kind", "flag": False}, + }, + } + with patch("importlib.import_module", side_effect=AssertionError("Stored type names are not imported")): + restored = deserialize_workflow_output(envelope) + assert type(restored) is AgentResponse and restored.value == {"type": "business.kind", "flag": False} + + +def test_existing_internal_pickle_contract_and_escaped_application_dictionary_are_unchanged() -> None: + from test_workflow_agent_contract_review import Answer + + response = _response(value=Answer(inputAnswer=42)) + internal = AgentExecutorResponse("A", response, full_conversation=response.messages) + encoded = serialize_value(internal) + assert "__pickled__" in encoded + restored = deserialize_value(encoded) + assert type(restored) is AgentExecutorResponse and isinstance(restored.agent_response.value, Answer) + assert "__pickled__" in serialize_value(response) + application = {"__pickled__": "literal", "__type__": "business", "_durable_agent_response": 99} + assert deserialize_value(serialize_value(application)) == application diff --git a/python/packages/durabletask/tests/test_workflow_protocol_review.py b/python/packages/durabletask/tests/test_workflow_protocol_review.py new file mode 100644 index 0000000..c10e7cd --- /dev/null +++ b/python/packages/durabletask/tests/test_workflow_protocol_review.py @@ -0,0 +1,383 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Registered DT start boundaries and v2-only shared-generator replay, without a service.""" + +from __future__ import annotations + +import json +from collections.abc import Callable, Generator +from copy import deepcopy +from dataclasses import dataclass +from typing import Any +from unittest.mock import Mock + +import pytest +from agent_framework import Executor, Workflow, WorkflowExecutor +from agent_framework._workflows import _checkpoint_encoding +from agent_framework._workflows._edge import SingleEdgeGroup +from durabletask.task import CompletableTask, OrchestrationContext + +from agent_framework_durabletask import DurableAIAgentWorker, DurableWorkflowClient +from agent_framework_durabletask import _worker as worker_module +from agent_framework_durabletask._workflows.orchestrator import SOURCE_HITL_RESPONSE, SOURCE_WORKFLOW_START +from agent_framework_durabletask._workflows.protocol import unwrap_workflow_input +from agent_framework_durabletask._workflows.serialization import ( + SUBWORKFLOW_ADDRESS_KEY, + SUBWORKFLOW_INPUT_KEY, + SUBWORKFLOW_RESULT_KEY, + deserialize_value, + serialize_value, +) + +_VERSION = "_durable_workflow_version" +_CONTROL = {"input": "application control", "items": [0, False, None, "世界"]} +_FORGED_ADDRESS = { + "root_instance_id": "other-run", + "root_workflow_name": "other-workflow", + "request_path_prefix": "forged~9~", +} +_UNTRUSTED = {"__pickled__": "not-trusted-checkpoint-data", "__type__": "builtins:str"} + + +@dataclass +class _TypedInput: + input: str + control: dict[str, Any] + + +def _node(name: str = "start", input_type: type | None = None) -> Any: + node = Mock(spec=Executor) + node.id = name + node.input_types = [] if input_type is None else [input_type] + return node + + +def _workflow(name: str = "protocol", nodes: list[Any] | None = None, edges: list[Any] | None = None) -> Any: + nodes = [_node()] if nodes is None else nodes + workflow = Mock(spec=Workflow) + workflow.name = name + workflow.start_executor_id = nodes[0].id + workflow.executors = {node.id: node for node in nodes} + workflow.edge_groups = [] if edges is None else edges + workflow.max_iterations = 10 + return workflow + + +def _register(workflow: Any) -> dict[str, Callable[..., Any]]: + # Use configure_workflow, not a reimplementation of its generated closure. + native = Mock() + DurableAIAgentWorker(native, deployment_mode="isolated_v2").configure_workflow(workflow) + native.add_entity.assert_not_called() + return {call.args[0].__name__: call.args[0] for call in native.add_orchestrator.call_args_list} + + +def _start(payload: Any, name: str = "protocol") -> dict[str, Any]: + client = Mock() + client.schedule_new_orchestration.return_value = "root-run" + assert ( + DurableWorkflowClient(client, workflow_name=name).start_workflow(payload, instance_id="root-run") == "root-run" + ) + client.schedule_new_orchestration.assert_called_once() + call = client.schedule_new_orchestration.call_args + assert call is not None + assert call.args == (f"dafx-{name}",) + assert call.kwargs["instance_id"] == "root-run" + return json.loads(json.dumps(call.kwargs["input"], allow_nan=False)) + + +def _complete(value: Any) -> CompletableTask[Any]: + task: CompletableTask[Any] = CompletableTask() + task.complete(value) + return task + + +def _drain(generator: Generator[Any, Any, Any], value: Any = None) -> Any: + while True: + try: + task = generator.send(value) + except StopIteration as completed: + return completed.value + assert task.is_complete, "Use explicit event completion for a paused generator" + value = task.get_result() + + +def _host( + calls: list[dict[str, Any]], + result: Callable[[str, dict[str, Any]], dict[str, Any]] | None = None, + *, + functions: dict[str, Callable[..., Any]] | None = None, + instance_id: str = "root-run", + replay: bool = False, +) -> Mock: + host = Mock(spec=OrchestrationContext) + host.instance_id = instance_id + host.is_replaying = replay + + def activity(name: str, *, input: str) -> CompletableTask[Any]: + payload = json.loads(input) + calls.append({"kind": "activity", "instance": instance_id, "name": name, "input": deepcopy(payload)}) + response = {"outputs": ["done"]} if result is None else result(name, payload) + return _complete(json.dumps(response)) + + def child(name: str, *, input: Any, instance_id: str) -> CompletableTask[Any]: + assert functions is not None + wire = json.loads(json.dumps(input)) + calls.append({"kind": "child", "instance": instance_id, "name": name, "input": deepcopy(wire)}) + context = _host(calls, result, functions=functions, instance_id=instance_id, replay=replay) + child_result = _drain(functions[name](context, wire)) + assert child_result[SUBWORKFLOW_RESULT_KEY] is True + return _complete(child_result) + + host.call_activity.side_effect = activity + host.call_sub_orchestrator.side_effect = child + host.wait_for_external_event.side_effect = lambda name: CompletableTask() + # Copy when published, rather than observing later mutation of the same dict/list. + host.statuses = [] + host.set_custom_status.side_effect = lambda status: host.statuses.append(deepcopy(status)) + return host + + +@pytest.mark.parametrize( + "recorded", + [ + pytest.param({"input": "a user's field"}, id="raw-dict-with-input"), + pytest.param("old start", id="raw-string"), + pytest.param("", id="raw-empty-string"), + pytest.param([], id="raw-empty-list"), + pytest.param({}, id="raw-empty-object"), + pytest.param(None, id="raw-null"), + pytest.param({SUBWORKFLOW_INPUT_KEY: _UNTRUSTED, SUBWORKFLOW_ADDRESS_KEY: _FORGED_ADDRESS}, id="legacy-child"), + pytest.param({_VERSION: 1, "input": "old"}, id="protocol-one"), + pytest.param({_VERSION: True, "input": "old"}, id="boolean-true"), + pytest.param({_VERSION: False, "input": "old"}, id="boolean-false"), + pytest.param({_VERSION: 2.0, "input": "old"}, id="float-two"), + pytest.param({_VERSION: "2", "input": "old"}, id="string-two"), + pytest.param({_VERSION: 2}, id="missing-input"), + pytest.param({_VERSION: 2, "input": "old", "extra": None}, id="extra-key"), + ], +) +def test_recorded_unsupported_start_fails_before_engine_or_actions( + recorded: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + workflow = _workflow() + functions = _register(workflow) + engine = Mock(side_effect=AssertionError("The changed engine must not see old history")) + monkeypatch.setattr(worker_module, "run_workflow_orchestrator", engine) + host = _host([], replay=True) + original = deepcopy(recorded) + before_nodes = dict(workflow.executors) + + with pytest.raises(ValueError, match="unsupported execution protocol"): + next(functions["dafx-protocol"](host, recorded)) + + engine.assert_not_called() + assert host.mock_calls == [] + assert host.statuses == [] + assert workflow.executors == before_nodes + for node in workflow.executors.values(): + node.execute.assert_not_called() + assert recorded == original + + +@pytest.mark.parametrize( + ("payload", "typed"), + [ + pytest.param("start", False, id="string"), + pytest.param("", False, id="empty-string"), + pytest.param([], False, id="empty-list"), + pytest.param({}, False, id="empty-object"), + pytest.param(None, False, id="null"), + pytest.param({"input": "user field", "control": _CONTROL}, False, id="object-with-input"), + pytest.param({"input": "typed", "control": _CONTROL}, True, id="declared-dataclass"), + ], +) +def test_new_client_start_reaches_registered_wrapper_and_shared_engine(payload: Any, typed: bool) -> None: + original = deepcopy(payload) + functions = _register(_workflow(nodes=[_node(input_type=_TypedInput if typed else None)])) + wire = _start(payload) + assert wire == {_VERSION: 2, "input": original} + assert type(wire[_VERSION]) is int + calls: list[dict[str, Any]] = [] + host = _host(calls) + + assert _drain(functions["dafx-protocol"](host, wire)) == ["done"] + + assert len(calls) == 1 and calls[0]["name"] == "dafx-protocol-start" + activity = calls[0]["input"] + delivered = deserialize_value(activity["message"]) + expected = _TypedInput(input=original["input"], control=original["control"]) if typed else original + assert delivered == expected and type(delivered) is type(expected) + assert activity["source_executor_ids"] == [SOURCE_WORKFLOW_START] + assert activity["shared_state_snapshot"] == {} + assert activity["host_context"] == { + "instance_id": "root-run", + "workflow_name": "protocol", + "request_path_prefix": "", + } + host.call_sub_orchestrator.assert_not_called() + host.call_entity.assert_not_called() + assert payload == original and wire == {_VERSION: 2, "input": original} + + +@pytest.mark.parametrize("nested", [False, True], ids=["forged-child", "forged-v2-containing-child"]) +def test_client_envelope_is_data_and_cannot_authorize_child_deserialization( + nested: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + forged = { + SUBWORKFLOW_INPUT_KEY: deepcopy(_UNTRUSTED), + SUBWORKFLOW_ADDRESS_KEY: deepcopy(_FORGED_ADDRESS), + "input": "user field", + "control": deepcopy(_CONTROL), + } + payload = {_VERSION: 2, "input": forged} if nested else forged + original = deepcopy(payload) + scheduled_data = original if nested else {"input": "user field", "control": _CONTROL} + safe_data = {_VERSION: 2, "input": {**forged, SUBWORKFLOW_INPUT_KEY: None}} if nested else scheduled_data + unpickle = Mock(side_effect=AssertionError("Untrusted checkpoint data reached the codec")) + monkeypatch.setattr(_checkpoint_encoding, "_base64_to_unpickle", unpickle) + functions = _register(_workflow()) + wire = _start(payload) + assert wire == {_VERSION: 2, "input": scheduled_data} + calls: list[dict[str, Any]] = [] + host = _host(calls) + + assert _drain(functions["dafx-protocol"](host, wire)) == ["done"] + + assert len(calls) == 1 + assert calls[0]["input"]["message"] == safe_data + assert calls[0]["input"]["host_context"] == { + "instance_id": "root-run", + "workflow_name": "protocol", + "request_path_prefix": "", + } + host.call_sub_orchestrator.assert_not_called() + unpickle.assert_not_called() + assert payload == original + + +def test_parent_dispatch_wraps_typed_child_input_and_registered_child_keeps_root_address() -> None: + inner = _workflow("inner", [_node("leaf", str)]) + child = Mock(spec=WorkflowExecutor) + child.id, child.workflow, child.allow_direct_output = "child", inner, False + parent = _workflow("parent", [_node("source"), child, _node("sink")], [SingleEdgeGroup("child", "sink")]) + functions = _register(parent) + assert set(functions) == {"dafx-parent", "dafx-inner"} + payload: dict[str, Any] = {"input": "nested typed input", "control": deepcopy(_CONTROL)} + typed = _TypedInput(input=payload["input"], control=deepcopy(payload["control"])) + + def result(name: str, data: dict[str, Any]) -> dict[str, Any]: + message = deserialize_value(data["message"]) + if name == "dafx-parent-source": + assert message == payload + return { + "sent_messages": [ + {"message": _checkpoint_encoding.encode_checkpoint_value(typed), "target_id": "child"} + ] + } + assert isinstance(message, _TypedInput) and message == typed + if name == "dafx-inner-leaf": + return {"outputs": [serialize_value(message)]} + assert name == "dafx-parent-sink" + return {"outputs": ["done"]} + + calls: list[dict[str, Any]] = [] + host = _host(calls, result, functions=functions) + assert _drain(functions["dafx-parent"](host, _start(payload, "parent"))) == ["done"] + assert [call["name"] for call in calls] == [ + "dafx-parent-source", + "dafx-inner", + "dafx-inner-leaf", + "dafx-parent-sink", + ] + dispatch = calls[1] + assert dispatch["instance"] == "root-run::child::0" + child_input = unwrap_workflow_input(dispatch["input"]) + assert dispatch["input"] == {_VERSION: 2, "input": child_input} + assert type(dispatch["input"][_VERSION]) is int + # Check typed semantics through the core codec without assuming a pickle byte layout. + decoded_child = _checkpoint_encoding.decode_checkpoint_value(child_input) + assert decoded_child == { + SUBWORKFLOW_INPUT_KEY: typed, + SUBWORKFLOW_ADDRESS_KEY: { + "root_instance_id": "root-run", + "root_workflow_name": "parent", + "request_path_prefix": "child~0~", + }, + } + assert type(decoded_child[SUBWORKFLOW_INPUT_KEY]) is _TypedInput + assert type(deserialize_value(calls[2]["input"]["message"])) is _TypedInput + assert calls[2]["input"]["host_context"] == { + "instance_id": "root-run", + "workflow_name": "parent", + "request_path_prefix": "child~0~", + } + assert calls[2]["input"]["source_executor_ids"] == [SOURCE_WORKFLOW_START] + assert calls[3]["input"]["source_executor_ids"] == ["child"] + assert calls[0]["input"]["message"] == payload + + +def test_v2_paused_hitl_replays_full_shared_generator_with_identical_dispatch_and_state() -> None: + """Cold generator replay of v2 only, not SDK history execution or old-history compatibility.""" + payload = {"input": "start", "control": deepcopy(_CONTROL)} + answer = {"input": "approved", "control": deepcopy(_CONTROL)} + + def result(name: str, data: dict[str, Any]) -> dict[str, Any]: + if data["source_executor_ids"] == [SOURCE_WORKFLOW_START]: + return { + "shared_state_updates": {"pending": payload}, + "pending_request_info_events": [ + { + "request_id": "approval", + "source_executor_id": "gate", + "data": payload, + "request_type": "builtins:dict", + "response_type": "builtins:dict", + } + ], + } + if name == "dafx-protocol-gate": + assert data["shared_state_snapshot"] == {"pending": payload} + assert deserialize_value(data["message"]) == { + "request_id": "approval", + "original_request": payload, + "response": answer, + "response_type": "builtins:dict", + } + return { + "shared_state_deletes": ["pending"], + "shared_state_updates": {"decision": answer}, + "sent_messages": [{"message": answer, "target_id": "sink"}], + } + assert name == "dafx-protocol-sink" + assert data["shared_state_snapshot"] == {"decision": answer} + assert data["message"] == answer + return {"outputs": ["done"]} + + wire = _start(payload) + executions = [] + for replay in (False, True): + functions = _register(_workflow(nodes=[_node("gate"), _node("sink")])) + calls: list[dict[str, Any]] = [] + host = _host(calls, result, replay=replay) + generator = functions["dafx-protocol"](host, deepcopy(wire)) + batch = next(generator) + assert batch.is_complete + waiting = generator.send(batch.get_result()) + assert not waiting.is_complete and len(calls) == 1 + if not replay: + assert host.statuses[-1]["state"] == "waiting_for_human_input" + assert host.statuses[-1]["pending_requests"]["approval"]["data"] == payload + waiting.complete(deepcopy(_UNTRUSTED)) + waiting_again = generator.send(waiting.get_result()) + assert not waiting_again.is_complete and len(calls) == 1 + waiting_again.complete(deepcopy(answer)) + assert _drain(generator, waiting_again.get_result()) == ["done"] + assert [call.args[0] for call in host.wait_for_external_event.call_args_list] == ["approval", "approval"] + assert len(calls) == 3 + assert calls[1]["input"]["source_executor_ids"] == [f"{SOURCE_HITL_RESPONSE}_approval"] + assert calls[2]["input"]["source_executor_ids"] == ["gate"] + if replay: + host.set_custom_status.assert_not_called() + executions.append(calls) + assert executions[0] == executions[1] + assert wire == {_VERSION: 2, "input": payload} diff --git a/python/packages/durabletask/tests/test_workflow_review_followup.py b/python/packages/durabletask/tests/test_workflow_review_followup.py new file mode 100644 index 0000000..e358756 --- /dev/null +++ b/python/packages/durabletask/tests/test_workflow_review_followup.py @@ -0,0 +1,433 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Occurrence ambiguity, serialized forwarding and declared HITL reconstruction.""" + +import json +from collections import defaultdict +from copy import deepcopy +from dataclasses import dataclass +from typing import Any, get_args +from unittest.mock import Mock, patch + +import pytest +from agent_framework import ( + AgentExecutor, + AgentExecutorResponse, + AgentResponse, + AgentSession, + Content, + Executor, + Message, + Workflow, + WorkflowContext, + WorkflowExecutor, + handler, + response_handler, +) +from agent_framework._types import ContentType +from pydantic import BaseModel + +from agent_framework_durabletask._workflows.activity import execute_workflow_activity +from agent_framework_durabletask._workflows.context import WorkflowOrchestrationContext +from agent_framework_durabletask._workflows.orchestrator import ( + SOURCE_HITL_RESPONSE, + TaskType, + _match_occurrences, + _prepare_activity_task, + _prepare_agent_task, + _prepare_subworkflow_task, + _process_activity_result, + _route_result_messages, + _WorkflowDeliveryLedger, +) +from agent_framework_durabletask._workflows.protocol import unwrap_workflow_input +from agent_framework_durabletask._workflows.serialization import ( + SUBWORKFLOW_INPUT_KEY, + deserialize_value, + reconstruct_to_type, + serialize_value, +) + + +class _Agent: + name = "target" + id = "target" + description = None + + def create_session(self, **kwargs: Any) -> AgentSession: + return AgentSession(**kwargs) + + async def run(self, messages: Any = None, **kwargs: Any) -> AgentResponse: + raise AssertionError("These tests schedule tasks without invoking a model") + + +def _agent(**kwargs: Any) -> AgentExecutor: + agent: Any = _Agent() + return AgentExecutor(agent, id="target", **kwargs) + + +def _host(instance_id: str = "review-run") -> Any: + host = Mock(spec=WorkflowOrchestrationContext) + host.instance_id = instance_id + host.prepare_activity_task.side_effect = lambda name, payload: payload + host.call_sub_orchestrator.side_effect = lambda name, payload, **kwargs: payload + return host + + +def _envelope(messages: list[Message], latest: list[Message]) -> AgentExecutorResponse: + return AgentExecutorResponse("producer", AgentResponse(messages=latest), list(messages)) + + +def _dispatch(host: Any, executor: AgentExecutor, source: Any, ledger: _WorkflowDeliveryLedger) -> tuple[Any, Any]: + _prepare_agent_task(host, executor, executor.id, source, "review", ledger) + call = host.prepare_agent_task.call_args + messages = json.loads(json.dumps(call.args[3])) + return messages, call.kwargs["context_message_ids"] + + +@pytest.mark.parametrize("latest_alias", [False, True]) +@pytest.mark.parametrize("selection", [[0], [0, 2]]) +def test_sparse_reused_alias_never_suppresses_a_previously_unsent_position( + latest_alias: bool, selection: list[int] +) -> None: + shared = Message("assistant", ["same"], message_id="opaque") + original = [shared, shared, Message("user", ["other"])] + source = _envelope(original, [shared] if latest_alias else []) + positions = list(selection) + target = _agent(context_mode="custom", context_filter=lambda values: [values[i] for i in positions]) + host, ledger = _host(), _WorkflowDeliveryLedger(instance_id="review-run") + + first, first_ids = _dispatch(host, target, source, ledger) + positions[0] = 1 + second, second_ids = _dispatch(host, target, source, ledger) + + assert first == [original[i].to_dict() for i in selection] + assert second == [shared.to_dict()] + assert len(second_ids) == 1 + assert set(first_ids).isdisjoint(second_ids) + assert shared.message_id == "opaque" + assert source.full_conversation[0] is source.full_conversation[1] is shared + + +@pytest.mark.parametrize("detached", [False, True]) +def test_whole_list_positions_keep_repeated_aliases_and_equal_detached_copies(detached: bool) -> None: + shared = Message("assistant", ["same"], message_id="opaque") + source = _envelope([shared, shared], [shared]) + host, ledger = _host(), _WorkflowDeliveryLedger(instance_id="review-run") + target = _agent(context_mode="custom", context_filter=lambda values: deepcopy(values) if detached else list(values)) + + first, ids = _dispatch(host, target, source, ledger) + repeated, repeated_ids = _dispatch(host, target, source, ledger) + + assert first == [shared.to_dict(), shared.to_dict()] + assert len(set(ids)) == 2 + assert repeated == repeated_ids == [] + + +def test_alias_reused_more_than_source_multiplicity_is_a_new_handoff_occurrence() -> None: + shared = Message("assistant", ["same"], message_id="opaque") + source = _envelope([shared], [shared]) + target = _agent(context_mode="custom", context_filter=lambda values: [values[0], values[0]]) + messages, ids = _dispatch(_host(), target, source, _WorkflowDeliveryLedger()) + assert messages == [shared.to_dict(), shared.to_dict()] + assert len(set(ids)) == 2 + + +def test_equal_but_wrong_position_aliases_are_not_a_whole_list_copy() -> None: + first = Message("user", ["same"]) + second = Message("user", ["same"]) + assert _match_occurrences([first, first], [first, second], ["first", "second"]) == ["first", None] + assert _match_occurrences([first, second], [first, second], ["first", "second"]) == ["first", "second"] + detached = deepcopy(first) + assert _match_occurrences([detached, detached], [first, second], ["first", "second"]) == [None, None] + + +def test_previously_unique_global_alias_does_not_collapse_a_later_repeated_history() -> None: + shared = Message("user", ["same"], message_id="opaque") + host, ledger = _host(), _WorkflowDeliveryLedger(instance_id="review-run") + _dispatch(host, _agent(), _envelope([shared], []), ledger) + repeated = _envelope([shared, shared], []) + sent, ids = _dispatch(host, _agent(), repeated, ledger) + assert sent == [shared.to_dict(), shared.to_dict()] + assert len(set(ids)) == 2 + + +class _Relay(Executor): + def __init__(self, mode: str = "unchanged") -> None: + super().__init__(id="relay") + self.mode = mode + + @handler + async def relay(self, message: AgentExecutorResponse, ctx: WorkflowContext[AgentExecutorResponse]) -> None: + if self.mode == "new-response": + # The same producer, response ID, application ID and text are not an event ID. + latest = deepcopy(message.agent_response.messages) + message = AgentExecutorResponse( + message.executor_id, + AgentResponse(messages=latest, response_id=message.agent_response.response_id), + [*message.full_conversation[: -len(latest)], *latest], + ) + elif self.mode == "replace-message": + # Even reusing the AgentResponse does not prove a replacement is forwarding. + latest = deepcopy(message.agent_response.messages) + message.agent_response.messages = latest + message.full_conversation = [*message.full_conversation[: -len(latest)], *latest] + elif self.mode == "changed-value": + message.agent_response.messages[-1].contents = [Content.from_text("changed")] + await ctx.send_message(message, target_id="target") + + +_ADDRESS = {"root_instance_id": "review-run", "root_workflow_name": "review", "request_path_prefix": ""} + + +@pytest.mark.parametrize("child_dispatch", [False, True]) +def test_failed_forwarding_dispatch_does_not_commit_provenance(child_dispatch: bool) -> None: + latest = Message("assistant", ["same"], message_id="opaque") + source = _envelope([latest], [latest]) + ledger = _WorkflowDeliveryLedger(instance_id="review-run") + host = _host() + host.prepare_activity_task.side_effect = OSError("prepare failed") + host.call_sub_orchestrator.side_effect = OSError("prepare failed") + child = Mock(spec=WorkflowExecutor) + child.workflow = Mock() + child.workflow.name = "inner" + with pytest.raises(OSError, match="prepare failed"): + if child_dispatch: + _prepare_subworkflow_task(host, child, source, "child", _ADDRESS, ledger) + else: + _prepare_activity_task(host, "relay", source, "producer", None, "review", _ADDRESS, ledger) + assert ledger == _WorkflowDeliveryLedger(instance_id="review-run") + assert not hasattr(source.agent_response, "_durable_workflow_forwarding") + assert latest.message_id == "opaque" + + +def _relay_result(host: Any, ledger: _WorkflowDeliveryLedger, source: Any, relay: _Relay) -> Any: + payload = _prepare_activity_task(host, relay.id, source, "producer", None, "review", _ADDRESS, ledger) + raw = execute_workflow_activity(relay, payload) + result = _process_activity_result(raw, relay.id, None, []) + result.source_message = source + return result + + +def _routed(result: Any, ledger: _WorkflowDeliveryLedger) -> AgentExecutorResponse: + workflow = Mock(spec=Workflow) + workflow.edge_groups = [] + pending: dict[str, list[tuple[Any, str]]] = {} + _route_result_messages(result, workflow, pending, defaultdict(dict), ledger) + response = pending["target"][0][0] + assert isinstance(response, AgentExecutorResponse) + return response + + +@pytest.mark.parametrize("mode", ["unchanged", "new-response", "replace-message", "changed-value"]) +@pytest.mark.parametrize("context_mode", ["full", "last_agent"]) +def test_real_activity_checkpoint_forwarding_distinguishes_new_identical_producer_events( + mode: str, context_mode: str +) -> None: + latest = Message("assistant", ["same"], message_id="opaque") + source = _envelope([Message("user", ["question"], message_id="question"), latest], [latest]) + before = [message.to_dict() for message in source.full_conversation] + response_before = source.agent_response.to_dict() + + def replay() -> tuple[Any, Any]: + host, ledger = _host(), _WorkflowDeliveryLedger(instance_id="review-run") + target = _agent(context_mode=context_mode) + _, first_ids = _dispatch(host, target, source, ledger) + result = _relay_result(host, ledger, source, _Relay(mode)) + forwarded = _routed(result, ledger) + assert forwarded is not source + assert forwarded.agent_response is not source.agent_response + assert forwarded.full_conversation[-1] is not latest + sent, sent_ids = _dispatch(host, target, forwarded, ledger) + if mode == "unchanged": + assert sent == sent_ids == [] + else: + expected = Message("assistant", ["changed"], message_id="opaque") if mode == "changed-value" else latest + assert sent == [expected.to_dict()] + assert len(sent_ids) == 1 + assert set(first_ids).isdisjoint(sent_ids) + return sent, sent_ids + + assert replay() == replay() + assert [message.to_dict() for message in source.full_conversation] == before + assert source.agent_response.to_dict() == response_before + assert not hasattr(source.agent_response, "_durable_workflow_forwarding") + + +def test_unchanged_relay_preserves_repeated_anonymous_history_positions() -> None: + shared = Message("user", ["same"]) + latest = Message("assistant", ["answer"]) + source = _envelope([shared, shared, latest], [latest]) + host, ledger = _host(), _WorkflowDeliveryLedger(instance_id="review-run") + _, first_ids = _dispatch(host, _agent(), source, ledger) + forwarded = _routed(_relay_result(host, ledger, source, _Relay()), ledger) + sent, sent_ids = _dispatch(host, _agent(), forwarded, ledger) + assert len(set(first_ids)) == 3 + assert sent == sent_ids == [] + + +@pytest.mark.parametrize("relay_only", [False, True]) +def test_child_keeps_inherited_prefix_receipts_but_scopes_fresh_equal_outputs(relay_only: bool) -> None: + shared = Message("user", ["same"], message_id="shared") + latest = Message("assistant", ["same"], message_id="opaque") + source = _envelope([shared, shared, latest], [latest]) + host, ledger = _host(), _WorkflowDeliveryLedger(instance_id="review-run") + _, original_ids = _dispatch(host, _agent(), source, ledger) + child = Mock(spec=WorkflowExecutor) + child.workflow = Mock() + child.workflow.name = "inner" + new_ids: list[str] = [] + for child_id in ["review-run::child::0", "review-run::child::1"]: + payload = _prepare_subworkflow_task(host, child, source, child_id, _ADDRESS, ledger) + child_input = unwrap_workflow_input(payload) + inherited = deserialize_value(child_input[SUBWORKFLOW_INPUT_KEY]) + assert isinstance(inherited, AgentExecutorResponse) + child_ledger = _WorkflowDeliveryLedger(instance_id=child_id) + if relay_only: + # A real child activity re-dispatch must retain the parent's witness. + result = _relay_result(_host(child_id), child_ledger, inherited, _Relay()) + else: + fresh = Message("assistant", ["same"], message_id="opaque") + produced = _envelope([*inherited.full_conversation, fresh], [fresh]) + result = _process_activity_result( + json.dumps({"sent_messages": [{"message": serialize_value(produced), "target_id": "target"}]}), + "child", + None, + [], + ) + result.source_message = source + result.child_instance_id = child_id + result.task_type = TaskType.SUBWORKFLOW + output = _routed(result, ledger) + sent, sent_ids = _dispatch(host, _agent(), output, ledger) + if relay_only: + assert sent == sent_ids == [] + else: + assert sent == [latest.to_dict()] + assert len(sent_ids) == 1 + assert set(original_ids + new_ids).isdisjoint(sent_ids) + new_ids.extend(sent_ids) + assert [message.message_id for message in source.full_conversation] == ["shared", "shared", "opaque"] + + +@dataclass +class _Request: + prompt: str + + +class _ValidatedReply(BaseModel): + approved: bool + + +class _HumanGate(Executor): + def __init__(self) -> None: + super().__init__(id="human-gate") + self.seen: list[tuple[str | None, Any]] = [] + + @handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + await ctx.request_info(_Request(message), Content, request_id="request-1") + + @response_handler + async def content_reply(self, original_request: _Request, response: Content, ctx: WorkflowContext) -> None: + self.seen.append((ctx.request_id, response)) + + @response_handler + async def message_reply(self, original_request: _Request, response: Message, ctx: WorkflowContext) -> None: + self.seen.append((ctx.request_id, response)) + + @response_handler + async def validated_reply( + self, original_request: _Request, response: _ValidatedReply, ctx: WorkflowContext + ) -> None: + self.seen.append((ctx.request_id, response)) + + +def _hitl_input(value: Any, response_type: type) -> str: + return json.dumps({ + "message": serialize_value({ + "request_id": "request-1", + "original_request": serialize_value(_Request("Review")), + "response": value, + "response_type": f"{response_type.__module__}:{response_type.__name__}", + }), + "source_executor_ids": [f"{SOURCE_HITL_RESPONSE}_request-1"], + }) + + +@pytest.mark.parametrize("reply_type", [Content, Message]) +def test_external_framework_reply_reconstructs_and_receives_request_id(reply_type: type) -> None: + payload: dict[str, Any] = { + "type": "image_generation_tool_result", + "outputs": [{"type": "untrusted.module:Class", "items": [{"type": "application_data"}]}], + "additional_properties": {"opaque": {"type": "untrusted.module:Class"}}, + } + if reply_type is Message: + payload = {"role": "user", "message_id": "application-id", "contents": [payload]} + before = deepcopy(payload) + executor = _HumanGate() + result = json.loads(execute_workflow_activity(executor, _hitl_input(payload, reply_type))) + assert result["pending_request_info_events"] == [] + assert len(executor.seen) == 1 + request_id, reply = executor.seen[0] + assert request_id == "request-1" + assert isinstance(reply, reply_type) + content = reply.contents[0] if isinstance(reply, Message) else reply + assert isinstance(content, Content) + assert content.outputs == [{"type": "untrusted.module:Class", "items": [{"type": "application_data"}]}] + assert content.additional_properties == {"opaque": {"type": "untrusted.module:Class"}} + assert payload == before + + +def test_request_id_is_preserved_for_an_already_supported_pydantic_reply() -> None: + executor = _HumanGate() + execute_workflow_activity(executor, _hitl_input({"approved": True}, _ValidatedReply)) + assert executor.seen == [("request-1", _ValidatedReply(approved=True))] + + +@pytest.mark.parametrize( + ("value", "reply_type"), + [ + pytest.param({"wrong_field": True}, _ValidatedReply, id="invalid-model"), + pytest.param("wrong-runtime-type", Content, id="unmatched-handler"), + pytest.param({"contents": []}, Message, id="missing-message-role"), + pytest.param({"type": ""}, Content, id="invalid-content-type"), + pytest.param({"__pickled__": "not-a-pickle", "__type__": "untrusted:Class"}, Content, id="markers"), + ], +) +def test_invalid_hitl_reply_fails_activity_instead_of_silently_completing(value: Any, reply_type: type) -> None: + executor = _HumanGate() + with pytest.raises((TypeError, ValueError)): + execute_workflow_activity(executor, _hitl_input(value, reply_type)) + assert executor.seen == [] + + +@pytest.mark.parametrize("reply_type", [Content, Message]) +def test_declared_framework_reconstruction_does_not_import_payload_type_names(reply_type: type) -> None: + payload: dict[str, Any] = {"type": "untrusted.module:Class", "additional_properties": {"type": "other:Class"}} + if reply_type is Message: + payload = {"role": "user", "contents": [payload]} + with patch("importlib.import_module", side_effect=AssertionError("Payload type names must remain data")): + restored = reconstruct_to_type(payload, reply_type) + assert isinstance(restored, reply_type) + + +def test_content_known_nested_envelopes_are_rebuilt_but_application_results_stay_dicts() -> None: + payload = { + "type": "function_approval_response", + "approved": True, + "function_call": {"type": "function_call", "call_id": "call", "name": "lookup", "arguments": "{}"}, + "result": {"type": "application_result", "items": [False, None, 0]}, + } + restored = reconstruct_to_type(payload, Content) + assert isinstance(restored, Content) + assert isinstance(restored.function_call, Content) + assert restored.function_call.call_id == "call" + assert restored.result == payload["result"] + + +@pytest.mark.parametrize("content_type", get_args(ContentType)) +def test_all_declared_core_content_kinds_reconstruct_without_a_copied_kind_allowlist(content_type: str) -> None: + restored = reconstruct_to_type({"type": content_type}, Content) + assert isinstance(restored, Content) + assert restored.type == content_type diff --git a/python/packages/durabletask/tests/test_workflow_semantics_review.py b/python/packages/durabletask/tests/test_workflow_semantics_review.py new file mode 100644 index 0000000..8e8ee10 --- /dev/null +++ b/python/packages/durabletask/tests/test_workflow_semantics_review.py @@ -0,0 +1,859 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Logical core conversations are independent of durable transport occurrences.""" + +from __future__ import annotations + +import json +from collections.abc import Callable, Mapping +from copy import deepcopy +from datetime import datetime, timezone +from typing import Any +from unittest.mock import AsyncMock, Mock +from uuid import UUID + +import pytest +from agent_framework import ( + AgentExecutor, + AgentExecutorRequest, + AgentExecutorResponse, + AgentResponse, + AgentSession, + Content, + Executor, + Message, + Workflow, + WorkflowBuilder, + WorkflowExecutor, +) +from agent_framework._workflows._edge import EdgeGroup, FanInEdgeGroup, FanOutEdgeGroup, SingleEdgeGroup +from durabletask.task import CompletableTask, OrchestrationContext + +from agent_framework_durabletask import AgentEntity, AgentEntityStateProviderMixin, RunRequest, serialize_agent_response +from agent_framework_durabletask._message_identity import message_identity +from agent_framework_durabletask._workflows.dt_context import DurableTaskWorkflowContext +from agent_framework_durabletask._workflows.orchestrator import ( + TaskMetadata, + TaskType, + _build_context_messages, + _prepare_agent_task, + _process_agent_response, + _WorkflowDeliveryLedger, + run_workflow_orchestrator, +) +from agent_framework_durabletask._workflows.serialization import ( + SUBWORKFLOW_RESULT_KEY, + deserialize_value, + serialize_value, +) + + +class _Agent: + name = "stub" + id = "stub" + description = None + + def __init__(self, response: AgentResponse | None = None) -> None: + self.response = response if response is not None else AgentResponse(messages=[]) + self.inputs: list[list[dict[str, Any]]] = [] + + def create_session(self, **kwargs: Any) -> AgentSession: + return AgentSession(**kwargs) + + async def run(self, messages: list[Message], **kwargs: Any) -> AgentResponse: + self.inputs.append(_wire(messages)) + return self.response + + +def _agent(name: str, response: AgentResponse | None = None, **kwargs: Any) -> AgentExecutor: + stub: Any = _Agent(response) + return AgentExecutor(stub, id=name, **kwargs) + + +def _wire(messages: list[Message]) -> list[dict[str, Any]]: + return [message.to_dict() for message in messages] + + +def _envelope( + messages: list[Message], producer: str = "source", latest: list[Message] | None = None +) -> AgentExecutorResponse: + return AgentExecutorResponse( + producer, AgentResponse(messages=messages[-1:] if latest is None else latest), messages + ) + + +def _activity(name: str) -> Any: + node = Mock(spec=Executor) + node.id = name + node.input_types = [str] + return node + + +def _workflow(nodes: list[Any], edges: list[EdgeGroup]) -> Any: + workflow = Mock(spec=Workflow) + workflow.name = "review" + workflow.start_executor_id = nodes[0].id + workflow.executors = {node.id: node for node in nodes} + workflow.edge_groups = edges + workflow.max_iterations = 30 + return workflow + + +def _send(messages: list[Any], target: str | None = None, *, wait: bool = False) -> dict[str, Any]: + result: dict[str, Any] = { + "sent_messages": [{"message": serialize_value(message), "target_id": target} for message in messages] + } + if wait: + result["pending_request_info_events"] = [ + {"request_id": "approval", "source_executor_id": "gate", "data": "review"} + ] + return result + + +class _Host: + supports_event_streaming = False + current_utc_datetime = datetime(2026, 1, 1, tzinfo=timezone.utc) + + def __init__( + self, + responses: Mapping[str, AgentResponse | dict[str, Any]] | None = None, + *, + activities: dict[str, list[dict[str, Any]]] | None = None, + children: list[Any] | None = None, + instance_id: str = "run", + is_replaying: bool = False, + ) -> None: + self.instance_id = instance_id + self.is_replaying = is_replaying + self.calls: list[dict[str, Any]] = [] + self.responses = responses or {} + self.activities = {key: iter(values) for key, values in (activities or {}).items()} + self.children = iter(children or []) + self.child_ids: list[str | None] = [] + self.waits: list[str] = [] + self.batches: list[int] = [] + self.fail_prepare = False + + def prepare_agent_task( + self, + executor_id: str, + message: str, + orchestration_instance_id: str, + context_messages: list[dict[str, Any]] | None = None, + context_message_ids: list[str] | None = None, + ) -> AgentResponse | dict[str, Any]: + assert (context_messages is None) == (context_message_ids is None) + if context_messages is not None: + assert context_message_ids is not None + assert len(context_messages) == len(context_message_ids) + self.calls.append( + json.loads( + json.dumps( + { + "executor": executor_id, + "instance": orchestration_instance_id, + "message": message, + "contextMessages": context_messages, + "contextMessageIds": context_message_ids, + }, + allow_nan=False, + ) + ) + ) + if self.fail_prepare: + raise OSError("prepare failed") + return self.responses.get(executor_id, AgentResponse(messages=[Message("assistant", ["approved"])])) + + def prepare_activity_task(self, activity_name: str, input_json: str) -> str: + return json.dumps(next(self.activities[json.loads(input_json)["executor_id"]])) + + def call_sub_orchestrator(self, name: str, input: Any, instance_id: str | None = None) -> Any: + self.child_ids.append(instance_id) + return next(self.children) + + def task_all(self, tasks: list[Any]) -> list[Any]: + self.batches.append(len(tasks)) + return tasks + + def task_any(self, tasks: list[Any]) -> Any: + raise AssertionError("These workflows do not race tasks") + + def set_custom_status(self, status: Any) -> None: + pass + + def wait_for_external_event(self, name: str) -> str: + self.waits.append(name) + return "approved" + + def create_timer(self, fire_at: datetime) -> Any: + raise AssertionError("These workflows have no timers") + + def new_uuid(self) -> str: + raise AssertionError("Message identity must not require UUIDs") + + def cancel_task(self, task: Any) -> None: + raise AssertionError("These workflows do not cancel tasks") + + def get_task_result(self, task: Any) -> Any: + return task + + +def _run(host: _Host, workflow: Any, message: Any = "start") -> Any: + generator = run_workflow_orchestrator(host, workflow, message) + result: Any = None + while True: + try: + result = generator.send(result) + except StopIteration as completed: + return completed.value + + +def _turn( + host: _Host, executor: AgentExecutor, message: Any, ledger: _WorkflowDeliveryLedger +) -> tuple[dict[str, Any], AgentExecutorResponse]: + metadata = TaskMetadata(executor.id, message, "source", TaskType.AGENT) + result = _prepare_agent_task(host, executor, executor.id, message, "review", ledger, metadata) + response = _process_agent_response(result, executor.id, message, ledger, metadata).output_message + assert response is not None + return host.calls[-1], response + + +def _ids(call: dict[str, Any]) -> list[str]: + ids = call["contextMessageIds"] + assert isinstance(ids, list) + assert len(ids) == len(call["contextMessages"]) + assert all(isinstance(value, str) and value for value in ids) + return ids + + +def _texts(call: dict[str, Any]) -> list[str]: + return [Message.from_dict(message).text for message in call["contextMessages"]] + + +async def _core_turn(executor: AgentExecutor, message: Any) -> AgentExecutorResponse: + context = Mock() + context.source_executor_ids = ["source"] + context.is_streaming.return_value = False + context.get_state.return_value = {} + context.send_message = AsyncMock() + context.yield_output = AsyncMock() + if isinstance(message, AgentExecutorResponse): + await executor.from_response(message, context) + elif isinstance(message, str): + await executor.from_str(message, context) + elif isinstance(message, AgentExecutorRequest): + await executor.run(message, context) + elif isinstance(message, Message): + await executor.from_message(message, context) + else: + await executor.from_messages(message, context) + return context.send_message.call_args.args[0] + + +def _redact(messages: list[Message]) -> list[Message]: + selected = Message.from_dict(next(message for message in messages if message.message_id == "opaque").to_dict()) + selected.contents = [Content.from_text("redacted")] + return [selected, Message("system", ["summary"], additional_properties={"_is_summary": True})] + + +@pytest.mark.parametrize("mode", ["full", "last_agent", "custom"]) +@pytest.mark.parametrize("serialized", [False, True]) +async def test_three_hops_match_real_core_selected_cache_and_actual_response(mode: str, serialized: bool) -> None: + responses = { + "A": AgentResponse(messages=[Message("assistant", ["secret"], message_id="opaque")]), + "B": AgentResponse( + messages=[ + Message( + "assistant", [Content.from_function_call("call", "lookup", arguments={"id": 7})], message_id="call" + ), + Message("tool", [Content.from_function_result("call", result={"answer": 42})], message_id="tool"), + Message( + "assistant", + [Content.from_uri("https://example.com/image.png", media_type="image/png")], + message_id="picture", + additional_properties={"label": "original"}, + ), + ], + response_id="original-response", + additional_properties={"provider": {"opaque": True}}, + ), + "C": AgentResponse(messages=[]), + } + options: dict[str, Any] = {"context_mode": mode, "context_filter": _redact if mode == "custom" else None} + core_a = await _core_turn(_agent("A", responses["A"]), "question") + core_b = await _core_turn(_agent("B", responses["B"], **options), core_a) + core_c = await _core_turn(_agent("C", responses["C"]), core_b) + before = {name: response.to_dict() for name, response in responses.items()} + seen: list[AgentExecutorResponse] = [] + + def capture(response: AgentExecutorResponse) -> bool: + seen.append(response) + return True + + workflow = _workflow( + [_agent("A"), _agent("B", **options), _agent("C")], + [SingleEdgeGroup("A", "B"), SingleEdgeGroup("B", "C", condition=capture)], + ) + payloads = { + f"review-{name}": serialize_agent_response(response) if serialized else response + for name, response in responses.items() + } + host = _Host(payloads) + assert _run(host, workflow, "question") == [] + assert host.calls[1]["contextMessages"] == _wire(core_b.full_conversation[: -len(responses["B"].messages)]) + assert host.calls[2]["contextMessages"] == _wire(core_c.full_conversation) + assert seen[0].agent_response.to_dict() == responses["B"].to_dict() + assert _wire(seen[0].full_conversation) == _wire(core_b.full_conversation) + if not serialized: + assert seen[0].agent_response is responses["B"] + if mode == "custom": + assert "secret" not in json.dumps(host.calls[2]) + assert _texts(host.calls[2])[:2] == ["redacted", "summary"] + assert {name: response.to_dict() for name, response in responses.items()} == before + + +def test_serialized_response_tolerates_unknown_delivery_fields_without_mutating_payload() -> None: + response = AgentResponse( + messages=[Message("assistant", ["approved"], message_id="opaque")], + response_id="actual-response", + additional_properties={"provider": {"opaque": True}}, + ) + payload = serialize_agent_response(response) + payload["future_delivery_metadata"] = {"type": "provider_extension", "opaque": [1]} + before = deepcopy(payload) + _, outgoing = _turn(_Host({"review-A": payload}), _agent("A"), "question", _WorkflowDeliveryLedger()) + assert outgoing.agent_response.to_dict() == response.to_dict() + assert outgoing.full_conversation[-1].message_id == "opaque" + assert payload == before + + +@pytest.mark.parametrize("application_id", ["opaque", "wf_source_0", "wf:external:" + "a" * 64]) +def test_later_filter_observes_original_application_ids_not_transport_namespaces(application_id: str) -> None: + original = Message("assistant", ["approved"], message_id=application_id, additional_properties={"nested": [1]}) + upstream = _envelope([original]) + host, ledger = _Host(), _WorkflowDeliveryLedger() + _, outgoing = _turn(host, _agent("B"), upstream, ledger) + executor = _agent( + "C", + context_mode="custom", + context_filter=lambda messages: [m for m in messages if m.message_id == application_id], + ) + call, _ = _turn(host, executor, outgoing, ledger) + assert call["contextMessages"] == [original.to_dict()] + assert _ids(call) == _ids(host.calls[0]) + assert _build_context_messages(executor, outgoing) == [original.to_dict()] + assert outgoing.full_conversation[0] is original + + +def test_selected_full_context_not_just_delta_becomes_the_next_conversation() -> None: + messages = [Message("user", [str(i)], message_id=f"app-{i}") for i in range(4)] + upstream = _envelope(messages) + host, ledger = _Host(), _WorkflowDeliveryLedger() + first_executor = _agent("B", context_mode="custom", context_filter=lambda values: [values[1], values[3]]) + _turn(host, first_executor, upstream, ledger) + next_executor = _agent("B", context_mode="custom", context_filter=lambda values: [values[2], values[0], values[3]]) + call, outgoing = _turn(host, next_executor, upstream, ledger) + assert _texts(call) == ["2", "0"] + assert _wire(outgoing.full_conversation[:-1]) == _wire([messages[2], messages[0], messages[3]]) + repeated, _ = _turn(host, _agent("B"), upstream, ledger) + assert repeated["contextMessages"] == _ids(repeated) == [] + assert repeated["message"] == "" + downstream, _ = _turn(host, _agent("C"), outgoing, ledger) + assert _texts(downstream) == ["2", "0", "3", "approved"] + + +def test_detached_copy_updates_keep_occurrence_but_changed_fingerprint_is_delivered() -> None: + original = Message("assistant", ["secret"], message_id="opaque") + upstream = _envelope([original]) + host, ledger = _Host(), _WorkflowDeliveryLedger() + first, _ = _turn(host, _agent("B"), upstream, ledger) + redacted = _agent("B", context_mode="custom", context_filter=lambda messages: _redact(messages)[:1]) + changed, outgoing = _turn(host, redacted, upstream, ledger) + repeated, _ = _turn(host, redacted, upstream, ledger) + assert _ids(changed) == _ids(first) + assert _texts(changed) == ["redacted"] + assert changed["contextMessages"][0]["message_id"] == "opaque" + assert repeated["contextMessages"] == _ids(repeated) == [] + assert message_identity(original) != message_identity(outgoing.full_conversation[0]) + assert original.text == "secret" + + +@pytest.mark.parametrize("detached", [False, True]) +def test_source_selection_order_and_copies_have_parallel_occurrence_ids(detached: bool) -> None: + messages = [Message("user", [str(i)]) for i in range(4)] + source = _envelope(messages, latest=[]) + + def projection(indices: list[int]) -> Callable[[list[Message]], list[Message]]: + return lambda values: [Message.from_dict(values[i].to_dict()) if detached else values[i] for i in indices] + + host, ledger = _Host(), _WorkflowDeliveryLedger() + first, _ = _turn(host, _agent("B", context_mode="custom", context_filter=projection([1, 3])), source, ledger) + second, _ = _turn(host, _agent("B", context_mode="custom", context_filter=projection([2, 0, 3])), source, ledger) + other, _ = _turn(host, _agent("C", context_mode="custom", context_filter=projection([3, 1])), source, ledger) + assert _texts(first) == ["1", "3"] + assert _texts(second) == ["2", "0"] + assert _ids(other) == list(reversed(_ids(first))) + assert len(set(_ids(first) + _ids(second))) == 4 + assert all(message.message_id is None for message in messages) + + +def test_synthesized_detached_messages_are_handoff_scoped_even_with_equal_ids() -> None: + source = _envelope([Message("user", ["source"])]) + executor = _agent( + "B", + context_mode="custom", + context_filter=lambda _: [ + Message("system", ["summary"], message_id="summary"), + Message("system", ["summary"], message_id="summary"), + ], + ) + + def replay() -> list[dict[str, Any]]: + host, ledger = _Host(), _WorkflowDeliveryLedger() + for _ in range(2): + _turn(host, executor, deepcopy(source), ledger) + return host.calls + + calls = replay() + assert len(set(_ids(calls[0]) + _ids(calls[1]))) == 4 + assert calls == replay() + assert all(message["message_id"] == "summary" for call in calls for message in call["contextMessages"]) + + +@pytest.mark.parametrize("kind", ["raw", "anonymous", "same-id"]) +@pytest.mark.parametrize("pause", [False, True]) +def test_independent_producer_events_do_not_collide_across_sequential_or_hitl_dispatch(kind: str, pause: bool) -> None: + values: list[Any] = ( + ["same prompt", "same prompt"] + if kind == "raw" + else [ + _envelope([Message("assistant", ["approved"], message_id="opaque" if kind == "same-id" else None)]) + for _ in range(2) + ] + ) + activities = { + "gate": [_send(values[:1], "A", wait=True), _send(values[1:], "A")] if pause else [_send(values, "A")] + } + workflow = _workflow([_activity("gate"), _agent("A"), _agent("B")], [SingleEdgeGroup("A", "B")]) + live, replay = _Host(activities=activities), _Host(activities=activities, is_replaying=True) + assert _run(live, workflow) == _run(replay, workflow) == [] + assert live.calls == replay.calls + consumers = [call for call in live.calls if call["executor"] == "review-B"] + assert len(consumers) == 2 + assert _texts(consumers[0]) == _texts(consumers[1]) == ["same prompt" if kind == "raw" else "approved", "approved"] + assert set(_ids(consumers[0])).isdisjoint(_ids(consumers[1])) + assert live.waits == (["approval"] if pause else []) + + +def test_child_invocations_scope_equal_logical_ids_at_the_child_boundary() -> None: + child = Mock(spec=WorkflowExecutor) + child.id = "child" + child.workflow = Mock(name="inner") + child.workflow.name = "inner" + child.allow_direct_output = False + outputs = [_envelope([Message("assistant", ["approved"], message_id="wf_inner_0")], "inner") for _ in range(2)] + children = [ + {SUBWORKFLOW_RESULT_KEY: True, "outputs": [serialize_value(output)], "events": []} for output in outputs + ] + workflow = _workflow([_activity("gate"), child, _agent("B")], [SingleEdgeGroup("child", "B")]) + activities = {"gate": [_send(["same", "same"], "child")]} + host, replay = _Host(activities=activities, children=children), _Host(activities=activities, children=children) + assert _run(host, workflow) == _run(replay, workflow) == [] + assert host.child_ids == ["run::child::0", "run::child::1"] + assert host.calls == replay.calls + assert _texts(host.calls[0]) == _texts(host.calls[1]) == ["approved"] + assert set(_ids(host.calls[0])).isdisjoint(_ids(host.calls[1])) + assert [call["contextMessages"][0]["message_id"] for call in host.calls] == ["wf_inner_0"] * 2 + + +def test_activity_forwarded_copies_reuse_source_positions_but_new_output_is_an_event() -> None: + original = Message("user", ["question"], message_id="question") + source = _envelope([original, Message("assistant", ["approved"], message_id="opaque")], "A") + host, ledger = _Host(), _WorkflowDeliveryLedger(instance_id="run") + first, _ = _turn(host, _agent("B"), source, ledger) + transformed = _envelope( + [Message.from_dict(original.to_dict()), Message("assistant", ["approved"], message_id="opaque")], "A" + ) + # This is the association routing makes between an activity's input and output. + ledger.identify(transformed, source) + second, _ = _turn(host, _agent("B"), transformed, ledger) + assert _texts(first) == ["question", "approved"] + assert _texts(second) == ["approved"] + assert set(_ids(first)).isdisjoint(_ids(second)) + + +def test_fanin_keeps_logical_selected_order_and_deduplicates_transport_per_target() -> None: + responses = {"review-A": AgentResponse(messages=[Message("assistant", ["secret"], message_id="opaque")])} + seen: list[AgentExecutorResponse] = [] + + def capture(response: AgentExecutorResponse) -> bool: + seen.append(response) + return True + + workflow = _workflow( + [ + _agent("A"), + _agent("left", context_mode="custom", context_filter=_redact), + _agent("right", context_mode="custom", context_filter=_redact), + _agent("join"), + _agent("end"), + ], + [ + FanOutEdgeGroup("A", ["left", "right"]), + FanInEdgeGroup(["left", "right"], "join"), + SingleEdgeGroup("join", "end", condition=capture), + ], + ) + host = _Host(responses) + assert _run(host, workflow) == [] + joined = next(call for call in host.calls if call["executor"] == "review-join") + assert _texts(joined) == ["redacted", "summary", "approved", "summary", "approved"] + assert [message.text for message in seen[0].full_conversation] == [ + "redacted", + "summary", + "approved", + "redacted", + "summary", + "approved", + "approved", + ] + assert "secret" not in json.dumps(joined) + assert len(_ids(joined)) == 5 + + +@pytest.mark.parametrize("kind", ["string", "message", "messages", "mixed", "request"]) +async def test_all_core_input_handlers_keep_all_messages_and_contents(kind: str) -> None: + message = Message( + "tool", + [Content.from_function_result("call", result={"data": [0, False, None]})], + message_id="opaque", + additional_properties={"source": "app"}, + ) + inputs: dict[str, Any] = { + "string": "hello", + "message": message, + "messages": [message, message], + "mixed": ["hello", message], + "request": AgentExecutorRequest([message, message]), + } + core = await _core_turn(_agent("A"), inputs[kind]) + host = _Host({"review-A": AgentResponse(messages=[])}) + workflow = _workflow([_agent("A"), _agent("B")], [SingleEdgeGroup("A", "B")]) + assert _run(host, workflow, inputs[kind]) == [] + assert host.calls[1]["contextMessages"] == _wire(core.full_conversation) + assert len(set(_ids(host.calls[1]))) == len(core.full_conversation) + assert message.message_id == "opaque" + + +@pytest.mark.parametrize("value", [None, [], AgentExecutorRequest([])]) +def test_empty_inputs_are_explicit_empty_context_not_text_fallback(value: Any) -> None: + host = _Host() + assert _run(host, _workflow([_agent("A")], []), value) == [] + assert host.calls[0]["contextMessages"] == _ids(host.calls[0]) == [] + assert host.calls[0]["message"] == "" + + +@pytest.mark.parametrize("pause", [False, True]) +@pytest.mark.parametrize("prior_run", [False, True]) +def test_cache_only_requests_schedule_nothing_and_flush_all_messages_on_next_run(pause: bool, prior_run: bool) -> None: + cached = AgentExecutorRequest([Message("system", ["rules"], message_id="rules"), Message("user", ["draft"])], False) + prefix: list[Any] = ["first"] if prior_run else [] + batches = ( + [_send([*prefix, cached], "A", wait=True), _send(["answer"], "A")] + if pause + else [_send([*prefix, cached, "answer"], "A")] + ) + host = _Host(activities={"gate": batches}) + workflow = _workflow([_activity("gate"), _agent("A")], []) + assert _run(host, workflow) == [] + assert len(host.calls) == 1 + int(prior_run) + assert _texts(host.calls[-1]) == ["rules", "draft", "answer"] + assert len(set(_ids(host.calls[-1]))) == 3 + assert host.waits == (["approval"] if pause else []) + + +def test_cache_only_workflow_does_not_yield_a_model_task() -> None: + host = _Host() + assert _run(host, _workflow([_agent("A")], []), AgentExecutorRequest([Message("user", ["later"])], False)) == [] + assert host.calls == host.batches == [] + + +@pytest.mark.parametrize("failure", ["prepare", "serialization", "filter"]) +def test_failed_preparation_does_not_consume_delivery_or_occurrence_ordinals(failure: str) -> None: + source = _envelope([Message("user", ["question"])]) + host, ledger = _Host(), _WorkflowDeliveryLedger() + + def projection(messages: list[Message]) -> list[Message]: + if failure == "filter": + raise ValueError("filter failed") + if failure == "serialization": + return [Message("system", ["summary"], additional_properties={"bad": float("nan")})] + return messages + + host.fail_prepare = failure == "prepare" + with pytest.raises((TypeError, ValueError, OSError)): + _turn(host, _agent("B", context_mode="custom", context_filter=projection), source, ledger) + assert ledger == _WorkflowDeliveryLedger() + host.fail_prepare = False + call, _ = _turn(host, _agent("B"), source, ledger) + clean, _ = _turn(_Host(), _agent("B"), deepcopy(source), _WorkflowDeliveryLedger()) + assert call == clean + + +@pytest.mark.parametrize("mode", ["full", "last_agent", "custom"]) +def test_empty_projection_never_leaks_unselected_context_into_next_hop(mode: str) -> None: + secret = Message("assistant", ["secret"]) + source = _envelope([] if mode == "full" else [secret], latest=[] if mode == "last_agent" else [secret]) + executor = _agent("B", context_mode=mode, context_filter=(lambda _: []) if mode == "custom" else None) + host, ledger = _Host(), _WorkflowDeliveryLedger() + first, outgoing = _turn(host, executor, source, ledger) + assert first["contextMessages"] == _ids(first) == [] + assert first["message"] == "" + second, _ = _turn(host, _agent("C"), outgoing, ledger) + assert _texts(second) == ["approved"] + assert "secret" not in json.dumps(second) + + +def test_growing_conversations_send_only_new_messages_and_parallel_ids() -> None: + host, ledger = _Host(), _WorkflowDeliveryLedger() + source: Any = "question" + for _ in range(80): + _, source = _turn(host, _agent("A"), source, ledger) + call, _ = _turn(host, _agent("B"), source, ledger) + assert _texts(call) == ["approved"] + assert len(_ids(call)) == 1 + assert set(call) == {"executor", "instance", "message", "contextMessages", "contextMessageIds"} + assert len(json.dumps(call)) < 500 + assert len(json.dumps(_wire(source.full_conversation))) > 5 * len(json.dumps(call)) + + +@pytest.mark.parametrize("sequential", [False, True]) +@pytest.mark.parametrize("status", ["error", "already_completed"]) +def test_terminal_results_stop_routing_before_next_pending_model_call(sequential: bool, status: str) -> None: + host = _Host(activities={"gate": [_send(["first", "second", "must not run"], "A")]}) + workflow = _workflow([_activity("gate"), _agent("A"), _agent("B")], [SingleEdgeGroup("A", "B")]) + generator = run_workflow_orchestrator(host, workflow, "start") + yielded = generator.send(next(generator)) + if sequential: + generator.send(yielded) + error = AgentResponse(messages=[], additional_properties={"durable_status": status}).to_dict() + error["unknown_field"] = {"private": "must not deserialize"} + with pytest.raises(RuntimeError, match="expired durable response|terminal runtime error"): + generator.send(error if sequential else [error]) + assert [call["executor"] for call in host.calls] == ["review-A"] * (2 if sequential else 1) + + +def test_application_identity_survives_activity_serialization() -> None: + original = Message("assistant", ["approved"], message_id="opaque", additional_properties={"app": [1]}) + restored = deserialize_value(json.loads(json.dumps(serialize_value(_envelope([original]))))) + host, ledger = _Host(), _WorkflowDeliveryLedger() + call, outgoing = _turn(host, _agent("B"), restored, ledger) + assert call["contextMessages"] == [original.to_dict()] + assert outgoing.full_conversation[0].to_dict() == original.to_dict() + + +def test_exact_whole_list_copy_preserves_repeated_anonymous_positions() -> None: + source = _envelope([Message("user", ["same"]), Message("user", ["same"])], latest=[]) + host, ledger = _Host(), _WorkflowDeliveryLedger() + first, _ = _turn(host, _agent("B"), source, ledger) + copied, _ = _turn( + host, _agent("B", context_mode="custom", context_filter=lambda messages: deepcopy(messages)), source, ledger + ) + assert len(set(_ids(first))) == 2 + assert copied["contextMessages"] == _ids(copied) == [] + + +def test_last_agent_uses_output_occurrence_when_same_object_is_also_input() -> None: + shared = Message("assistant", ["same"], message_id="opaque") + host = _Host({"review-A": AgentResponse(messages=[shared])}) + ledger = _WorkflowDeliveryLedger() + _, outgoing = _turn(host, _agent("A"), AgentExecutorRequest([shared]), ledger) + full, _ = _turn(host, _agent("B"), outgoing, ledger) + latest, _ = _turn(host, _agent("C", context_mode="last_agent"), outgoing, ledger) + assert len(set(_ids(full))) == 2 + assert _ids(latest) == _ids(full)[1:] + assert shared.message_id == "opaque" + + +def test_workflow_instance_scopes_occurrences_without_changing_logical_messages() -> None: + workflow = _workflow([_agent("A"), _agent("B")], [SingleEdgeGroup("A", "B")]) + first, second = _Host(instance_id="first"), _Host(instance_id="second") + assert _run(first, workflow) == _run(second, workflow) == [] + assert first.calls[1]["contextMessages"] == second.calls[1]["contextMessages"] + assert set(_ids(first.calls[1])).isdisjoint(_ids(second.calls[1])) + + +def test_projection_can_exclude_non_json_source_messages() -> None: + excluded = Message("assistant", ["secret"], additional_properties={"invalid": float("nan")}) + selected = Message("user", ["selected"]) + source = _envelope([excluded, selected]) + executor = _agent( + "B", context_mode="custom", context_filter=lambda messages: [Message.from_dict(messages[-1].to_dict())] + ) + call, _ = _turn(_Host(), executor, source, _WorkflowDeliveryLedger()) + assert call["contextMessages"] == [selected.to_dict()] + assert len(_ids(call)) == 1 + assert "secret" not in json.dumps(call) + + +def test_empty_string_retains_its_user_message_without_falling_back_to_none() -> None: + host = _Host() + assert _run(host, _workflow([_agent("A")], []), "") == [] + assert host.calls[0]["contextMessages"] == [Message("user", [""]).to_dict()] + assert len(_ids(host.calls[0])) == 1 + assert host.calls[0]["message"] == "" + + +def test_reused_output_alias_retains_only_two_ambiguity_witnesses() -> None: + shared = Message("assistant", ["same"], message_id="opaque") + host = _Host({"review-A": AgentResponse(messages=[shared])}) + ledger = _WorkflowDeliveryLedger() + calls: list[dict[str, Any]] = [] + for _ in range(12): + _, output = _turn(host, _agent("A"), "question", ledger) + call, _ = _turn(host, _agent("B", context_mode="last_agent"), output, ledger) + calls.append(call) + assert len(ledger.aliases[id(shared)][1]) == 2 + assert len({identity for call in calls for identity in _ids(call)}) == len(calls) + assert all(call["contextMessages"] == [shared.to_dict()] for call in calls) + assert shared.message_id == "opaque" + + +class _JsonStateProvider(AgentEntityStateProviderMixin): + def __init__(self, name: str) -> None: + self.name = name + self.raw: dict[str, Any] = {} + + def _get_state_dict(self) -> dict[str, Any]: + return json.loads(json.dumps(self.raw, allow_nan=False)) + + def _set_state_dict(self, state: dict[str, Any]) -> None: + self.raw = json.loads(json.dumps(state, allow_nan=False)) + + def _get_session_id_from_entity(self) -> str: + return "adapter-run" + + def _get_entity_name_from_entity(self) -> str: + return self.name + + +@pytest.mark.parametrize("adapter", ["dt", "af"]) +@pytest.mark.parametrize("mode", ["full", "last_agent", "custom"]) +async def test_real_adapter_three_hops_preserve_selected_context_through_receiver(adapter: str, mode: str) -> None: + responses = { + "A": AgentResponse(messages=[Message("assistant", ["secret"], message_id="opaque")]), + "B": AgentResponse( + messages=[ + Message( + "assistant", + [Content.from_function_call("call", "lookup", arguments={"id": 7})], + message_id="same-id", + ), + Message("tool", [Content.from_function_result("call", result={"answer": 42})], message_id="same-id"), + Message( + "assistant", + [Content.from_uri("https://example.com/image.png", media_type="image/png")], + message_id="same-id", + additional_properties={"label": "original"}, + ), + ], + response_id="actual-response", + additional_properties={"provider": {"opaque": True}}, + ), + "C": AgentResponse(messages=[]), + } + options: dict[str, Any] = {"context_mode": mode, "context_filter": _redact if mode == "custom" else None} + core_a = await _core_turn(_agent("A", responses["A"]), "question") + core_b = await _core_turn(_agent("B", responses["B"], **options), core_a) + core_c = await _core_turn(_agent("C", responses["C"]), core_b) + expected_b = _wire(core_b.full_conversation[: -len(responses["B"].messages)]) + expected_c = _wire(core_c.full_conversation) + original_responses = {name: response.to_dict() for name, response in responses.items()} + agents: dict[str, Any] = {name: _Agent(response) for name, response in responses.items()} + providers = {name: _JsonStateProvider(name) for name in agents} + entities = {name: AgentEntity(agent, state_provider=providers[name]) for name, agent in agents.items()} + observed: list[AgentExecutorResponse] = [] + + def capture(response: AgentExecutorResponse) -> bool: + observed.append(response) + return True + + a, b, c = _agent("A"), _agent("B", **options), _agent("C") + workflow = ( + WorkflowBuilder(name="review", start_executor=a, output_from=[c]) + .add_edge(a, b) + .add_edge(b, c, condition=capture) + .build() + ) + if adapter == "dt": + native = Mock(spec=OrchestrationContext) + children: list[Any] = [CompletableTask() for _ in agents] + context: Any = DurableTaskWorkflowContext(native) + else: + df = pytest.importorskip("azure.durable_functions") + af_context = pytest.importorskip("agent_framework_azurefunctions._workflow_af_context") + from azure.durable_functions.models.actions.NoOpAction import NoOpAction + from azure.durable_functions.models.Task import AtomicTask + + native = Mock(spec=df.DurableOrchestrationContext) + children = [AtomicTask(index, NoOpAction()) for index in range(len(agents))] + native.task_all.side_effect = lambda tasks: tasks + context = af_context.AzureFunctionsWorkflowContext(native) + native.instance_id = "adapter-run" + native.is_replaying = False + native.current_utc_datetime = datetime(2026, 1, 1, tzinfo=timezone.utc) + native.new_uuid.side_effect = [str(UUID(int=index + 1)) for index in range(len(agents))] + native.call_entity.side_effect = children + orchestration = run_workflow_orchestrator(context, workflow, "question") + yielded = next(orchestration) + wires: list[dict[str, Any]] = [] + for index, name in enumerate(agents): + assert native.call_entity.call_count == index + 1 + entity_id, operation, payload = native.call_entity.call_args.args + expected_name = f"dafx-review-{name}".lower() if adapter == "dt" else f"dafx-review-{name}" + assert (entity_id.entity if adapter == "dt" else entity_id.name) == expected_name + assert entity_id.key == "adapter-run" + assert operation == "run" + wire = json.loads(json.dumps(payload, allow_nan=False)) + wires.append(wire) + response = await entities[name].run(wire) + result = json.loads(json.dumps(serialize_agent_response(response), allow_nan=False)) + assert result == original_responses[name] + if adapter == "dt": + children[index].complete(result) + completed_results = context.get_task_result(yielded) + else: + children[index].set_value(is_error=False, value=result) + completed_results = [context.get_task_result(task) for task in yielded] + assert len(completed_results) == 1 + assert completed_results[0].to_dict() == result + if name == "C": + with pytest.raises(StopIteration) as completed: + orchestration.send(completed_results) + # C is the designated output executor, even when its original response has no messages. + final_outputs = [deserialize_value(output) for output in completed.value.value] + assert len(final_outputs) == 1 + assert isinstance(final_outputs[0], AgentResponse) + assert final_outputs[0].to_dict() == original_responses["C"] + else: + yielded = orchestration.send(completed_results) + + for index, (name, expected) in enumerate([("B", expected_b), ("C", expected_c)], start=1): + wire = wires[index] + assert wire["contextMessages"] == expected + assert len(wire["contextMessageIds"]) == len(expected) + assert RunRequest.from_dict(wire).context_message_ids == wire["contextMessageIds"] + assert agents[name].inputs == [expected] + receipts = providers[name].raw["data"]["ingestedMessages"] + assert receipts == { + identity: [message_identity(Message.from_dict(message))] + for identity, message in zip(wire["contextMessageIds"], expected, strict=True) + } + assert wires[2]["contextMessageIds"][: len(expected_b)] == wires[1]["contextMessageIds"] + assert len(set(wires[2]["contextMessageIds"][-3:])) == 3 + assert [message.get("message_id") for message in wires[2]["contextMessages"][-3:]] == ["same-id"] * 3 + assert observed[0].agent_response.to_dict() == original_responses["B"] + assert _wire(observed[0].full_conversation) == expected_c + assert {name: response.to_dict() for name, response in responses.items()} == original_responses + if mode == "custom": + assert "secret" not in json.dumps(wires[2]) diff --git a/python/pyproject.toml b/python/pyproject.toml index 3aba926..879cd2f 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -47,6 +47,8 @@ dev = [ ] test = [ "azure-monitor-opentelemetry", + # Validates that persisted entity state matches the shared cross-language schema. + "jsonschema", "mcp[ws]", "redis", # Model provider packages used by the sample workers that the integration diff --git a/python/samples/08_workflow/worker.py b/python/samples/08_workflow/worker.py index eda7890..244f56d 100644 --- a/python/samples/08_workflow/worker.py +++ b/python/samples/08_workflow/worker.py @@ -149,7 +149,7 @@ def create_workflow() -> Workflow: email_sender = EmailSenderExecutor(id="email_sender") return ( - WorkflowBuilder(name=WORKFLOW_NAME, start_executor=spam_agent) + WorkflowBuilder(name=WORKFLOW_NAME, start_executor=spam_agent, output_from=[spam_handler, email_sender]) .add_switch_case_edge_group( spam_agent, [ diff --git a/python/samples/09_workflow_hitl/worker.py b/python/samples/09_workflow_hitl/worker.py index ccd1c53..226f93c 100644 --- a/python/samples/09_workflow_hitl/worker.py +++ b/python/samples/09_workflow_hitl/worker.py @@ -286,7 +286,7 @@ def create_workflow() -> Workflow: publish_executor = PublishExecutor() return ( - WorkflowBuilder(name=WORKFLOW_NAME, start_executor=input_router) + WorkflowBuilder(name=WORKFLOW_NAME, start_executor=input_router, output_from=[publish_executor]) .add_edge(input_router, content_analyzer_agent) .add_edge(content_analyzer_agent, content_analyzer_executor) .add_edge(content_analyzer_executor, human_review_executor) diff --git a/python/samples/10_workflow_streaming/worker.py b/python/samples/10_workflow_streaming/worker.py index 8ecc4c1..cb29856 100644 --- a/python/samples/10_workflow_streaming/worker.py +++ b/python/samples/10_workflow_streaming/worker.py @@ -86,7 +86,7 @@ def create_workflow() -> Workflow: publish = PublishExecutor(id="publish") return ( - WorkflowBuilder(start_executor=writer_agent) + WorkflowBuilder(start_executor=writer_agent, output_from=[publish]) .add_edge(writer_agent, reviewer_agent) .add_edge(reviewer_agent, publish) .build() diff --git a/python/samples/11_subworkflow/worker.py b/python/samples/11_subworkflow/worker.py index 2e996d2..7bf306d 100644 --- a/python/samples/11_subworkflow/worker.py +++ b/python/samples/11_subworkflow/worker.py @@ -134,7 +134,7 @@ def create_inner_workflow(chat_client: FoundryChatClient) -> Workflow: sentiment_formatter = SentimentFormatterExecutor(id="sentiment_formatter") return ( - WorkflowBuilder(name=INNER_WORKFLOW_NAME, start_executor=sentiment_agent) + WorkflowBuilder(name=INNER_WORKFLOW_NAME, start_executor=sentiment_agent, output_from=[sentiment_formatter]) .add_edge(sentiment_agent, sentiment_formatter) .build() ) @@ -152,7 +152,7 @@ def create_workflow() -> Workflow: reporter = ReporterExecutor(id="reporter") return ( - WorkflowBuilder(name=OUTER_WORKFLOW_NAME, start_executor=intake) + WorkflowBuilder(name=OUTER_WORKFLOW_NAME, start_executor=intake, output_from=[reporter]) .add_edge(intake, sentiment_sub) .add_edge(sentiment_sub, reporter) .build() diff --git a/python/samples/12_subworkflow_hitl/worker.py b/python/samples/12_subworkflow_hitl/worker.py index 18f1bd2..0e4b24d 100644 --- a/python/samples/12_subworkflow_hitl/worker.py +++ b/python/samples/12_subworkflow_hitl/worker.py @@ -161,7 +161,7 @@ async def handle_approval_response( def create_inner_workflow() -> Workflow: """Build the inner ``human_review`` workflow (a single HITL gate).""" review_gate = ReviewGateExecutor() - return WorkflowBuilder(name=INNER_WORKFLOW_NAME, start_executor=review_gate).build() + return WorkflowBuilder(name=INNER_WORKFLOW_NAME, start_executor=review_gate, output_from=[review_gate]).build() # ============================================================================ @@ -212,7 +212,7 @@ def create_workflow() -> Workflow: publish = PublishExecutor() return ( - WorkflowBuilder(name=OUTER_WORKFLOW_NAME, start_executor=intake) + WorkflowBuilder(name=OUTER_WORKFLOW_NAME, start_executor=intake, output_from=[publish]) .add_edge(intake, review_sub) .add_edge(review_sub, publish) .build() diff --git a/python/samples/13_conversation_compaction/.env.example b/python/samples/13_conversation_compaction/.env.example new file mode 100644 index 0000000..30f5c34 --- /dev/null +++ b/python/samples/13_conversation_compaction/.env.example @@ -0,0 +1,5 @@ +# Azure AI Foundry project endpoint URL, e.g. https://your-project.services.ai.azure.com/api/projects/your-project +FOUNDRY_PROJECT_ENDPOINT= + +# Model deployment name in your Foundry project +FOUNDRY_MODEL= diff --git a/python/samples/13_conversation_compaction/README.md b/python/samples/13_conversation_compaction/README.md new file mode 100644 index 0000000..1f874c2 --- /dev/null +++ b/python/samples/13_conversation_compaction/README.md @@ -0,0 +1,121 @@ +# Conversation Compaction with Durable Agents + +Shows an agent whose conversation history is **persisted durably** and **compacted as it grows**, +using the same configuration you would write for in-process Agent Framework. + +## What this demonstrates + +The agent is built with a plain `InMemoryHistoryProvider` and a `CompactionProvider`: + +```python +history = InMemoryHistoryProvider(skip_excluded=True) +compaction = CompactionProvider( + after_strategy=SlidingWindowStrategy(keep_last_groups=4), + history_source_id=history.source_id, +) +agent = Agent( + client=..., + name="Historian", + default_options={"store": False}, + context_providers=[history, compaction], +) +``` + +Registering that agent with the durable runtime changes nothing about how you configure it, but: + +- **History becomes durable.** The runtime swaps the in-memory provider for a durable-backed one, + preserving its `source_id` and storage flags. The provider owns transcript appends according to + `store_inputs`, `store_outputs`, `store_context_messages`, and `store_context_from`. This sample's + stored inputs and outputs survive worker restarts in the agent's durable entity. +- **Compaction state is persisted.** Annotations produced by the strategy are stored alongside the + messages and are available on later turns. +- **The history window stays small.** Only the history groups the strategy keeps are sent to the + model on the next turn. Individual messages can still be large. + +With this sample's storage flags and explicit `retention="keep_all", max_state_bytes=None`, +compaction does not delete the local transcript. Original responses are also retained temporarily +in `responseMailbox`, keyed by correlation id, independently of the model's compacted history. +`completedCorrelations` records completion even after response delivery expires. + +### Retention and state budgets + +Compaction selects model context. `retention` controls eager deletion of eligible compaction +exclusions. `max_state_bytes` independently controls pressure eviction of local transcript groups. +Set these on `DurableAIAgentWorker` or override them with `add_agent`. + +| Mode | Behavior | +| --- | --- | +| `keep_all` (default) | Does not eagerly delete compaction exclusions. An explicitly configured byte budget can still evict eligible transcript groups. | +| `follow_compaction` | Eagerly deletes eligible exclusions from local durable history, protecting system messages and the newest/current exchange. It does not enable a byte budget. | + +`max_state_bytes=None` is the default and disables pressure eviction, not the backend's size limit. +On the standalone DTS worker, `max_state_bytes="backend_limit"` resolves to 1,048,576 bytes (1 MiB). +An explicit positive integer is also accepted. Azure Functions cannot infer its backend limit, so +it requires an integer to enable pressure eviction and rejects `"backend_limit"`. + +For example, to retain compaction exclusions until state pressure requires eviction, use +`retention="keep_all", max_state_bytes=1_048_576, high_watermark=0.85, low_watermark=0.70`. +Use `retention="follow_compaction"` to opt into eager pruning as well. The watermark defaults are +`0.85` and `0.70`, with `0 < low_watermark < high_watermark <= 1`. Pressure eviction starts at the +high watermark and aims for the low watermark, or the protected state size if that is larger. + +The budget measures the whole serialized entity, including live mailbox payloads, completion +receipts, session state, and metadata. Pressure eviction preserves protected state and evicts whole +atomic transcript groups. If protected state alone reaches the high watermark, the operation fails with +`StateCapacityError` rather than discarding responses still owed to callers. Size a budget for those +delivery obligations and the backend limit. A burst of turns can fill a small budget even after +transcript pruning. Do not shorten delivery expiry just to make a demo fit. + +Neither retention mode provides unlimited capacity. Completion receipts persist until entity +deletion, and mailbox payloads expire independently of transcript retention. Retention does not +prune an external store or service-managed history. + +### Client-side vs service-managed history + +Compaction only applies to history the **client** owns. When Foundry or the Responses API owns a +turn's history, the durable history provider neither loads nor appends a local transcript. The +entity still persists session state, response delivery payloads, and completion receipts, not a +second conversation record. Existing local history is not erased when ownership changes. + +Ownership is resolved for each run from its `store` option, then the agent's `default_options`, +then the client's default. This sample sets `store=False` so the client-side history provider and +compaction control model context. + +## Running the sample + +1. Start the Durable Task Scheduler emulator: + + ```bash + docker run -d --name dts-emulator -p 8080:8080 -p 8082:8082 mcr.microsoft.com/dts/dts-emulator:latest + ``` + +2. Copy `.env.example` to `.env` and set `FOUNDRY_PROJECT_ENDPOINT` and `FOUNDRY_MODEL`. + +3. Sign in for `AzureCliCredential`: + + ```bash + az login + ``` + +4. Install dependencies and start the worker: + + ```bash + pip install -r requirements.txt + python worker.py + ``` + +5. In another terminal, run the client: + + ```bash + python client.py + ``` + +## What to look for + +The client runs a multi-turn conversation and then asks the agent to recall a fact from a **recent** +turn. The fact should still be in the retained window. + +A sliding window leaves older turns out of model context, so the model may no longer recall their +facts. With the sample's `keep_all` and disabled pressure budget, those messages remain in local +durable storage, marked as excluded. Opting into eager pruning or a byte budget can delete eligible +history. Summarization is an alternative when older details need to stay in context. diff --git a/python/samples/13_conversation_compaction/client.py b/python/samples/13_conversation_compaction/client.py new file mode 100644 index 0000000..0c62a0f --- /dev/null +++ b/python/samples/13_conversation_compaction/client.py @@ -0,0 +1,102 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Client that exercises a durable agent whose history is compacted as it grows. + +Runs a multi-turn conversation against the ``Historian`` agent hosted by ``worker.py`` and +checks recall within a sliding history window. This is not a state-capacity stress test. +""" + +import logging +import os + +from agent_framework_durabletask import DurableAIAgentClient +from azure.identity import AzureCliCredential +from dotenv import load_dotenv +from durabletask.azuremanaged.client import DurableTaskSchedulerClient + +load_dotenv() + +logging.basicConfig(level=logging.WARNING) +logger = logging.getLogger(__name__) + +# Turns that fill the conversation before recall is tested. +FILLER_TURNS = [ + "Name a color.", + "Name a country.", + "Name a fruit.", + "Name a musical instrument.", +] + +CODENAME = "BLUEHERON" + + +def get_client( + taskhub: str | None = None, endpoint: str | None = None, log_handler: logging.Handler | None = None +) -> DurableAIAgentClient: + """Create a configured DurableAIAgentClient. + + Args: + taskhub: Task hub name (defaults to TASKHUB env var or "default") + endpoint: Scheduler endpoint (defaults to ENDPOINT env var or "http://localhost:8080") + log_handler: Optional logging handler for client logging + + Returns: + Configured DurableAIAgentClient instance + """ + taskhub_name = taskhub or os.getenv("TASKHUB", "default") + endpoint_url = endpoint or os.getenv("ENDPOINT", "http://localhost:8080") + + credential = None if endpoint_url == "http://localhost:8080" else AzureCliCredential() + + dts_client = DurableTaskSchedulerClient( + host_address=endpoint_url, + secure_channel=endpoint_url != "http://localhost:8080", + taskhub=taskhub_name, + token_credential=credential, + log_handler=log_handler, + ) + + return DurableAIAgentClient(dts_client) + + +def run_client(agent_client: DurableAIAgentClient) -> None: + """Run a multi-turn conversation against the compacting agent. + + Args: + agent_client: The durable agent client to use. + """ + agent = agent_client.get_agent("Historian") + session = agent.create_session() + + print("Running a multi-turn conversation...\n") + + for turn in FILLER_TURNS: + response = agent.run(turn, session=session) + print(f"[user] {turn}") + print(f"[agent] {response.text}\n") + + fact = f"My project codename is {CODENAME}." + print(f"[user] {fact}") + print(f"[agent] {agent.run(fact, session=session).text}\n") + + question = "What is my project codename? Reply with just the codename." + answer = agent.run(question, session=session) + print(f"[user] {question}") + print(f"[agent] {answer.text}\n") + + if CODENAME.lower() in answer.text.lower(): + print("The agent recalled the fact from its recent history window.") + else: + print("The agent did not recall the recent fact. Inspect its response for errors.") + + +def main() -> None: + """Client entry point.""" + try: + run_client(get_client()) + except Exception as e: + logger.exception(f"Error during agent interaction: {e}") + + +if __name__ == "__main__": + main() diff --git a/python/samples/13_conversation_compaction/requirements.txt b/python/samples/13_conversation_compaction/requirements.txt new file mode 100644 index 0000000..ea73f71 --- /dev/null +++ b/python/samples/13_conversation_compaction/requirements.txt @@ -0,0 +1,13 @@ +# Agent Framework packages +# To use the deployed version, uncomment the lines below and comment out the local installation lines +# agent-framework-foundry +# agent-framework-durabletask + +# Local installation (for development and testing) +# Each package must be listed explicitly because pip doesn't resolve uv workspace sources. +# Without explicit entries, pip would fetch transitive dependencies from PyPI instead of local source. +agent-framework-foundry>=1.10.1,<2 # Foundry support from PyPI +-e ../../packages/durabletask # Local Durable Task package under development + +# Azure authentication +azure-identity \ No newline at end of file diff --git a/python/samples/13_conversation_compaction/sample.py b/python/samples/13_conversation_compaction/sample.py new file mode 100644 index 0000000..5e1ee06 --- /dev/null +++ b/python/samples/13_conversation_compaction/sample.py @@ -0,0 +1,49 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Conversation Compaction Sample - Durable Task Integration (Combined Worker + Client) + +Runs both the worker and client in a single process. The worker is started first to +register the compacting agent, then the client drives a multi-turn conversation. + +Prerequisites: +- Set FOUNDRY_PROJECT_ENDPOINT and FOUNDRY_MODEL +- Sign in with Azure CLI for AzureCliCredential authentication +- Durable Task Scheduler must be running (e.g., using Docker) + +To run this sample: + python sample.py +""" + +import logging + +from client import get_client, run_client # pyrefly: ignore[missing-import] +from dotenv import load_dotenv +from worker import get_worker, setup_worker # pyrefly: ignore[missing-import] + +# Configure logging (must be after imports to override their basicConfig) +logging.basicConfig(level=logging.INFO, force=True) +logger = logging.getLogger(__name__) + + +def main(): + """Main entry point - runs both worker and client in single process.""" + silent_handler = logging.NullHandler() + + dts_worker = get_worker(log_handler=silent_handler) + with dts_worker: + setup_worker(dts_worker) + dts_worker.start() + logger.debug("Worker started and listening for requests...") + + agent_client = get_client(log_handler=silent_handler) + try: + run_client(agent_client) + except Exception as e: + logger.exception(f"Error during agent interaction: {e}") + + logger.debug("Sample completed. Worker shutting down...") + + +if __name__ == "__main__": + load_dotenv() + main() diff --git a/python/samples/13_conversation_compaction/worker.py b/python/samples/13_conversation_compaction/worker.py new file mode 100644 index 0000000..986b682 --- /dev/null +++ b/python/samples/13_conversation_compaction/worker.py @@ -0,0 +1,155 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Worker hosting an agent whose conversation history is compacted as it grows. + +The agent is configured exactly as it would be for in-process Agent Framework: an +``InMemoryHistoryProvider`` plus a ``CompactionProvider``. Registering it with the durable +runtime transparently swaps the history provider for a durable-backed one, so: + +- the history provider persists the inputs and outputs selected by its storage flags, +- that client-owned history lives in the agent's durable entity and survives restarts, +- the compaction strategy still runs, and its annotations are persisted alongside the + messages for later turns, +- only the history groups compaction keeps are sent to the model on the next turn. + +No durable-specific configuration is required on the agent itself. + +The sample keeps the default ``retention="keep_all"`` and ``max_state_bytes=None``. Compaction +limits the number of history groups sent to the model, not total entity size or message size. +Pruning exclusions and pressure eviction are separate opt-ins. Neither provides unlimited capacity. + +Compaction applies to history the client owns. On a service-owned turn the durable history +provider neither loads nor appends a local transcript. The entity keeps session state, original +responses in its delivery mailbox, and completion receipts. This sample sets ``store=False`` +so the history provider owns the model's context instead of Foundry's service-managed history. + +Prerequisites: +- Set FOUNDRY_PROJECT_ENDPOINT and FOUNDRY_MODEL +- Sign in with Azure CLI for AzureCliCredential authentication +- Start a Durable Task Scheduler (e.g., using Docker) +""" + +import asyncio +import logging +import os + +from agent_framework import Agent, CompactionProvider, InMemoryHistoryProvider, SlidingWindowStrategy +from agent_framework.foundry import FoundryChatClient +from agent_framework_durabletask import DurableAIAgentWorker +from azure.identity import AzureCliCredential +from azure.identity.aio import AzureCliCredential as AsyncAzureCliCredential +from dotenv import load_dotenv +from durabletask.azuremanaged.worker import DurableTaskSchedulerWorker + +# Load environment variables from .env file +load_dotenv() + +# Configure logging +logging.basicConfig(level=logging.WARNING) +logger = logging.getLogger(__name__) + +# Keep only the most recent turns in the model's context. Deliberately small so the +# effect is easy to observe in a short sample conversation. +KEEP_LAST_GROUPS = 4 + + +def create_historian_agent() -> Agent: + """Create an agent that recalls facts within a sliding history window. + + Returns: + Agent: The configured Historian agent. + """ + # A plain in-memory history provider: the durable runtime replaces it with a + # durable-backed provider at registration, preserving this ``source_id`` so the + # compaction provider below stays wired to it. + history = InMemoryHistoryProvider(skip_excluded=True) + + compaction = CompactionProvider( + after_strategy=SlidingWindowStrategy(keep_last_groups=KEEP_LAST_GROUPS), + history_source_id=history.source_id, + ) + + return Agent( + client=FoundryChatClient( + project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + model=os.environ["FOUNDRY_MODEL"], + credential=AsyncAzureCliCredential(), + ), + name="Historian", + instructions=( + "You are a concise assistant. Answer in one short sentence. " + "When the user tells you a fact, remember it and repeat it exactly when asked." + ), + # Keep the conversation client-side so the history provider (and therefore compaction) + # owns the model's context. + default_options={"store": False}, + context_providers=[history, compaction], + ) + + +def get_worker( + taskhub: str | None = None, endpoint: str | None = None, log_handler: logging.Handler | None = None +) -> DurableTaskSchedulerWorker: + """Create a configured DurableTaskSchedulerWorker. + + Args: + taskhub: Task hub name (defaults to TASKHUB env var or "default") + endpoint: Scheduler endpoint (defaults to ENDPOINT env var or "http://localhost:8080") + log_handler: Optional logging handler for worker logging + + Returns: + Configured DurableTaskSchedulerWorker instance + """ + taskhub_name = taskhub or os.getenv("TASKHUB", "default") + endpoint_url = endpoint or os.getenv("ENDPOINT", "http://localhost:8080") + + credential = None if endpoint_url == "http://localhost:8080" else AzureCliCredential() + + return DurableTaskSchedulerWorker( + host_address=endpoint_url, + secure_channel=endpoint_url != "http://localhost:8080", + taskhub=taskhub_name, + token_credential=credential, + log_handler=log_handler, + ) + + +def setup_worker(worker: DurableTaskSchedulerWorker) -> DurableAIAgentWorker: + """Register the compacting agent with the durable worker. + + Args: + worker: The DurableTaskSchedulerWorker instance + + Returns: + DurableAIAgentWorker with agents registered + """ + # Keep compacted-out history by default. To delete eligible exclusions, choose + # retention="follow_compaction". Pressure eviction is independent: opt in with + # max_state_bytes="backend_limit" (1 MiB on DTS) or a positive integer budget. + # Budget for live response payloads, receipts and session state as well as history. + agent_worker = DurableAIAgentWorker(worker, retention="keep_all", max_state_bytes=None) + + agent = create_historian_agent() + agent_worker.add_agent(agent) + + logger.debug(f"✓ Registered agent: {agent.name}") + return agent_worker + + +async def main(): + """Main entry point for the worker process.""" + worker = get_worker() + setup_worker(worker) + + logger.info("Worker is ready and listening for requests...") + + try: + worker.start() + while True: + await asyncio.sleep(1) + except KeyboardInterrupt: + logger.debug("Worker shutdown initiated") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/samples/14_external_history_redis/.env.example b/python/samples/14_external_history_redis/.env.example new file mode 100644 index 0000000..b89a793 --- /dev/null +++ b/python/samples/14_external_history_redis/.env.example @@ -0,0 +1,8 @@ +# Azure AI Foundry project endpoint URL, e.g. https://your-project.services.ai.azure.com/api/projects/your-project +FOUNDRY_PROJECT_ENDPOINT= + +# Model deployment name in your Foundry project +FOUNDRY_MODEL= + +# Redis connection string used by the external history provider +REDIS_CONNECTION_STRING=redis://localhost:6379 diff --git a/python/samples/14_external_history_redis/README.md b/python/samples/14_external_history_redis/README.md new file mode 100644 index 0000000..e8cf445 --- /dev/null +++ b/python/samples/14_external_history_redis/README.md @@ -0,0 +1,93 @@ +# External Conversation History (Redis) with Durable Agents + +Shows an agent whose conversation history lives in a **user-chosen external store** rather than in +durable entity state, using the same configuration you would write for in-process Agent Framework. + +## What this demonstrates + +The agent is built with an ordinary `HistoryProvider` that happens to be backed by Redis: + +```python +history = RedisHistoryProvider("redis://localhost:6379") +agent = Agent( + client=..., + name="Archivist", + default_options={"store": False}, + context_providers=[history], +) +``` + +Registering that agent with the durable runtime changes nothing about how you configure it: + +- **Your provider stays active for this client-owned run.** The exact built-in + `InMemoryHistoryProvider` is swapped for a durable-backed one (see + [13_conversation_compaction](../13_conversation_compaction)), but this Redis provider keeps its + hooks and storage. On a service-owned run, the inactive primary is wrapped to suppress load/store + hooks. Use a distinct store-only sink to audit both branches. +- **It receives a stable session id.** The durable entity creates a fresh session per operation but + gives it the entity's own session id, so the provider reads and writes the same key every turn. + Without that, an externally keyed store would start a new conversation on each turn. +- **The provider owns transcript writes.** Core calls it according to `store_inputs`, + `store_outputs`, `store_context_messages`, and `store_context_from`. This sample uses the default + input/output storage flags and `store=False` so the provider supplies model context. +- **Delivery is separate from history.** Fresh durable entity state has an empty + `conversationHistory`, not metadata-only exchange envelopes or a local transcript mirror. It + stores session state, original responses in `responseMailbox` by correlation id, and completion + evidence in `completedCorrelations`. Delivery payloads expire independently of Redis history. + +The [Redis provider](redis_history_provider.py) is deliberately small, using a list read and a blind +`RPUSH`. It is not an exactly-once storage implementation. If Redis accepts an append but the durable +operation is interrupted before its local state commits, retrying can append the same messages +again. A stable session id does not prevent that. A production provider needs its own idempotency +policy for external writes. + +Portable durable `reset` is unsupported for this external provider. Clearing its history requires a +provider-owned operation and coordination with the caller. The sample does not implement one. + +The default `retention="keep_all"` and `max_state_bytes=None` do not prune Redis and do not enable +local pressure eviction. `follow_compaction` or an explicit local byte budget does not manage Redis +retention either. External storage does not give the entity unlimited capacity. Live response +payloads, session state, and completion receipts still need space, and completion receipts persist +until entity deletion. Existing local history from before an ownership change is not erased. + +## Running the sample + +1. Start the Durable Task Scheduler emulator and Redis: + + ```bash + docker run -d --name dts-emulator -p 8080:8080 -p 8082:8082 mcr.microsoft.com/dts/dts-emulator:latest + docker run -d --name redis -p 6379:6379 redis:latest + ``` + +2. Copy `.env.example` to `.env` and set `FOUNDRY_PROJECT_ENDPOINT` and `FOUNDRY_MODEL`. + +3. Sign in for `AzureCliCredential`: + + ```bash + az login + ``` + +4. Install dependencies and start the worker: + + ```bash + pip install -r requirements.txt + python worker.py + ``` + +5. In another terminal, run the client: + + ```bash + python client.py + ``` + +## What to look for + +The client states a fact and then asks for it back in a later turn. The agent answers correctly, +which is only possible if Redis served the earlier turn back into the model's context. The durable +runtime itself never replays history for this agent. + +To see it directly, inspect the Redis key while the sample runs: + +```bash +docker exec -it redis redis-cli KEYS 'durable_sample:history:*' +``` diff --git a/python/samples/14_external_history_redis/client.py b/python/samples/14_external_history_redis/client.py new file mode 100644 index 0000000..f4b3071 --- /dev/null +++ b/python/samples/14_external_history_redis/client.py @@ -0,0 +1,88 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Client that exercises a durable agent whose history lives in Redis. + +Runs a multi-turn conversation against the ``Archivist`` agent hosted by ``worker.py`` and shows +that a user-chosen external store keeps the conversation going under the durable runtime. +""" + +import logging +import os + +from agent_framework_durabletask import DurableAIAgentClient +from azure.identity import AzureCliCredential +from dotenv import load_dotenv +from durabletask.azuremanaged.client import DurableTaskSchedulerClient + +load_dotenv() + +logging.basicConfig(level=logging.WARNING) +logger = logging.getLogger(__name__) + +FACT = "My library card number is 4417." + + +def get_client( + taskhub: str | None = None, endpoint: str | None = None, log_handler: logging.Handler | None = None +) -> DurableAIAgentClient: + """Create a configured DurableAIAgentClient. + + Args: + taskhub: Task hub name (defaults to TASKHUB env var or "default") + endpoint: Scheduler endpoint (defaults to ENDPOINT env var or "http://localhost:8080") + log_handler: Optional logging handler for client logging + + Returns: + Configured DurableAIAgentClient instance + """ + taskhub_name = taskhub or os.getenv("TASKHUB", "default") + endpoint_url = endpoint or os.getenv("ENDPOINT", "http://localhost:8080") + + credential = None if endpoint_url == "http://localhost:8080" else AzureCliCredential() + + dts_client = DurableTaskSchedulerClient( + host_address=endpoint_url, + secure_channel=endpoint_url != "http://localhost:8080", + taskhub=taskhub_name, + token_credential=credential, + log_handler=log_handler, + ) + + return DurableAIAgentClient(dts_client) + + +def run_client(agent_client: DurableAIAgentClient) -> None: + """Run a multi-turn conversation served from the external Redis store. + + Args: + agent_client: The durable agent client to use. + """ + agent = agent_client.get_agent("Archivist") + session = agent.create_session() + + print("Running a multi-turn conversation backed by Redis...\n") + + print(f"[user] {FACT}") + print(f"[agent] {agent.run(FACT, session=session).text}\n") + + question = "What is my library card number? Reply with just the number." + answer = agent.run(question, session=session) + print(f"[user] {question}") + print(f"[agent] {answer.text}\n") + + if "4417" in answer.text: + print("The agent recalled the fact, so Redis served the prior turn back to the model.") + else: + print("The agent did not recall the fact - check that Redis is reachable.") + + +def main() -> None: + """Client entry point.""" + try: + run_client(get_client()) + except Exception as e: + logger.exception(f"Error during agent interaction: {e}") + + +if __name__ == "__main__": + main() diff --git a/python/samples/14_external_history_redis/redis_history_provider.py b/python/samples/14_external_history_redis/redis_history_provider.py new file mode 100644 index 0000000..b1bf0f0 --- /dev/null +++ b/python/samples/14_external_history_redis/redis_history_provider.py @@ -0,0 +1,108 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""A minimal Redis-backed history provider. + +This is an ordinary Agent Framework ``HistoryProvider`` - nothing about it is durable-specific. +It demonstrates reading messages for a session id and appending the new messages selected by +the provider's storage flags. + +The durable runtime leaves providers like this alone: the user chose where their conversation +lives. It does not add a local transcript mirror or make Redis writes exactly-once. This provider +blindly appends, so retrying an interrupted operation after Redis accepted its write can duplicate +messages. It supplies no provider-owned clear operation, so portable durable reset is unsupported. +""" + +from collections.abc import Sequence +from typing import Any + +import redis.asyncio as aioredis +from agent_framework import HistoryProvider, Message + + +class RedisHistoryProvider(HistoryProvider): + """Stores conversation history in a Redis list, one entry per message. + + Messages are keyed by session id, so the same session id must be used on every turn for the + conversation to continue - which is exactly what the durable entity guarantees. + Writes are intentionally not idempotent. A production provider needs its own retry and + clearing policy; a stable session key alone does not deduplicate interrupted appends. + """ + + DEFAULT_SOURCE_ID = "redis_history" + + def __init__( + self, + redis_url: str, + *, + source_id: str = DEFAULT_SOURCE_ID, + key_prefix: str = "durable_sample:history", + ) -> None: + """Create a Redis-backed history provider. + + Args: + redis_url: Redis connection URL, for example ``redis://localhost:6379``. + source_id: Unique identifier for this provider instance. + key_prefix: Prefix for the Redis keys this provider owns. + """ + super().__init__(source_id) + self.key_prefix = key_prefix + self._client: aioredis.Redis = aioredis.from_url(redis_url, decode_responses=True) + + def _key(self, session_id: str | None) -> str: + """Build the Redis key holding the history for a session. + + Args: + session_id: The session ID to build a key for. + + Returns: + The Redis key for this session's history. + """ + return f"{self.key_prefix}:{session_id or 'default'}" + + async def get_messages( + self, session_id: str | None, *, state: dict[str, Any] | None = None, **kwargs: Any + ) -> list[Message]: + """Read this session's messages from Redis, oldest first. + + Args: + session_id: The session ID to retrieve messages for. + state: Unused, since this provider keeps nothing in session state. + **kwargs: Additional arguments (unused). + + Returns: + The stored messages in chronological order. + """ + stored: list[str] = await self._client.lrange(self._key(session_id), 0, -1) + return [Message.from_json(entry) for entry in stored] + + async def save_messages( + self, + session_id: str | None, + messages: Sequence[Message], + *, + state: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + """Append messages to this session's Redis list. + + This intentionally uses blind RPUSH. Repeating a save repeats its messages, including + when a durable operation is interrupted after this write but before its local commit. + + Args: + session_id: The session ID to store messages for. + messages: The messages to persist. + state: Unused, since this provider keeps nothing in session state. + **kwargs: Additional arguments (unused). + """ + if not messages: + return + await self._client.rpush(self._key(session_id), *[message.to_json() for message in messages]) + + async def aclose(self) -> None: + """Close the Redis connection pool. + + A provider that opens a connection should offer a way to give it back. Without this the + pool stays open until the process exits, which is survivable in a sample but shows up as + unclosed-connection warnings and is the wrong thing to copy into an application. + """ + await self._client.aclose() diff --git a/python/samples/14_external_history_redis/requirements.txt b/python/samples/14_external_history_redis/requirements.txt new file mode 100644 index 0000000..ffb066a --- /dev/null +++ b/python/samples/14_external_history_redis/requirements.txt @@ -0,0 +1,13 @@ +# Agent Framework packages +# To use the deployed version, uncomment the lines below and comment out the local installation lines +# agent-framework-foundry +# agent-framework-durabletask + +# Local installation (for development and testing) +# Each package must be listed explicitly because pip doesn't resolve uv workspace sources. +# Without explicit entries, pip would fetch transitive dependencies from PyPI instead of local source. +agent-framework-foundry>=1.10.1,<2 # Foundry support from PyPI +-e ../../packages/durabletask # Local Durable Task package under development + +# External history store used by this sample +redis diff --git a/python/samples/14_external_history_redis/sample.py b/python/samples/14_external_history_redis/sample.py new file mode 100644 index 0000000..2a9e52f --- /dev/null +++ b/python/samples/14_external_history_redis/sample.py @@ -0,0 +1,49 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""External History (Redis) Sample - Durable Task Integration (Combined Worker + Client) + +Runs both the worker and client in a single process. The worker is started first to register +the Redis-backed agent, then the client drives a multi-turn conversation. + +Prerequisites: +- Set FOUNDRY_PROJECT_ENDPOINT and FOUNDRY_MODEL +- Sign in with Azure CLI for AzureCliCredential authentication +- Durable Task Scheduler and Redis must be running (e.g., using Docker) + +To run this sample: + python sample.py +""" + +import logging + +from client import get_client, run_client # pyrefly: ignore[missing-import] +from dotenv import load_dotenv +from worker import get_worker, setup_worker # pyrefly: ignore[missing-import] + +# Configure logging (must be after imports to override their basicConfig) +logging.basicConfig(level=logging.INFO, force=True) +logger = logging.getLogger(__name__) + + +def main(): + """Main entry point - runs both worker and client in single process.""" + silent_handler = logging.NullHandler() + + dts_worker = get_worker(log_handler=silent_handler) + with dts_worker: + setup_worker(dts_worker) + dts_worker.start() + logger.debug("Worker started and listening for requests...") + + agent_client = get_client(log_handler=silent_handler) + try: + run_client(agent_client) + except Exception as e: + logger.exception(f"Error during agent interaction: {e}") + + logger.debug("Sample completed. Worker shutting down...") + + +if __name__ == "__main__": + load_dotenv() + main() diff --git a/python/samples/14_external_history_redis/worker.py b/python/samples/14_external_history_redis/worker.py new file mode 100644 index 0000000..27cbd81 --- /dev/null +++ b/python/samples/14_external_history_redis/worker.py @@ -0,0 +1,144 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Worker hosting an agent whose conversation history lives in Redis, not in durable state. + +The agent is configured exactly as it would be for in-process Agent Framework: a history +provider the user chose (here Redis) is passed as a context provider. Registering it with the +durable runtime requires no changes: + +- the runtime **leaves the provider alone** - the user picked where their conversation lives, +- it hands the provider the entity's **stable** session id on every turn, so history continues + across turns and across worker restarts, +- the provider owns transcript appends according to its storage flags, +- durable state stores session state, original responses in a correlation-keyed delivery mailbox, + and completion receipts, not a local transcript mirror. + +This minimal provider blindly appends to Redis. An interrupted operation retried after Redis +accepted the append can duplicate messages. Durable execution does not make that external write +exactly-once. Portable reset is unsupported for this provider and requires provider-owned clearing. + +Contrast with ``13_conversation_compaction``, where an in-memory provider is transparently +swapped for a durable-backed one. + +Prerequisites: +- Set FOUNDRY_PROJECT_ENDPOINT and FOUNDRY_MODEL +- Sign in with Azure CLI for AzureCliCredential authentication +- Start a Durable Task Scheduler and a Redis instance (e.g., using Docker) +""" + +import asyncio +import logging +import os + +from agent_framework import Agent +from agent_framework.foundry import FoundryChatClient +from agent_framework_durabletask import DurableAIAgentWorker +from azure.identity import AzureCliCredential +from azure.identity.aio import AzureCliCredential as AsyncAzureCliCredential +from dotenv import load_dotenv +from durabletask.azuremanaged.worker import DurableTaskSchedulerWorker +from redis_history_provider import RedisHistoryProvider # pyrefly: ignore[missing-import] + +# Load environment variables from .env file +load_dotenv() + +# Configure logging +logging.basicConfig(level=logging.WARNING) +logger = logging.getLogger(__name__) + +# Providers holding a Redis connection pool, closed when the worker stops. +_open_history_providers: list[RedisHistoryProvider] = [] + + +def create_archivist_agent() -> Agent: + """Create an agent whose history is stored in Redis. + + Returns: + Agent: The configured Archivist agent. + """ + history = RedisHistoryProvider(os.getenv("REDIS_CONNECTION_STRING", "redis://localhost:6379")) + # Kept so the worker can hand the connection pool back on the way out. + _open_history_providers.append(history) + + return Agent( + client=FoundryChatClient( + project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + model=os.environ["FOUNDRY_MODEL"], + credential=AsyncAzureCliCredential(), + ), + name="Archivist", + instructions=( + "You are a concise assistant. Answer in one short sentence. " + "When the user tells you a fact, remember it and repeat it exactly when asked." + ), + # Keep the conversation client-side so the history provider owns the model's context. + default_options={"store": False}, + context_providers=[history], + ) + + +def get_worker( + taskhub: str | None = None, endpoint: str | None = None, log_handler: logging.Handler | None = None +) -> DurableTaskSchedulerWorker: + """Create a configured DurableTaskSchedulerWorker. + + Args: + taskhub: Task hub name (defaults to TASKHUB env var or "default") + endpoint: Scheduler endpoint (defaults to ENDPOINT env var or "http://localhost:8080") + log_handler: Optional logging handler for worker logging + + Returns: + Configured DurableTaskSchedulerWorker instance + """ + taskhub_name = taskhub or os.getenv("TASKHUB", "default") + endpoint_url = endpoint or os.getenv("ENDPOINT", "http://localhost:8080") + + credential = None if endpoint_url == "http://localhost:8080" else AzureCliCredential() + + return DurableTaskSchedulerWorker( + host_address=endpoint_url, + secure_channel=endpoint_url != "http://localhost:8080", + taskhub=taskhub_name, + token_credential=credential, + log_handler=log_handler, + ) + + +def setup_worker(worker: DurableTaskSchedulerWorker) -> DurableAIAgentWorker: + """Register the Redis-backed agent with the durable worker. + + Args: + worker: The DurableTaskSchedulerWorker instance + + Returns: + DurableAIAgentWorker with agents registered + """ + agent_worker = DurableAIAgentWorker(worker) + + agent = create_archivist_agent() + agent_worker.add_agent(agent) + + logger.debug(f"✓ Registered agent: {agent.name}") + return agent_worker + + +async def main(): + """Main entry point for the worker process.""" + worker = get_worker() + setup_worker(worker) + + logger.info("Worker is ready and listening for requests...") + + try: + worker.start() + while True: + await asyncio.sleep(1) + except KeyboardInterrupt: + logger.debug("Worker shutdown initiated") + finally: + for history in _open_history_providers: + await history.aclose() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/samples/README.md b/python/samples/README.md index 81a95ef..c738d39 100644 --- a/python/samples/README.md +++ b/python/samples/README.md @@ -2,9 +2,94 @@ This directory contains samples for durable agent hosting using the Durable Task Scheduler. These samples demonstrate the worker-client architecture pattern, enabling distributed agent execution with persistent conversation state. +## PR #59 prototype scope + +This is an integrated reference for [ADR PR #88](https://github.com/microsoft/agent-framework-durable-extension/pull/88), +not the final implementation PR. After ADR approval, the agreed changes will be split into stacked +implementation PRs. [PR #59](https://github.com/microsoft/agent-framework-durable-extension/pull/59) +remains the prototype until that stack lands. Its APIs and deployment choices are provisional. + +The prototype's version-2 runtime requires `deployment_mode="isolated_v2"` on `DurableAIAgentWorker`, +`AgentFunctionApp` and the standalone Functions entity factory, or +`DURABLE_AGENTS_DEPLOYMENT_MODE=isolated_v2` when the argument is omitted/`None`. Configure the sample +host environment accordingly. This is operator acknowledgement, not proof of isolation. Use a +separate hub/deployment with compatible workers and clients. Old workers and workflow histories, +including paused legacy HITL, must stay on the old engine. + +Only `2.0.0` entity state is writable. Legacy and supported future-minor state is read-only. +Names are unchanged, so using an old `@name@key` on an empty new hub is not migration. Explicit +backend migration needs an empty, separately addressed destination and authorized ownership transfer. +Scalar legacy ingestion cursors require a complete accepted-message journal, including evicted +inputs. If that journal is unavailable, keep the session on the old engine. Migration does not move +workflow history or reconstruct missing original responses. + +The workflow client, generated start routes and child dispatch wrap new starts with protocol version +2. Native custom schedulers must use public `wrap_workflow_input` for new instances. Old/raw starts +reject before revised actions execute. Rewrapping old starts is not history migration. + +## Prototype validation + +The published baseline is +[prototype commit 9b4550d](https://github.com/microsoft/agent-framework-durable-extension/commit/9b4550d). +The results below were recorded locally for the outcome, media, failure-boundary and telemetry +follow-up. They are not the current remote CI status or a claim of release readiness. See +[PR #59 checks](https://github.com/microsoft/agent-framework-durable-extension/pull/59/checks) +for remote results. + +| Local check | Result | +| --- | --- | +| Python 3.13 / core 1.16 | 3,427 passed, zero skipped | +| Python 3.13 / real cached core 1.13 | 3,427 passed, zero skipped | +| Python 3.10 / core 1.16 | 3,427 passed, zero skipped | +| Media retention units | 30 passed, six content kinds across four policies plus six protected-floor cases | +| Cancellation/failure units and Functions consumers | 11 + 3 passed | +| Retention OTel units | 20 passed | +| Completion-outcome units | 52 passed, including formatted and unformatted acceptance-only regressions | +| Existing consumer parameterizations | Eight additional cases passed | +| Direct DTS integration suite | 45 passed in 354.65 seconds, prior 42 plus three new cases | +| Azure Functions integration suite | 45 passed in 649.44 seconds, prior 43 plus two media cases | +| Ruff lint/format, Pyright, MyPy, offline lock and both package builds | Passed | + +The focused unit counts are subsets of each 3,427-test run, not additional tests. Media cases cover +inline PNG, inline text files, image URIs, hosted files, mixed binary/text tool results and large +tool payloads. They check all retention/budget combinations, JSON cold reload, exact subsequent +model input, atomic tool groups, protected floors and staged deletion measurements. Failure cases +cover cancellation at provider/model/retention boundaries, warm rollback, lost write acknowledgement +and provider failure combined with rejected error persistence and bounded polling. Caller polling +cancellation does not cancel the entity. Outcome tests cover retained success/failure, unknown +legacy receipts, strict migration and rejection of fresh acceptance-only completion records. + +The three new direct tests use real DTS persistence and process restarts with a deterministic +`BaseChatClient`, not Foundry. Two exercise PNG and inline-file pressure with persisted-state +readback, exact next model input and matching truncation/OTel counts. The third hard-kills a worker +before commit, observes repeated simulated external effects on retry, then kills after confirmed +scheduler readback and verifies duplicate suppression. These are not live graceful execution +cancellation tests. The two new Functions cases use the production entity handler and actual +`DurableEntityContext` with Azure Storage via Azurite. They verify PNG/inline-file pressure, +persisted JSON, a restarted host, exact subsequent model input and staged OTel measurements. +Inline media bytes dominate the live pressure cases, rather than text padding alone. The existing +42 direct tests and 43 Functions tests remain text-based and include Foundry-backed scenarios. + +The Functions rerun required the local test Azurite setting `--skipApiVersionCheck`. The initial +36 failures and seven passes were caused by unsupported Storage API `2026-02-06`, not product +changes. The corrected final run passed all 45 tests. Coverage percentage was not remeasured here. + +Mutation checks reject disabled pressure/eager pruning, lost content metadata, missing rollback, +missing binding cleanup and missing telemetry. Live DTS cases fail when pressure is disabled. +Functions cases fail on actual stored byte size when the configured budget is deliberately inflated. +Restored runs pass. Mutation changes stayed in fresh process memory or generated temporary test apps. + +Graceful-shutdown-specific host behavior and hosted-model media acceptance are not established by +these tests. Remaining release validation includes actual scheduler-limit/offload behavior as those +capabilities are enabled, exact Pydantic 2.11 runtime validation (artifact downloads +remain blocked), and shared reader/writer, client, replay and rollback compatibility. The reduced +live budget is not a scheduler-limit test. No compiled C# or cross-runtime schema acceptance is +claimed. Existing .NET readers and legacy workflow histories remain incompatible with the revised +contract. These gaps do not replace or defer the ADR's required validation. + ## Import convention -These samples import the durable hosting types **directly from the extension packages** — +These samples import the durable hosting types **directly from the extension packages**, `agent_framework_durabletask` and `agent_framework_azurefunctions`: ```python @@ -15,7 +100,7 @@ from agent_framework_azurefunctions import AgentFunctionApp For backward compatibility these entry-point types are also re-exported from `agent_framework.azure` in the core `agent-framework` package, so existing `from agent_framework.azure import ...` code keeps working. **New and updated samples should use -the direct package imports shown above** — the canonical, self-contained path for this repo — +the direct package imports shown above**, the self-contained path for this repo, rather than routing through the `agent_framework.azure` shim. ## Quick Prerequisites Checklist @@ -70,6 +155,55 @@ az account show - **[11_subworkflow](11_subworkflow/)**: Compose workflows by embedding an inner `Workflow` as a node via `WorkflowExecutor`. On the durable host the inner workflow runs as its own child orchestration, and a single `configure_workflow` call registers both. - **[12_subworkflow_hitl](12_subworkflow_hitl/)**: A human-in-the-loop pause that lives **inside a sub-workflow**. The nested request surfaces to the client with a qualified request id (`{executor}~{ordinal}~{requestId}`) behind a single top-level addressing surface. +These workflow samples and their Azure Functions counterparts explicitly set +`WorkflowBuilder(output_from=[...])` to the executors that produce their final results. +For composed workflows, the inner workflow selects the result forwarded to its parent, +and the outer workflow selects its final report or publication message. Agent responses +still travel along the graph edges but are not additional results in these samples. +For a workflow intended to return agent responses, include those agents in `output_from` +or use `output_from="all"`. + +### Conversation History + +History providers own transcript writes according to their storage flags. External and +service-managed history do not get a local transcript mirror. The entity keeps response delivery +payloads and completion receipts separately from model history. Retention defaults to `keep_all` +with `max_state_bytes=None`. Eager pruning and pressure eviction are separate opt-ins, not a promise +of unlimited capacity. + +Only exact built-in in-memory providers are substituted. Subclasses retain custom hooks and session +transcripts in the protected floor, outside durable transcript eviction. Service-owned runs +intentionally suppress both load and store hooks on the inactive primary, including per-call hooks. +That differs from core 1.16 behavior. Use a distinct store-only sink to audit both service/client +branches. Do not assume universal unchanged-hook semantics or retry-safe external effects. + +Workflow delta transport uses parallel occurrence IDs, not public `Message.message_id` rewrites. +The full selected logical conversation and all response messages remain available for downstream +projection. Private forwarding provenance is internal checkpoint data, not application metadata. +Typed/cache-only requests, agent approval/HITL and output-designated agents are supported locally. + +Delivery expires logically even while an idle entity retains its payload. New runs, duplicate runs, +reset and backend `expire_responses` clean expired payloads without erasing completion receipts. +Idle physical cleanup needs an application-owned schedule or explicit backend signal/manual +operation. No public HTTP/MCP cleanup endpoint is generated. Receipts can exhaust capacity, and +entity commits do not provide a distributed transaction or exactly-once external tool execution. + +New receipts retain invocation success/failure after payload expiry. Expired lookup exposes +`durable_outcome` as `succeeded`, `failed` or `unknown`, without changing original response payloads. +An older unknown receipt still prevents reruns. Strict migration can require trustworthy outcomes, +but a possibly pruned legacy transcript without an error is not proof of success. The default +legacy-compatible path retains duplicate protection. Fresh acceptance-only responses do not record +completion, and fire-and-forget acceptance remains distinct from completion. + +The shared [retention telemetry](../packages/durabletask/README.md#retention-telemetry) measures staged +deletion, not committed deletion. A `set_state` return or failure leaves commit status unknown. +Pair separate persisted readback with subsequent model input. Applications own SDK/exporter setup. +`"backend_limit"` remains a non-normative Python-only Scheduler convenience, outside the portable +`None` or positive-integer budget contract and without assumed shared-review agreement. + +- **[13_conversation_compaction](13_conversation_compaction/)**: Compact client-owned history with `InMemoryHistoryProvider` and `CompactionProvider`. Keep excluded history by default and choose transcript pruning or a state budget independently. +- **[14_external_history_redis](14_external_history_redis/)**: Use an ordinary Redis history provider with a stable session id and no local transcript mirror. The minimal blind-append provider documents interrupted-retry duplicates and unsupported portable reset. + ### Azure Functions Hosting These samples host workflows and agents on Azure Durable Functions (`func start`) instead of the worker-client model above. Each has its own setup steps in its README, and shared environment setup lives in [azure_functions/README.md](azure_functions/README.md). @@ -87,6 +221,7 @@ These samples host workflows and agents on Azure Durable Functions (`func start` - **[azure_functions/11_workflow_parallel](azure_functions/11_workflow_parallel/)**: Parallel execution of executors and agents in an Azure Durable Functions workflow. - **[azure_functions/12_workflow_hitl](azure_functions/12_workflow_hitl/)**: The workflow human-in-the-loop pattern on Azure Durable Functions, with the reviewer notified from inside the workflow via `WorkflowHitlContext`. - **[azure_functions/13_subworkflow_hitl](azure_functions/13_subworkflow_hitl/)**: A human-in-the-loop pause inside a sub-workflow on Azure Durable Functions, exposed through a single top-level respond surface. +- **[azure_functions/14_conversation_compaction](azure_functions/14_conversation_compaction/)**: Compact client-owned history on Azure Functions with independent retention and explicit byte-budget options. The Functions counterpart to [13_conversation_compaction](13_conversation_compaction/). ## Running the Samples @@ -96,7 +231,7 @@ These samples are designed to be run locally in a cloned repository. The following prerequisites are required to run the samples: -- [Python 3.9 or later](https://www.python.org/downloads/) +- [Python 3.10 or later](https://www.python.org/downloads/), `agent-framework-core>=1.13.0,<2` and `pydantic>=2.11,<3` - [Azure CLI](https://learn.microsoft.com/cli/azure/install-azure-cli) installed and authenticated (`az login`) - [Microsoft Foundry project](https://learn.microsoft.com/azure/foundry/how-to/create-projects) with a deployed model, configured through `FOUNDRY_PROJECT_ENDPOINT` and `FOUNDRY_MODEL` (gpt-4o-mini or better is recommended) - [Durable Task Scheduler](https://learn.microsoft.com/azure/azure-functions/durable/durable-task-scheduler/develop-with-durable-task-scheduler) (local emulator or Azure-hosted) diff --git a/python/samples/azure_functions/09_workflow_shared_state/function_app.py b/python/samples/azure_functions/09_workflow_shared_state/function_app.py index 2d50f6d..e781b5d 100644 --- a/python/samples/azure_functions/09_workflow_shared_state/function_app.py +++ b/python/samples/azure_functions/09_workflow_shared_state/function_app.py @@ -220,7 +220,9 @@ def _create_workflow() -> Workflow: # False -> submit_to_email_assistant -> email_assistant_agent -> finalize_and_send # True -> handle_spam return ( - WorkflowBuilder(name="email_triage_shared_state", start_executor=store_email) + WorkflowBuilder( + name="email_triage_shared_state", start_executor=store_email, output_from=[handle_spam, finalize_and_send] + ) .add_edge(store_email, spam_detection_agent) .add_edge(spam_detection_agent, to_detection_result) .add_edge(to_detection_result, submit_to_email_assistant, condition=get_condition(False)) diff --git a/python/samples/azure_functions/10_workflow_no_shared_state/README.md b/python/samples/azure_functions/10_workflow_no_shared_state/README.md index 0b7e8cf..c02566e 100644 --- a/python/samples/azure_functions/10_workflow_no_shared_state/README.md +++ b/python/samples/azure_functions/10_workflow_no_shared_state/README.md @@ -109,8 +109,7 @@ Email sent: Hi, Thank you for the reminder about the sprint planning meeting tom ```python workflow = ( - WorkflowBuilder() - .set_start_executor(spam_agent) + WorkflowBuilder(name="email_triage", start_executor=spam_agent, output_from=[spam_handler, email_sender]) .add_switch_case_edge_group( spam_agent, [ diff --git a/python/samples/azure_functions/10_workflow_no_shared_state/function_app.py b/python/samples/azure_functions/10_workflow_no_shared_state/function_app.py index 6b7e4cf..3f9572e 100644 --- a/python/samples/azure_functions/10_workflow_no_shared_state/function_app.py +++ b/python/samples/azure_functions/10_workflow_no_shared_state/function_app.py @@ -182,7 +182,7 @@ def _create_workflow() -> Workflow: # Build workflow return ( - WorkflowBuilder(name="email_triage", start_executor=spam_agent) + WorkflowBuilder(name="email_triage", start_executor=spam_agent, output_from=[spam_handler, email_sender]) .add_switch_case_edge_group( spam_agent, [ diff --git a/python/samples/azure_functions/11_workflow_parallel/function_app.py b/python/samples/azure_functions/11_workflow_parallel/function_app.py index 3438b8c..715bed2 100644 --- a/python/samples/azure_functions/11_workflow_parallel/function_app.py +++ b/python/samples/azure_functions/11_workflow_parallel/function_app.py @@ -348,7 +348,7 @@ def _create_workflow() -> Workflow: # Build workflow with parallel patterns return ( - WorkflowBuilder(name="parallel_review", start_executor=input_router) + WorkflowBuilder(name="parallel_review", start_executor=input_router, output_from=[final_report_executor]) # Pattern 1: Fan-out to two executors (run in parallel) .add_fan_out_edges( source=input_router, diff --git a/python/samples/azure_functions/12_workflow_hitl/function_app.py b/python/samples/azure_functions/12_workflow_hitl/function_app.py index 865c608..95e5c9d 100644 --- a/python/samples/azure_functions/12_workflow_hitl/function_app.py +++ b/python/samples/azure_functions/12_workflow_hitl/function_app.py @@ -477,7 +477,7 @@ def _create_workflow() -> Workflow: # Side-branch: human_review_executor -> notify_executor emails the reviewer a respond # link (built from WorkflowHitlContext) in the same superstep, before the pause. return ( - WorkflowBuilder(name="content_moderation", start_executor=input_router) + WorkflowBuilder(name="content_moderation", start_executor=input_router, output_from=[publish_executor]) .add_edge(input_router, content_analyzer_agent) .add_edge(content_analyzer_agent, content_analyzer_executor) .add_edge(content_analyzer_executor, human_review_executor) diff --git a/python/samples/azure_functions/13_subworkflow_hitl/function_app.py b/python/samples/azure_functions/13_subworkflow_hitl/function_app.py index 1b6e0f9..dc7e505 100644 --- a/python/samples/azure_functions/13_subworkflow_hitl/function_app.py +++ b/python/samples/azure_functions/13_subworkflow_hitl/function_app.py @@ -237,7 +237,11 @@ def create_inner_workflow() -> Workflow: notify = NotifyExecutor() # Side-branch: review_gate -> notify builds the qualified respond URL in the same # superstep that raises the request, before the inner workflow pauses. - return WorkflowBuilder(name=INNER_WORKFLOW_NAME, start_executor=review_gate).add_edge(review_gate, notify).build() + return ( + WorkflowBuilder(name=INNER_WORKFLOW_NAME, start_executor=review_gate, output_from=[review_gate]) + .add_edge(review_gate, notify) + .build() + ) # ============================================================================ @@ -288,7 +292,7 @@ def _create_workflow() -> Workflow: publish = PublishExecutor() return ( - WorkflowBuilder(name=OUTER_WORKFLOW_NAME, start_executor=intake) + WorkflowBuilder(name=OUTER_WORKFLOW_NAME, start_executor=intake, output_from=[publish]) .add_edge(intake, review_sub) .add_edge(review_sub, publish) .build() diff --git a/python/samples/azure_functions/14_conversation_compaction/README.md b/python/samples/azure_functions/14_conversation_compaction/README.md new file mode 100644 index 0000000..414b28a --- /dev/null +++ b/python/samples/azure_functions/14_conversation_compaction/README.md @@ -0,0 +1,113 @@ +# Conversation Compaction Sample (Python) + +This sample demonstrates hosting an agent whose conversation history is **persisted durably** and +**compacted as it grows**, using the same configuration you would write for in-process Agent +Framework. It is the Azure Functions counterpart to the standalone +[`13_conversation_compaction`](../../13_conversation_compaction) sample. + +## Key Concepts Demonstrated + +- Configuring compaction the ordinary core way, an `InMemoryHistoryProvider` plus a + `CompactionProvider`, with **no durable-specific configuration on the agent**. +- The durable runtime swapping the in-memory provider for a durable-backed one at registration, + preserving its `source_id` and storage flags. The provider owns appends according to + `store_inputs`, `store_outputs`, `store_context_messages`, and `store_context_from`. +- Compaction annotations being persisted alongside the stored messages for later turns. +- Limiting the number of history groups sent to the model, not the size of individual messages + or the whole entity. + +```python +history = InMemoryHistoryProvider(skip_excluded=True) +compaction = CompactionProvider( + after_strategy=SlidingWindowStrategy(keep_last_groups=4), + history_source_id=history.source_id, +) +agent = Agent( + client=..., + name="Historian", + default_options={"store": False}, + context_providers=[history, compaction], +) + +app = AgentFunctionApp( + agents=[agent], + enable_health_check=True, + retention="keep_all", + max_state_bytes=None, +) +``` + +This sample stores inputs and outputs and keeps them with explicit `retention="keep_all"` and +`max_state_bytes=None`, which are also the host defaults. Original responses live independently +in the correlation-keyed `responseMailbox` until delivery expiry. `completedCorrelations` keeps +completion evidence after those payloads expire. + +### Retention and state budgets + +`retention="keep_all"` disables eager deletion of compaction exclusions. It does not disable an +explicit byte budget. `retention="follow_compaction"` prunes eligible exclusions from local durable +history, protecting system messages and the newest/current exchange, but does not enable pressure +eviction by itself. + +`max_state_bytes=None` disables pressure eviction, not the backend's capacity limit. To opt in, +pass a positive integer chosen for your backend and workload, for example +`max_state_bytes=1_048_576, high_watermark=0.85, low_watermark=0.70`. These watermark values are the +defaults and must satisfy `0 < low_watermark < high_watermark <= 1`. The budget is independent of +`retention` and can be used with either mode. Configure them on `AgentFunctionApp` or override +them with `add_agent`. + +Functions cannot infer the backend's limit and rejects `max_state_bytes="backend_limit"`. That +option resolves to 1,048,576 bytes (1 MiB) only on the standalone DTS worker. + +Pressure eviction measures the whole serialized entity. It starts at the high watermark and aims +for the low watermark, or the protected state size if larger. Live responses, completion receipts, +session state, metadata, and protected transcript groups all need space. If the protected state +reaches the high watermark, the operation fails with `StateCapacityError` rather than deleting +responses still owed to callers. A burst of turns can fill a small budget even after history is +pruned. Do not shorten delivery expiry to force the sample to fit. + +Neither mode gives unlimited capacity. Completion receipts persist until entity deletion, and +mailbox expiry is separate from transcript retention. These settings do not prune an external +store or service-managed history. + +### Client-side vs service-managed history + +Compaction only applies to history the **client** owns. On a service-owned turn, the durable history +provider neither loads nor appends a local transcript. Session state, response delivery payloads, +and completion receipts are still persisted, not a second conversation record. Switching ownership +does not erase existing local history. + +The runtime resolves ownership from the run's `store` option, then the agent's `default_options`, +then the client's default. This sample sets `store=False` so client-side history and compaction +control model context rather than Foundry's service-managed history. + +## Prerequisites + +Follow the common setup steps in `../README.md` to install tooling, configure Foundry +credentials, and install the Python dependencies for this sample. This sample uses +`FOUNDRY_PROJECT_ENDPOINT` and `FOUNDRY_MODEL`. + +## Running the Sample + +Send several turns using the **same** session id so they form one conversation. `demo.http` contains +a ready-made sequence, and the equivalent with `curl` is: + +```bash +curl -X POST http://localhost:7071/api/agents/Historian/run \ + -H "Content-Type: application/json" \ + -d '{"message": "My project codename is BLUEHERON.", "session_id": "compaction-demo-001"}' + +curl -X POST http://localhost:7071/api/agents/Historian/run \ + -H "Content-Type: application/json" \ + -d '{"message": "What is my project codename? Reply with just the codename.", "session_id": "compaction-demo-001"}' +``` + +## What to look for + +The agent answers correctly from a **recent** turn while older turns fall outside the retained +window. + +A sliding window leaves older turns out of model context, so the model may no longer recall their +facts. With the sample's `keep_all` and disabled pressure budget, those messages remain in local +durable storage, marked as excluded. Opting into eager pruning or a byte budget can delete eligible +history. Summarization is an alternative when older details need to stay in context. diff --git a/python/samples/azure_functions/14_conversation_compaction/demo.http b/python/samples/azure_functions/14_conversation_compaction/demo.http new file mode 100644 index 0000000..e273795 --- /dev/null +++ b/python/samples/azure_functions/14_conversation_compaction/demo.http @@ -0,0 +1,63 @@ +### Conversation Compaction Sample Interactions +@baseUrl = http://localhost:7071 +@agentName = Historian +@agentRoute = {{baseUrl}}/api/agents/{{agentName}} +@healthRoute = {{baseUrl}}/api/health +@sessionId = compaction-demo-001 + +### Health Check +GET {{healthRoute}} + +### Turn 1 - filler +POST {{agentRoute}}/run +Content-Type: application/json + +{ + "message": "Name a color.", + "session_id": "{{sessionId}}" +} + +### Turn 2 - filler +POST {{agentRoute}}/run +Content-Type: application/json + +{ + "message": "Name a country.", + "session_id": "{{sessionId}}" +} + +### Turn 3 - filler +POST {{agentRoute}}/run +Content-Type: application/json + +{ + "message": "Name a fruit.", + "session_id": "{{sessionId}}" +} + +### Turn 4 - filler +POST {{agentRoute}}/run +Content-Type: application/json + +{ + "message": "Name a musical instrument.", + "session_id": "{{sessionId}}" +} + +### Turn 5 - state the fact to recall later +POST {{agentRoute}}/run +Content-Type: application/json + +{ + "message": "My project codename is BLUEHERON.", + "session_id": "{{sessionId}}" +} + +### Turn 6 - the fact is inside the retained window, so it is answered +POST {{agentRoute}}/run +Content-Type: application/json + +{ + "message": "What is my project codename? Reply with just the codename.", + "session_id": "{{sessionId}}" +} diff --git a/python/samples/azure_functions/14_conversation_compaction/function_app.py b/python/samples/azure_functions/14_conversation_compaction/function_app.py new file mode 100644 index 0000000..18a8559 --- /dev/null +++ b/python/samples/azure_functions/14_conversation_compaction/function_app.py @@ -0,0 +1,92 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Host an agent whose conversation history is compacted as it grows, inside Azure Functions. + +The agent is configured exactly as it would be for in-process Agent Framework: an +``InMemoryHistoryProvider`` plus a ``CompactionProvider``. Registering it with +``AgentFunctionApp`` transparently swaps the history provider for a durable-backed one, so +the provider persists the inputs and outputs selected by its storage flags in the agent's durable +entity. Compaction annotations are persisted alongside those messages. Only the history groups +compaction keeps are sent to the model on the next turn. + +This is the Azure Functions counterpart to the standalone ``13_conversation_compaction`` sample. + +Compaction applies to history the client owns. On a service-owned turn the durable history +provider neither loads nor appends a local transcript. The entity still persists session state, +original responses in its delivery mailbox, and completion receipts. This sample sets +``store=False`` so the history provider owns model context instead of the service. + +The sample explicitly keeps ``retention="keep_all"`` and ``max_state_bytes=None``. A sliding +window limits history groups, not message size or total state. Neither pruning nor an optional +pressure budget provides unlimited capacity. + +Prerequisites: set `FOUNDRY_PROJECT_ENDPOINT`, `FOUNDRY_MODEL`, and sign in +with Azure CLI before starting the Functions host.""" + +import os +from typing import Any + +from agent_framework import Agent, CompactionProvider, InMemoryHistoryProvider, SlidingWindowStrategy +from agent_framework.foundry import FoundryChatClient +from agent_framework_azurefunctions import AgentFunctionApp +from azure.identity.aio import AzureCliCredential +from dotenv import load_dotenv + +load_dotenv() + +# Keep only the most recent turns in the model's context. Deliberately small so the +# effect is easy to observe in a short sample conversation. +KEEP_LAST_GROUPS = 4 + + +# 1. Instantiate the agent the ordinary core way - no durable-specific configuration. +def _create_agent() -> Any: + """Create the Historian agent.""" + # A plain in-memory history provider: the durable runtime replaces it with a + # durable-backed provider at registration, preserving this ``source_id`` so the + # compaction provider below stays wired to it. + history = InMemoryHistoryProvider(skip_excluded=True) + + compaction = CompactionProvider( + after_strategy=SlidingWindowStrategy(keep_last_groups=KEEP_LAST_GROUPS), + history_source_id=history.source_id, + ) + + return Agent( + client=FoundryChatClient( + project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + model=os.environ["FOUNDRY_MODEL"], + credential=AzureCliCredential(), + ), + name="Historian", + instructions=( + "You are a concise assistant. Answer in one short sentence. " + "When the user tells you a fact, remember it and repeat it exactly when asked." + ), + # Keep the conversation client-side so the history provider (and therefore compaction) + # owns the model's context. + default_options={"store": False}, + context_providers=[history, compaction], + ) + + +# 2. Register the agent with AgentFunctionApp so Azure Functions exposes the required triggers. +# Choose retention="follow_compaction" to prune eligible exclusions. Independently, set a +# positive integer max_state_bytes to enable pressure eviction. Functions cannot resolve +# "backend_limit". Allow space for live responses, completion receipts and session state. +app = AgentFunctionApp( + agents=[_create_agent()], + enable_health_check=True, + max_poll_retries=50, + retention="keep_all", + max_state_bytes=None, +) + +""" +Expected behavior when posting several turns with the same `session_id`: + +- each turn uses the recent history groups kept by compaction, +- the number of history groups sent to the model stops growing once the sliding window fills, +- this configuration keeps stored inputs and outputs, with compacted-out messages marked excluded, +- original responses and completion receipts are separate from the compacted local transcript. +""" diff --git a/python/samples/azure_functions/14_conversation_compaction/host.json b/python/samples/azure_functions/14_conversation_compaction/host.json new file mode 100644 index 0000000..9e7fd87 --- /dev/null +++ b/python/samples/azure_functions/14_conversation_compaction/host.json @@ -0,0 +1,12 @@ +{ + "version": "2.0", + "extensionBundle": { + "id": "Microsoft.Azure.Functions.ExtensionBundle", + "version": "[4.*, 5.0.0)" + }, + "extensions": { + "durableTask": { + "hubName": "%TASKHUB_NAME%" + } + } +} diff --git a/python/samples/azure_functions/14_conversation_compaction/local.settings.json.template b/python/samples/azure_functions/14_conversation_compaction/local.settings.json.template new file mode 100644 index 0000000..1d8bc82 --- /dev/null +++ b/python/samples/azure_functions/14_conversation_compaction/local.settings.json.template @@ -0,0 +1,11 @@ +{ + "IsEncrypted": false, + "Values": { + "FUNCTIONS_WORKER_RUNTIME": "python", + "AzureWebJobsStorage": "UseDevelopmentStorage=true", + "DURABLE_TASK_SCHEDULER_CONNECTION_STRING": "Endpoint=http://localhost:8080;TaskHub=default;Authentication=None", + "TASKHUB_NAME": "default", + "FOUNDRY_PROJECT_ENDPOINT": "", + "FOUNDRY_MODEL": "" + } +} diff --git a/python/samples/azure_functions/14_conversation_compaction/requirements.txt b/python/samples/azure_functions/14_conversation_compaction/requirements.txt new file mode 100644 index 0000000..07296cd --- /dev/null +++ b/python/samples/azure_functions/14_conversation_compaction/requirements.txt @@ -0,0 +1,17 @@ +# Agent Framework packages +# To use the deployed version, uncomment the lines below and comment out the local installation lines +# agent-framework-foundry +# agent-framework-azurefunctions + +# Local installation (for development and testing) +# Each package must be listed explicitly because pip doesn't resolve uv workspace sources. +# Without explicit entries, pip would fetch transitive dependencies from PyPI instead of local source. +agent-framework-foundry>=1.10.1,<2 # Foundry support from PyPI (pulls in core) +-e ../../../packages/durabletask # Durable Task support - dependency of azurefunctions +-e ../../../packages/azurefunctions # Azure Functions integration - the main package for this sample + +# Azure authentication +azure-identity + +# Local environment loading +python-dotenv diff --git a/python/uv.lock b/python/uv.lock index 19e5653..202cb8d 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -109,6 +109,7 @@ test = [ { name = "agent-framework-foundry", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "agent-framework-openai", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "azure-monitor-opentelemetry", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "jsonschema", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "mcp", extra = ["ws"], marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "redis", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, ] @@ -141,6 +142,7 @@ test = [ { name = "agent-framework-foundry", specifier = ">=1.10.1,<2" }, { name = "agent-framework-openai", specifier = ">=1.10.1,<2" }, { name = "azure-monitor-opentelemetry" }, + { name = "jsonschema" }, { name = "mcp", extras = ["ws"] }, { name = "redis" }, ] @@ -153,6 +155,8 @@ dependencies = [ { name = "agent-framework-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "durabletask", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "durabletask-azuremanaged", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "opentelemetry-api", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "pydantic", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "python-dateutil", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, ] @@ -166,6 +170,8 @@ requires-dist = [ { name = "agent-framework-core", specifier = ">=1.13.0,<2" }, { name = "durabletask", specifier = ">=1.5.0,<2" }, { name = "durabletask-azuremanaged", specifier = ">=1.4.0,<2" }, + { name = "opentelemetry-api", specifier = ">=1.39.0,<2" }, + { name = "pydantic", specifier = ">=2.11,<3" }, { name = "python-dateutil", specifier = ">=2.8.0,<3" }, ] diff --git a/schemas/durable-agent-entity-state.json b/schemas/durable-agent-entity-state.json index 53ac064..fb024d8 100644 --- a/schemas/durable-agent-entity-state.json +++ b/schemas/durable-agent-entity-state.json @@ -1,10 +1,12 @@ { "$schema": "https://json-schema.org/draft/2020-12/schema", "$id": "https://github.com/microsoft/agent-framework-durable-extension/schemas/durable-agent-entity-state.json", + "description": "Durable agent state. Version 2 separates response delivery and completion evidence from the mutable model transcript. Legacy version 1 layouts remain readable. Readers preserve unknown properties, including nested entry, message, content and usage properties, when writing state back.", "$defs": { "usage": { "type": "object", "description": "Token usage statistics.", + "additionalProperties": true, "properties": { "inputTokenCount": { "type": "integer" }, "outputTokenCount": { "type": "integer" }, @@ -67,7 +69,7 @@ "uri": { "type": "string", "description": "The URI." }, "mediaType": { "type": "string", "description": "The media type of the URI." } }, - "required": ["$type", "uri", "mediaType"] + "required": ["$type", "uri"] }, "usageContent": { "type": "object", @@ -94,7 +96,7 @@ "$type": { "type": "string", "const": "functionCall" }, "callId": { "type": "string", "description": "The identifier of the function being called." }, "name": { "type": "string", "description": "The name of the function being called." }, - "arguments": { "type": "object", "description": "The arguments provided to the function call." } + "arguments": { "type": ["object", "string"], "description": "The arguments provided to the function call, either a mapping or the original core argument string." } }, "required": ["$type", "callId", "name"] }, @@ -118,6 +120,21 @@ "required": ["$type", "content"] }, "chatContentItem": { + "type": "object", + "additionalProperties": true, + "properties": { + "extensionData": { + "type": "object", + "properties": { + "coreContent": { + "type": "object", + "description": "Canonical core fields not represented by this content subtype's existing shared-schema fields. Does not duplicate text, URI, arguments or result. Function results retain canonical nested items here alongside their legacy result representation. Known shared-schema fields remain authoritative when edited.", + "additionalProperties": true + } + }, + "additionalProperties": true + } + }, "oneOf": [ { "$ref": "#/$defs/dataContent" }, { "$ref": "#/$defs/errorContent" }, @@ -129,20 +146,42 @@ { "$ref": "#/$defs/textContent" }, { "$ref": "#/$defs/textReasoningContent" }, { "$ref": "#/$defs/uriContent" }, - { "$ref": "#/$defs/unknownContent" } + { "$ref": "#/$defs/unknownContent" }, + { + "type": "object", + "properties": { + "$type": { + "type": "string", + "minLength": 1, + "not": { "enum": ["data", "error", "functionCall", "functionResult", "hostedFile", "hostedVectorStore", "usage", "text", "reasoning", "uri", "unknown"] } + } + }, + "required": ["$type"], + "additionalProperties": true + } ] }, "chatMessage": { "type": "object", + "additionalProperties": true, "description": "Single chat message exchanged with the agent.", "properties": { "authorName": { "type": "string", "description": "The name of the author of the message." }, - "role": { "type": "string", "enum": ["user", "assistant", "system", "tool"] }, + "role": { "type": "string", "description": "Core message role, including user, assistant, system, developer and tool." }, "contents": { "type": "array", + "description": "Model transcript content, which may be changed or removed by compaction and retention. Empty content remains valid for legacy records, but version 2 does not require contentless request or response mirrors when another provider owns history. Version 2 delivers responses from responseMailbox, never by reconstructing them from this transcript.", "items": { "$ref": "#/$defs/chatContentItem" } }, - "createdAt": { "type": "string", "format": "date-time", "description": "When this message was created (RFC 3339)." } + "createdAt": { "type": "string", "format": "date-time", "description": "When this message was created (RFC 3339)." }, + "messageId": { + "type": "string", + "description": "Stable identity for this message. Context management reconciles its results back onto stored messages by this value, so an implementation that drops it on round-trip silently loses compaction state. Assigned by the runtime when the producer left it unset." + }, + "extensionData": { + "type": "object", + "description": "Message-level metadata, carrying the annotations context management writes (for example exclusion markers and summary markers). Must round-trip: discarding it loses compaction state rather than failing loudly." + } }, "required": ["role"] }, @@ -153,6 +192,8 @@ }, "conversationEntry": { "type": "object", + "description": "Fields shared by the known conversation entry kinds, discriminated by $type. Additional properties on known entries must survive a read/write round-trip. Future entry kinds are preserved separately as opaque entries rather than interpreted as requests or responses.", + "additionalProperties": true, "properties": { "createdAt": { "type": "string", "format": "date-time", "description": "When this exchange was created (RFC 3339)." }, "correlationId": { "type": "string", "description": "An optional correlation ID to group related exchanges." }, @@ -164,6 +205,7 @@ { "$ref": "#/$defs/conversationEntry" } ], "description": "The request (i.e. prompt) sent to the agent.", + "required": ["$type"], "properties": { "$type": { "type": "string", "const": "request" }, "orchestrationId": { @@ -184,7 +226,8 @@ "allOf": [ { "$ref": "#/$defs/conversationEntry" } ], - "description": "The response received from the agent.", + "description": "A response in the mutable model transcript, not an immutable delivery result. Legacy version 1 polling reads this entry. Version 2 polling reads responseMailbox instead, including when transcript content is compacted, retained elsewhere or removed.", + "required": ["$type"], "properties": { "$type": { "type": "string", "const": "response" }, "usage": { @@ -192,24 +235,195 @@ } } }, + "agentErrorResponse": { + "allOf": [ + { "$ref": "#/$defs/conversationEntry" } + ], + "description": "A turn that failed, never replayed to the model as conversation. The distinction is carried by $type rather than a transient flag. Legacy version 1 polling can return this entry; version 2 delivers the error through responseMailbox independently of transcript retention.", + "required": ["$type"], + "properties": { + "$type": { "type": "string", "const": "errorResponse" }, + "usage": { + "$ref": "#/$defs/usage" + } + } + }, + "compaction": { + "allOf": [ + { "$ref": "#/$defs/conversationEntry" } + ], + "description": "A message produced by context compaction, such as a summary replacing the turns it stands in for. Part of the model's transcript and positioned in conversation order, but it answers no request, so it carries no correlation ID and is never returned to a caller polling for a response.", + "required": ["$type"], + "properties": { + "$type": { "type": "string", "const": "compaction" } + } + }, + "opaqueConversationEntry": { + "type": "object", + "description": "An entry from a future writer. Preserve the entire object unchanged, but do not replay it to the model or return it as a response. This branch excludes every known discriminator so malformed known entries cannot bypass their typed contracts.", + "properties": { + "$type": { + "type": "string", + "minLength": 1, + "not": { "enum": ["request", "response", "errorResponse", "compaction"] } + } + }, + "required": ["$type"], + "additionalProperties": true + }, + "coreContent": { + "type": "object", + "description": "Inline core Content.to_dict() JSON. Uses the core type discriminator and snake_case fields, not the transcript's $type conversion. Content metadata and nested content remain part of the delivery snapshot.", + "properties": { + "type": { "type": "string", "minLength": 1 } + }, + "required": ["type"], + "additionalProperties": true + }, + "coreMessage": { + "type": "object", + "description": "Inline core Message.to_dict() JSON, including author, identity and message metadata.", + "properties": { + "type": { "type": "string", "const": "message" }, + "role": { "type": "string" }, + "contents": { "type": "array", "items": { "$ref": "#/$defs/coreContent" } }, + "author_name": { "type": "string" }, + "message_id": { "type": "string" }, + "additional_properties": { "type": "object" } + }, + "required": ["role", "contents"], + "additionalProperties": true + }, + "coreAgentResponse": { + "type": "object", + "description": "An independent, inline base-response JSON snapshot restored through the version-aware durable response loader, not a transcript entry, JSON-encoded string or storage reference. Unknown envelope fields remain in storage and are filtered only for the consumer. Raw SDK representations are not persisted.", + "properties": { + "type": { "type": "string", "const": "agent_response" }, + "_durable_response_version": { "type": "integer", "const": 1, "description": "Durable response envelope version. Absent on legacy inline snapshots." }, + "_durable_value_by_name": { "type": "boolean", "description": "Validate the retained structured value by field name rather than serialization alias when true." }, + "messages": { "type": "array", "items": { "$ref": "#/$defs/coreMessage" } }, + "response_id": { "type": "string" }, + "agent_id": { "type": "string" }, + "created_at": { "type": "string", "format": "date-time" }, + "finish_reason": { "type": "string" }, + "usage_details": { "type": "object", "additionalProperties": { "type": ["integer", "null"] } }, + "value": { "description": "The structured result in JSON form, when present. Capture it independently of mutable transcript text, including when core keeps the value outside to_dict()." }, + "continuation_token": { "type": "object", "description": "Opaque core continuation metadata, when present on the recorded response." }, + "additional_properties": { "type": "object" } + }, + "required": ["type", "messages"], + "additionalProperties": true + }, + "deliveryTimestamp": { + "type": "string", + "format": "date-time", + "description": "An RFC 3339 delivery timestamp with an explicit offset. The pattern enforces its shape even when optional format checking is unavailable.", + "pattern": "^[0-9]{4}-[0-9]{2}-[0-9]{2}[Tt][0-9]{2}:[0-9]{2}:[0-9]{2}(?:\\.[0-9]+)?(?:[Zz]|[+-][0-9]{2}:[0-9]{2})$", + "not": { "pattern": "\\s" } + }, + "responseMailboxEntry": { + "type": "object", + "description": "A retained response delivery snapshot. Expiry removes only this payload, not its completedCorrelations receipt. An expired or removed payload must never be reconstructed from conversationHistory.", + "properties": { + "response": { "$ref": "#/$defs/coreAgentResponse" }, + "createdAt": { "$ref": "#/$defs/deliveryTimestamp", "description": "When this delivery snapshot was recorded. For legacy conversion, this is the conversion time, not evidence of the original completion time." }, + "expiresAt": { "$ref": "#/$defs/deliveryTimestamp", "description": "Exclusive end of the configured response delivery window." } + }, + "required": ["response", "createdAt", "expiresAt"], + "additionalProperties": true + }, + "completedCorrelation": { + "type": "object", + "description": "Completion evidence retained until entity deletion, independently of response delivery expiry and transcript retention. A receipt without a live mailbox payload yields already_completed with response_expired, not a new execution or a transcript fallback.", + "properties": { + "completedAt": { "$ref": "#/$defs/deliveryTimestamp", "description": "When completion evidence was recorded. For a legacy receipt this is the conversion time, not a recovered historical timestamp." }, + "legacy": { "type": "boolean", "description": "True for a snapshot converted from surviving version 1 transcript data. Such a snapshot receives a fresh delivery grace window but is not claimed to be the immutable original response." } + }, + "required": ["completedAt"], + "additionalProperties": true + }, "data": { "type": "object", - "description": "The durable agent's state data.", + "description": "The durable agent's state data. Unknown properties must survive a read/write round-trip, including a version 1 to version 2 writer upgrade.", + "additionalProperties": true, "properties": { "conversationHistory": { "type": "array", - "description": "Ordered list of conversation entries.", - "items": { "$ref": "#/$defs/conversationEntry" } + "description": "Ordered model transcript entries when the durable runtime owns history. Every entry declares its kind through $type. Known kinds retain their typed contracts; future kinds are preserved opaquely and excluded from model context. This transcript is not the version 2 response delivery store and need not mirror externally owned history.", + "items": { + "oneOf": [ + { "$ref": "#/$defs/agentRequest" }, + { "$ref": "#/$defs/agentResponse" }, + { "$ref": "#/$defs/agentErrorResponse" }, + { "$ref": "#/$defs/compaction" }, + { "$ref": "#/$defs/opaqueConversationEntry" } + ] + } + }, + "session": { + "type": "object", + "description": "Opaque, owner-managed serialized session state, including provider state and service-issued conversation identifiers. .NET and Python use different shapes. Preserve owner state without interpreting provider names or treating every message-shaped slice as a duplicate of conversationHistory. Ownership determines which runtime working buffers are excluded; external provider state is not owned by transcript retention." + }, + "responseMailbox": { + "type": "object", + "description": "Version 2 delivery payloads keyed by correlation ID, separate from the model transcript. Each payload is an independent core AgentResponse JSON snapshot.", + "additionalProperties": { "$ref": "#/$defs/responseMailboxEntry" } + }, + "completedCorrelations": { + "type": "object", + "description": "Completion receipts keyed by correlation ID. Preserve them after mailbox expiry to distinguish a completed request from one that has never completed.", + "additionalProperties": { "$ref": "#/$defs/completedCorrelation" } + }, + "ingestedMessages": { + "type": "object", + "description": "Exact ingestion evidence keyed by message ID, retained independently of transcript content. Lists hold hashes of the complete delivered core message representation, so changed content with the same ID is distinguishable. Null is a legacy known-ID marker, not a claim about which content hashes were previously delivered.", + "additionalProperties": { + "oneOf": [ + { "type": "array", "items": { "type": "string" } }, + { "type": "null" } + ] + } + }, + "ingestedPositions": { + "type": "object", + "description": "Legacy version 1 scalar maximum positions keyed by workflow executor ID. Retained for legacy readers only. A scalar maximum does not prove delivery of skipped, sparse or evicted positions and cannot be automatically migrated into version 2 exact receipts without recorded delivery evidence. Non-empty legacy cursors require an explicit version-gated migration, rejected without state changes when that evidence is unavailable.", + "deprecated": true, + "additionalProperties": { "type": "integer", "minimum": 0 } + }, + "truncation": { + "type": "object", + "description": "Recorded retention loss from this entity's transcript, represented by a counter and timestamps rather than an unbounded list. This does not describe mailbox expiry or changes to externally owned history. Absence means no transcript eviction has been recorded, not that the entity owns a complete conversation.", + "properties": { + "evictedMessageCount": { + "type": "integer", + "minimum": 1, + "description": "Total messages retention has evicted over the life of this conversation." + }, + "firstEvictedAt": { + "type": "string", + "format": "date-time", + "description": "When this conversation first became incomplete." + }, + "lastEvictedAt": { + "type": "string", + "format": "date-time", + "description": "When retention last removed anything." + } + }, + "required": ["evictedMessageCount", "firstEvictedAt", "lastEvictedAt"] } } } }, "type": "object", + "additionalProperties": true, "properties": { "schemaVersion": { "type": "string", - "description": "Semantic version of this state schema. By convention, this should be the first property.", - "pattern": "^\\d+\\.\\d+\\.\\d+$" + "description": "Semantic version of the persisted layout. New writers emit 2.0.0. Readers also accept legacy 1.x layouts and preserve the read version until an explicit writer upgrade. Unknown major versions and missing versions must fail rather than reset existing state. By convention, this is the first property.", + "default": "2.0.0", + "pattern": "^[12]\\.[0-9]+\\.[0-9]+$", + "not": { "pattern": "\\s" } }, "data": { "$ref": "#/$defs/data" } },