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/decisions/0032-durable-thread-compaction.md b/docs/decisions/0032-durable-thread-compaction.md new file mode 100644 index 0000000..3aa4553 --- /dev/null +++ b/docs/decisions/0032-durable-thread-compaction.md @@ -0,0 +1,418 @@ +--- +status: proposed +contact: ahmedmuhsin +date: 2026-07-27 +deciders: +consulted: +informed: +--- + +# Thread Compaction for Durable Agents and Workflows + +> **How to read this.** The decision comes first. The later sections record the Python prototype +> and the gaps it exposed. **.NET is not implemented yet.** +> +> **Naming.** .NET's `ChatHistoryProvider` and Python's `HistoryProvider` are the same concept. The +> decision sections use the .NET name, and the implementation sections use the Python one. + +## Context and Problem Statement + +Long-running **durable** agents and workflows accumulate conversation history in durable +storage and replay it on every turn. Durable agents persist a full `ConversationHistory` in +entity state (`AgentEntity` → `DurableAgentState`). Durable workflows persist inter-executor +messages (`AgentExecutor.full_conversation`) as checkpointed envelopes. An in-memory agent keeps its +history in process RAM, where it disappears when the process recycles. Durable history instead +survives restarts and is reloaded on later turns. + +It helps to separate **three distinct pressures**, because they have different owners. + +| Pressure | What bounds it | Same in core? | Owner | +| --- | --- | --- | --- | +| **Context window**, the model's max input per call | the model | **Yes**, identical in core and durable | Compaction (in-run filter) | +| **Token cost / latency**, resending history each turn | tokens billed / round-trip | **Yes**, same mechanism | Compaction (in-run filter) | +| **Storage capacity**, the cumulative persisted state | backend state-size limit | **No**, durable-only | Backend offload and durable retention | + +The first two are per-operation and identical in both runtimes. The third is cumulative. +`ConversationHistory` is one blob appended to every turn and re-persisted whole, so it is bounded by +what the backend will store, and the two backends fail differently. **Durable Task Scheduler caps a +message at 1 MB.** The Azure Storage backend has no hard cap, because it compresses anything over +45 KB into a `-largemessages` blob, but it pays for size in CPU, I/O and memory. So one +backend stops working at the limit and the other degrades toward it, while a core process is bounded +only by RAM and resets on restart. + +**Storage capacity is an infrastructure concern, not a context-window concern.** It is relieved +first by raising the ceiling where blob offload is available and only then by deleting. A tool for +bounding what the model reads is not a tool for bounding what the backend holds. + +Core MAF already has a compaction system ([ADR-0019](https://github.com/microsoft/agent-framework/blob/main/docs/decisions/0019-python-context-compaction-strategy.md), +.NET `Microsoft.Agents.AI.Compaction`, Python `agent_framework._compaction`) with **two hooks**. + +1. **In-run filter.** A `CompactionProvider` (`AIContextProvider`) / `compaction_strategy` runs + before each model call. It is **non-lossy**. It filters the projection sent to the model and + stores incremental group state in the `AgentSession.StateBag`, leaving the underlying store + untouched. This hook works with **any** history provider, since it acts on the messages already + loaded into the invocation context. +2. **Store reducer.** **Lossily** rewrites the stored conversation, applying the same strategies at + the store instead of at the model call. Unlike the in-run filter, this hook is tied to a specific + storage mechanism in both languages. .NET exposes an `IChatReducer` on `InMemoryChatHistoryProvider` + only (bridged from any strategy by `strategy.AsChatReducer()`), and Python's + `CompactionProvider.after_strategy` reads the messages out of session state. Neither offers it to + a provider backed by anything else - see "Core Interface Gaps" below. + +The durable layer benefited from **neither**, because `AgentEntity` **bypassed the history +provider**: it created a fresh session per operation (so the StateBag - and any history provider +store or reducer in it - was discarded) and fed `ConversationHistory` directly as input messages. +Both the in-run filter's incremental state and the store reducer were thrown away every turn. + +The goal is **configuration parity**: a user's core compaction config must carry over to a durable +entity or workflow **unchanged**, reusing the same strategies and hooks on the durable runtime, +without a parallel durable compaction API. Storage retention is a separate deployment policy because +it has no meaning for an in-memory agent. + +**How should core compaction be reused on the durable runtime, in both .NET and Python, so the same +agent configuration bounds model input and the persisted store can be bounded separately?** + +## Decision Drivers + +- **Configuration parity.** The same core compaction config (strategies, `CompactionProvider`, + `IChatReducer`) must apply unchanged when moving core → durable entity → durable workflow. No + parallel durable-only API. +- **Reuse existing core hooks.** Do not reinvent triggers, strategies or grouping. Reuse the in-run + filter and the store reducer. +- **Separate storage capacity from context management.** Bound the model input with compaction + (parity with core), raise backend capacity where possible, and use observable deletion only as a + fallback. +- **Deleting is a last resort, and never silent.** Entity state is a state bag, not an immutable + system of record, so deleting from it is legitimate. But deletion should happen only when capacity + demands it, should remove no more than capacity demands, and should always be observable. +- **Determinism and idempotency.** Durable entity operations can be retried, so a lossy reducer + (especially LLM summarization) must not corrupt or diverge persisted state across retries. +- **Message-list correctness.** Preserve atomic groups (assistant tool-call plus tool-result, and + reasoning pairings) so the model input stays valid. +- **Cover both surfaces.** Durable agents **and** durable workflows, in **both** languages. Core's + compaction system is agent-level, so a workflow agent node inherits it unchanged. The conversation + chained *between* nodes is governed by `AgentExecutor`'s `context_mode` / `context_filter` seam, + which is a plain callable rather than the compaction system. That difference is real and is called + out rather than papered over. +- **Defer when the model provider owns the conversation.** When the chat client keeps history on the + service, the client holds no history for core compaction. "Service" here means the model provider, + not the durable entity. The entity's own durable record still follows its retention policy. + +## Considered Options + +- **Option 1, in-run filter only (rejected).** Reuse the agent's core compaction without changing + durable history. This bounds model input but not persisted state. +- **Option 2, bespoke pre-write compaction in the agent entity.** Add durable-specific code that + compacts `ConversationHistory` inside the entity operation before checkpoint. Rejected because it + duplicates core's store-reducer behavior. +- **Option 3, on-storage maintenance compaction (deferred).** Compact persisted history from a + separate entity operation. This may suit expensive summarization but does not prevent in-turn + growth. +- **Option 4, workflow context projection (chosen).** Honor `AgentExecutor.context_mode` and + `context_filter` for the `full_conversation` chained between executors. +- **Option 5, auto-derive a durable store reducer (rejected as default).** Derive a lossy reducer + from a configured in-run strategy. The explicit equivalent is `follow_compaction`. The default + must also protect agents with no compaction strategy. +- **Option 6, durable store as a `ChatHistoryProvider` (chosen).** Back the durable entity's + persisted conversation with a core `ChatHistoryProvider` implementation, so both core hooks apply + on the durable runtime from the user's unchanged configuration. The in-run filter runs in the + agent pipeline (L1), and a user-configured reducer or strategy can bound the durable store (L2, + opt-in). External history providers also rejoin the context pipeline, but the entity still keeps + its own conversation record, so they do not currently remove the need for retention. +- **Option 7, offload large payloads to blob storage (chosen where available).** Raise the ceiling + instead of reducing content, using the Durable Task Scheduler [large payload + extension](https://learn.microsoft.com/azure/durable-task/scheduler/durable-task-scheduler-large-payloads). + Non-lossy, and the same technique the Azure Storage backend has always used internally. + +## Decision Outcome + +Chosen option: **Option 6, express durable conversation storage as a core `ChatHistoryProvider`**, +combined with workflow context projection (Option 4). The two solve different surfaces. + +| Surface | Mechanism | Behavior | +| --- | --- | --- | +| **L1, agent context** | The user's configured core `CompactionProvider` / `compaction_strategy` | Non-lossy projection of model input. The same agent configuration works durably. | +| **L2, eager store pruning** | Core compaction annotations plus `retention="follow_compaction"` | Opt-in deletion of messages the user's strategy excluded. Python only today because .NET is blocked by duplicated compaction state (gap 4). | +| **L3, workflow context** | Existing `AgentExecutor.context_mode` / `context_filter` projection | Controls `full_conversation` passed between executors. This is not a core compaction hook. | +| **Capacity fallback** | Durable retention | Bounds entity state independently of whether compaction is configured. | + +This gives agent-level **configuration parity**, not byte-for-byte parity in every workflow cycle. +Durable workflow nodes intentionally deduplicate repeated upstream context before persisting it. The +L3 section explains the measured difference. + +Capacity is handled in this order: raise the ceiling non-lossily where blob offload is available, +honor an explicit `follow_compaction` choice, then evict under pressure. An exclusion normally means +only "do not send this to the model". It means "delete this" only under `follow_compaction`. + +### Retention + +| Mode | Behavior | +| --- | --- | +| `keep_all` | Never delete. The entity may reach the backend limit and fail. | +| `auto` **(default)** | Delete only under storage pressure, targeting the low watermark. | +| `follow_compaction` | Delete whatever compaction excluded every turn, then use the same pressure eviction as `auto` if the remaining state is still too large. | + +**How pressure eviction works.** After the turn is recorded and before the state is persisted, the +entity measures its serialized state. `auto` uses only this path. `follow_compaction` uses it after +eager pruning. Below the high watermark, nothing happens. Above it, the entity targets the low +watermark using detached message copies with existing exclusions cleared and +`TokenBudgetComposedStrategy(strategies=[])`. Clearing exclusions makes the budget reflect what is +stored, while the empty strategy list bypasses the user's context policy and uses core's +deterministic oldest-group fallback. System messages, atomic tool groups, and the newest exchange +are protected. The entity remeasures after each pass and logs what it removes. + +`max_state_bytes` defaults to `1_048_576`. High and low watermarks of `0.85` and `0.70` provide +hysteresis and room for estimation error. Measuring a 1 MB prototype state took about 8 ms. + +**Why not simply reduce the store by default.** A default-on reducer only helps agents that already +configured compaction, because nothing else marks messages excludable. It would leave every other +configuration exposed while changing behavior for only a subset of users. + +**Why not rely on blob offload alone.** It raises the ceiling roughly tenfold and does not remove it. +It is also unreachable on the Durable Functions Python path today (gap 6). Retention is therefore +the fallback that works on every host. + +**Service-managed model context** is outside compaction scope, mirroring ADR-0019. When the model +provider owns the conversation, the client holds no history to compact. Entity retention still +applies. See "Service-managed conversations". + +**Why workflows largely come "for free."** Durable workflow agent execution +(`DurableExecutorDispatcher.ExecuteAgentAsync`) runs an agent through the same +`DurableAIAgent → AgentEntity → inner agent` path as standalone durable agents, so **L1, L2 and +retention are inherited by workflow agent executors**. The workflow's own `full_conversation` between +executors does not pass through the agent, so it needs the separate **L3** hook. + +### Consequences + +- **Configuration parity.** Existing agent compaction configuration works durably without changing + the agent. Retention does not choose the current model projection. +- **Broad capacity protection.** Because pressure eviction lives in the entity, it covers external + providers, service-managed agents, and agents with no context pipeline. A single oversized newest + exchange can still fail because the current result is never evicted. +- **Proportionate deletion.** Under `auto`, the budget decides how much to remove. The user's context + strategy does not. +- **Larger entity change.** The history-provider design must preserve response polling and the + entity's conversation record. +- **Python-only eager pruning.** .NET would duplicate the transcript if it persisted current + compaction state (gap 4), so a .NET implementation could use pressure retention but not L2 yet. +- **Core workarounds.** Python must publish and reconcile a working buffer because core binds + store-side compaction to session state (gaps 1 and 2). +- **Threshold behavior.** `auto` changes behavior only near the capacity limit. This is less uniform + than always pruning, but it avoids changing unaffected conversations. + +### Validation + +**Done (Python).** Unit tests cover provider substitution, annotation round-trips, synthetic summary +insertion and reconciliation, all retention modes, session persistence, and workflow projection and +deduplication. The retention test drives a real agent through twenty turns against a reduced budget. +`keep_all` is the control proving the same run exceeds it. Scheduler integration covers persisted +annotations and message ids, external-provider session identity, schema conformance, and downstream +workflow context. + +**Outstanding.** Not covered yet. + +- Retention crossing the real scheduler limit against a live backend, rather than a reduced budget + in process. +- The .NET realization and its schema parity (gap 3), and the .NET compaction-state blocker (gap 4). +- Blob offload (Option 7) against a real scheduler. It remains unreachable through Durable Functions + Python 1.x and the 2.x preview (gap 6). +- Idempotency of an LLM-based reducer across simulated entity retries. + +## Cross-Cutting Design Details + +- Honor the user's configured reducer trigger. Durable registration must not change compaction + cadence. +- Reuse core grouping so tool-call/result and reasoning groups remain atomic. +- Pressure eviction is deterministic and uses the estimator tokenizer without a model call. Any + future LLM reducer must give summaries stable identities and be tested across retries. +- The durable history provider belongs in `AgentEntity`. Workflow projection belongs at the + existing `AgentExecutor.context_mode` / `context_filter` seam. + +## Core Interface Gaps for Pluggable History Providers + +Prototyping the Python `DurableHistoryProvider` surfaced places where the current contracts assume a +*session-state-backed* history provider. They are recorded here because they affect **any** provider +whose store is not session state (Cosmos, Valkey, durable), not just this one. The prototype works +around them, but the cleaner fix is upstream. + +1. **Store-side compaction is bound to session state rather than to the provider.** `CompactionProvider` + has two hooks, but only `before_strategy` works with any provider because it acts on invocation + context. `after_strategy` mutates `session.state[history_source_id]["messages"]` and assumes that + mutation rewrites storage. External providers can therefore bound model input but cannot use + core to rewrite their stores. .NET similarly exposes `IChatReducer` only on + `InMemoryChatHistoryProvider`. + + *Workaround:* the durable provider publishes a working buffer under the session-state key core + expects. *Upstream fix:* put store-rewrite compaction on the provider abstraction. + +2. **`save_messages()` is append-only.** The other half of the same open question. It receives only + new messages, so changes to existing messages and inserted summaries have no path back to storage. + *Workaround:* the durable provider reconciles its working buffer **by `message_id`** during + `after_run`. *Upstream fix:* add an explicit replace/flush operation alongside append. + +3. **Message-level metadata was not persisted (durable schema).** Python wrote + `extension_data` asymmetrically, so annotations disappeared on round-trip. This is fixed. The + shared schema now declares `messageId` and `extensionData` as round-trip-required, describes + `session` as runtime-discriminated, and has a conformance test. A validator would not previously + have rejected these fields because the schema permits extra properties. The defect was an + under-declared contract. + + .NET still loses `ChatMessage.AdditionalProperties` and `MessageId` in + `FromChatMessage`/`ToChatMessage`. Its `[JsonExtensionData]` property is only an overflow bucket for + unmapped JSON. `ChatMessage.MessageId` **does** exist in the pinned package and needs mapping. + Exclusions themselves live on `CompactionMessageGroup.IsExcluded`, while + `AdditionalProperties` carries the `_is_summary` marker, so mapping both fields is necessary but + not sufficient for .NET compaction parity. + +4. **.NET compaction state cannot be persisted without duplicating the transcript.** This is the + blocker behind "L2 is Python-only today". `CompactionProvider.State` is documented as living in + `AgentSession.StateBag`, holds `List`, and each group serializes its full + `ChatMessage` objects. Every run rewrites it wholesale. That leaves three unappealing choices for a + durable provider that also persists the session: + + | Choice | Consequence | + | --- | --- | + | Persist the session | The conversation is stored twice, in `ConversationHistory` and again in the state bag, so entity state roughly doubles instead of being bounded | + | Omit the provider state | Exclusions and summaries are discarded and summarization can re-run | + | Return only included messages | `CompactionMessageIndex.Update()` sees a trimmed front and rebuilds from scratch, losing the incremental state | + + None of this is inherent to the history-provider approach. It resolves if core can persist + lightweight compaction metadata keyed by `MessageId` rather than whole message copies. Until then + a .NET implementation can bound entity state through the retention path, which is independent of + `CompactionProvider`. + +5. **Provider cadence splits under per-service-call persistence.** With + `require_per_service_call_history_persistence=True`, history providers run per model call while + `CompactionProvider` remains once per run. Compaction can then annotate after the last history + flush, delaying persistence until the next flush. Only `HarnessAgent` enables this today, so the + gap is latent. + +6. **Blob offload is unreachable on the Durable Functions Python path.** Not a core gap but an + upstream gap. Version 1.x, which this package pins (`>=1.3.1,<2`), does not use the durabletask + SDK, so it has no Python-side payload-store seam. The 2.x previews (`2.0.0b1` and `2.0.0b2`, Python + 3.13+) depend on `durabletask>=1.9.0`, but `DurableFunctionsWorker` and + `DurableFunctionsClient` do not expose the base types' `payload_store` parameter. The direct + durabletask path does. *Upstream fix:* expose `payload_store` on both Functions types. + +Two more core gaps are described where they matter: the process-local state-type registry under +session persistence, and the lack of a public resolved history-ownership decision under +service-managed conversations. + +## L3 Realization: Workflow Context Parity + +In-process workflows give a downstream `AgentExecutor` the upstream conversation through +`AgentExecutorResponse.full_conversation`, governed by `context_mode` (`full` | `last_agent` | +`custom` + `context_filter`). The durable orchestrator previously flattened that to the **last +message's text**, so a downstream agent lost everything earlier nodes produced. + +Agent-level compaction needs no workflow-specific work: `AgentExecutor` passes its own session to +`agent.run()`, so the agent's `CompactionProvider` runs normally. Inter-executor context has no core +compaction hook. Durable instead honors the existing `context_mode` and invokes `context_filter` for +`custom` mode, then sends the projection as `RunRequest.context_messages`. Those messages become part +of the request entry and are visible to agent-level compaction. + +Cycles need deduplication because a node receives the accumulated upstream conversation again on +each visit. The orchestrator stamps each forwarded message as `wf_{executor}_{position}`. The entity +stores the highest ingested position per executor and drops older positions, keeping the newest +message as input when everything repeats. Per-executor watermarks are required because fan-out +branches can share a position. + +Stored-id comparison is insufficient: retention removes old ids, after which a cycle would re-ingest +exactly what was evicted and oscillate instead of converging. The small position map survives +deletion. Once content is evicted, the node no longer sees it. Re-ingesting it would defeat +retention. + +This intentionally differs from core in one measured case. On the third visit of a `full`-mode +cycle, core in-process sends 11 messages with repeated context while durable sends 8 after dedup. In +`last_agent` mode they are identical. Each durable node also keeps history keyed by workflow instance +and executor, so its memory survives restarts independently of the workflow envelope. + +## Zero-Configuration Registration + +Registration must not require edits to an agent that already works in core. The entity therefore +substitutes history at construction time. It shallow-copies the agent when substitution is needed, +so the caller's instance remains unchanged. + +| User configured | Durable behavior | +| --- | --- | +| Nothing | Inject a durable history provider, using the `source_id` core's auto-injected provider would have, so default-wired compaction still resolves. No compaction by default (same as core). | +| `InMemoryHistoryProvider` (± compaction) | Replace with the durable provider, **preserving `source_id` and `skip_excluded`** so any attached `CompactionProvider` keeps working untouched. | +| Cosmos / Redis / file / custom provider | **Leave alone.** The user chose where their conversation lives, and durable still supplies execution durability. | +| Service-managed history | **Leave alone.** The model service owns the conversation. Decided by core's precedence, explicit `store` first and then the client's `STORES_BY_DEFAULT`. | +| Agent without the core context pipeline | **Leave alone.** Falls back to replaying persisted history. | + +Preserving `source_id` is the load-bearing detail. `CompactionProvider` locates history through +`history_source_id` (default `"in_memory"`), so a provider swapped in under the same id is invisible +to the rest of the configuration. An explicit `DurableHistoryProvider` takes precedence. + +### Entity Context Ownership + +1. **Who supplies conversation context?** If the agent exposes core's context-provider pipeline, + the providers do, so the entity passes a session and delivers **only the new messages**. This + holds whether history lives in durable state, an external store, or the model service. +2. **Who bounds entity state?** Retention does, for every configuration, because the entity records + the conversation even when another provider owns model context. + +The entity therefore replays its own persisted history in exactly one case, an agent that does not +expose the context pipeline. Passing a session re-engages external providers and core's in-run +filter. It does not let core rewrite an external store (gap 1). The session id is derived from the +full entity identity, name plus key, so workflow nodes cannot share an external-provider key. + +### The session is persisted, not just its conversation id + +Providers use session state for data that must survive turns, including pending approvals. Because +the entity creates a session per operation, it persists the **whole serialized session** rather than +selecting fields. Two details prevent duplication and type loss: + +- The service-issued conversation id needs no bespoke field of its own - it is already part of + `AgentSession.to_dict()`. +- The durable history provider's own slice is **excluded** before persisting. It is derived from + `conversationHistory`, so storing it would duplicate the transcript. + +Restore applies the stored state onto a session created by the agent's own `create_session()`, so +the agent's session type is preserved. Core's state-type registry is process-local, so the entity +pre-registers serializable types already loaded in the process before restore. Pydantic state remains +a core gap because broad subclass discovery would be collision-prone. + +### Service-managed conversations + +When the model service stores the conversation, it identifies the thread with an id. The entity +creates a fresh session per operation, so that id is **persisted in durable state and restored on +the next turn** as part of the serialized session. Without it, every turn would start a new thread. + +Whether the service owns history is decided with **core's precedence, not the client class alone**. +An explicit `store` in the agent's options wins, and only when it is unset does the client's +`STORES_BY_DEFAULT` apply. This matters because clients that store by default (such as the Responses +API) are routinely put back into client-side mode with `store=False`. Consulting only +`STORES_BY_DEFAULT` would leave such an agent with a plain in-memory provider that the durable +runtime never persists, silently losing the conversation between turns. + +Core resolves this rule inside `Agent._run` and does not expose the result, so this layer +**re-derives it** and can drift if core changes. *Upstream fix:* expose the resolved decision. The +integration sample covers `store=False` against a store-by-default client. + +### Retention is a deployment policy, not agent configuration + +Compaction annotates, it does not delete. Deletion is configured at **registration** (an app-level +default with a per-agent override) rather than on the agent, so the agent definition stays portable: +the same agent runs in-memory where retention has no meaning. `auto` applies no context policy, but +deletion necessarily shortens future available history. Only `follow_compaction` treats a compaction +exclusion as permission to delete. Under `auto`, the storage budget alone chooses what is removed. + +## Out of Scope: Entity Lifetime + +Idle TTL and cleanup bound how many abandoned entities remain. They do not bound an actively used +entity because each interaction extends its lifetime. Cross-language TTL parity is a separate +decision. + +## More Information + +- Builds on [ADR-0019](https://github.com/microsoft/agent-framework/blob/main/docs/decisions/0019-python-context-compaction-strategy.md) (context compaction strategy), + which defines the in-run / pre-write / on-existing-storage compaction points and the atomic-group + constraint. +- Core reference mechanisms reused: `CompactionProvider` (in-run filter), `InMemoryChatHistoryProvider` + + `IChatReducer` (store reducer), `strategy.AsChatReducer()` bridge, and the existing external + `ChatHistoryProvider` implementations (`CosmosChatHistoryProvider`, `ValkeyChatHistoryProvider`). +- Relevant durable code: `AgentEntity` and `DurableAgentState` (durable agents), + `DurableExecutorDispatcher.ExecuteAgentAsync` (durable workflow agent execution), and + `AgentExecutor` (`context_mode` / `context_filter`, `full_conversation`). diff --git a/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py b/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py index dc90d13..5271786 100644 --- a/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py +++ b/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py @@ -44,6 +44,11 @@ execute_workflow_activity, plan_workflow_registration, ) +from agent_framework_durabletask._retention import ( + DEFAULT_MAX_STATE_BYTES, + DEFAULT_RETENTION, + RetentionMode, +) from agent_framework_durabletask._workflows.naming import ( SUBWORKFLOW_REQUEST_SEPARATOR, split_subworkflow_request_id, @@ -244,6 +249,9 @@ 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: int = DEFAULT_MAX_STATE_BYTES, ): """Initialize the AgentFunctionApp. @@ -263,6 +271,16 @@ 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 retention: Default conversation retention for agents hosted by this app, including + agents inside hosted workflows. ``auto`` deletes only under storage pressure, + ``keep_all`` never deletes and lets the entity fail at the backend limit, and + ``follow_compaction`` first deletes what compaction excluded, then uses the same + pressure eviction as ``auto`` if that is not enough. ``add_agent`` can override it + per agent. + :param max_state_bytes: Budget for serialized entity state. + :param workflow_retention: Retention for agent nodes inside hosted workflows. When None, + ``retention`` applies. Worth setting separately, since a workflow node's entity lives + for one orchestration while a standalone agent's can live indefinitely. :note: If no agents are provided, they can be added later using :meth:`add_agent`. """ @@ -283,6 +301,9 @@ def __init__( 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 | None = workflow_retention + self._max_state_bytes = max_state_bytes try: retries = int(max_poll_retries) @@ -419,6 +440,7 @@ def _register_workflow_primitives(self, workflow: Workflow) -> None: agent_executor.agent, callback=self.default_callback, entity_id=workflow_scoped_executor_id(workflow.name, agent_executor.id), + retention=self._workflow_retention, ) for executor in plan.activity_executors: # Set up a Functions activity trigger for each non-agent executor, scoped @@ -826,6 +848,7 @@ def add_agent( enable_mcp_tool_trigger: bool | None = None, *, entity_id: str | None = None, + retention: RetentionMode | None = None, ) -> None: """Add an agent to the function app after initialization. @@ -842,6 +865,7 @@ 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. Raises: ValueError: If the agent doesn't have a 'name' attribute. @@ -890,9 +914,15 @@ def add_agent( ) effective_callback = callback or self.default_callback + effective_retention: RetentionMode = self._retention if retention is None else retention self._setup_agent_functions( - agent, registration_name, effective_callback, effective_enable_http_endpoint, effective_enable_mcp_endpoint + agent, + registration_name, + effective_callback, + effective_enable_http_endpoint, + effective_enable_mcp_endpoint, + retention=effective_retention, ) logger.debug(f"[AgentFunctionApp] Agent '{registration_name}' added successfully") @@ -937,6 +967,8 @@ def _setup_agent_functions( callback: AgentResponseCallbackProtocol | None, enable_http_endpoint: bool, enable_mcp_tool_trigger: bool, + *, + retention: RetentionMode = DEFAULT_RETENTION, ) -> None: """Set up the HTTP trigger, entity, and MCP tool trigger for a specific agent. @@ -946,6 +978,7 @@ 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. """ logger.debug(f"[AgentFunctionApp] Setting up functions for agent '{agent_name}'...") @@ -956,7 +989,7 @@ 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) if enable_mcp_tool_trigger: agent_description = agent.description @@ -1098,6 +1131,8 @@ def _setup_agent_entity( agent: SupportsAgentRun, agent_name: str, callback: AgentResponseCallbackProtocol | None, + *, + retention: RetentionMode = DEFAULT_RETENTION, ) -> None: """Register the durable entity responsible for agent state. @@ -1105,6 +1140,7 @@ 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. """ # Use the prefixed entity name for both registration and function naming entity_name_with_prefix = AgentSessionId.to_entity_name(agent_name) @@ -1117,7 +1153,7 @@ def entity_function(context: df.DurableEntityContext) -> None: - run_agent: (Deprecated) Execute the agent with a message - reset: Clear conversation history """ - entity_handler = create_agent_entity(agent, callback) + entity_handler = create_agent_entity(agent, callback, retention=retention) entity_handler(context) # Set function name for Azure Functions (used in function.json generation) diff --git a/python/packages/azurefunctions/agent_framework_azurefunctions/_entities.py b/python/packages/azurefunctions/agent_framework_azurefunctions/_entities.py index 83ad50a..46de237 100644 --- a/python/packages/azurefunctions/agent_framework_azurefunctions/_entities.py +++ b/python/packages/azurefunctions/agent_framework_azurefunctions/_entities.py @@ -21,6 +21,7 @@ AgentResponseCallbackProtocol, run_agent_coroutine, ) +from agent_framework_durabletask._retention import DEFAULT_MAX_STATE_BYTES, DEFAULT_RETENTION, RetentionMode logger = logging.getLogger("agent_framework.azurefunctions") @@ -47,10 +48,16 @@ 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, + *, + retention: RetentionMode = DEFAULT_RETENTION, + max_state_bytes: int = DEFAULT_MAX_STATE_BYTES, ) -> Callable[[df.DurableEntityContext], None]: """Factory function to create an agent entity class. @@ -58,6 +65,12 @@ 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: + retention: How much of the conversation durable state may discard. ``auto`` deletes only + under storage pressure, ``keep_all`` never deletes, and ``follow_compaction`` first + deletes what compaction excluded, then uses pressure eviction if needed. + max_state_bytes: Budget for serialized entity state. + Returns: Entity function configured with the agent """ @@ -69,7 +82,13 @@ 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=max_state_bytes, + ) operation = context.operation_name 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_af_context.py b/python/packages/azurefunctions/agent_framework_azurefunctions/_workflow_af_context.py index eaf99a5..96fe027 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,20 @@ 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, + ) -> Any: + return build_agent_task( + AzureFunctionsAgentExecutor(self._context), + executor_id, + message, + orchestration_instance_id, + context_messages, + ) def prepare_activity_task(self, activity_name: str, input_json: str) -> Any: orchestration_context: Any = self._context @@ -103,3 +111,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/test_01_single_agent.py b/python/packages/azurefunctions/tests/integration_tests/test_01_single_agent.py index ff0e425..940189d 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 @@ -93,13 +93,13 @@ 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.""" + """History must accumulate *and* reach the model on later turns.""" 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}, ) assert response1.status_code in [200, 202] @@ -107,13 +107,17 @@ def test_conversation_continuity(self) -> None: data1 = response1.json() assert data1["message_count"] == 2 # Initial + reply - # Second message in same session + # Second message in same session; only answerable from persisted history. response2 = self.helper.post_json( - f"{self.base_url}/run", {"message": "What about San Francisco?", "session_id": session_id} + f"{self.base_url}/run", + {"message": "What is my favorite animal? Reply with just the animal name.", "session_id": session_id}, ) assert response2.status_code == 200 data2 = response2.json() assert data2["message_count"] == 4 + assert "axolotl" in str(data2["response"]).lower(), ( + f"Agent lost conversation context across turns. Got: {data2['response']!r}" + ) else: # In async mode, we can't easily test message count # Just verify we can make multiple calls 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..4d07c6e --- /dev/null +++ b/python/packages/azurefunctions/tests/integration_tests/test_14_conversation_compaction.py @@ -0,0 +1,83 @@ +# 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 +configuration, mirroring the standalone durabletask coverage. + +The function app is automatically started by the test fixture. + +Prerequisites: +- Azure OpenAI credentials configured (see packages/azurefunctions/tests/integration_tests/.env.example) +- Azurite or Azure Storage account configured + +Usage: + uv run pytest packages/azurefunctions/tests/integration_tests/test_14_conversation_compaction.py -v +""" + +import uuid + +import pytest + +# 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}) + assert response.status_code in [200, 202] + return response.json() + + 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() diff --git a/python/packages/azurefunctions/tests/test_app.py b/python/packages/azurefunctions/tests/test_app.py index 15fff90..8560ad3 100644 --- a/python/packages/azurefunctions/tests/test_app.py +++ b/python/packages/azurefunctions/tests/test_app.py @@ -269,7 +269,7 @@ 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="auto") assert app._agent_metadata["OverrideAgent"].http_endpoint_enabled is True def test_agent_override_disables_http_route_when_app_enabled(self) -> None: @@ -286,7 +286,7 @@ 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="auto") assert app._agent_metadata["DisabledOverride"].http_endpoint_enabled is False def test_multiple_apps_independent(self) -> None: diff --git a/python/packages/durabletask/agent_framework_durabletask/__init__.py b/python/packages/durabletask/agent_framework_durabletask/__init__.py index a3e2727..7e91d30 100644 --- a/python/packages/durabletask/agent_framework_durabletask/__init__.py +++ b/python/packages/durabletask/agent_framework_durabletask/__init__.py @@ -50,10 +50,11 @@ ) from ._entities import AgentEntity, AgentEntityStateProviderMixin from ._executors import DurableAgentExecutor +from ._history_provider import DurableHistoryBinding, DurableHistoryProvider 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 ._shim import DurableAIAgent, build_agent_task from ._worker import DurableAIAgentWorker from ._workflows.activity import execute_workflow_activity from ._workflows.client import DurableWorkflowClient @@ -159,6 +160,8 @@ def __dir__() -> list[str]: "DurableAgentStateUriContent", "DurableAgentStateUsage", "DurableAgentStateUsageContent", + "DurableHistoryBinding", + "DurableHistoryProvider", "DurableStateFields", "DurableTaskWorkflowContext", "DurableWorkflowClient", @@ -166,6 +169,7 @@ def __dir__() -> list[str]: "WorkflowOrchestrationContext", "WorkflowRegistrationPlan", "__version__", + "build_agent_task", "collect_hosted_workflows", "deserialize_workflow_output", "ensure_response_format", diff --git a/python/packages/durabletask/agent_framework_durabletask/_constants.py b/python/packages/durabletask/agent_framework_durabletask/_constants.py index 9e48b51..0e25e02 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_constants.py +++ b/python/packages/durabletask/agent_framework_durabletask/_constants.py @@ -131,6 +131,16 @@ 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" + + # Highest chained-conversation position ingested from each workflow executor. Survives + # retention, which identity-based duplicate detection cannot. + INGESTED_POSITIONS: Final[str] = "ingestedPositions" + 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..10e74db 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_durable_agent_state.py +++ b/python/packages/durabletask/agent_framework_durabletask/_durable_agent_state.py @@ -326,25 +326,42 @@ 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: Highest chained-conversation position taken from each workflow + executor. A workflow re-sends the whole conversation on every visit, and comparing + against stored ids stops working once retention deletes any of them, so the mark is + kept separately. 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 extension_data: dict[str, Any] | None 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, ) -> 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: Highest chained-conversation position taken from each workflow + executor, used to recognize context this entity has already recorded """ self.conversation_history = conversation_history or [] self.extension_data = extension_data + self.session = session + self.ingested_positions = ingested_positions def to_dict(self) -> dict[str, Any]: result: dict[str, Any] = { @@ -352,6 +369,10 @@ def to_dict(self) -> dict[str, Any]: } if self.extension_data is not None: result[DurableStateFields.EXTENSION_DATA] = self.extension_data + if self.session is not None: + result[DurableStateFields.SESSION] = self.session + if self.ingested_positions: + result[DurableStateFields.INGESTED_POSITIONS] = self.ingested_positions return result @classmethod @@ -359,6 +380,8 @@ def from_dict(cls, data_dict: dict[str, Any]) -> DurableAgentStateData: return 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), ) @@ -392,7 +415,7 @@ class DurableAgentState: """ # Durable Agent Schema version - SCHEMA_VERSION: str = "1.1.0" + SCHEMA_VERSION: str = "1.2.0" data: DurableAgentStateData schema_version: str = SCHEMA_VERSION @@ -611,10 +634,18 @@ def from_dict(cls, data: dict[str, Any]) -> DurableAgentStateRequest: @staticmethod def from_run_request(request: RunRequest) -> DurableAgentStateRequest: + # A workflow may deliver the upstream conversation instead of a single message. + if request.context_messages: + messages = [ + DurableAgentStateMessage.from_chat_message(Message.from_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), @@ -722,13 +753,19 @@ 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 def __init__( @@ -738,11 +775,13 @@ 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 def to_dict(self) -> dict[str, Any]: @@ -763,6 +802,10 @@ def to_dict(self) -> dict[str, Any]: result[DurableStateFields.CREATED_AT] = self.created_at.isoformat() if self.author_name is not None: result[DurableStateFields.AUTHOR_NAME] = self.author_name + if self.message_id is not None: + result[DurableStateFields.MESSAGE_ID] = self.message_id + if self.extension_data: + result[DurableStateFields.EXTENSION_DATA] = self.extension_data return result @classmethod @@ -775,6 +818,7 @@ def from_dict(cls, data: dict[str, Any]) -> DurableAgentStateMessage: 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), ) @@ -820,6 +864,7 @@ def from_chat_message(chat_message: Message) -> 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, + message_id=getattr(chat_message, "message_id", None), extension_data=dict(chat_message.additional_properties) if chat_message.additional_properties else None, ) @@ -841,6 +886,9 @@ 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 diff --git a/python/packages/durabletask/agent_framework_durabletask/_entities.py b/python/packages/durabletask/agent_framework_durabletask/_entities.py index 63f7098..13db529 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_entities.py +++ b/python/packages/durabletask/agent_framework_durabletask/_entities.py @@ -5,18 +5,22 @@ from __future__ import annotations import inspect +import json import logging import warnings +from collections.abc import Sequence from datetime import datetime, timezone from typing import Any, cast from agent_framework import ( AgentResponse, AgentResponseUpdate, + AgentSession, Content, Message, ResponseStream, SupportsAgentRun, + register_state_type, ) from durabletask.entities import DurableEntity @@ -28,10 +32,73 @@ DurableAgentStateRequest, DurableAgentStateResponse, ) +from ._history_provider import ( + DurableHistoryBinding, + DurableHistoryProvider, + bind_durable_history, + ensure_durable_history, + unbind_durable_history, +) from ._models import RunRequest +from ._retention import ( + DEFAULT_MAX_STATE_BYTES, + DEFAULT_RETENTION, + RetentionMode, + enforce_budget, + prunes_excluded, +) +from ._workflows.naming import parse_workflow_message_id 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() + + +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 +130,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`.""" @@ -117,10 +207,16 @@ def __init__( callback: AgentResponseCallbackProtocol | None = None, *, state_provider: AgentEntityStateProviderMixin, + retention: RetentionMode = DEFAULT_RETENTION, + max_state_bytes: int = DEFAULT_MAX_STATE_BYTES, ) -> None: - self.agent = agent + # 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 = max_state_bytes logger.debug("[AgentEntity] Initialized with agent type: %s", type(agent).__name__) @@ -169,18 +265,49 @@ async def run( logger.debug("[AgentEntity.run] Received SessionId %s Message: %s", session_id, run_request) state_request = DurableAgentStateRequest.from_run_request(run_request) + if run_request.context_messages: + state_request.messages = self._drop_already_stored(state_request.messages) 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 - ] + durable_history = self._find_durable_history_provider() + uses_context_pipeline = self._has_context_pipeline() + binding_token = ( + bind_durable_history( + DurableHistoryBinding(state_provider=self._state_provider, correlation_id=correlation_id) + ) + if durable_history is not None + else None + ) - run_kwargs: dict[str, Any] = {"messages": chat_messages, "options": options} + try: + 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 + # exactly one place and core providers work unchanged on the durable runtime. + session = self._create_session() + chat_messages = [ + replayable_message + for m in state_request.messages + if (replayable_message := self._to_replayable_message(m)) is not None + ] + 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 = [ + 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 = {"messages": chat_messages, "options": options} agent_run_response: AgentResponse = await self._invoke_agent( run_kwargs=run_kwargs, @@ -191,6 +318,8 @@ async def run( state_response = DurableAgentStateResponse.from_run_response(correlation_id, agent_run_response) self.state.data.conversation_history.append(state_response) + self._capture_session(session) + await self._enforce_retention() self.persist_state() return agent_run_response @@ -209,10 +338,193 @@ async def run( 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) + await self._enforce_retention() self.persist_state() return error_response + finally: + if binding_token is not None: + unbind_durable_history(binding_token) + + async def _enforce_retention(self) -> None: + """Bound durable state before it is persisted, unless the caller asked to keep everything. + + This lives on the entity rather than the history provider because the entity records the + conversation in every configuration, including external providers, service-managed agents + and agents with no context pipeline. Those are exactly the cases with no other mitigation. + """ + if self._retention == "keep_all": + return + await enforce_budget(self.state, max_state_bytes=self._max_state_bytes) + + 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)) + + 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. + + The durable history provider's own slice is dropped before persisting: it is derived from + ``conversation_history`` on every turn, so storing it would duplicate the transcript and + let the copy drift from the record of truth. It is removed *before* serializing rather + than after, because that slice holds the working message buffer and its position index, + and serializing the whole transcript only to discard it is pure waste. + + Provider state is arbitrary, so the payload is checked before it replaces the last good + one. Core neither raises nor warns on a value it cannot serialize, it passes the live + object through, and the entity state provider serializes eagerly. An unusable payload + would therefore fail the save, and fail it again from the error handler, masking whatever + the agent actually returned. + """ + 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 + 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) + except (TypeError, ValueError) as exc: + logger.warning( + "[AgentEntity] Session state could not be serialized and was not persisted, so the " + "previous turn's state is kept. A context provider is holding a value that is not " + "JSON-compatible: %s", + exc, + ) + return + self.state.data.session = payload + + def _drop_already_stored(self, messages: list[DurableAgentStateMessage]) -> list[DurableAgentStateMessage]: + """Filter out chained conversation this entity has already recorded. + + A workflow node that runs more than once (for example in a cycle) receives the whole + upstream conversation each time. Without filtering it re-records all of it on every visit. + + Filtering is by **position**, not by stored identity. The obvious check, "is this id + already in my history", stops working the moment retention evicts anything: those ids + leave the comparison set, the orchestrator re-sends them because its own conversation is + never evicted, and the entity re-records exactly what was deleted. That oscillates instead + of settling. A high-water mark per producing executor is unaffected by deletion, and is + per executor rather than global because a fan-out gives two branches the same position. + + Messages without a workflow id fall back to the identity check, which is enough for them + because nothing re-delivers them. + + The final message is always kept so the agent still receives an input. + """ + ingested = dict(self.state.data.ingested_positions or {}) + seen: dict[str, int] = {} + kept: list[DurableAgentStateMessage] = [] + + known_ids = { + stored.message_id + for entry in self.state.data.conversation_history + for stored in entry.messages + if stored.message_id + } + + for message in messages: + marker = parse_workflow_message_id(message.message_id) + if marker is not None: + executor, position = marker + seen[executor] = max(seen.get(executor, -1), position) + if position <= ingested.get(executor, -1): + continue + elif message.message_id and message.message_id in known_ids: + continue + kept.append(message) + + for executor, position in seen.items(): + ingested[executor] = max(ingested.get(executor, -1), position) + if ingested: + self.state.data.ingested_positions = ingested + + if not kept and messages: + # Keep the newest message so the agent still has an input, but drop the id it shares + # with the copy already in history. Two stored messages under one id collide in the + # compaction position map, so annotations and pruning would target the wrong one. + repeated = messages[-1] + repeated.message_id = None + return [repeated] + 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()." + ) + session: Any = create_session(session_id=self._state_provider.core_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 + @staticmethod def _to_replayable_message(message: DurableAgentStateMessage) -> Message | None: """Convert persisted history into a message safe to replay into chat clients.""" @@ -372,3 +684,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..90f3b54 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_executors.py +++ b/python/packages/durabletask/agent_framework_durabletask/_executors.py @@ -155,11 +155,20 @@ 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, ) -> RunRequest: """Create a RunRequest from message and options.""" correlation_id = self.generate_unique_id() @@ -179,6 +188,8 @@ def get_run_request( wait_for_response=wait_for_response, correlation_id=correlation_id, options=opts, + context_messages=context_messages, + orchestration_id=self._orchestration_id(), ) def _create_acceptance_response(self, correlation_id: str) -> AgentResponse: @@ -449,23 +460,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..75f72ba --- /dev/null +++ b/python/packages/durabletask/agent_framework_durabletask/_history_provider.py @@ -0,0 +1,495 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""A core ``HistoryProvider`` backed by durable entity state. + +This lets the durable runtime plug into the Agent Framework context-provider pipeline +instead of managing conversation history itself. Because the agent's own history +provider supplies context, core compaction (``CompactionProvider``) works unchanged and +its annotations are persisted alongside the messages in durable entity state - a single +stored copy, no side-car session blob. + +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 +from typing import TYPE_CHECKING, Any, cast + +from agent_framework import HistoryProvider, InMemoryHistoryProvider, Message, SupportsAgentRun + +from ._durable_agent_state import DurableAgentStateEntry, DurableAgentStateMessage, DurableAgentStateResponse + +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 + """Correlation id of the in-flight request, whose entry is excluded from loaded history.""" + + +_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): + """History provider whose store is the durable entity's conversation history. + + The durable entity remains the writer of record for requests and responses, so this + provider does not append messages itself (``store_inputs``/``store_outputs`` are off). + What it does provide is: + + * **load** - flattens persisted conversation history into ``Message`` objects, restoring + any compaction annotations that were stored with them. + * **flush** - writes annotations that compaction applied during the run back into the + persisted messages, so compaction state survives across entity operations. + + 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. This is **lossy** and opt-in - it is what actually bounds the + size of persisted state. + """ + + DEFAULT_SOURCE_ID = "durable_history" + + def __init__( + self, + source_id: str | None = None, + *, + skip_excluded: bool = True, + prune_excluded: bool = False, + ) -> None: + """Initialize the durable history provider. + + Args: + source_id: Unique identifier for this provider instance. + skip_excluded: Omit compaction-excluded messages from loaded context. + prune_excluded: Physically delete excluded messages from durable storage + on flush. Lossy, so it is disabled by default. + """ + super().__init__( + source_id=source_id or self.DEFAULT_SOURCE_ID, + load_messages=True, + # The durable entity owns appends to conversation history. + store_inputs=False, + store_outputs=False, + ) + 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.""" + yield from replayable_entries( + binding.state_provider.state.data.conversation_history, + correlation_id=binding.correlation_id, + ) + + @staticmethod + def _synthetic_message_id(entry: DurableAgentStateEntry, index: int) -> str: + """Build an id for a stored message that arrived without one. + + 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 unique within the conversation history. + """ + # 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() + return f"durable_{entry.json_type.value}_{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 = 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, + ) + + 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 [] + + loaded: list[Message] = [] + id_map: dict[str, tuple[DurableAgentStateEntry, int]] = {} + for entry, index in self._replayable_entries(binding): + stored = entry.messages[index] + message = self._to_message(stored) + if message is None: + continue + if not message.message_id: + # Give every loaded message a stable identity so compaction results can be + # reconciled back onto durable state on flush. + message.message_id = self._synthetic_message_id(entry, index) + stored.message_id = message.message_id + loaded.append(message) + id_map[message.message_id] = (entry, index) + + 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: + """No-op: the durable entity appends requests and responses to its own state.""" + return + + @staticmethod + def _is_service_managed(session: Any) -> bool: + """Return whether the conversation is stored by the model service, not by us.""" + return bool(getattr(session, "service_session_id", None)) + + 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.""" + if self._is_service_managed(session): + logger.debug("[DurableHistoryProvider] Session is service-managed, skipping durable history load.") + return + await super().before_run(agent=agent, session=session, context=context, state=state) + + async def after_run( + self, + *, + agent: Any, + session: Any, + context: Any, + state: dict[str, Any], + ) -> None: + """Flush compaction annotations from the working buffer into durable state.""" + if self._is_service_managed(session): + return + self.flush(state) + + def flush(self, state: dict[str, Any]) -> None: + """Persist compaction results back into 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. + + Args: + state: The provider-scoped session state holding the working buffer. + """ + binding = current_durable_history_binding() + if binding is None: + return + + raw_buffer = state.get(WORKING_BUFFER_KEY) + raw_positions = state.get(POSITIONS_KEY) + if not isinstance(raw_buffer, list) or not isinstance(raw_positions, dict): + return + buffer = cast("list[Message]", raw_buffer) + stored_by_id = cast("dict[str, tuple[DurableAgentStateEntry, int]]", raw_positions) + + pruned: list[tuple[DurableAgentStateEntry, DurableAgentStateMessage]] = [] + # 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: + annotations = dict(message.additional_properties) if message.additional_properties else None + position = stored_by_id.get(message.message_id) if message.message_id else None + + if position is None: + inserted = self._insert_new_message(binding, message, after=last_known) + if inserted is not None: + # The insertion pushed everything after it in that entry along by one, so the + # recorded positions have to move too or later updates land on the wrong message. + self._shift_positions(stored_by_id, inserted) + last_known = inserted + continue + + entry, index = position + stored = entry.messages[index] + stored.extension_data = annotations + last_known = position + if self.prune_excluded and annotations and annotations.get(EXCLUDED_KEY): + pruned.append((entry, stored)) + + if pruned: + self._prune(binding, pruned) + + binding.state_provider.persist_state() + + @staticmethod + def _shift_positions( + stored_by_id: dict[str, tuple[DurableAgentStateEntry, int]], + inserted: tuple[DurableAgentStateEntry, int], + ) -> None: + """Move recorded positions that an insertion pushed further along their entry. + + Args: + stored_by_id: Recorded ``message_id`` to position mapping, updated in place. + inserted: The entry and index the new message was inserted at. + """ + entry, index = inserted + for message_id, (stored_entry, stored_index) in list(stored_by_id.items()): + if stored_entry is entry and stored_index >= index: + stored_by_id[message_id] = (stored_entry, stored_index + 1) + + @staticmethod + def _insert_new_message( + binding: DurableHistoryBinding, + message: Message, + *, + after: tuple[DurableAgentStateEntry, int] | None, + ) -> tuple[DurableAgentStateEntry, int] | None: + """Persist a message that compaction produced (for example a summary).""" + stored = DurableAgentStateMessage.from_chat_message(message) + if after is not None: + entry, index = after + entry.messages.insert(index + 1, stored) + return entry, index + 1 + + history = binding.state_provider.state.data.conversation_history + if not history: + return None + first = history[0] + first.messages.insert(0, stored) + return first, 0 + + @staticmethod + def _prune( + binding: DurableHistoryBinding, + pruned: list[tuple[DurableAgentStateEntry, DurableAgentStateMessage]], + ) -> None: + """Physically remove excluded messages (and any entries left empty). + + Removal is by identity rather than index, since insertions earlier in this flush may have + moved messages within their entry. + """ + prune_messages(binding.state_provider.state.data.conversation_history, pruned) + + +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. + + Shared by the history provider and by retention, so both agree on which stored messages are + real conversation rather than bookkeeping. + + Args: + history: The entity's conversation history. + correlation_id: The in-flight request, which is delivered as run input rather than history. + + Yields: + Each replayable message as its owning entry and its index within that entry. + """ + for entry in history: + if isinstance(entry, DurableAgentStateResponse) and entry.is_error: + 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: + """Physically remove the given messages, and any entries left empty. + + 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. + """ + for entry, stored in pruned: + for index, candidate in enumerate(entry.messages): + if candidate is stored: + del entry.messages[index] + break + + remaining = [entry for entry in history if entry.messages] + if len(remaining) != len(history): + history[:] = remaining + + +def _service_stores_history(agent: Any) -> bool: + """Return whether the service keeps conversation history for this agent. + + Mirrors core's precedence: an explicit ``store`` in the agent's default options wins, and only + when it is 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 with ``store=False``, + in which case durable history is what makes the conversation survive. + """ + 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 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. + + The rules mirror what core would do, so behavior stays predictable: + + * **No history provider** - a :class:`DurableHistoryProvider` is added. It uses the same + ``source_id`` core's auto-injected provider would have, so a ``CompactionProvider`` left on + its defaults still finds it. + * **In-memory history** - replaced by a :class:`DurableHistoryProvider` carrying the *same* + ``source_id`` and ``skip_excluded``, so any compaction wired to it keeps working untouched. + * **Any other history provider** (Cosmos, Redis, file, custom) - left alone. The user chose + where their conversation lives, and durable still provides execution durability. + * **Service-managed history** - left alone. The model service owns the conversation. + * **Agents without the core context pipeline** - 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. It only affects providers this function creates. + + Returns: + The agent to run, either unchanged or a shallow copy with durable-backed history. + """ + providers = getattr(agent, "context_providers", None) + if not isinstance(providers, (list, tuple)): + return agent + + if _service_stores_history(agent): + logger.debug( + "[DurableHistoryProvider] Agent %s stores history service-side, leaving providers unchanged.", + getattr(agent, "name", type(agent).__name__), + ) + 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 the source_id core's auto-injected provider would use so default-wired + # compaction keeps resolving. + updated = [ + DurableHistoryProvider( + source_id=InMemoryHistoryProvider.DEFAULT_SOURCE_ID, + prune_excluded=prune_excluded, + ), + *provider_list, + ] + elif isinstance(existing, InMemoryHistoryProvider): + replacement = DurableHistoryProvider( + source_id=existing.source_id, + skip_excluded=existing.skip_excluded, + prune_excluded=prune_excluded, + ) + 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 + + try: + clone = copy.copy(agent) + clone.context_providers = updated # type: ignore[attr-defined] + except Exception: + logger.warning( + "[DurableHistoryProvider] Could not attach durable history to agent %s, " + "falling back to replaying persisted history.", + getattr(agent, "name", type(agent).__name__), + exc_info=True, + ) + return agent + + return clone diff --git a/python/packages/durabletask/agent_framework_durabletask/_models.py b/python/packages/durabletask/agent_framework_durabletask/_models.py index e8eabca..f6ac97d 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,7 @@ 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 def __init__( self, @@ -134,6 +140,7 @@ 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, ) -> None: self.message = message self.correlation_id = correlation_id @@ -145,6 +152,7 @@ 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 {} + self.context_messages = context_messages @staticmethod def coerce_role(value: str | None) -> str: @@ -158,7 +166,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 +181,8 @@ 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: + result["contextMessages"] = self.context_messages return result @classmethod @@ -200,6 +210,8 @@ 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") + context_messages = cast("list[dict[str, Any]]", raw_context) if isinstance(raw_context, list) else None return cls( message=data.get("message", ""), @@ -212,6 +224,7 @@ 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, ) 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..ed41bcd --- /dev/null +++ b/python/packages/durabletask/agent_framework_durabletask/_retention.py @@ -0,0 +1,216 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Bounding durable entity state so an agent does not simply stop working at the backend limit. + +Retention is a **capacity** concern, deliberately separate from compaction. Compaction decides what +the model should read. Retention decides what durable state can afford to hold. An exclusion made +for token cost is not consent to delete the record, so the two never share a decision. + +See ADR 0032, "Retention". +""" + +from __future__ import annotations + +import json +import logging +from typing import Literal, cast + +from agent_framework import ( + CharacterEstimatorTokenizer, + Message, + TokenBudgetComposedStrategy, +) + +from ._durable_agent_state import ( + DurableAgentState, + DurableAgentStateEntry, + DurableAgentStateMessage, +) +from ._history_provider import EXCLUDED_KEY, prune_messages, replayable_entries + +logger = logging.getLogger("agent_framework.durabletask") + +RetentionMode = Literal["keep_all", "auto", "follow_compaction"] +"""How much of the conversation durable state is allowed to discard. + +``keep_all`` + Never delete. The entity may reach the backend limit and fail. The honest choice when the + complete record matters more than availability. +``auto`` + Delete only under storage pressure, and only down to the low watermark. The default. +``follow_compaction`` + Delete whatever compaction excluded every turn, then use the same pressure eviction as + ``auto`` if the remaining state is still too large. +""" + +DEFAULT_RETENTION: RetentionMode = "auto" + +DEFAULT_MAX_STATE_BYTES = 1_048_576 +"""The Durable Task Scheduler message limit. Raise it when large payload offload is configured.""" + +HIGH_WATERMARK = 0.85 +"""Fraction of the budget that triggers eviction. + +Below 0.9 because the budget is approximate twice over, once in the byte-to-token estimate and once +because a message's non-text content is not counted when calibrating that estimate. +""" + +LOW_WATERMARK = 0.70 +"""Fraction of the budget to evict down to. + +The gap from the high watermark is hysteresis. Evicting to just under the trigger would evict again +on every subsequent turn. +""" + +_BYTES_PER_TOKEN = 4 +"""Matches ``CharacterEstimatorTokenizer``, which is a flat 4 characters per token.""" + +_MAX_PASSES = 3 +"""Eviction re-measures rather than trusting the estimate, but must not loop indefinitely.""" + + +def prunes_excluded(retention: RetentionMode) -> bool: + """Whether compaction exclusions should be deleted as they are made.""" + return retention == "follow_compaction" + + +async def enforce_budget(state: DurableAgentState, *, max_state_bytes: int = DEFAULT_MAX_STATE_BYTES) -> int: + """Evict oldest conversation groups when persisted state approaches the backend limit. + + The measurement is exact rather than estimated. Serializing state at the 1 MB limit costs a few + milliseconds against a turn dominated by a model call, and ``to_dict()`` already runs on every + persist, so the incremental cost is small and only paid once per turn. + + Args: + state: The entity state, modified in place. + + Keyword Args: + max_state_bytes: The budget for serialized state. + + Returns: + How many messages were removed. Zero is the common case. + """ + high = int(max_state_bytes * HIGH_WATERMARK) + size = _serialized_size(state) + if size < high: + return 0 + + history = state.data.conversation_history + target = int(max_state_bytes * LOW_WATERMARK) + removed: list[str] = [] + + for attempt in range(_MAX_PASSES): + # Tighten on each pass, since the byte-to-token conversion is a heuristic and a first + # attempt can land short of the target. + evicted = await _evict_once(history, serialized_size=size, target_bytes=target >> attempt) + if not evicted: + break + removed.extend(evicted) + size = _serialized_size(state) + if size < high: + break + + if removed: + logger.warning( + "[Retention] Durable state passed %d bytes of a %d budget, so %d message(s) were " + "evicted oldest-first (%s .. %s), leaving %d bytes. Set retention='keep_all' to " + "disable this, or raise max_state_bytes if large payload offload is enabled.", + high, + max_state_bytes, + len(removed), + removed[0], + removed[-1], + size, + ) + elif size >= high: + logger.error( + "[Retention] Durable state is %d bytes against a %d budget and nothing could be " + "evicted. The newest exchange is never evicted, so a single turn larger than the " + "budget cannot be resolved by retention.", + size, + max_state_bytes, + ) + return len(removed) + + +def _serialized_size(state: DurableAgentState) -> int: + """Measure the state exactly as it will be persisted.""" + return len(json.dumps(state.to_dict())) + + +async def _evict_once( + history: list[DurableAgentStateEntry], + *, + serialized_size: int, + target_bytes: int, +) -> list[str]: + """Run one eviction pass, returning the ids of the messages removed. + + Core already knows how to drop oldest groups to a budget while preserving system messages and + keeping tool-call groups whole, so that judgement is borrowed rather than reimplemented. + """ + candidates: list[Message] = [] + origins: list[tuple[DurableAgentStateEntry, DurableAgentStateMessage]] = [] + protected = _newest_exchange(history) + for entry, index in replayable_entries(history): + if entry in protected: + # Never evict the exchange that just happened. Core's budget fallback will drop + # everything if the budget demands it, and losing the current turn would break + # response polling and discard the result the caller is waiting for. + continue + stored = entry.messages[index] + message = cast("Message", stored.to_chat_message()) + # The budget is computed over *included* messages, so a user's own compaction exclusions + # would make an over-budget conversation look empty. Clearing them here makes the budget + # reflect what is stored. This is a detached copy, so the persisted annotation is untouched. + message.additional_properties.pop(EXCLUDED_KEY, None) + candidates.append(message) + origins.append((entry, stored)) + + if not candidates: + return [] + + strategy = TokenBudgetComposedStrategy( + token_budget=_token_budget(candidates, serialized_size=serialized_size, target_bytes=target_bytes), + tokenizer=CharacterEstimatorTokenizer(), + # No strategies, so this goes straight to core's deterministic oldest-group eviction. + # Passing the user's strategy would satisfy the budget immediately under early stop, and + # everything it had excluded for context reasons would then be deleted. + strategies=[], + ) + await strategy(candidates) + + evicted = [ + (position, origins[position]) + for position, message in enumerate(candidates) + if message.additional_properties.get(EXCLUDED_KEY) + ] + if not evicted: + return [] + prune_messages(history, [origin for _, origin in evicted]) + return [candidates[position].message_id or "" for position, _ in evicted] + + +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. + """ + if not history: + return [] + newest = history[-1].correlation_id + if newest is None: + return [history[-1]] + return [entry for entry in history if entry.correlation_id == newest] + + +def _token_budget(candidates: list[Message], *, serialized_size: int, target_bytes: int) -> int: + """Convert a byte budget into the token budget the strategy expects. + + Serialized state is larger than the text it contains, because of keys, escaping, ids and + annotations. Rather than assume an overhead constant, the ratio is measured from the state in + hand, so a conversation of long prose and one full of tool-call metadata are both handled. + """ + content_chars = sum(len(message.text or "") for message in candidates) + ratio = (content_chars / serialized_size) if serialized_size else 1.0 + return max(int(target_bytes * ratio) // _BYTES_PER_TOKEN, 1) diff --git a/python/packages/durabletask/agent_framework_durabletask/_shim.py b/python/packages/durabletask/agent_framework_durabletask/_shim.py index e6e9f5d..084163c 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_shim.py +++ b/python/packages/durabletask/agent_framework_durabletask/_shim.py @@ -16,13 +16,43 @@ from agent_framework._types import AgentRunInputs from ._executors import DurableAgentExecutor -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, +) -> 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. + + 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) + + class DurableAgentProvider(ABC, Generic[TaskT]): """Abstract provider for constructing durable agent proxies. @@ -92,6 +122,7 @@ 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, ) -> TaskT: """Execute the agent via the injected provider. @@ -103,6 +134,9 @@ 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. Note: This method overrides SupportsAgentRun.run() with a different return type: @@ -122,9 +156,13 @@ def run( # type: ignore[override] raise ValueError("DurableAIAgent does not support streaming mode (stream must be False)") message_str = 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 else {} run_request = self._executor.get_run_request( message=message_str, options=options, + **extra, ) return self._executor.run_durable_agent( diff --git a/python/packages/durabletask/agent_framework_durabletask/_worker.py b/python/packages/durabletask/agent_framework_durabletask/_worker.py index 3eed81f..81b301c 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_worker.py +++ b/python/packages/durabletask/agent_framework_durabletask/_worker.py @@ -19,6 +19,7 @@ from ._async_bridge import run_agent_coroutine from ._callbacks import AgentResponseCallbackProtocol from ._entities import AgentEntity, DurableTaskEntityStateProvider +from ._retention import DEFAULT_MAX_STATE_BYTES, DEFAULT_RETENTION, RetentionMode from ._workflows.activity import execute_workflow_activity from ._workflows.dt_context import DurableTaskWorkflowContext from ._workflows.naming import ( @@ -78,15 +79,26 @@ def __init__( self, worker: TaskHubGrpcWorker, callback: AgentResponseCallbackProtocol | None = None, + *, + retention: RetentionMode = DEFAULT_RETENTION, + max_state_bytes: int = DEFAULT_MAX_STATE_BYTES, ): """Initialize the worker wrapper. Args: worker: The durabletask worker instance to wrap callback: Optional callback for agent response notifications + retention: Default conversation retention for registered agents. ``auto`` deletes only + under storage pressure, ``keep_all`` never deletes and lets the entity fail at the + backend limit, and ``follow_compaction`` first deletes what compaction excluded, + then uses the same pressure eviction as ``auto`` if that is not enough. + max_state_bytes: Budget for serialized entity state. Raise it when large payload + offload is configured on the worker and client. """ self._worker = worker self._callback = callback + self._retention: RetentionMode = retention + self._max_state_bytes = max_state_bytes self._registered_agents: dict[str, SupportsAgentRun] = {} self._workflows: dict[str, Workflow] = {} # Every workflow whose orchestration has been registered (top-level plus nested @@ -102,6 +114,7 @@ def add_agent( callback: AgentResponseCallbackProtocol | None = None, *, entity_id: str | None = None, + retention: RetentionMode | None = None, ) -> None: """Register an agent with the worker. @@ -115,6 +128,7 @@ 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. Raises: ValueError: If the agent doesn't have a name or is already registered @@ -137,7 +151,14 @@ def add_agent( 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) + effective_retention: RetentionMode = self._retention if retention is None else retention + entity_class = self.__create_agent_entity( + agent, + effective_callback, + entity_id=registration_name, + retention=effective_retention, + max_state_bytes=self._max_state_bytes, + ) # Register the entity class with the worker # The worker.add_entity method takes a class @@ -195,6 +216,8 @@ def configure_workflow( self, workflow: Workflow, callback: AgentResponseCallbackProtocol | None = None, + *, + retention: RetentionMode | None = None, ) -> None: """Register a :class:`Workflow` for automatic orchestration. @@ -220,6 +243,9 @@ 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. Raises: ValueError: If the workflow (or a nested sub-workflow) name is missing, @@ -266,12 +292,13 @@ def configure_workflow( for hosted in hosted_workflows: if hosted.name.casefold() in self._registered_orchestrations: continue - self._register_single_workflow(hosted, callback) + self._register_single_workflow(hosted, callback, retention) def _register_single_workflow( self, workflow: Workflow, callback: AgentResponseCallbackProtocol | None, + retention: RetentionMode | None = None, ) -> None: """Register one workflow's durable primitives (no recursion into sub-workflows). @@ -291,7 +318,7 @@ def _register_single_workflow( 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) # Register non-agent executors as durable activities, scoped by workflow name. # WorkflowExecutor nodes are intentionally not registered as activities: their @@ -356,6 +383,8 @@ def __create_agent_entity( callback: AgentResponseCallbackProtocol | None = None, *, entity_id: str | None = None, + retention: RetentionMode = DEFAULT_RETENTION, + max_state_bytes: int = DEFAULT_MAX_STATE_BYTES, ) -> type[DurableTaskEntityStateProvider]: """Factory function to create a DurableEntity class configured with an agent. @@ -368,6 +397,8 @@ 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: Budget for serialized entity state. Returns: A new DurableEntity subclass configured for this agent @@ -385,6 +416,8 @@ def __init__(self) -> None: agent=agent, callback=callback, state_provider=self, + retention=retention, + max_state_bytes=max_state_bytes, ) logger.debug( "[ConfiguredAgentEntity] Initialized entity for agent: %s (entity name: %s)", diff --git a/python/packages/durabletask/agent_framework_durabletask/_workflows/context.py b/python/packages/durabletask/agent_framework_durabletask/_workflows/context.py index d757d00..d129148 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_workflows/context.py +++ b/python/packages/durabletask/agent_framework_durabletask/_workflows/context.py @@ -73,13 +73,21 @@ 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, + ) -> 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. 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..5ed23d0 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,20 @@ 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, + ) -> Any: + return build_agent_task( + self._executor, + executor_id, + message, + orchestration_instance_id, + context_messages, + ) 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..2dce5f3 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_workflows/orchestrator.py +++ b/python/packages/durabletask/agent_framework_durabletask/_workflows/orchestrator.py @@ -51,8 +51,10 @@ 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, ) @@ -231,7 +233,19 @@ def build_agent_executor_response( if isinstance(previous_message, AgentExecutorResponse) and previous_message.full_conversation: full_conversation.extend(previous_message.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), + ) + ) + # Core leaves message_id unset, and a node that runs more than once receives this + # conversation again every time. Without an id the entity cannot tell the repeat from new + # input, so it re-records the whole conversation on each visit and state grows without bound. + # The position is fixed once a message joins the conversation and the orchestrator rebuilds + # the same sequence on replay, so deriving the id from it is both unique and replay-safe. + assistant_message.message_id = workflow_message_id(executor_id, len(full_conversation)) full_conversation.append(assistant_message) return AgentExecutorResponse( @@ -246,8 +260,38 @@ def build_agent_executor_response( # ============================================================================ +def _build_context_messages(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 conversation to forward (for example the first + node in a workflow, which receives the raw input instead). + """ + if not isinstance(message, AgentExecutorResponse): + return None + + mode = getattr(executor, "_context_mode", "full") + if mode == "last_agent": + selected = list(message.agent_response.messages) if message.agent_response else [] + elif mode == "custom": + context_filter = getattr(executor, "_context_filter", None) + if context_filter is None: + return None + selected = list(context_filter(list(message.full_conversation))) + else: + selected = list(message.full_conversation) + + if not selected: + return None + return [m.to_dict() for m in selected] + + def _prepare_agent_task( ctx: WorkflowOrchestrationContext, + executor: AgentExecutor, executor_id: str, message: Any, workflow_name: str, @@ -259,10 +303,14 @@ 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. + + Any upstream conversation is forwarded as context messages so a downstream agent sees + what earlier nodes produced, matching in-process workflow behavior. """ message_content = _extract_message_content(message) + context_messages = _build_context_messages(executor, message) scoped_id = workflow_scoped_executor_id(workflow_name, executor_id) - return ctx.prepare_agent_task(scoped_id, message_content, ctx.instance_id) + return ctx.prepare_agent_task(scoped_id, message_content, ctx.instance_id, context_messages) def _prepare_activity_task( @@ -945,7 +993,13 @@ def _prepare_all_tasks( 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) + task = _prepare_agent_task( + ctx, + cast(AgentExecutor, workflow.executors[first_msg[0]]), + first_msg[0], + first_msg[1], + workflow.name, + ) all_tasks.append(task) task_metadata_list.append( TaskMetadata( @@ -1159,7 +1213,13 @@ def publish_live_status( # Phase 3: Process sequential 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) + task = _prepare_agent_task( + ctx, + cast(AgentExecutor, workflow.executors[executor_id]), + executor_id, + message, + workflow.name, + ) agent_response: AgentResponse = yield task logger.debug("Agent %s sequential response completed", executor_id) 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..2d0be5d 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 @@ -69,6 +69,31 @@ 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.dts_client.schedule_new_orchestration( + orchestrator=workflow_orchestrator_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( 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..54777a6 --- /dev/null +++ b/python/packages/durabletask/tests/integration_tests/test_13_dt_conversation_compaction.py @@ -0,0 +1,190 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Integration tests for durable conversation compaction. + +Covers the behavior an agent gets by simply being registered with the durable runtime: + +- history is persisted in the agent's durable entity and reaches the model on later turns, +- the configured compaction strategy runs and its annotations are persisted, so compaction + state survives entity state serialization rather than being recomputed each turn, +- the full conversation record is retained in storage even though the model sees less. +""" + +import json +from pathlib import Path +from typing import Any, Protocol + +import pytest +from durabletask.entities import EntityInstanceId + +from agent_framework_durabletask import DurableAgentState, DurableAIAgentClient + +# 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: + """Compaction runs durably without any durable-specific agent 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() + + 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, because message ids are assigned when history is first loaded rather + # than when it is written. After one turn there is nothing to load and nothing to stamp. + 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. + + This is what stops compaction from being recomputed on every turn, and it only works + because message-level metadata and ids are persisted with the conversation. + """ + 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" + + # Reconciling compaction results across turns relies on stable ids, so every message + # the provider has processed must carry one. (The newest turn is annotated on the + # following load, so it is not required to have an id yet.) + assert all(m.message_id for m in annotated), "annotated messages must carry stable message ids" + + def test_full_record_is_retained(self) -> None: + """Compaction bounds what the model sees; it does not delete the record by default.""" + agent = self.agent_client.get_agent("Historian") + session = agent.create_session() + + turns = KEEP_LAST_GROUPS + 3 + for index in range(turns): + assert agent.run(f"Name city number {index + 1}.", session=session) is not None + + state = self._read_state(session.durable_session_id) + + # One request entry and one response entry per turn: nothing was pruned. + assert len(state.data.conversation_history) == turns * 2 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..7c2323e --- /dev/null +++ b/python/packages/durabletask/tests/integration_tests/test_14_dt_external_history_redis.py @@ -0,0 +1,144 @@ +# 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. + +The last point is the load-bearing one: 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. +""" + +import os +from typing import Any, Protocol + +import pytest +import redis.asyncio as aioredis + +from agent_framework_durabletask import DurableAgentState, DurableAIAgentClient + + +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_durable_state_still_records_the_conversation(self) -> None: + """Durable state remains the audit record even when history lives elsewhere.""" + agent = self.agent_client.get_agent("Archivist") + session = agent.create_session() + + assert agent.run("Note that the archive opens at nine.", session=session) is not None + + state = self._read_state(session.durable_session_id) + assert state.data.conversation_history, "expected the entity to record the conversation" + + 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/test_durable_agent_state.py b/python/packages/durabletask/tests/test_durable_agent_state.py index d3a36c9..3c78a81 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 == "1.2.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"] == "1.2.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) 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..fee1bc1 --- /dev/null +++ b/python/packages/durabletask/tests/test_durable_history_autoswap.py @@ -0,0 +1,310 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Tests for automatic durable history backing (ADR-0032). + +A user should be able to take an agent that already works in core, register it with the +durable runtime, and get durable conversation history with no configuration change. +These tests cover the substitution rules and confirm the user's agent is never mutated. +""" + +from typing import Any + +from agent_framework import ( + Agent, + HistoryProvider, + InMemoryHistoryProvider, + Message, +) + +from agent_framework_durabletask import AgentEntity, AgentEntityStateProviderMixin, DurableHistoryProvider +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 _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): + def __init__(self, *, session_id: str = "autoswap-session") -> None: + self._session_id = session_id + 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 _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_is_left_alone(self) -> None: + agent = _agent(_ServiceStoringClient()) + + prepared = ensure_durable_history(agent) + + assert prepared is agent + assert not _history_providers(prepared) + + 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_keeps_history_with_the_service(self) -> None: + agent = _agent(default_options={"store": True}) + + prepared = ensure_durable_history(agent) + + assert prepared is agent + assert not _history_providers(prepared) + + 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: + agent = _agent() + + prepared = ensure_durable_history(agent) + + assert _history_providers(prepared)[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 + + def test_other_retention_modes_do_not_prune_on_write(self) -> None: + """Only ``follow_compaction`` treats a compaction exclusion as consent to delete.""" + for mode in ("auto", "keep_all"): + entity = AgentEntity(_agent(), state_provider=_InMemoryStateProvider(), retention=mode) + + assert _history_providers(entity.agent)[0].prune_excluded is False, 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 + + +class TestServiceManagedSessions: + """Service-backed agents let the service own the conversation.""" + + 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" 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..f3a9bc0 --- /dev/null +++ b/python/packages/durabletask/tests/test_durable_history_provider.py @@ -0,0 +1,664 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Tests for :class:`DurableHistoryProvider` (ADR-0032 Option 6). + +The provider makes durable entity state the store behind core's ``HistoryProvider`` +interface, so conversation history is persisted exactly once and core compaction +plugs in unchanged. +""" + +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, +) + +from agent_framework_durabletask import ( + AgentEntity, + AgentEntityStateProviderMixin, + DurableHistoryProvider, +) + +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") -> None: + self._session_id = session_id + 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 durable SDK serializes entity state as it is set, so a value it cannot encode + # surfaces here rather than later. Mirrored so tests see the same failure the host does. + json.dumps(state) + self._state_dict = state + + 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: + """Durable entity state is the single store behind core's HistoryProvider.""" + + async def test_history_is_stored_once(self) -> None: + """Messages live only in conversation history, never duplicated into the session blob.""" + 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. + from agent_framework_durabletask import DurableAgentState + + 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. + from agent_framework_durabletask import DurableAgentState + + 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}" + + async def test_service_managed_session_is_skipped(self) -> None: + """When the model service owns the conversation, the provider must not participate.""" + from types import SimpleNamespace + + 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() + token = bind_durable_history(DurableHistoryBinding(state_provider=provider)) + try: + state: dict[str, Any] = {} + context = SimpleNamespace(session_id="s", extend_messages=lambda *_: None) + service_session = SimpleNamespace(service_session_id="svc-123", state={}) + + await history.before_run(agent=None, session=service_session, context=context, state=state) + # Nothing was loaded, so no working buffer was published. + assert "messages" not in state + + # Flushing is likewise a no-op and must not raise. + await history.after_run(agent=None, session=service_session, context=context, state=state) + 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_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" + + async def test_unserializable_provider_state_does_not_break_the_turn(self) -> None: + """Core passes a value it cannot serialize straight through, without raising or warning. + + Assigning that to entity state fails the save, and the error handler saves again with the + same payload, so the second failure escapes and masks whatever the agent returned. The + payload is checked first instead, keeping the last good session. + """ + + class _UnserializableProvider(ContextProvider): + def __init__(self) -> None: + super().__init__("unserializable") + + async def after_run(self, *, agent: Any, session: Any, context: Any, state: dict[str, Any]) -> None: + state["handle"] = object() + + provider = _InMemoryStateProvider() + agent = _agent([InMemoryHistoryProvider(), _UnserializableProvider()]) + entity = _make_entity(agent, provider) + + await _run_turns(entity, ["first"]) + + stored = provider._get_state_dict() + assert stored["data"].get("session") is None, "an unusable session payload was persisted" + # The turn still completed and the conversation was recorded. + assert len(entity.state.data.conversation_history) == 2 + json.dumps(stored) 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..293d2e5 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,13 @@ 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, + ) -> 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_retention.py b/python/packages/durabletask/tests/test_retention.py new file mode 100644 index 0000000..0034d89 --- /dev/null +++ b/python/packages/durabletask/tests/test_retention.py @@ -0,0 +1,337 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Tests for retention (ADR-0032, "Retention"). + +Retention bounds durable entity state so an agent does not simply stop working when it reaches the +backend limit. It is a capacity concern and deliberately separate from compaction: an exclusion made +for token cost is not consent to delete the record. +""" + +import json +from collections.abc import AsyncIterator +from datetime import datetime, timezone +from typing import Any, cast + +from agent_framework import ( + Agent, + BaseChatClient, + ChatResponse, + ChatResponseUpdate, + Content, + Message, + ResponseStream, +) + +from agent_framework_durabletask import ( + AgentEntity, + AgentEntityStateProviderMixin, + DurableAgentState, + DurableAgentStateMessage, + DurableAgentStateRequest, + DurableAgentStateResponse, +) +from agent_framework_durabletask._retention import ( + HIGH_WATERMARK, + LOW_WATERMARK, + 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 entity 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. + """ + state = DurableAgentState() + now = datetime.now(tz=timezone.utc) + marked = 0 + for index in range(turns): + request = DurableAgentStateRequest( + correlation_id=f"c{index}", + created_at=now, + 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=now, + 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: + assert prunes_excluded("follow_compaction") is True + assert prunes_excluded("auto") is False + assert prunes_excluded("keep_all") is False + + def test_the_default_is_auto(self) -> None: + """Which is the deliberate behavior change: previously nothing bounded storage.""" + from agent_framework_durabletask import DurableAIAgentWorker + + assert DurableAIAgentWorker(cast(Any, object()))._retention == "auto" + + +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 = _message_ids(state) + + removed = await enforce_budget(state, max_state_bytes=BUDGET) + + assert removed == 0 + assert _message_ids(state) == 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 * HIGH_WATERMARK, "eviction did not get back under the trigger" + + 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)[-1] == "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_keep_all_is_the_caller_s_decision(self) -> None: + """``keep_all`` is enforced by the entity, so the budget helper itself always acts.""" + state = _state(turns=60) + + assert await enforce_budget(state, max_state_bytes=BUDGET) > 0 + + +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" + + +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: + """Core's fallback will drop everything if asked, which would lose the result being polled.""" + state = _state(turns=1, chars=BUDGET * 2) + + removed = await enforce_budget(state, max_state_bytes=BUDGET) + + assert removed == 0 + 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.extend(_state(turns=1, chars=BUDGET * 2).data.conversation_history) + + await enforce_budget(state, max_state_bytes=BUDGET) + + assert _message_ids(state)[-2:] == ["u0", "a0"], "the newest exchange must survive" + + +class TestStateShape: + """Eviction must leave durable state usable.""" + + async def test_empty_entries_are_removed(self) -> None: + state = _state(turns=60) + + await enforce_budget(state, max_state_bytes=BUDGET) + + 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. + json.dumps(state) + self._state_dict = 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) + + replies: list[str] = [] + for turn in range(self.TURNS): + result = await entity.run({"message": f"question {turn}", "correlationId": f"corr-{turn}"}) + replies.append(result.text) + return provider, replies + + async def test_state_stays_bounded_across_many_turns(self) -> None: + provider, _ = await self._drive(max_state_bytes=self.LIMIT) + assert len(json.dumps(provider._get_state_dict())) <= self.LIMIT + + 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 0 < len(state.data.conversation_history) < 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) + kept = len(DurableAgentState.from_dict(provider._get_state_dict()).data.conversation_history) + assert 0 < kept < self.TURNS * 2 + + async def test_keep_all_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=self.LIMIT) + assert len(json.dumps(provider._get_state_dict())) > self.LIMIT 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..ffc631a --- /dev/null +++ b/python/packages/durabletask/tests/test_state_schema.py @@ -0,0 +1,123 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""The persisted state must match the shared cross-language schema. + +``schemas/durable-agent-entity-state.json`` is the contract between the Python and .NET hosting +layers. Nothing enforced it before, so fields this runtime persisted (``messageId`` and +``extensionData``, both load-bearing for context management) went undeclared and a .NET +implementer reading the schema would not have known to round-trip them. +""" + +import json +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import pytest +from agent_framework import Message + +from agent_framework_durabletask import ( + DurableAgentState, + DurableAgentStateMessage, + DurableAgentStateRequest, + DurableAgentStateResponse, +) + +jsonschema = pytest.importorskip("jsonschema") + +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).validate(payload) + + +def _populated_state() -> DurableAgentState: + """Build state exercising every field this runtime persists.""" + now = datetime.now(tz=timezone.utc) + request = DurableAgentStateRequest( + correlation_id="c0", + created_at=now, + messages=[ + DurableAgentStateMessage.from_chat_message( + Message(role="user", contents=["hello"], message_id="wf_input_0") + ) + ], + ) + response = DurableAgentStateResponse( + correlation_id="c0", + created_at=now, + messages=[ + DurableAgentStateMessage.from_chat_message( + Message(role="assistant", contents=["hi"], message_id="wf_writer_1") + ) + ], + ) + # 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_positions = {"input": 0, "writer": 1} + 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_the_ingestion_watermark_is_declared(schema: dict[str, Any]) -> None: + """It is how a repeated workflow node recognizes context it already recorded.""" + assert "ingestedPositions" in schema["$defs"]["data"]["properties"] + + +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_positions == {"input": 0, "writer": 1} 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..3ad74c4 --- /dev/null +++ b/python/packages/durabletask/tests/test_workflow_context_parity.py @@ -0,0 +1,370 @@ +# 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 + +from agent_framework import ( + AgentExecutor, + AgentExecutorResponse, + AgentResponse, + Message, +) + +from agent_framework_durabletask import ( + AgentEntity, + AgentEntityStateProviderMixin, + DurableAgentState, + DurableAgentStateRequest, + RunRequest, +) +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")] + entity.state.data.conversation_history.append( + DurableAgentStateRequest.from_run_request(self._request(first, "corr-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_keeps_last_message(self) -> None: + """The agent must always receive at least one input message. + + The kept copy loses its id, because storing two messages under one id would collide in the + compaction position map and send annotations or pruning to the wrong stored message. + """ + provider = _InMemoryStateProvider() + entity = AgentEntity(_stub_agent(), state_provider=provider) + + messages = [Message(role="user", contents=["hello"], message_id="m0")] + entity.state.data.conversation_history.append( + DurableAgentStateRequest.from_run_request(self._request(messages, "corr-0")) + ) + + entry = DurableAgentStateRequest.from_run_request(self._request(messages, "corr-1")) + entry.messages = entity._drop_already_stored(entry.messages) + + assert len(entry.messages) == 1 + assert entry.messages[0].message_id is None + assert entry.messages[0].to_chat_message().text == "hello" + + 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}" + + +class TestWorkflowConversationIdentity: + """Messages the workflow itself builds must carry ids, or a repeated node cannot spot them. + + Core leaves ``message_id`` unset, and the entity's duplicate check treats a message without one + as new. An unstamped conversation therefore defeats the check entirely, and a node in a cycle + re-records the whole conversation on every visit. + """ + + 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") + + marks = entity.state.data.ingested_positions or {} + assert set(marks) == {"input", "A", "B"}, f"expected a mark per producing executor, got {marks}" + + 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_positions == entity.state.data.ingested_positions + + +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 survives 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], + ) + + restored = RunRequest.from_dict(request.to_dict()) + + assert restored.context_messages is not None + assert len(restored.context_messages) == 1 + + 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 "contextMessages" not in request.to_dict() 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/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..b8a5733 --- /dev/null +++ b/python/samples/13_conversation_compaction/README.md @@ -0,0 +1,100 @@ +# 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` so the compaction provider stays wired to it. Conversation state lives + in the agent's durable entity and survives worker restarts. +- **Compaction state is persisted.** Annotations produced by the strategy are stored alongside the + messages, so compaction is not recomputed from scratch on every turn. +- **Context stays bounded.** Only the messages the strategy keeps are sent to the model, so a long + conversation does not grow the per-turn context without limit. + +The full conversation remains in durable storage, and compaction bounds what the *model* sees. + +### Retention: what durable storage is allowed to discard + +Compaction and retention answer different questions. Compaction decides what the model should read. +Retention decides what durable state can afford to hold, and an exclusion made to save tokens is not +consent to delete the record. Set it at registration with `add_agent(agent, retention=...)`, or +app-wide on the worker. + +| Mode | Behavior | +| --- | --- | +| `auto` (default) | Deletes only when state approaches the backend limit, and only enough to get back under it. Nothing changes for a conversation that never gets close. | +| `keep_all` | Never deletes. The entity may reach the limit and fail. Choose this when the complete record matters more than staying available. | +| `follow_compaction` | Deletes whatever compaction excluded every turn. If that does not free enough space, it also uses the same pressure eviction as `auto`. | + +`auto` exists because the alternative is an agent that simply stops working mid-conversation, with +no warning. It evicts oldest-first, keeps system messages and tool-call groups intact, never touches +the exchange that just completed, and logs what it removed. + +### Client-side vs service-managed history + +Compaction only applies to history the **client** owns. When a chat client keeps the conversation on +the service (Foundry and the Responses API both do so by default), the service owns the model's +context, the durable entity keeps the transcript purely as a record, and the durable history provider +stays out of the way. This sample sets `store=False` so history is client-side and compaction has +something to compact. + +## 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, which it answers correctly. + +The trade-off is the point of the sample: a sliding window keeps context bounded by *dropping* older +turns from what the model sees, so facts from long-past turns are genuinely no longer available to +the model. Those messages are **not deleted**. They remain in durable storage, marked as excluded, +so the conversation record stays complete and auditable. Choose a strategy accordingly: use +summarization if old details must survive in the model's context, and a sliding window when only +recent context matters. diff --git a/python/samples/13_conversation_compaction/client.py b/python/samples/13_conversation_compaction/client.py new file mode 100644 index 0000000..2203104 --- /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 +shows that the conversation keeps working while the model's context stays bounded. +""" + +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("Recent context was retained while the conversation stayed compacted.") + else: + print("The codename fell outside the retained window.") + + +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..be6b523 --- /dev/null +++ b/python/samples/13_conversation_compaction/worker.py @@ -0,0 +1,147 @@ +# 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: + +- conversation history is persisted in the agent's durable entity and survives restarts, +- the compaction strategy still runs, and its annotations are persisted alongside the + messages, so compaction state is not recomputed on every turn, +- only the messages compaction keeps are sent to the model, bounding context growth. + +No durable-specific configuration is required on the agent itself. + +Note on service-managed conversations: compaction applies to history the *client* owns. When a +chat client keeps the conversation on the service (Foundry and the Responses API both do so by +default), the service owns the model's context and the durable entity keeps the full transcript +purely as a record. This sample therefore sets ``store=False`` so history is client-side and +compaction has something to compact. + +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 remembers facts while its context stays bounded. + + 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 + """ + agent_worker = DurableAIAgentWorker(worker) + + 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..75cb2f9 --- /dev/null +++ b/python/samples/14_external_history_redis/README.md @@ -0,0 +1,75 @@ +# 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 is left alone.** An `InMemoryHistoryProvider` is swapped for a durable-backed one + (see [13_conversation_compaction](../13_conversation_compaction)), but a provider + you chose deliberately is never substituted. You picked where the conversation lives. +- **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. +- **Execution is still durable.** Retries, restarts, and orchestration guarantees are unchanged, and + durable state still records the conversation for audit. + +`redis_history_provider.py` is deliberately small, roughly "read a list, append to a list", to show +how little a bring-your-own-store provider needs. The same shape applies to Cosmos DB, a file, or any +other backend. + +## 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..539af2b --- /dev/null +++ b/python/samples/14_external_history_redis/redis_history_provider.py @@ -0,0 +1,93 @@ +# 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 is included in the sample rather than imported from a package to keep the sample dependency +free and to show exactly how little a "bring your own store" provider needs: read the messages +for a session id, append new ones. + +The durable runtime leaves providers like this alone: the user chose where their conversation +lives, so durable supplies execution durability and stays out of the way of storage. +""" + +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. + """ + + 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. + + 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]) 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..cfe0b5e --- /dev/null +++ b/python/samples/14_external_history_redis/worker.py @@ -0,0 +1,130 @@ +# 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, +- durable state still records the conversation for audit, and execution stays durable. + +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__) + + +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")) + + 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") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/samples/README.md b/python/samples/README.md index 81a95ef..11277c0 100644 --- a/python/samples/README.md +++ b/python/samples/README.md @@ -70,6 +70,10 @@ 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. +### Conversation History +- **[13_conversation_compaction](13_conversation_compaction/)**: Persist conversation history durably and compact it as it grows. An agent configured the ordinary core way (`InMemoryHistoryProvider` + `CompactionProvider`) gets durable-backed history automatically, with compaction annotations persisted alongside the messages. +- **[14_external_history_redis](14_external_history_redis/)**: Keep conversation history in a store you chose (Redis here) instead of durable state. The durable runtime leaves your provider alone and hands it the entity's stable session id, so it continues the conversation across turns and restarts. + ### 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 +91,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/)**: Persist conversation history durably and compact it as it grows, on Azure Functions. The Functions counterpart to [13_conversation_compaction](13_conversation_compaction/). ## Running the Samples 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..98c4449 --- /dev/null +++ b/python/samples/azure_functions/14_conversation_compaction/README.md @@ -0,0 +1,78 @@ +# 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` so the compaction provider stays wired to it. +- Compaction annotations being persisted alongside the messages, so compaction state is not + recomputed from scratch on every turn. +- Context growth being bounded: only the messages the strategy keeps are sent to the model. + +```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) +``` + +The full conversation remains in durable storage, and compaction bounds what the *model* sees. To +also delete what compaction excluded, use +`AgentFunctionApp(..., retention="follow_compaction")`. It deletes exclusions every turn, then +uses the same pressure eviction as `auto` if the remaining state is still too large. + +### Client-side vs service-managed history + +Compaction only applies to history the **client** owns. When a chat client keeps the conversation on +the service (Foundry and the Responses API both do so by default), the service owns the model's +context, the durable entity keeps the transcript purely as a record, and the durable history provider +stays out of the way. This sample sets `store=False` so history is client-side and compaction has +something to compact. + +## 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. + +The trade-off is the point of the sample: a sliding window keeps context bounded by *dropping* older +turns from what the model sees, so facts from long-past turns are genuinely no longer available to +the model. Those messages are **not deleted**. They remain in durable storage, marked as excluded, +so the conversation record stays complete and auditable. Choose a strategy accordingly: use +summarization if old details must survive in the model's context, and a sliding window when only +recent context matters. 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..715eb4a --- /dev/null +++ b/python/samples/azure_functions/14_conversation_compaction/function_app.py @@ -0,0 +1,81 @@ +# 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 +history is persisted in the agent's durable entity, the compaction strategy still runs, and its +annotations are persisted alongside the messages. Only the messages compaction keeps are sent to +the model, bounding context growth. + +This is the Azure Functions counterpart to the standalone ``13_conversation_compaction`` sample. + +Note on service-managed conversations: compaction applies to history the *client* owns. When a +chat client keeps the conversation on the service (Foundry and the Responses API both do so by +default), the service owns the model's context and the durable entity keeps the full transcript +purely as a record. This sample therefore sets ``store=False`` so history is client-side and +compaction has something to compact. + +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. +# Set retention="follow_compaction" here to delete compacted-out messages immediately, with +# pressure eviction as a fallback if the remaining state is still too large. +app = AgentFunctionApp(agents=[_create_agent()], enable_health_check=True, max_poll_retries=50) + +""" +Expected behavior when posting several turns with the same `session_id`: + +- every turn is answered with the earlier turns in context, +- the model's context stops growing once the sliding window fills, +- the durable entity keeps the whole conversation, with compacted-out messages marked excluded. +""" 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/schemas/durable-agent-entity-state.json b/schemas/durable-agent-entity-state.json index 53ac064..5f6f907 100644 --- a/schemas/durable-agent-entity-state.json +++ b/schemas/durable-agent-entity-state.json @@ -142,7 +142,15 @@ "type": "array", "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"] }, @@ -200,6 +208,15 @@ "type": "array", "description": "Ordered list of conversation entries.", "items": { "$ref": "#/$defs/conversationEntry" } + }, + "session": { + "type": "object", + "description": "Serialized agent session carried between turns, holding the per-provider state bag and any service-issued conversation id. The shape is the hosting runtime's own and is deliberately not fixed here: .NET serializes conversationId plus stateBag, Python serializes session_id, service_session_id and state. Treat it as opaque and discriminate on the properties present. The agent's own history provider slice is excluded, since conversationHistory is the record of truth." + }, + "ingestedPositions": { + "type": "object", + "description": "Highest chained-conversation position this entity has taken from each workflow executor, keyed by executor id. A workflow re-sends the whole conversation on every visit, so this is what lets a repeated node recognize context it already recorded. Kept separately from the messages because retention may delete them.", + "additionalProperties": { "type": "integer", "minimum": 0 } } } }