From 2b26c2f75576cffad13ebcb03fd2fc11a6407746 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 24 Jul 2026 16:30:17 +0200 Subject: [PATCH 1/8] Python: Move session persistence into core Move SessionStore and durable msgspec-backed storage into core, restore sessions in Foundry Responses hosting with per-user isolation, and document the serialization design. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3e5c81ad-75e8-4e92-a883-8bbd676c6c8c --- ...0033-python-session-store-serialization.md | 236 ++++++ .../agent-framework-py-release/SKILL.md | 5 + python/AGENTS.md | 7 + python/PACKAGE_STATUS.md | 55 +- python/packages/core/AGENTS.md | 5 +- .../packages/core/agent_framework/__init__.py | 4 + .../core/agent_framework/__init__.pyi | 4 + .../core/agent_framework/_feature_stage.py | 40 +- .../core/agent_framework/_harness/_memory.py | 6 +- .../core/agent_framework/_sessions.py | 791 ++++++++++++++---- python/packages/core/pyproject.toml | 1 + .../core/tests/core/test_feature_stage.py | 17 + .../core/tests/core/test_harness_memory.py | 6 +- .../packages/core/tests/core/test_sessions.py | 377 ++++++++- python/packages/foundry_hosting/README.md | 87 ++ .../__init__.py | 9 +- .../_responses.py | 365 ++++++-- .../_session_isolation.py | 118 +++ .../foundry_hosting/tests/test_responses.py | 628 ++++++++++++-- python/packages/hosting-responses/README.md | 4 +- python/packages/hosting-telegram/README.md | 7 +- .../_parsing.py | 8 +- .../tests/hosting_telegram/test_parsing.py | 6 +- python/packages/hosting/README.md | 45 +- .../agent_framework_hosting/__init__.py | 2 - .../hosting/agent_framework_hosting/_state.py | 88 +- .../hosting/tests/hosting/test_state.py | 95 +-- .../simple_context_provider.py | 16 +- .../conversations/file_history_provider.py | 18 +- ...story_provider_conversation_persistence.py | 18 +- .../samples/04-hosting/af-hosting/README.md | 2 +- .../af-hosting/local_responses/README.md | 6 +- .../af-hosting/local_responses/app.py | 11 +- .../af-hosting/local_telegram/README.md | 4 +- python/uv.lock | 50 ++ 35 files changed, 2589 insertions(+), 552 deletions(-) create mode 100644 docs/decisions/0033-python-session-store-serialization.md create mode 100644 python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_isolation.py diff --git a/docs/decisions/0033-python-session-store-serialization.md b/docs/decisions/0033-python-session-store-serialization.md new file mode 100644 index 0000000000..4f06334c90 --- /dev/null +++ b/docs/decisions/0033-python-session-store-serialization.md @@ -0,0 +1,236 @@ +--- +status: proposed +contact: eavanvalkenburg +date: 2026-07-24 +deciders: eavanvalkenburg, chetantoshnival, taochenosu, moonbox3, giles17 +--- + +# Python session storage and serialization + +## Context and Problem Statement + +Python does not have a broadly shared session-store API in +`agent-framework-core`. The alpha `agent-framework-hosting` package has a small process-local `SessionStore`, but that +type is hosting-specific, in-memory only, and unavailable to packages such as Foundry Hosting without taking a +dependency on the hosting helper package. + +The alpha implementation is a prototype, not a compatibility constraint. This decision may replace its location, +names, method shape, and behavior if another design is preferable. + +The existing file-backed persistence surfaces solve narrower problems: + +- `FileHistoryProvider` stores conversation `Message` records, not complete `AgentSession` snapshots; +- `FileCheckpointStorage` stores workflow checkpoints; and +- the Responses provider stores protocol history, but not Agent Framework runtime state carried in + `AgentSession.state`. + +`AgentSession.to_dict()` / `from_dict()` already provide a dictionary snapshot shape. Session state may contain +framework or application-defined objects, and `register_state_type` provides dynamic type restoration, but the +registration and collision behavior is not yet strong enough to serve as a durable, cold-start persistence contract. + +The framework therefore needs to decide: + +- where a reusable in-memory and file-backed session store belongs; +- how a complete `AgentSession` should be serialized atomically and validated; +- how custom nested state types are registered and restored after process restart; and +- how to provide the required readable JSON format while leaving room for an optional optimized binary format. + +## Decision Drivers + +### Session-store ownership and API + +- Make session storage reusable by core, hosting, and provider packages without creating dependency cycles. +- Keep the smallest public API that supports in-memory use, durable implementations, and application-defined stores. +- Define the minimum async operations required for lookup, replacement, and deletion. +- Decide explicitly whether reads return shared instances or independent snapshots suitable for branching. +- Simpler is better + +### Serialization and type restoration + +- Provide readable JSON serialization as a required capability. +- Treat an optimized binary format as a nice-to-have only when the chosen JSON implementation supports it without a + separate state model or substantial additional complexity. +- Perform one typed encode and decode operation per file write/read. +- Preserve existing dynamic application registration of nested state types. +- Fail before persistence when an object cannot be restored after a cold start. +- Keep the existing serialized `{"type": "", ...}` representation compatible. + +## Decision 1: Session-store ownership and API shape + +### Keep `SessionStore` in `agent-framework-hosting` + +- Good: keeps the abstraction local to app-owned hosting scenarios. +- Bad: Foundry Hosting and other packages cannot reuse it without depending on the hosting helper package. +- Bad: a generic session snapshot store is not inherently or only a web-hosting concern. +- Bad: durable implementations would either be duplicated or placed in an unrelated package. + +### Add an abstract store plus separate in-memory and file implementations + +For example, define a `SessionStore` protocol/ABC with `InMemorySessionStore` and `FileSessionStore`. + +- Good: clearly separates the contract from implementations. +- Good: implementation names state their storage behavior explicitly. +- Neutral: follows a familiar repository/adapter pattern. +- Bad: introduces an additional public type and rename for a three-method experimental API. +- Bad: callers must choose an implementation even for the default in-memory case. +- Bad: the abstraction adds little value while every implementation still needs the same method overrides. + +### Move the concrete store to core and use it as the overridable base + +Move `SessionStore` to `agent-framework-core`, retain its in-memory behavior, and implement `FileSessionStore` by +overriding the same async methods. + +- Good: one public type is both the useful default and the extension point. +- Good: existing custom stores can continue subclassing and overriding `get` / `set` / `delete`. +- Good: core and provider packages can share the API without depending on hosting helpers. +- Good: `FileSessionStore` remains a focused subclass while the base stays free of file-system concerns. +- Bad: the class name does not explicitly say "in memory" when used without overrides. + +## Decision 2: Serialization and type restoration + +Once a file-backed store exists, it needs an on-disk format and a reliable way to reconstruct the complete +`AgentSession`, including nested framework and application-defined state. This decision is independent of where the +store API lives or whether that API is abstract or concrete. + +The alternatives below compare top-level snapshot validation, JSON encoding/decoding cost, and how each option +interacts with the dynamic custom-state registry. Binary storage is not a primary selection criterion. + +### Considered options + +The standard-library and optimized-JSON options are not mutually exclusive. A store can default to `json` while +accepting caller-supplied `dumps` / `loads` callables for `orjson` or another compatible implementation. This is the +pre-msgspec `FileHistoryProvider` design; those hooks remain only as a deprecated compatibility path. + +### Standard library `json` + +- Good: no additional dependency and familiar readable output. +- Good: accepts the existing dictionary snapshots without a schema. +- Good: can remain the fallback/default behind pluggable `dumps` / `loads`. +- Neutral: custom state restoration still requires the framework registry. +- Bad: slower encoding and decoding than optimized native implementations. +- Bad: provides no typed snapshot validation during file reads. + +### Optimized drop-in JSON libraries such as `orjson` + +- Good: substantially faster JSON encoding and decoding than the standard library. +- Good: can preserve the existing dictionary-oriented snapshot and custom `dumps` / `loads` shape. +- Good: can be an opt-in codec without making the optimized package a framework dependency. +- Neutral: returns bytes when encoding, which the file stores can already handle. +- Neutral: custom state restoration still requires the framework registry. +- Bad: remains an untyped top-level decode; the framework must separately validate the session snapshot shape. +- Bad: choosing one drop-in implementation as a core dependency adds a dependency without providing typed construction. + +### Pydantic `model_dump` / `model_validate` + +- Good: Pydantic is already a core dependency. +- Good: a typed session snapshot model can validate top-level fields and provide `model_dump_json` / + `model_validate_json` for file serialization. +- Good: validation errors include useful field paths. +- Neutral: the dynamic `state` field remains `dict[str, Any]`, so custom nested state restoration still requires the + framework registry. +- Neutral: the public `AgentSession` does not need to become a Pydantic model; an internal snapshot model can bridge it. +- Bad: benchmarked encode/decode includes model construction and dumping overhead on every operation. +- Bad: core dependency on Pydantic run the risk of us not being able to use different versions or users of the framework being unable to upgrade or having additional extra code dealing with major version bumps in Pydantic. + +### msgspec typed/tagged unions only + +- Good: msgspec owns validation and reconstruction end to end. +- Neutral: works well for a closed set of framework-owned `msgspec.Struct` types. +- Bad: every external type must be known when the decoder schema is constructed; dynamic registration is lost. + +### msgspec codecs plus an explicit dynamic registry + +- Good: one typed file encode/decode and dynamic nested custom types. +- Good: it satisfies the required readable JSON format. +- Neutral: the same typed snapshot can also support optional MessagePack as a low-cost implementation detail. +- Good: the registry can enforce stable IDs, codec completeness, and collision handling. +- Neutral: a single state-payload hook still recursively applies registry codecs. +- Bad: msgspec cannot infer dynamic types from JSON without the framework's type tags. + +## Benchmark Evidence + +A benchmark using a large `AgentSession` with 2,000 `Message` objects stored through +`InMemoryHistoryProvider`, nested standard dictionaries, registered custom classes, and registered Pydantic models +measured the complete `AgentSession.to_dict()` / codec / `AgentSession.from_dict()` path. + +| Codec | File size | Encode median (ms) | Decode median (ms) | Round-trip median (ms) | Disk round-trip median (ms) | +| --- | ---: | ---: | ---: | ---: | ---: | +| Standard library JSON | 1.57 MiB | 33.503 | 14.316 | 55.261 | 75.226 | +| orjson | 1.57 MiB | 25.808 | 11.754 | 39.398 | 63.319 | +| Pydantic JSON | 1.57 MiB | 28.330 | 18.344 | 53.522 | 77.096 | +| msgspec JSON | 1.57 MiB | 26.019 | 11.379 | **38.060** | 62.230 | +| msgspec MessagePack | **1.45 MiB** | **25.134** | **11.201** | 38.512 | **58.112** | + +The JSON encodings produced the same 1.57 MiB file size. msgspec JSON had the best median JSON round-trip latency, +slightly ahead of orjson, while also supporting typed top-level decoding. Pydantic validation added measurable decode +and disk-round-trip overhead without eliminating the dynamic state registry. + +MessagePack reduced file size to 92.2% of JSON (about 7.8% smaller) and produced the best encode, decode, and disk +round-trip medians. Its in-memory round-trip median was effectively tied with msgspec JSON. This supports offering it +as a nice-to-have, but it is not required to justify choosing msgspec for JSON. + +These results are workload- and machine-dependent, but they validate the architectural choice: + +- use msgspec JSON as the readable default; +- optionally offer msgspec MessagePack when storage size or disk latency matters; +- retain the explicit registry for dynamic custom state in both formats; +- do not add orjson solely for a small JSON performance difference without typed decoding; and +- do not use Pydantic as the file codec when its validation overhead does not replace the registry. + +## Decision Outcome + +### Decision 1: Move the concrete overridable store to core + +`SessionStore` moves to `agent-framework-core` as an experimental public API. It remains a concrete in-memory store and +the default used by `AgentState` in the `hosting` package. Its async `get`, `set`, and `delete` methods remain overridable for custom storage +implementations. + +`FileSessionStore` subclasses `SessionStore` and provides durable atomic file persistence. No separate +`InMemorySessionStore`, protocol, or ABC is introduced. `agent-framework-hosting` consumes the core type and no longer +owns or re-exports `SessionStore` (this will be a breaking change in the `hosting` package). + +Session-store keys use one restricted contract suitable for the built-in file implementation: at most 128 ASCII +letters, digits, `-`, and `_`. Logical `AgentSession.session_id` values outside a store are not globally constrained by +this file-oriented key contract. + +### Decision 2: Use msgspec codecs plus an explicit dynamic registry + +Chosen option: **msgspec codecs plus an explicit dynamic registry**. + +`FileSessionStore` uses a typed internal `msgspec.Struct` snapshot with reusable JSON and MessagePack encoders/decoders. +JSON is the required and default format. Because msgspec can reuse the same typed snapshot and registry hooks, +`serialization_format="msgpack"` is also exposed as an optional compact binary convenience. The complete state +dictionary is wrapped in one custom field; its encode/decode hooks recursively translate explicitly registered types +to and from the existing tagged mappings in either format. + +The dependency floor is `msgspec>=0.20.0` because core supports Python 3.14 and msgspec added Python 3.14 support in +version 0.20.0. + +The public `AgentSession` remains a normal framework class. The msgspec Struct is an internal persistence DTO rather +than the inheritance base for runtime sessions. A targeted comparison against the dictionary-based msgspec path found +that the internal Struct improved JSON encode latency by about 9% and JSON round-trip latency by about 4%, without +changing file size. + +`register_state_type` requires explicit registration, supports stable type IDs and optional codecs, rejects collisions, +and provides defaults for `to_dict` / `from_dict` classes and explicitly registered Pydantic models. Unknown persisted +type IDs remain raw dictionaries for compatibility. + +`FileHistoryProvider` also adds msgspec JSON as its default JSON Lines codec. It supports the same explicit +`serialization_format="msgpack"` choice using length-prefixed append-only MessagePack records. Its existing `dumps` / +`loads` extension points remain temporarily for JSON compatibility, emit `DeprecationWarning` when supplied, and do +not apply to MessagePack. New code uses the built-in msgspec codecs. + +## Follow-up Work + +Audit the remaining file-backed stores to determine whether they benefit from the same typed msgspec treatment and +optional JSON / MessagePack formats. `FileCheckpointStorage` is the first candidate because it persists large, +structured workflow state and currently uses JSON plus custom checkpoint value encoding. + +`MemoryContextProvider` is another candidate because its file-backed path combines `MemoryFileStore` state with +transcript files and still exposes `history_dumps` / `history_loads` passthroughs to the deprecated +`FileHistoryProvider` codec hooks. + +The follow-up should measure real framework payloads before changing formats, preserve compatibility or define a clear +migration path for existing files, and consider whether each store needs readable JSON, compact binary storage, append +semantics, or atomic whole-file replacement. Other candidates include file-backed todo state, but each should be +evaluated independently rather than adopting msgspec by default solely for consistency. diff --git a/python/.github/skills/agent-framework-py-release/SKILL.md b/python/.github/skills/agent-framework-py-release/SKILL.md index af78284901..cfad652af7 100644 --- a/python/.github/skills/agent-framework-py-release/SKILL.md +++ b/python/.github/skills/agent-framework-py-release/SKILL.md @@ -39,6 +39,8 @@ If the user states target versions or a date explicitly, use exactly what they s ## Non-negotiable rules +- **Release workflow owns the CHANGELOG**: individual feature/fix PRs do not edit + `python/CHANGELOG.md`; this workflow creates the entries centrally from merged changes. - **CHANGELOG-driven bumps**: only packages mentioned in the new CHANGELOG section get version bumps. Exceptions: root follows core (==pin); user-opted cohort bump on betas. - **Follow `python-package-management` for package lifecycle and versioning rules** — do not duplicate those rules in this release workflow. @@ -136,6 +138,9 @@ Root `agent-framework` is touched when `python/pyproject.toml`, `python/agent_fr ### 4. Draft CHANGELOG entries (THIS DRIVES THE BUMP LIST) +Individual feature/fix PRs intentionally do not add CHANGELOG entries. During release preparation, derive this section +centrally from the merged PRs since the last released tag. + Locate `## [Unreleased]` and the top existing release header. INSERT a new section between them. **New section structure:** diff --git a/python/AGENTS.md b/python/AGENTS.md index b99f0ec7da..c423aadcc8 100644 --- a/python/AGENTS.md +++ b/python/AGENTS.md @@ -37,6 +37,13 @@ team norms from a single conversation without explicit confirmation. match the feature-lifecycle stages documented in the `python-feature-lifecycle` skill. +## Changelog Ownership + +- Individual feature, fix, documentation, and dependency PRs must not update + `python/CHANGELOG.md`. +- CHANGELOG entries and release sections are assembled centrally during the + Python release-preparation workflow. + ## Pull Request Description Guidance When preparing a PR description: diff --git a/python/PACKAGE_STATUS.md b/python/PACKAGE_STATUS.md index a7dfa47b90..da187a62b8 100644 --- a/python/PACKAGE_STATUS.md +++ b/python/PACKAGE_STATUS.md @@ -80,12 +80,61 @@ listed below. helper checks defined in `agent_framework/_evaluation.py` - `agent-framework-foundry`: `FoundryEvals`, `evaluate_traces`, and `evaluate_foundry_target` -#### `SKILLS` +#### `FILE_HISTORY` -- `agent-framework-core`: exported skills APIs from `agent_framework`, including `Skill`, - `SkillResource`, `SkillScript`, `SkillScriptRunner`, and `SkillsProvider` from +- `agent-framework-core`: `FileHistoryProvider` from `agent_framework/_sessions.py` + +#### `FIDES` + +- `agent-framework-core`: security labeling, content indirection, policy enforcement, and secure MCP + APIs from `agent_framework/security.py`, including `IntegrityLabel`, `ConfidentialityLabel`, + `ContentLabel`, `ContentVariableStore`, `SecureAgentConfig`, and `SecureMCPToolProxy` + +#### `FOUNDRY_TOOLS` + +- `agent-framework-foundry`: released-service tool helpers on `FoundryChatClient`, currently + `get_bing_grounding_tool` and `get_azure_ai_search_tool` + +#### `FOUNDRY_PREVIEW_TOOLS` + +- `agent-framework-foundry`: preview-service tool helpers on `FoundryChatClient`, including Bing + Custom Search, SharePoint, Fabric, Memory Search, Computer Use, Browser Automation, and A2A + +#### `FUNCTIONAL_WORKFLOWS` + +- `agent-framework-core`: functional workflow APIs from + `agent_framework/_workflows/_functional.py`, including `RunContext`, `step`, + `FunctionalWorkflow`, `workflow`, and `FunctionalWorkflowAgent` + +#### `HARNESS` + +- `agent-framework-core`: experimental harness APIs for background agents, file access, looping, + memory, and file-backed todo storage under `agent_framework/_harness/` + +#### `MCP_LONG_RUNNING_TASKS` + +- `agent-framework-core`: `MCPTaskOptions` from `agent_framework/_mcp.py` + +#### `MCP_SKILLS` + +- `agent-framework-core`: `MCPSkillResource`, `MCPSkill`, and `MCPSkillsSource` from `agent_framework/_skills.py` +#### `PROGRESSIVE_TOOLS` + +- `agent-framework-core`: `FunctionInvocationContext.add_tools` and + `FunctionInvocationContext.remove_tools` from `agent_framework/_middleware.py` + +#### `SESSION_STORE` + +- `agent-framework-core`: `SessionStore` and `FileSessionStore` from + `agent_framework/_sessions.py` + +#### `TO_PROMPT_AGENT` + +- `agent-framework-foundry`: `to_prompt_agent` from + `agent_framework_foundry/_to_prompt_agent.py` + ### Release-candidate features There are currently no feature-level `rc` APIs. diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 7ebfaa59f4..4529a19de9 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -74,13 +74,16 @@ agent_framework/ ### Sessions (`_sessions.py`) - **`AgentSession`** - Manages conversation state and session metadata +- **`SessionStore`** - Experimental in-memory `session_id -> AgentSession` snapshot store; reads return independent copies +- **`FileSessionStore`** - Experimental msgspec file-backed session snapshot store with atomic last-writer-wins updates; JSON is the default and `serialization_format="msgpack"` enables binary MessagePack +- **`register_state_type`** - Explicitly registers custom `AgentSession.state` classes with stable type IDs and optional mapping codecs; registration must happen before persistence/restoration, and conflicting IDs are rejected - **`ServiceSessionId`** - Mapping alias for structured service-owned continuation handles used in `AgentSession.service_session_id` - **`SessionContext`** - Context object for session-scoped data during agent runs. `extend_messages(...)` can attach ordered, deduplicated `origin_session_ids` attribution when a provider injects content from other sessions. - **`ContextProvider`** - Base class for context providers (RAG, memory systems) - **`HistoryProvider`** - Base class for conversation history storage - **`InMemoryHistoryProvider`** - Built-in session-state history provider for local runs -- **`FileHistoryProvider`** - JSON Lines file-backed history provider storing one file per session with one message record per line +- **`FileHistoryProvider`** - Experimental append-only file-backed history provider; msgspec JSON Lines is the default and `serialization_format="msgpack"` uses length-prefixed binary MessagePack records. Custom `dumps`/`loads` remain as deprecated JSON-only compatibility hooks and emit `DeprecationWarning` when supplied. ### Skills (`_skills.py`) diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index ae09f7610b..62a385aa8d 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -194,11 +194,13 @@ "AgentSession", "ContextProvider", "FileHistoryProvider", + "FileSessionStore", "HistoryProvider", "InMemoryHistoryProvider", "MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY", "MessageInjectionMiddleware", "ServiceSessionId", + "SessionStore", "SessionContext", "enqueue_messages", "register_state_type", @@ -451,6 +453,7 @@ "FileMemoryProvider", "FileSearchMatch", "FileSearchResult", + "FileSessionStore", "FileSkill", "FileSkillScript", "FileSkillsSource", @@ -516,6 +519,7 @@ "SelectiveToolCallCompactionStrategy", "ServiceSessionId", "SessionContext", + "SessionStore", "SingleEdgeGroup", "Skill", "SkillFrontmatter", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index d4547535e8..cf9a7ce004 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -154,11 +154,13 @@ from ._sessions import ( AgentSession, ContextProvider, FileHistoryProvider, + FileSessionStore, HistoryProvider, InMemoryHistoryProvider, MessageInjectionMiddleware, ServiceSessionId, SessionContext, + SessionStore, enqueue_messages, register_state_type, ) @@ -418,6 +420,7 @@ __all__ = [ "FileMemoryProvider", "FileSearchMatch", "FileSearchResult", + "FileSessionStore", "FileSkill", "FileSkillScript", "FileSkillsSource", @@ -483,6 +486,7 @@ __all__ = [ "SelectiveToolCallCompactionStrategy", "ServiceSessionId", "SessionContext", + "SessionStore", "SingleEdgeGroup", "Skill", "SkillFrontmatter", diff --git a/python/packages/core/agent_framework/_feature_stage.py b/python/packages/core/agent_framework/_feature_stage.py index 39f9c74f4b..7f6c26ffc0 100644 --- a/python/packages/core/agent_framework/_feature_stage.py +++ b/python/packages/core/agent_framework/_feature_stage.py @@ -61,6 +61,7 @@ class ExperimentalFeature(str, Enum): MCP_LONG_RUNNING_TASKS = "MCP_LONG_RUNNING_TASKS" MCP_SKILLS = "MCP_SKILLS" PROGRESSIVE_TOOLS = "PROGRESSIVE_TOOLS" + SESSION_STORE = "SESSION_STORE" TO_PROMPT_AGENT = "TO_PROMPT_AGENT" @@ -140,6 +141,15 @@ def _get_descriptor_callable(obj: Any) -> Callable[..., Any]: return cast(Callable[..., Any], obj.__func__) +def _is_internal_framework_subclass(args: tuple[Any, ...]) -> bool: + """Return whether an experimental subclass is defined inside an Agent Framework package.""" + subclass = next((arg for arg in args if isinstance(arg, type)), None) + if subclass is None: + return False + module = subclass.__module__ + return module == "agent_framework" or module.startswith(("agent_framework.", "agent_framework_")) + + def _is_protocol_class(obj: Any) -> bool: return isinstance(obj, type) and bool(getattr(obj, "_is_protocol", False)) @@ -326,28 +336,30 @@ def __new__(cls: type[Any], /, *args: Any, **kwargs: Any) -> Any: @functools.wraps(original_init_subclass_func) def bound_init_subclass_wrapper(*args: Any, **kwargs: Any) -> Any: - _warn_on_feature_use( - stage=stage, - feature_id=feature_id, - object_name=object_name, - category=category, - ) + if not _is_internal_framework_subclass(args): + _warn_on_feature_use( + stage=stage, + feature_id=feature_id, + object_name=object_name, + category=category, + ) return original_init_subclass_func(*args, **kwargs) experimental_class.__init_subclass__ = classmethod(bound_init_subclass_wrapper) # type: ignore[assignment] else: @functools.wraps(original_init_subclass) - def init_subclass_wrapper(*args: Any, **kwargs: Any) -> Any: - _warn_on_feature_use( - stage=stage, - feature_id=feature_id, - object_name=object_name, - category=category, - ) + def init_subclass_wrapper(subclass: type[Any], /, *args: Any, **kwargs: Any) -> Any: + if not _is_internal_framework_subclass((subclass, *args)): + _warn_on_feature_use( + stage=stage, + feature_id=feature_id, + object_name=object_name, + category=category, + ) return original_init_subclass(*args, **kwargs) - experimental_class.__init_subclass__ = init_subclass_wrapper + experimental_class.__init_subclass__ = classmethod(init_subclass_wrapper) # type: ignore[assignment] return cast(FeatureStageT, experimental_class) diff --git a/python/packages/core/agent_framework/_harness/_memory.py b/python/packages/core/agent_framework/_harness/_memory.py index 04d80cf049..cc30512ab3 100644 --- a/python/packages/core/agent_framework/_harness/_memory.py +++ b/python/packages/core/agent_framework/_harness/_memory.py @@ -976,8 +976,10 @@ def __init__( consolidation_client: Optional chat client override used only for consolidation so the cleanup pass can use a cheaper or faster model than the main agent client. history_message_filter: Optional callback that can rewrite or drop messages before transcript save. - history_dumps: Callable used to serialize transcript JSONL. - history_loads: Callable used to deserialize transcript JSONL and state JSON. + history_dumps: Deprecated callback forwarded to ``FileHistoryProvider.dumps``. + Omit it to use msgspec JSON. + history_loads: Deprecated callback forwarded to ``FileHistoryProvider.loads``. + Omit it to use msgspec JSON. """ super().__init__( source_id=source_id, diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index 9a6dab9a7c..f57a2cb23d 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -7,6 +7,8 @@ - ContextProvider: Base class for context providers - HistoryProvider: Base class for history storage providers - AgentSession: Lightweight session state container +- SessionStore: In-memory session snapshot storage +- FileSessionStore: msgspec JSON file-backed session snapshot storage - InMemoryHistoryProvider: Built-in in-memory history provider - FileHistoryProvider: Built-in JSON Lines file history provider """ @@ -15,17 +17,21 @@ import asyncio import copy -import json import logging +import math +import os import threading import uuid +import warnings import weakref from abc import abstractmethod from base64 import urlsafe_b64encode from collections.abc import AsyncIterable, Awaitable, Callable, Iterable, Mapping, Sequence -from contextlib import suppress +from dataclasses import dataclass from pathlib import Path -from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias, TypeGuard, cast +from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, TypeGuard, TypeVar, cast + +import msgspec from ._feature_stage import ExperimentalFeature, experimental from ._middleware import ChatContext, ChatMiddleware @@ -48,22 +54,79 @@ logger = logging.getLogger("agent_framework") -# Registry of known types for state deserialization -_STATE_TYPE_REGISTRY: dict[str, type] = {} MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY: str = "message_injection.pending_messages" _MESSAGE_INJECTION_LOCK = threading.Lock() JsonDumps: TypeAlias = Callable[[Any], str | bytes] JsonLoads: TypeAlias = Callable[[str | bytes], Any] ServiceSessionId: TypeAlias = Mapping[str, Any] +StateT = TypeVar("StateT") +StateEncoder: TypeAlias = Callable[[Any], Mapping[str, Any]] +StateDecoder: TypeAlias = Callable[[Mapping[str, Any]], Any] +_STATE_SCALAR_TYPES = (str, int, float, bool, type(None)) +_WINDOWS_RESERVED_FILE_STEMS: frozenset[str] = frozenset({ + "CON", + "PRN", + "AUX", + "NUL", + "COM1", + "COM2", + "COM3", + "COM4", + "COM5", + "COM6", + "COM7", + "COM8", + "COM9", + "LPT1", + "LPT2", + "LPT3", + "LPT4", + "LPT5", + "LPT6", + "LPT7", + "LPT8", + "LPT9", +}) + + +_DEFAULT_JSON_ENCODER = msgspec.json.Encoder() +_DEFAULT_JSON_DECODER = msgspec.json.Decoder() +_DEFAULT_MSGPACK_ENCODER = msgspec.msgpack.Encoder() +_DEFAULT_MSGPACK_DECODER = msgspec.msgpack.Decoder() +_JSON_FILE_EXTENSION = ".json" +_JSON_LINES_FILE_EXTENSION = ".jsonl" +_MSGPACK_FILE_EXTENSION = ".msgpack" + + +def _default_json_dumps(value: Any) -> bytes: + return _DEFAULT_JSON_ENCODER.encode(value) -def _default_json_dumps(value: Any) -> str: - return json.dumps(value, ensure_ascii=False) +def _default_json_loads(value: str | bytes) -> Any: + return _DEFAULT_JSON_DECODER.decode(value) -def _default_json_loads(value: str | bytes) -> Any: - return json.loads(value) +def _is_literal_session_file_stem_safe(session_id: str) -> bool: + """Return whether a session ID can be used directly as a filename stem.""" + if ( + not session_id + or session_id.startswith(".") + or session_id.endswith((" ", ".")) + or session_id.upper() in _WINDOWS_RESERVED_FILE_STEMS + ): + return False + if any(ord(character) < 32 for character in session_id): + return False + return all(character.isalnum() or character in "._-" for character in session_id) + + +def _session_file_stem(session_id: str, *, encoded_prefix: str) -> str: + """Return a safe filename stem for an opaque session ID.""" + if _is_literal_session_file_stem_safe(session_id): + return session_id + encoded_session_id = urlsafe_b64encode(session_id.encode("utf-8")).decode("ascii").rstrip("=") + return f"{encoded_prefix}{encoded_session_id}" def _deduplicate_origin_session_ids(origin_session_ids: Iterable[str]) -> list[str]: @@ -89,80 +152,204 @@ def _is_single_middleware( return not _is_middleware_sequence(middleware) -def register_state_type(cls: type) -> None: - """Register a type for automatic deserialization in session state. +@dataclass(frozen=True, slots=True) +class _StateTypeRegistration: + cls: type[Any] + type_id: str + encoder: StateEncoder + decoder: StateDecoder + + +_STATE_TYPE_REGISTRY: dict[str, _StateTypeRegistration] = {} +_STATE_CLASS_REGISTRY: dict[type[Any], _StateTypeRegistration] = {} + - Call this for any custom type (including Pydantic models) that you store - in ``session.state`` and want to survive ``to_dict()`` / ``from_dict()`` - round-trips. Types with ``to_dict``/``from_dict`` methods or Pydantic - ``BaseModel`` subclasses are handled automatically. +def _resolve_state_type_id(cls: type[Any], type_id: str | None) -> str: + """Resolve the stable serialized ID for a registered session-state type.""" + if type_id is not None: + resolved = type_id + elif callable(identifier := getattr(cls, "_get_type_identifier", None)): + resolved = identifier() + elif isinstance(identifier := getattr(cls, "TYPE", None), str): + resolved = identifier + else: + resolved = cls.__name__.lower() + if not isinstance(resolved, str) or not resolved: + raise ValueError("State type identifier must be a non-empty string.") + return resolved + + +def _default_state_encoder(cls: type[Any]) -> StateEncoder: + """Create the default encoder for one explicitly registered state type.""" + if callable(getattr(cls, "to_dict", None)): + + def encode_to_dict(value: Any) -> Mapping[str, Any]: + payload = value.to_dict() + if not isinstance(payload, Mapping): + raise TypeError(f"{cls.__name__}.to_dict() must return a mapping.") + return cast(Mapping[str, Any], payload) + + return encode_to_dict + + from pydantic import BaseModel + + if issubclass(cls, BaseModel): + + def encode_pydantic(value: Any) -> Mapping[str, Any]: + return cast(Mapping[str, Any], value.model_dump()) + + return encode_pydantic + + raise ValueError( + f"State type {cls.__name__!r} must define to_dict()/from_dict(), be a Pydantic model, " + "or provide encoder and decoder callbacks." + ) - The type identifier defaults to ``cls.__name__.lower()`` but can be - overridden by defining a ``_get_type_identifier`` classmethod. - Note: - Pydantic models are auto-registered on first serialization, but - pre-registering ensures deserialization works even if the model - hasn't been serialized in this process yet (e.g. cold-start restore). +def _default_state_decoder(cls: type[Any]) -> StateDecoder: + """Create the default decoder for one explicitly registered state type.""" + if callable(getattr(cls, "from_dict", None)): + + def decode_from_dict(payload: Mapping[str, Any]) -> Any: + return cls.from_dict(dict(payload)) + + return decode_from_dict + + from pydantic import BaseModel + + if issubclass(cls, BaseModel): + + def decode_pydantic(payload: Mapping[str, Any]) -> Any: + return cls.model_validate({key: value for key, value in payload.items() if key != "type"}) + + return decode_pydantic + + raise ValueError( + f"State type {cls.__name__!r} must define to_dict()/from_dict(), be a Pydantic model, " + "or provide encoder and decoder callbacks." + ) + + +def register_state_type( + cls: type[StateT], + *, + type_id: str | None = None, + encoder: Callable[[StateT], Mapping[str, Any]] | None = None, + decoder: Callable[[Mapping[str, Any]], StateT] | None = None, +) -> None: + """Register a type for automatic deserialization in session state. + + Registration is explicit so persisted sessions can be restored after a + process restart. The type identifier is resolved from ``type_id``, then + ``_get_type_identifier()``, then ``TYPE``, and finally the lowercased class + name. Classes implementing ``to_dict`` / ``from_dict`` and Pydantic models + receive default codecs; other classes must provide both callbacks. + + Call this at module level immediately after defining the custom state class. + Importing that module then registers the type before any persisted session + is loaded, without requiring the related context provider to be instantiated + first. Args: cls: The type to register. - """ - type_id: str = getattr(cls, "_get_type_identifier", lambda: cls.__name__.lower())() - _STATE_TYPE_REGISTRY[type_id] = cls - -# Keep internal alias for framework use -_register_state_type = register_state_type + Keyword Args: + type_id: Stable identifier stored in the serialized ``type`` field. + encoder: Optional callback that converts an instance to a mapping. + decoder: Optional callback that reconstructs an instance from a mapping. + Raises: + ValueError: If the registration is incomplete or conflicts with an existing type. + """ + resolved_type_id = _resolve_state_type_id(cls, type_id) + existing_for_class = _STATE_CLASS_REGISTRY.get(cls) + if existing_for_class is not None: + if existing_for_class.type_id != resolved_type_id: + raise ValueError( + f"State type {cls.__name__!r} is already registered as {existing_for_class.type_id!r}, " + f"not {resolved_type_id!r}." + ) + if encoder is not None and existing_for_class.encoder is not encoder: + raise ValueError(f"State type {cls.__name__!r} is already registered with a different encoder.") + if decoder is not None and existing_for_class.decoder is not decoder: + raise ValueError(f"State type {cls.__name__!r} is already registered with a different decoder.") + return -def _serialize_value(value: Any) -> Any: - """Serialize a single value, handling objects with to_dict() and Pydantic models.""" - if hasattr(value, "to_dict") and callable(value.to_dict): - return value.to_dict() - # Pydantic BaseModel support — import lazily to avoid hard dep at module level - with suppress(ImportError): - from pydantic import BaseModel + existing_for_id = _STATE_TYPE_REGISTRY.get(resolved_type_id) + if existing_for_id is not None: + raise ValueError( + f"State type identifier {resolved_type_id!r} is already registered for {existing_for_id.cls.__name__!r}." + ) - if isinstance(value, BaseModel): - data = value.model_dump() - type_id: str = getattr(value.__class__, "_get_type_identifier", lambda: value.__class__.__name__.lower())() - data["type"] = type_id - # Auto-register for round-trip deserialization - _STATE_TYPE_REGISTRY.setdefault(type_id, value.__class__) - return data + if (encoder is None) != (decoder is None): + raise ValueError("State type encoder and decoder must be provided together.") + resolved_encoder = cast(StateEncoder, encoder) if encoder is not None else _default_state_encoder(cls) + resolved_decoder = cast(StateDecoder, decoder) if decoder is not None else _default_state_decoder(cls) + registration = _StateTypeRegistration( + cls=cls, + type_id=resolved_type_id, + encoder=resolved_encoder, + decoder=resolved_decoder, + ) + _STATE_TYPE_REGISTRY[resolved_type_id] = registration + _STATE_CLASS_REGISTRY[cls] = registration + + +def _serialize_value(value: Any, *, path: str) -> Any: + """Serialize one session-state value through the explicit type registry.""" + value_type = cast(type[Any], type(value)) + registration = _STATE_CLASS_REGISTRY.get(value_type) + if registration is not None: + payload = registration.encoder(value) + payload_type = payload.get("type") + if payload_type is not None and payload_type != registration.type_id: + raise ValueError( + f"State encoder for {registration.cls.__name__!r} returned type {payload_type!r}; " + f"expected {registration.type_id!r}." + ) + serialized = { + str(key): _serialize_value(item, path=f"{path}.{key}") for key, item in payload.items() if key != "type" + } + serialized["type"] = registration.type_id + return serialized + if isinstance(value, float) and not math.isfinite(value): + raise ValueError(f"Session state value at {path} must be a finite float.") + if value_type in _STATE_SCALAR_TYPES: + return value if isinstance(value, list): - return [_serialize_value(item) for item in value] # pyright: ignore[reportUnknownVariableType] - if isinstance(value, dict): - return {str(k): _serialize_value(v) for k, v in value.items()} # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] - return value + return [_serialize_value(item, path=f"{path}[{index}]") for index, item in enumerate(cast(list[Any], value))] + if isinstance(value, tuple): + return [ + _serialize_value(item, path=f"{path}[{index}]") for index, item in enumerate(cast(tuple[Any, ...], value)) + ] + if isinstance(value, Mapping): + return { + str(key): _serialize_value(item, path=f"{path}.{key}") + for key, item in cast(Mapping[Any, Any], value).items() + } + raise TypeError( + f"Session state value at {path} has unregistered type {type(value).__name__!r}; " + "call register_state_type() before persisting the session." + ) def _deserialize_value(value: Any) -> Any: """Deserialize a single value, restoring registered types.""" - if isinstance(value, dict) and "type" in value: - type_id = str(value["type"]) # pyright: ignore[reportUnknownArgumentType] - cls = _STATE_TYPE_REGISTRY.get(type_id) - if cls is not None: - if hasattr(cls, "from_dict"): - return cls.from_dict(value) # type: ignore[union-attr] - # Pydantic BaseModel support - with suppress(ImportError): - from pydantic import BaseModel - - if issubclass(cls, BaseModel): - data: dict[str, Any] = {str(k): v for k, v in value.items() if k != "type"} # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] - return cls.model_validate(data) if isinstance(value, list): - return [_deserialize_value(item) for item in value] # pyright: ignore[reportUnknownVariableType] - if isinstance(value, dict): - return {str(k): _deserialize_value(v) for k, v in value.items()} # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] + return [_deserialize_value(item) for item in cast(list[Any], value)] + if isinstance(value, Mapping): + raw_mapping = {str(key): item for key, item in cast(Mapping[Any, Any], value).items()} + if "type" in raw_mapping: + registration = _STATE_TYPE_REGISTRY.get(str(raw_mapping["type"])) + if registration is not None: + return registration.decoder(raw_mapping) + return {key: _deserialize_value(item) for key, item in raw_mapping.items()} return value def _serialize_state(state: dict[str, Any]) -> dict[str, Any]: """Deep-serialize a state dict, converting SerializationProtocol objects to dicts.""" - return {k: _serialize_value(v) for k, v in state.items()} + return {key: _serialize_value(value, path=f"state.{key}") for key, value in state.items()} def _deserialize_state(state: dict[str, Any]) -> dict[str, Any]: @@ -170,8 +357,52 @@ def _deserialize_state(state: dict[str, Any]) -> dict[str, Any]: return {k: _deserialize_value(v) for k, v in state.items()} +class _SessionStatePayload: + """Wrapper that routes the complete dynamic state mapping through msgspec hooks.""" + + __slots__ = ("value",) + + def __init__(self, value: dict[str, Any]) -> None: + self.value = value + + def serialize(self) -> dict[str, Any]: + """Serialize the wrapped state through the explicit registry.""" + return _serialize_state(self.value) + + +class _SessionSnapshot(msgspec.Struct): + """Typed on-disk representation of one AgentSession.""" + + type: Literal["session"] + session_id: str + service_session_id: str | dict[str, Any] | None + state: _SessionStatePayload + + +def _session_snapshot_enc_hook(value: Any) -> Any: + """Encode the complete dynamic state payload for msgspec.""" + if isinstance(value, _SessionStatePayload): + return value.serialize() + raise NotImplementedError(f"Objects of type {type(value).__name__!r} are not supported.") + + +def _session_snapshot_dec_hook(target_type: type[Any], value: Any) -> Any: + """Restore the complete dynamic state payload for msgspec.""" + if target_type is _SessionStatePayload: + if not isinstance(value, Mapping): + raise TypeError("Session state payload must be a mapping.") + return _SessionStatePayload(_deserialize_state(dict(cast(Mapping[str, Any], value)))) + raise NotImplementedError(f"Objects of type {target_type.__name__!r} are not supported.") + + +_SESSION_SNAPSHOT_ENCODER = msgspec.json.Encoder(enc_hook=_session_snapshot_enc_hook) +_SESSION_SNAPSHOT_DECODER = msgspec.json.Decoder(_SessionSnapshot, dec_hook=_session_snapshot_dec_hook) +_SESSION_SNAPSHOT_MSGPACK_ENCODER = msgspec.msgpack.Encoder(enc_hook=_session_snapshot_enc_hook) +_SESSION_SNAPSHOT_MSGPACK_DECODER = msgspec.msgpack.Decoder(_SessionSnapshot, dec_hook=_session_snapshot_dec_hook) + + # Register known types -_register_state_type(Message) +register_state_type(Message) class SessionContext: @@ -421,6 +652,17 @@ class ContextProvider: Context providers participate in the context engineering pipeline, adding context before model invocation and processing responses after. + Provider-scoped ``state`` is stored inside :attr:`AgentSession.state` and + may be persisted by a session store. Standard JSON-native Python values + (``None``, booleans, integers, finite floats, strings, lists, tuples, and + mappings) require no registration. If a provider stores an instance of a + custom class or Pydantic model, the application must call + :func:`register_state_type` for that class at module level, immediately + after its definition. Importing the module then registers the type before a + persisted session is loaded, even when the related context provider has not + been instantiated yet. Framework-owned state types such as :class:`Message` + are registered by Agent Framework. + Attributes: source_id: Unique identifier for this provider instance (required). Used for message/tool attribution so other providers can filter. @@ -489,6 +731,16 @@ class HistoryProvider(ContextProvider): The default ``before_run``/``after_run`` handle loading and storing based on configuration flags. Override them for custom behavior. + Normal :class:`Message` history requires no custom registration because + Agent Framework registers ``Message`` for session persistence. If a history + provider stores any other custom class or Pydantic model in its + provider-scoped ``state`` or elsewhere in :attr:`AgentSession.state`, the + application must call :func:`register_state_type` at module level + immediately after defining that class. This ensures importing the provider + module registers its state types before session restoration and before the + provider itself is instantiated. Prefer JSON-native Python values when a + custom runtime type is not needed after restoration. + Attributes: load_messages: Whether to load messages before invocation (default True). When False, the agent skips calling ``before_run`` entirely. @@ -1007,9 +1259,9 @@ def session_id(self) -> str: def to_dict(self) -> dict[str, Any]: """Serialize session to a plain dict for storage/transfer. - Values in ``state`` that implement ``SerializationProtocol`` (i.e. have - ``to_dict``/``from_dict``) are serialized automatically. Built-in types - (str, int, float, bool, None, list, dict) are kept as-is. + Custom values in ``state`` must be registered with + :func:`register_state_type`. Built-in JSON scalar, list, tuple, and + mapping values are serialized recursively. """ return { "type": "session", @@ -1039,6 +1291,218 @@ def from_dict(cls, data: dict[str, Any]) -> AgentSession: return session +@experimental(feature_id=ExperimentalFeature.SESSION_STORE) +class SessionStore: + """In-memory storage for Agent Framework session snapshots. + + The store maps an opaque caller-selected ID to an :class:`AgentSession`. + Reads return independent working copies so one continuation does not mutate + another stored snapshot. The store has no eviction; callers that need TTLs, + durable storage, or distributed coordination should provide another + implementation with the same async methods. + + Session-store IDs are limited to 128 characters containing only ASCII + letters, digits, ``-``, and ``_``. This restricted shape is defense in + depth for storage backends; custom implementations must still use + parameterized queries and their backend's normal key-handling protections. + """ + + MAX_SESSION_ID_LENGTH: ClassVar[int] = 128 + + def __init__(self) -> None: + """Create an empty session store.""" + self._sessions: dict[str, AgentSession] = {} + + @staticmethod + def validate_session_id(session_id: str) -> None: + """Validate an ID for use with a session store. + + Args: + session_id: Session-store ID to validate. + + Raises: + ValueError: If the ID is empty or contains characters other than + ASCII letters, digits, ``-``, and ``_``. + """ + if not isinstance(session_id, str) or not session_id: + raise ValueError("session_id must be a non-empty string") + if len(session_id) > SessionStore.MAX_SESSION_ID_LENGTH: + raise ValueError(f"session_id must be at most {SessionStore.MAX_SESSION_ID_LENGTH} characters") + if not all(character.isascii() and (character.isalnum() or character in "-_") for character in session_id): + raise ValueError("session_id must contain only ASCII letters, digits, '-' and '_'") + + async def get(self, session_id: str) -> AgentSession | None: + """Return a copy of the stored session, or ``None`` when absent. + + Args: + session_id: Opaque caller-selected session ID. + + Returns: + An independent copy of the stored session, or ``None``. + + Raises: + ValueError: If ``session_id`` is empty or contains unsupported characters. + """ + SessionStore.validate_session_id(session_id) + session = self._sessions.get(session_id) + return copy.deepcopy(session) if session is not None else None + + async def set(self, session_id: str, session: AgentSession) -> None: + """Store ``session`` under ``session_id``, replacing any existing entry. + + Args: + session_id: Opaque caller-selected session ID. + session: The session to store. + + Raises: + ValueError: If ``session_id`` is empty or contains unsupported characters. + """ + SessionStore.validate_session_id(session_id) + self._sessions[session_id] = session + + async def delete(self, session_id: str) -> None: + """Delete the stored session, if present. + + Args: + session_id: Opaque caller-selected session ID. + + Raises: + ValueError: If ``session_id`` is empty or contains unsupported characters. + """ + SessionStore.validate_session_id(session_id) + self._sessions.pop(session_id, None) + + +@experimental(feature_id=ExperimentalFeature.SESSION_STORE) +class FileSessionStore(SessionStore): + """File-backed storage for Agent Framework session snapshots. + + Each session is stored as one JSON or MessagePack file beneath + ``storage_path``. JSON is the default; pass ``serialization_format="msgpack"`` + for a compact binary representation. + Writes use a unique sibling temporary file followed by :func:`os.replace` + so readers never observe a partially written snapshot. Concurrent writers + use last-writer-wins semantics. + + The complete snapshot is encoded and decoded in one call through a typed + :mod:`msgspec` codec. The dynamic state mapping is routed through the + explicit :func:`register_state_type` registry by one state-payload hook. + + Security posture: + Persisted session snapshots are stored as plaintext JSON or binary + MessagePack on the local filesystem. + Treat ``storage_path`` as trusted application storage, not as a secret + store. Restricted store keys and resolved-path validation help prevent + path traversal via ``session_id``, but they do not encrypt file contents + or coordinate concurrent updates across tasks, processes, or hosts. + Atomic replacement prevents partial writes, while concurrent writers + still use last-writer-wins semantics. Use OS-level file permissions, + trusted directories, and carefully review what session state is allowed + to be persisted. + """ + + _ENCODED_SESSION_PREFIX: ClassVar[str] = "~session-" + + def __init__( + self, + storage_path: str | Path, + *, + serialization_format: Literal["json", "msgpack"] = "json", + ) -> None: + """Initialize file-backed session storage. + + Args: + storage_path: Directory where session snapshot files are stored. + + Keyword Args: + serialization_format: ``"json"`` (default) for readable JSON files + or ``"msgpack"`` for binary MessagePack files. + + Raises: + ValueError: If ``serialization_format`` is unsupported. + """ + if serialization_format not in ("json", "msgpack"): + raise ValueError("serialization_format must be 'json' or 'msgpack'") + self.storage_path = Path(storage_path) + self.storage_path.mkdir(parents=True, exist_ok=True) + self._storage_root = self.storage_path.resolve() + self.serialization_format = serialization_format + if serialization_format == "json": + self._encoder = _SESSION_SNAPSHOT_ENCODER + self._decoder = _SESSION_SNAPSHOT_DECODER + self._file_extension = _JSON_FILE_EXTENSION + else: + self._encoder = _SESSION_SNAPSHOT_MSGPACK_ENCODER + self._decoder = _SESSION_SNAPSHOT_MSGPACK_DECODER + self._file_extension = _MSGPACK_FILE_EXTENSION + + async def get(self, session_id: str) -> AgentSession | None: + """Load a session snapshot, or return ``None`` when it does not exist.""" + file_path = self._session_file_path(session_id) + + def _read() -> AgentSession | None: + try: + serialized = file_path.read_bytes() + except FileNotFoundError: + return None + try: + snapshot = self._decoder.decode(serialized) + session = AgentSession( + session_id=snapshot.session_id, + service_session_id=snapshot.service_session_id, + ) + session.state = snapshot.state.value + return session + except (msgspec.DecodeError, TypeError, ValueError) as exc: + raise ValueError(f"Failed to deserialize session from '{file_path}'.") from exc + + return await asyncio.to_thread(_read) + + async def set(self, session_id: str, session: AgentSession) -> None: + """Persist a session snapshot atomically.""" + file_path = self._session_file_path(session_id) + if type(session) is not AgentSession: + raise TypeError( + "FileSessionStore supports AgentSession instances only; " + "custom AgentSession subclasses require a custom SessionStore." + ) + service_session_id = session.service_session_id + serialized_service_session_id = ( + dict(service_session_id) if isinstance(service_session_id, Mapping) else service_session_id + ) + snapshot = _SessionSnapshot( + type="session", + session_id=session.session_id, + service_session_id=serialized_service_session_id, + state=_SessionStatePayload(session.state), + ) + serialized = self._encoder.encode(snapshot) + + def _write() -> None: + temp_path = file_path.with_name(f".{file_path.name}.{uuid.uuid4().hex}.tmp") + try: + temp_path.write_bytes(serialized) + os.replace(temp_path, file_path) + finally: + temp_path.unlink(missing_ok=True) + + await asyncio.to_thread(_write) + + async def delete(self, session_id: str) -> None: + """Delete a persisted session snapshot, if present.""" + file_path = self._session_file_path(session_id) + await asyncio.to_thread(file_path.unlink, missing_ok=True) + + def _session_file_path(self, session_id: str) -> Path: + """Resolve the contained snapshot path for ``session_id``.""" + SessionStore.validate_session_id(session_id) + file_stem = _session_file_stem(session_id, encoded_prefix=self._ENCODED_SESSION_PREFIX) + file_path = (self._storage_root / f"{file_stem}{self._file_extension}").resolve() + if not file_path.is_relative_to(self._storage_root): + raise ValueError(f"Session path escaped storage directory: {session_id!r}") + return file_path + + class InMemoryHistoryProvider(HistoryProvider): """Built-in history provider that stores messages in session.state. @@ -1120,16 +1584,17 @@ async def save_messages( @experimental(feature_id=ExperimentalFeature.FILE_HISTORY) class FileHistoryProvider(HistoryProvider): - """File-backed history provider that stores one JSON Lines file per session. + """File-backed history provider that stores one append-only file per session. - Each persisted message is written as a single JSON object per line. The - provider does not serialize full session snapshots into the file. By default - it uses the standard library ``json`` module, but callers can inject - alternative ``dumps`` and ``loads`` callables compatible with the JSON - Lines format. + JSON Lines is the default: each message is one JSON object per line. + Pass ``serialization_format="msgpack"`` to store length-prefixed binary + MessagePack records instead. Both formats use :mod:`msgspec` by default. + The custom ``dumps`` and ``loads`` constructor arguments remain available + for JSON Lines compatibility but are deprecated. Security posture: - Persisted history is stored as plaintext JSONL on the local filesystem. + Persisted history is stored as plaintext JSONL or binary MessagePack on + the local filesystem. Treat ``storage_path`` as trusted application storage, not as a secret store. Encoded fallback filenames and resolved-path validation help prevent path traversal via ``session_id``, but they do not encrypt file @@ -1140,36 +1605,13 @@ class FileHistoryProvider(HistoryProvider): DEFAULT_SOURCE_ID: ClassVar[str] = "file_history" DEFAULT_SESSION_FILE_STEM: ClassVar[str] = "default" - FILE_EXTENSION: ClassVar[str] = ".jsonl" _FILE_LOCK_STRIPE_COUNT: ClassVar[int] = 64 _ENCODED_SESSION_PREFIX: ClassVar[str] = "~session-" + _MSGPACK_RECORD_HEADER_BYTES: ClassVar[int] = 4 + _MAX_MSGPACK_RECORD_BYTES: ClassVar[int] = 64 * 1024 * 1024 _FILE_WRITE_LOCKS: ClassVar[tuple[threading.Lock, ...]] = tuple( threading.Lock() for _ in range(_FILE_LOCK_STRIPE_COUNT) ) - _WINDOWS_RESERVED_FILE_STEMS: ClassVar[frozenset[str]] = frozenset({ - "CON", - "PRN", - "AUX", - "NUL", - "COM1", - "COM2", - "COM3", - "COM4", - "COM5", - "COM6", - "COM7", - "COM8", - "COM9", - "LPT1", - "LPT2", - "LPT3", - "LPT4", - "LPT5", - "LPT6", - "LPT7", - "LPT8", - "LPT9", - }) def __init__( self, @@ -1182,6 +1624,7 @@ def __init__( store_context_from: set[str] | None = None, store_outputs: bool = True, skip_excluded: bool = False, + serialization_format: Literal["json", "msgpack"] = "json", dumps: JsonDumps | None = None, loads: JsonLoads | None = None, ) -> None: @@ -1199,11 +1642,31 @@ def __init__( store_outputs: Whether to store response messages. skip_excluded: When True, ``get_messages`` omits messages whose ``additional_properties["_excluded"]`` is truthy. - dumps: Callable that serializes a message payload dict to JSON text - or UTF-8 bytes. The returned JSON must fit on a single line. - loads: Callable that deserializes JSON text or bytes back to a - message payload dict. + serialization_format: ``"json"`` (default) for JSON Lines or + ``"msgpack"`` for length-prefixed binary MessagePack records. + dumps: Deprecated. Callable that serializes a message payload dict + to single-line JSON text or UTF-8 bytes. Omit this argument to + use the built-in msgspec codec. + loads: Deprecated. Callable that deserializes JSON text or bytes + back to a message payload dict. Omit this argument to use the + built-in msgspec codec. + + Raises: + ValueError: If the format is unsupported or custom JSON codecs are + supplied with MessagePack. """ + if serialization_format not in ("json", "msgpack"): + raise ValueError("serialization_format must be 'json' or 'msgpack'") + if serialization_format == "msgpack" and (dumps is not None or loads is not None): + raise ValueError("Custom dumps and loads are supported only with serialization_format='json'") + if dumps is not None or loads is not None: + warnings.warn( + "The FileHistoryProvider constructor arguments `dumps` and `loads` are deprecated and will be " + "removed in a future version. Omit them to use the built-in msgspec codec selected by " + "`serialization_format`.", + DeprecationWarning, + stacklevel=2, + ) super().__init__( source_id=source_id, load_messages=load_messages, @@ -1216,6 +1679,8 @@ def __init__( self.storage_path.mkdir(parents=True, exist_ok=True) self._storage_root = self.storage_path.resolve() self.skip_excluded = skip_excluded + self.serialization_format = serialization_format + self._file_extension = _JSON_LINES_FILE_EXTENSION if serialization_format == "json" else _MSGPACK_FILE_EXTENSION self.dumps = dumps or _default_json_dumps self.loads = loads or _default_json_loads self._async_write_locks_by_loop: weakref.WeakKeyDictionary[ @@ -1230,7 +1695,7 @@ async def get_messages( state: dict[str, Any] | None = None, **kwargs: Any, ) -> list[Message]: - """Retrieve messages from the session's JSON Lines file.""" + """Retrieve messages from the session's history file.""" del state, kwargs file_path = self._session_file_path(session_id) async_lock = self._session_async_write_lock(file_path) @@ -1241,31 +1706,9 @@ def _read_messages() -> list[Message]: if not file_path.exists(): return [] - messages: list[Message] = [] - with file_path.open(encoding="utf-8") as file_handle: - for line_number, line in enumerate(file_handle, start=1): - serialized = line.strip() - if not serialized: - continue - try: - payload = self.loads(serialized) - except (TypeError, ValueError) as exc: - raise ValueError( - f"Failed to deserialize history line {line_number} from '{file_path}'." - ) from exc - if not isinstance(payload, Mapping): - raise ValueError( - f"History line {line_number} in '{file_path}' did not deserialize to a mapping." - ) - - try: - message = Message.from_dict(dict(cast(Mapping[str, Any], payload))) - except ValueError as exc: - raise ValueError( - f"History line {line_number} in '{file_path}' is not a valid Message payload." - ) from exc - messages.append(message) - return messages + if self.serialization_format == "json": + return self._read_json_messages(file_path) + return self._read_msgpack_messages(file_path) async with async_lock: messages = await asyncio.to_thread(_read_messages) @@ -1281,7 +1724,7 @@ async def save_messages( state: dict[str, Any] | None = None, **kwargs: Any, ) -> None: - """Append messages to the session's JSON Lines file.""" + """Append messages to the session's history file.""" del state, kwargs if not messages: return @@ -1291,14 +1734,77 @@ async def save_messages( file_lock = self._session_write_lock(file_path) def _append_messages() -> None: - with file_lock, file_path.open("a", encoding="utf-8") as file_handle: - for message in messages: - file_handle.write(f"{self._serialize_message(message)}\n") + with file_lock: + if self.serialization_format == "json": + with file_path.open("a", encoding="utf-8") as file_handle: + for message in messages: + file_handle.write(f"{self._serialize_json_message(message)}\n") + return + with file_path.open("ab") as file_handle: + for message in messages: + serialized = _DEFAULT_MSGPACK_ENCODER.encode(message.to_dict()) + file_handle.write(len(serialized).to_bytes(self._MSGPACK_RECORD_HEADER_BYTES, "big")) + file_handle.write(serialized) async with async_lock: await asyncio.to_thread(_append_messages) - def _serialize_message(self, message: Message) -> str: + def _read_json_messages(self, file_path: Path) -> list[Message]: + """Read JSON Lines messages from ``file_path``.""" + messages: list[Message] = [] + with file_path.open(encoding="utf-8") as file_handle: + for line_number, line in enumerate(file_handle, start=1): + serialized = line.strip() + if not serialized: + continue + try: + payload = self.loads(serialized) + except (TypeError, ValueError) as exc: + raise ValueError(f"Failed to deserialize history line {line_number} from '{file_path}'.") from exc + messages.append(self._parse_message_payload(payload, file_path=file_path, record_number=line_number)) + return messages + + def _read_msgpack_messages(self, file_path: Path) -> list[Message]: + """Read length-prefixed MessagePack records from ``file_path``.""" + messages: list[Message] = [] + with file_path.open("rb") as file_handle: + record_number = 0 + while True: + header = file_handle.read(self._MSGPACK_RECORD_HEADER_BYTES) + if not header: + return messages + record_number += 1 + if len(header) != self._MSGPACK_RECORD_HEADER_BYTES: + raise ValueError(f"History record {record_number} in '{file_path}' has a truncated length header.") + record_length = int.from_bytes(header, "big") + if record_length <= 0 or record_length > self._MAX_MSGPACK_RECORD_BYTES: + raise ValueError( + f"History record {record_number} in '{file_path}' has invalid length {record_length}." + ) + serialized = file_handle.read(record_length) + if len(serialized) != record_length: + raise ValueError(f"History record {record_number} in '{file_path}' is truncated.") + try: + payload = _DEFAULT_MSGPACK_DECODER.decode(serialized) + except msgspec.DecodeError as exc: + raise ValueError( + f"Failed to deserialize history record {record_number} from '{file_path}'." + ) from exc + messages.append(self._parse_message_payload(payload, file_path=file_path, record_number=record_number)) + + @staticmethod + def _parse_message_payload(payload: Any, *, file_path: Path, record_number: int) -> Message: + """Validate and reconstruct one stored Message payload.""" + if not isinstance(payload, Mapping): + raise ValueError(f"History record {record_number} in '{file_path}' did not deserialize to a mapping.") + try: + return Message.from_dict(dict(cast(Mapping[str, Any], payload))) + except ValueError as exc: + raise ValueError( + f"History record {record_number} in '{file_path}' is not a valid Message payload." + ) from exc + + def _serialize_json_message(self, message: Message) -> str: """Serialize a message payload to a single JSON Lines record.""" serialized = self.dumps(message.to_dict()) if isinstance(serialized, bytes): @@ -1314,7 +1820,7 @@ def _serialize_message(self, message: Message) -> str: def _session_file_path(self, session_id: str | None) -> Path: """Resolve the on-disk history file path for a session.""" - file_path = (self._storage_root / f"{self._session_file_stem(session_id)}{self.FILE_EXTENSION}").resolve() + file_path = (self._storage_root / f"{self._session_file_stem(session_id)}{self._file_extension}").resolve() if not file_path.is_relative_to(self._storage_root): raise ValueError(f"Session history path escaped storage directory: {session_id!r}") return file_path @@ -1322,11 +1828,7 @@ def _session_file_path(self, session_id: str | None) -> Path: def _session_file_stem(self, session_id: str | None) -> str: """Return the filename stem for a session.""" raw_session_id = session_id or self.DEFAULT_SESSION_FILE_STEM - if self._is_literal_session_file_stem_safe(raw_session_id): - return raw_session_id - - encoded_session_id = urlsafe_b64encode(raw_session_id.encode("utf-8")).decode("ascii").rstrip("=") - return f"{self._ENCODED_SESSION_PREFIX}{encoded_session_id or self.DEFAULT_SESSION_FILE_STEM}" + return _session_file_stem(raw_session_id, encoded_prefix=self._ENCODED_SESSION_PREFIX) def _session_async_write_lock(self, file_path: Path) -> asyncio.Lock: """Return the event-loop-local async lock for a session history file.""" @@ -1350,13 +1852,4 @@ def _lock_index(cls, file_path: Path) -> int: @classmethod def _is_literal_session_file_stem_safe(cls, session_id: str) -> bool: """Return whether the session ID can be used directly as a filename stem.""" - if ( - not session_id - or session_id.startswith(".") - or session_id.endswith((" ", ".")) - or session_id.upper() in cls._WINDOWS_RESERVED_FILE_STEMS - ): - return False - if any(ord(character) < 32 for character in session_id): - return False - return all(character.isalnum() or character in "._-" for character in session_id) + return _is_literal_session_file_stem_safe(session_id) diff --git a/python/packages/core/pyproject.toml b/python/packages/core/pyproject.toml index ac535dc693..e2ea3bdc1a 100644 --- a/python/packages/core/pyproject.toml +++ b/python/packages/core/pyproject.toml @@ -23,6 +23,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ + "msgspec>=0.20.0,<1", "typing-extensions>=4.15.0,<5", "pydantic>=2,<3", "python-dotenv>=1,<2", diff --git a/python/packages/core/tests/core/test_feature_stage.py b/python/packages/core/tests/core/test_feature_stage.py index ba05281d1f..bf8be941ea 100644 --- a/python/packages/core/tests/core/test_feature_stage.py +++ b/python/packages/core/tests/core/test_feature_stage.py @@ -176,6 +176,23 @@ def do(self) -> int: assert Concrete().do() == 1 +def test_experimental_internal_framework_subclass_does_not_warn() -> None: + @experimental(feature_id=AlternateExperimentalFeature.EXPERIMENTAL_FEATURE) # type: ignore[arg-type] # ty: ignore[invalid-argument-type] + class ExperimentalBase: + pass + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + internal_subclass = type( + "InternalSubclass", + (ExperimentalBase,), + {"__module__": "agent_framework._internal_test"}, + ) + + assert caught == [] + assert issubclass(internal_subclass, ExperimentalBase) + + def test_experimental_runtime_checkable_protocol_keeps_protocol_runtime_checks() -> None: with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") diff --git a/python/packages/core/tests/core/test_harness_memory.py b/python/packages/core/tests/core/test_harness_memory.py index f5b2733e57..a23f9e58d9 100644 --- a/python/packages/core/tests/core/test_harness_memory.py +++ b/python/packages/core/tests/core/test_harness_memory.py @@ -194,11 +194,7 @@ async def test_memory_file_store_writes_topics_index_state_and_transcripts(tmp_p source_id=DEFAULT_MEMORY_SOURCE_ID, )["sessions_since_consolidation"] == ["session-1"] - history_provider = FileHistoryProvider( - store.get_transcripts_directory(session, source_id=DEFAULT_MEMORY_SOURCE_ID), - dumps=lambda value: json.dumps(value, separators=(",", ":"), sort_keys=True), - loads=json.loads, - ) + history_provider = FileHistoryProvider(store.get_transcripts_directory(session, source_id=DEFAULT_MEMORY_SOURCE_ID)) await history_provider.save_messages( session.session_id, [ diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py index c4f4061d42..43179dfdd0 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -4,11 +4,14 @@ import json import threading import time -from collections.abc import Awaitable, Callable, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass from pathlib import Path from typing import Any +import msgspec import pytest +from typing_extensions import Self from agent_framework import ( AgentContext, @@ -17,12 +20,15 @@ ContextProvider, ExperimentalFeature, FileHistoryProvider, + FileSessionStore, HistoryProvider, InMemoryHistoryProvider, Message, SessionContext, + SessionStore, agent_middleware, chat_middleware, + register_state_type, ) from agent_framework._sessions import LOCAL_HISTORY_CONVERSATION_ID, is_local_history_conversation_id from agent_framework.exceptions import MiddlewareException @@ -571,6 +577,325 @@ def test_from_dict_missing_state(self) -> None: assert session.state == {} +class _RegisteredSessionState: + def __init__(self, value: str) -> None: + self.value = value + + @classmethod + def _get_type_identifier(cls) -> str: + return "registered_session_state" + + def to_dict(self) -> dict[str, Any]: + return {"type": self._get_type_identifier(), "value": self.value} + + @classmethod + def from_dict(cls, value: dict[str, Any]) -> "_RegisteredSessionState": + return cls(value=str(value["value"])) + + +register_state_type(_RegisteredSessionState) + + +class TestStateTypeRegistry: + def test_same_registration_is_idempotent(self) -> None: + register_state_type(_RegisteredSessionState) + + def test_type_constant_is_used_as_default_identifier(self) -> None: + class TypeConstantState: + TYPE = "type_constant_state_test" + + def __init__(self, value: str) -> None: + self.value = value + + def to_dict(self) -> dict[str, Any]: + return {"value": self.value} + + @classmethod + def from_dict(cls, value: dict[str, Any]) -> Self: + return cls(str(value["value"])) + + register_state_type(TypeConstantState) + session = AgentSession(session_id="type-constant") + session.state["value"] = TypeConstantState("ok") + + restored = AgentSession.from_dict(session.to_dict()) + + assert isinstance(restored.state["value"], TypeConstantState) + assert restored.state["value"].value == "ok" + + def test_custom_encoder_and_decoder_support_plain_classes(self) -> None: + @dataclass + class CallbackState: + value: str + + def encode(value: CallbackState) -> dict[str, Any]: + return {"value": value.value} + + def decode(value: Mapping[str, Any]) -> CallbackState: + return CallbackState(str(value["value"])) + + register_state_type( + CallbackState, + type_id="callback_state_test", + encoder=encode, + decoder=decode, + ) + session = AgentSession(session_id="callback") + session.state["value"] = CallbackState("ok") + + restored = AgentSession.from_dict(session.to_dict()) + + assert isinstance(restored.state["value"], CallbackState) + assert restored.state["value"].value == "ok" + + def test_explicit_pydantic_registration_round_trips(self) -> None: + from pydantic import BaseModel + + class PydanticState(BaseModel): + value: str + + register_state_type(PydanticState, type_id="pydantic_state_test") + session = AgentSession(session_id="pydantic") + session.state["value"] = PydanticState(value="ok") + + restored = AgentSession.from_dict(session.to_dict()) + + assert isinstance(restored.state["value"], PydanticState) + assert restored.state["value"].value == "ok" + + def test_conflicting_type_identifier_is_rejected(self) -> None: + class FirstState: + def to_dict(self) -> dict[str, Any]: + return {} + + @classmethod + def from_dict(cls, value: dict[str, Any]) -> Self: + del value + return cls() + + class SecondState(FirstState): + pass + + register_state_type(FirstState, type_id="collision_state_test") + + with pytest.raises(ValueError, match="already registered"): + register_state_type(SecondState, type_id="collision_state_test") + + def test_unregistered_object_fails_with_state_path(self) -> None: + class UnregisteredState: + pass + + session = AgentSession(session_id="unregistered") + session.state["nested"] = [{"value": UnregisteredState()}] + + with pytest.raises(TypeError, match=r"state\.nested\[0\]\.value.*UnregisteredState"): + session.to_dict() + + def test_unknown_type_tag_remains_a_raw_dict(self) -> None: + session = AgentSession.from_dict({ + "session_id": "unknown", + "state": {"value": {"type": "future_state_type", "nested": {"count": 1}}}, + }) + + assert session.state["value"] == { + "type": "future_state_type", + "nested": {"count": 1}, + } + + def test_registered_child_tag_does_not_hijack_message_payload(self) -> None: + @dataclass + class TextState: + value: str + + register_state_type( + TextState, + type_id="text", + encoder=lambda value: {"value": value.value}, + decoder=lambda value: TextState(str(value["value"])), + ) + session = AgentSession(session_id="message") + session.state["message"] = Message(role="user", contents=["hello"]) + + restored = AgentSession.from_dict(session.to_dict()) + + assert isinstance(restored.state["message"], Message) + assert restored.state["message"].text == "hello" + + @pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf")]) + def test_non_finite_float_is_rejected(self, value: float) -> None: + session = AgentSession(session_id="float") + session.state["value"] = value + + with pytest.raises(ValueError, match="finite float"): + session.to_dict() + + +class TestSessionStore: + def test_is_marked_experimental(self) -> None: + for store_type in (SessionStore, FileSessionStore): + assert store_type.__feature_stage__ == "experimental" # type: ignore[attr-defined, union-attr] # ty: ignore[unresolved-attribute] + assert store_type.__feature_id__ == ExperimentalFeature.SESSION_STORE.value # type: ignore[attr-defined, union-attr] # ty: ignore[unresolved-attribute] + assert store_type.__doc__ is not None + assert ".. warning:: Experimental" in store_type.__doc__ + + async def test_get_returns_none_for_missing_id(self) -> None: + store = SessionStore() + + assert await store.get("session-1") is None + + async def test_set_then_get_returns_independent_copy(self) -> None: + store = SessionStore() + session = AgentSession(session_id="session-1") + session.state["nested"] = {"values": ["original"]} + + await store.set("session-1", session) + + stored = await store.get("session-1") + assert stored is not None + assert stored is not session + stored.state["nested"]["values"].append("changed") + + reread = await store.get("session-1") + assert reread is not None + assert reread.state["nested"]["values"] == ["original"] + + async def test_set_replaces_existing_entry(self) -> None: + store = SessionStore() + await store.set("session-1", AgentSession(session_id="first")) + await store.set("session-1", AgentSession(session_id="second")) + + stored = await store.get("session-1") + + assert stored is not None + assert stored.session_id == "second" + + async def test_delete_forgets_session(self) -> None: + store = SessionStore() + await store.set("session-1", AgentSession(session_id="session-1")) + + await store.delete("session-1") + + assert await store.get("session-1") is None + + @pytest.mark.parametrize( + "session_id", + [ + "", + "two words", + "tenant/user", + "tenant:user", + "session.id", + "' OR 1=1 --", + "café", + "line\nbreak", + "a" * 129, + ], + ) + async def test_invalid_session_id_raises(self, session_id: str) -> None: + store = SessionStore() + session = AgentSession() + + with pytest.raises(ValueError, match="session_id"): + await store.get(session_id) + with pytest.raises(ValueError, match="session_id"): + await store.set(session_id, session) + with pytest.raises(ValueError, match="session_id"): + await store.delete(session_id) + + +class TestFileSessionStore: + async def test_round_trips_session_across_store_instances(self, tmp_path: Path) -> None: + session = AgentSession(session_id="framework-session", service_session_id={"response_id": "resp-1"}) + session.state["nested"] = { + "messages": [Message(role="user", contents=["hello"])], + "custom": [_RegisteredSessionState("persisted")], + } + + await FileSessionStore(tmp_path).set("tenant_user-conversation", session) + restored = await FileSessionStore(tmp_path).get("tenant_user-conversation") + + assert restored is not None + assert restored is not session + assert restored.session_id == "framework-session" + assert restored.service_session_id == {"response_id": "resp-1"} + assert isinstance(restored.state["nested"]["messages"][0], Message) + assert restored.state["nested"]["messages"][0].text == "hello" + assert isinstance(restored.state["nested"]["custom"][0], _RegisteredSessionState) + assert restored.state["nested"]["custom"][0].value == "persisted" + + files = await asyncio.to_thread(lambda: list(tmp_path.iterdir())) + assert len(files) == 1 + assert files[0].parent == tmp_path + assert files[0].suffix == ".json" + + async def test_round_trips_binary_messagepack_session(self, tmp_path: Path) -> None: + store = FileSessionStore(tmp_path, serialization_format="msgpack") + session = AgentSession(session_id="binary-session") + session.state["nested"] = [_RegisteredSessionState("persisted")] + + await store.set("binary-session", session) + restored = await store.get("binary-session") + + assert restored is not None + assert isinstance(restored.state["nested"][0], _RegisteredSessionState) + assert restored.state["nested"][0].value == "persisted" + files = await asyncio.to_thread(lambda: list(tmp_path.iterdir())) + assert len(files) == 1 + assert files[0].suffix == ".msgpack" + assert not (await asyncio.to_thread(files[0].read_bytes)).startswith(b"{") + + async def test_rejects_custom_agent_session_subclass(self, tmp_path: Path) -> None: + class CustomAgentSession(AgentSession): + pass + + store = FileSessionStore(tmp_path) + + with pytest.raises(TypeError, match="AgentSession instances only"): + await store.set("custom-session", CustomAgentSession(session_id="custom-session")) + + async def test_missing_and_deleted_sessions_return_none(self, tmp_path: Path) -> None: + store = FileSessionStore(tmp_path) + assert await store.get("session-1") is None + + await store.set("session-1", AgentSession(session_id="session-1")) + await store.delete("session-1") + + assert await store.get("session-1") is None + + async def test_set_replaces_atomically_without_leaving_temp_files(self, tmp_path: Path) -> None: + store = FileSessionStore(tmp_path) + await store.set("session-1", AgentSession(session_id="first")) + await store.set("session-1", AgentSession(session_id="second")) + + stored = await store.get("session-1") + + assert stored is not None + assert stored.session_id == "second" + temp_files = await asyncio.to_thread(lambda: [*tmp_path.glob("*.tmp"), *tmp_path.glob(".*.tmp")]) + assert not temp_files + + async def test_corrupt_session_file_raises(self, tmp_path: Path) -> None: + store = FileSessionStore(tmp_path) + await store.set("session-1", AgentSession(session_id="session-1")) + session_file = await asyncio.to_thread(lambda: next(tmp_path.iterdir())) + await asyncio.to_thread(session_file.write_text, "{not-json", encoding="utf-8") + + with pytest.raises(ValueError, match="Failed to deserialize session"): + await store.get("session-1") + + @pytest.mark.parametrize("session_id", ["", "two words", "tenant/user", "session.id", "café"]) + async def test_invalid_session_id_raises(self, tmp_path: Path, session_id: str) -> None: + store = FileSessionStore(tmp_path) + session = AgentSession() + + with pytest.raises(ValueError, match="session_id"): + await store.get(session_id) + with pytest.raises(ValueError, match="session_id"): + await store.set(session_id, session) + with pytest.raises(ValueError, match="session_id"): + await store.delete(session_id) + + # --------------------------------------------------------------------------- # InMemoryHistoryProvider tests # --------------------------------------------------------------------------- @@ -677,6 +1002,42 @@ def test_is_marked_experimental(self) -> None: assert FileHistoryProvider.__doc__ is not None assert ".. warning:: Experimental" in FileHistoryProvider.__doc__ + def test_uses_msgspec_json_by_default(self, tmp_path: Path) -> None: + provider = FileHistoryProvider(tmp_path) + + serialized = provider.dumps({"text": "héllo"}) + + assert isinstance(serialized, bytes) + assert provider.loads(serialized) == {"text": "héllo"} + + async def test_stores_and_loads_length_prefixed_msgpack(self, tmp_path: Path) -> None: + provider = FileHistoryProvider(tmp_path, serialization_format="msgpack") + messages = [ + Message(role="user", contents=["hello"]), + Message(role="assistant", contents=["hi there"]), + ] + + await provider.save_messages("binary-history", messages) + loaded = await provider.get_messages("binary-history") + + assert [message.text for message in loaded] == ["hello", "hi there"] + session_file = provider._session_file_path("binary-history") + assert session_file.suffix == ".msgpack" + raw = await asyncio.to_thread(session_file.read_bytes) + first_record_length = int.from_bytes(raw[:4], "big") + assert first_record_length > 0 + assert raw[4 : 4 + first_record_length] == msgspec.msgpack.encode(messages[0].to_dict()) + + def test_msgpack_rejects_custom_json_codecs(self, tmp_path: Path) -> None: + with pytest.raises(ValueError, match="Custom dumps and loads"): + FileHistoryProvider(tmp_path, serialization_format="msgpack", dumps=json.dumps) + + def test_custom_json_codecs_are_deprecated(self, tmp_path: Path) -> None: + with pytest.warns(DeprecationWarning, match=r"dumps.*loads.*deprecated"): + FileHistoryProvider(tmp_path / "dumps", dumps=json.dumps) + with pytest.warns(DeprecationWarning, match=r"dumps.*loads.*deprecated"): + FileHistoryProvider(tmp_path / "loads", loads=json.loads) + async def test_stores_and_loads_messages(self, tmp_path: Path) -> None: from agent_framework import AgentResponse @@ -760,7 +1121,8 @@ def loads(payload: str | bytes) -> object: payload = payload.decode("utf-8") return json.loads(payload) - provider = FileHistoryProvider(tmp_path, dumps=dumps, loads=loads) + with pytest.warns(DeprecationWarning, match=r"dumps.*loads.*deprecated"): + provider = FileHistoryProvider(tmp_path, dumps=dumps, loads=loads) await provider.save_messages("custom-serializer", [Message(role="user", contents=["hello"])]) loaded = await provider.get_messages("custom-serializer") @@ -776,6 +1138,14 @@ async def test_invalid_jsonl_line_raises(self, tmp_path: Path) -> None: with pytest.raises(ValueError, match="Failed to deserialize history line 1"): await provider.get_messages("broken") + async def test_truncated_msgpack_record_raises(self, tmp_path: Path) -> None: + provider = FileHistoryProvider(tmp_path, serialization_format="msgpack") + session_file = provider._session_file_path("broken") + await asyncio.to_thread(session_file.write_bytes, (10).to_bytes(4, "big") + b"short") + + with pytest.raises(ValueError, match="record 1.*truncated"): + await provider.get_messages("broken") + async def test_missing_session_file_returns_empty_messages(self, tmp_path: Path) -> None: provider = FileHistoryProvider(tmp_path) @@ -819,7 +1189,8 @@ async def test_serializer_must_return_single_line_json(self, tmp_path: Path) -> def dumps(payload: object) -> str: return json.dumps(payload, indent=2) - provider = FileHistoryProvider(tmp_path, dumps=dumps) + with pytest.warns(DeprecationWarning, match=r"dumps.*loads.*deprecated"): + provider = FileHistoryProvider(tmp_path, dumps=dumps) with pytest.raises(ValueError, match="single-line JSON"): await provider.save_messages("pretty-json", [Message(role="user", contents=["hello"])]) diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index c0794fd4b0..93c243c540 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -1,3 +1,90 @@ # Foundry Hosting This package provides the integration of Agent Framework agents and workflows with the Foundry Agent Server, which can be hosted on Foundry infrastructure. + +`ResponsesHostServer` persists the Agent Framework `AgentSession` used by regular +agents in addition to the Responses provider's message history. By default it +uses core's experimental msgspec-backed `FileSessionStore` under `$HOME/.checkpoints/sessions` +when hosted and `{cwd}/.checkpoints/sessions` locally. The store is partitioned +by stable agent name, the platform user key, and the Responses conversation +partition. Foundry hashes that composite scope into a restricted-alphabet +`foundry_` session ID before it reaches the store. Pass `session_store=` to +use another `SessionStore` implementation. + +Workflow agents continue to use checkpoint storage; their checkpoints share the +same `$HOME/.checkpoints` or local `.checkpoints` root. + +## Custom authenticated session isolation + +`ResponsesHostServer` accepts an explicit sync or async +`session_isolation_key_resolver(request, context)` callable. The default uses +Foundry's trusted `context.platform_context.user_id_key`. Hosted requests fail +closed when no key is resolved; local requests may remain unscoped. + +File-backed state is partitioned under the validated identity: + +```text +.checkpoints/ + sessions/user-/.json + checkpoints/user-// + function-approvals/user-/approval_requests.json +``` + +The default store is `IsolationKeyScopedFileSessionStore`, a +`FileSessionStore` subclass. One store instance serves the server; each +`get`/`set`/`delete` operation resolves its path beneath the isolation directory +bound for that request. The server does not keep a store instance per user. +If `session_store=` is supplied with file-backed storage, it must already be an +`IsolationKeyScopedFileSessionStore`; an unscoped `FileSessionStore` is rejected. + +The resolver must return a stable, globally unique, portable path-safe identity. +The value is visible in the directory name, so return an application-owned +opaque ID rather than a sensitive claim when necessary. + +Python has no universal DI container or `IHttpContextAccessor`. Applications +that authenticate through ASGI middleware can explicitly bridge their verified +principal to the resolver with an application-owned `ContextVar`: + +```python +from contextvars import ContextVar +from typing import Any + +from agent_framework_foundry_hosting import ResponsesHostServer + + +current_subject: ContextVar[str | None] = ContextVar("current_subject", default=None) + + +class VerifiedSubjectMiddleware: + def __init__(self, app: Any) -> None: + self.app = app + + async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None: + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + # An outer authentication middleware must populate this trusted user. + user = scope.get("user") + subject = getattr(user, "identity", None) if getattr(user, "is_authenticated", False) else None + token = current_subject.set(subject) + try: + await self.app(scope, receive, send) + finally: + current_subject.reset(token) + + +def claims_isolation_key(request: Any, context: Any) -> str | None: + del request, context + return current_subject.get() + + +server = ResponsesHostServer( + agent, + session_isolation_key_resolver=claims_isolation_key, +) +app = VerifiedSubjectMiddleware(server) +``` + +Authentication middleware must run before `VerifiedSubjectMiddleware`. Do not +use an unverified caller-controlled header as the isolation key. diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/__init__.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/__init__.py index 9233f5a763..36e9027de7 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/__init__.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/__init__.py @@ -4,6 +4,7 @@ from ._invocations import InvocationsHostServer from ._responses import ResponsesHostServer +from ._session_isolation import IsolationKeyScopedFileSessionStore, ResponsesSessionIsolationKeyResolver from ._toolbox import FoundryToolbox try: @@ -11,4 +12,10 @@ except importlib.metadata.PackageNotFoundError: __version__ = "0.0.0" -__all__ = ["FoundryToolbox", "InvocationsHostServer", "ResponsesHostServer"] +__all__ = [ + "FoundryToolbox", + "InvocationsHostServer", + "IsolationKeyScopedFileSessionStore", + "ResponsesHostServer", + "ResponsesSessionIsolationKeyResolver", +] diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index 114ff4f005..d80c551e90 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -4,6 +4,7 @@ import asyncio import base64 +import hashlib import json import logging import os @@ -16,13 +17,17 @@ from typing import Literal, Protocol, cast from agent_framework import ( + AgentSession, ChatOptions, Content, ContextProvider, FileCheckpointStorage, + FileSessionStore, HistoryProvider, + InMemoryHistoryProvider, Message, RawAgent, + SessionStore, SupportsAgentRun, WorkflowAgent, ) @@ -115,9 +120,31 @@ from mcp import McpError from typing_extensions import Any +from ._session_isolation import ( + IsolationKeyScopedFileSessionStore, + ResolvedSessionIsolation, + ResponsesSessionIsolationKeyResolver, + platform_session_isolation_key_resolver, + resolve_session_isolation, +) + logger = logging.getLogger(__name__) _AZURE_RESPONSES_MESSAGE_ROLE_TYPE = f"{MessageRole.__module__}:{MessageRole.__qualname__}" +_CURRENT_RESPONSE_ID_BODY_LENGTH = 50 +_CURRENT_RESPONSE_ID_PARTITION_LENGTH = 18 +_LEGACY_RESPONSE_ID_BODY_LENGTH = 48 +_LEGACY_RESPONSE_ID_PARTITION_LENGTH = 16 +_HOSTED_SESSION_DATA_DIRECTORY_ENVIRONMENT_VARIABLE = "HOME" +_DEFAULT_HOSTED_SESSION_DATA_DIRECTORY = "/home/session" +_CHECKPOINT_DIRECTORY_NAME = ".checkpoints" +_SESSION_DIRECTORY_NAME = "sessions" +_WORKFLOW_CHECKPOINT_DIRECTORY_NAME = "checkpoints" +_FUNCTION_APPROVAL_DIRECTORY_NAME = "function-approvals" +_FUNCTION_APPROVAL_FILE_NAME = "approval_requests.json" +_HOSTED_RESPONSES_HISTORY_SOURCE_ID = "_foundry_responses_history" +_SESSION_ISOLATION_STATE_KEY = "foundry_hosting.session_isolation_fingerprint" +_UNSCOPED_SESSION_ISOLATION = ResolvedSessionIsolation(None, None, None) # region Approval Storage @@ -250,30 +277,129 @@ def _validate_path_segment(segment: str, *, kind: Literal["context id", "user id raise RuntimeError(f"Invalid {kind}: {segment!r}") -def _checkpoint_storage_for_context(root: str, context_id: str, *, user_id: str | None = None) -> FileCheckpointStorage: +def _is_usable_hosted_home_directory(home_directory: str | None) -> bool: + """Return whether ``home_directory`` can safely contain hosted durable state.""" + if home_directory is None or not home_directory.strip(): + return False + try: + resolved = Path(home_directory).expanduser().resolve() + except (OSError, RuntimeError): + return False + return resolved != Path(resolved.anchor) + + +def _resolve_durable_storage_root( + *, + is_hosted: bool, + home_directory: str | None, + current_directory: str, +) -> Path: + """Resolve the shared root for Foundry session snapshots and workflow checkpoints.""" + if is_hosted: + if home_directory is not None and _is_usable_hosted_home_directory(home_directory): + base_directory = Path(home_directory).expanduser() + else: + base_directory = Path(_DEFAULT_HOSTED_SESSION_DATA_DIRECTORY) + else: + base_directory = Path(current_directory) + return (base_directory / _CHECKPOINT_DIRECTORY_NAME).resolve() + + +def _response_id_partition(response_id: str | None) -> str | None: + """Extract the stable Responses SDK partition from an item/response ID.""" + if response_id is None or not response_id.strip(): + return None + _, separator, body = response_id.partition("_") + if not separator: + body = response_id + if len(body) == _CURRENT_RESPONSE_ID_BODY_LENGTH: + return body[:_CURRENT_RESPONSE_ID_PARTITION_LENGTH] + if len(body) == _LEGACY_RESPONSE_ID_BODY_LENGTH: + return body[-_LEGACY_RESPONSE_ID_PARTITION_LENGTH:] + return response_id + + +def _resolve_session_conversation_key(request: CreateResponse, context: ResponseContext) -> str: + """Resolve the stable conversation-level key used for AgentSession persistence.""" + conversation_key = ( + context.conversation_id + or _response_id_partition(request.previous_response_id) + or _response_id_partition(context.response_id) + ) + if conversation_key is None: + raise RuntimeError("A Responses session key could not be resolved.") + return conversation_key + + +def _conversation_object_id(conversation_key: str) -> str: + """Return a restricted store/file ID derived from the Responses conversation object.""" + try: + SessionStore.validate_session_id(conversation_key) + except ValueError: + return f"conversation_{hashlib.sha256(conversation_key.encode('utf-8')).hexdigest()}" + return conversation_key + + +def _custom_session_store_key( + agent: SupportsAgentRun, + object_id: str, + *, + isolation: ResolvedSessionIsolation, +) -> str: + """Create an isolation-scoped key for a non-file custom store.""" + key: dict[str, str] = {"object": object_id} + if agent.name: + key["agent"] = agent.name + if isolation.fingerprint: + key["isolation"] = isolation.fingerprint + canonical_key = json.dumps(key, ensure_ascii=False, separators=(",", ":"), sort_keys=True) + return f"foundry_{hashlib.sha256(canonical_key.encode('utf-8')).hexdigest()}" + + +def _stamp_or_validate_session_isolation( + session: AgentSession, + isolation: ResolvedSessionIsolation, +) -> None: + """Stamp a legacy/new session or reject a cross-isolation restore.""" + existing = session.state.get(_SESSION_ISOLATION_STATE_KEY) + expected = isolation.fingerprint + if existing is None: + if expected is not None: + session.state[_SESSION_ISOLATION_STATE_KEY] = expected + return + if existing != expected: + raise RuntimeError("Hosted session identity context mismatch.") + + +def _is_hosted_responses_history_sentinel(provider: ContextProvider) -> bool: + """Return whether ``provider`` is the host's no-op history-loading sentinel.""" + return ( + isinstance(provider, InMemoryHistoryProvider) + and provider.source_id == _HOSTED_RESPONSES_HISTORY_SOURCE_ID + and not provider.store_inputs + and not provider.store_context_messages + and not provider.store_outputs + ) + + +def _checkpoint_storage_for_context( + root: str | Path, + context_id: str, + *, + isolation: ResolvedSessionIsolation = _UNSCOPED_SESSION_ISOLATION, +) -> FileCheckpointStorage: """Build a ``FileCheckpointStorage`` for ``context_id`` rooted under ``root``. - When the platform supplies a per-user partition key (``user_id``, from the - ``x-agent-user-id`` header on container protocol v2), the per-conversation - checkpoint directory is nested under it: ``//``. - This isolates each tenant's workflow state so one user can never restore or - observe another user's checkpoint, even with a guessed or forged - ``context_id``. An absent (``None``) or empty ``user_id`` -- local - development or protocol v1 -- falls back to the unscoped - ``/`` layout. - - Both ``context_id`` and ``user_id`` are validated as single safe path - segments, and each resolved directory is verified to stay under its parent - before any directory is created on disk (CWE-22). + Scoped requests use ``/user-/``. Local + unscoped requests use ``/``. """ _validate_path_segment(context_id, kind="context id") base_path = Path(root).resolve() - if user_id: - _validate_path_segment(user_id, kind="user id") - user_path = (base_path / user_id).resolve() + if isolation.directory_segment: + user_path = (base_path / isolation.directory_segment).resolve() if not user_path.is_relative_to(base_path): - raise RuntimeError(f"Invalid user id: {user_id!r}") + raise RuntimeError(f"Invalid session isolation directory: {isolation.directory_segment!r}") base_path = user_path storage_path = (base_path / context_id).resolve() @@ -287,23 +413,17 @@ def _checkpoint_storage_for_context(root: str, context_id: str, *, user_id: str ) -def _approval_storage_path_for_user(base_path: str, user_id: str) -> str: - """Return the per-user approval storage file path under the base directory. - - Inserts the validated ``user_id`` as a directory segment between the base - directory and the file name (``//``), mirroring the - per-user checkpoint partitioning so one tenant can never read another - tenant's saved approval requests. The user id is validated as a single safe - path segment and the resulting directory is verified to stay under the base - directory before use (CWE-22). - """ - _validate_path_segment(user_id, kind="user id") - directory, filename = os.path.split(base_path) - base_dir = Path(directory or ".").resolve() - user_dir = (base_dir / user_id).resolve() - if not user_dir.is_relative_to(base_dir): - raise RuntimeError(f"Invalid user id: {user_id!r}") - return str(user_dir / filename) +def _approval_storage_path( + root: str | Path, + isolation: ResolvedSessionIsolation, +) -> str: + """Return the approval file path for one resolved isolation identity.""" + base_path = Path(root).resolve() + if isolation.directory_segment: + base_path = (base_path / isolation.directory_segment).resolve() + if not base_path.is_relative_to(Path(root).resolve()): + raise RuntimeError(f"Invalid session isolation directory: {isolation.directory_segment!r}") + return str(base_path / _FUNCTION_APPROVAL_FILE_NAME) # endregion Approval Storage @@ -387,10 +507,6 @@ def consent_url_from_error(exc: BaseException) -> list[ConsentError] | None: class ResponsesHostServer(ResponsesAgentServerHost): """A responses server host for an agent.""" - # TODO(@taochen): Allow a different checkpoint storage that stores checkpoints externally - CHECKPOINT_STORAGE_PATH = "/.checkpoints" - FUNCTION_APPROVAL_STORAGE_PATH = "/.function_approvals/approval_requests.json" - def __init__( self, agent: SupportsAgentRun, @@ -398,6 +514,8 @@ def __init__( prefix: str = "", options: ResponsesServerOptions | None = None, store: ResponseProviderProtocol | None = None, + session_store: SessionStore | None = None, + session_isolation_key_resolver: ResponsesSessionIsolationKeyResolver | None = None, **kwargs: Any, ) -> None: """Initialize a ResponsesHostServer. @@ -407,6 +525,13 @@ def __init__( prefix: The URL prefix for the server. options: Optional server options. store: Optional response store. + session_store: Optional Agent Framework session store. Defaults to + an :class:`IsolationKeyScopedFileSessionStore` under the Foundry + durable storage root. A caller-supplied file store must already + be an ``IsolationKeyScopedFileSessionStore``. + session_isolation_key_resolver: Optional sync/async callable that + resolves a stable authenticated identity for the request. + Defaults to the trusted Foundry platform user key. **kwargs: Additional keyword arguments. Note: @@ -420,6 +545,8 @@ def __init__( for provider in getattr(agent, "context_providers", []): if isinstance(provider, HistoryProvider) and provider.load_messages: + if _is_hosted_responses_history_sentinel(provider): + continue raise RuntimeError( "There shouldn't be a history provider with `load_messages=True` already present. " "History is managed by the hosting infrastructure." @@ -431,6 +558,12 @@ def __init__( provider.source_id, ) + self._storage_root = _resolve_durable_storage_root( + is_hosted=self.config.is_hosted, + home_directory=os.getenv(_HOSTED_SESSION_DATA_DIRECTORY_ENVIRONMENT_VARIABLE), + current_directory=os.getcwd(), + ) + self._session_isolation_key_resolver = session_isolation_key_resolver or platform_session_isolation_key_resolver self._is_workflow_agent = False self._checkpoint_storage_path = None if isinstance(agent, WorkflowAgent): @@ -439,27 +572,44 @@ def __init__( "There should not be a checkpoint storage already present in the workflow agent. " "The hosting infrastructure will manage checkpoints instead." ) - self._checkpoint_storage_path = ( - self.CHECKPOINT_STORAGE_PATH - if self.config.is_hosted - else os.path.join(os.getcwd(), self.CHECKPOINT_STORAGE_PATH.lstrip("/")) - ) + self._checkpoint_storage_path = str(self._storage_root / _WORKFLOW_CHECKPOINT_DIRECTORY_NAME) self._is_workflow_agent = True - self._agent = agent - self._approval_storage = ( - FileBasedFunctionApprovalStorage(self.FUNCTION_APPROVAL_STORAGE_PATH) - if self.config.is_hosted - else InMemoryFunctionApprovalStorage() - ) - # Per-user (multi-tenant) approval stores. Hosted file-based approval - # storage is partitioned by the platform per-user partition key so one - # tenant can never read another tenant's saved approval requests. - # Instances are cached so concurrent requests for the same user share one - # lock, preserving serialized read-modify-write on the JSON file. Local - # (in-memory) dev and protocol v1 (no user id) keep the single shared - # ``self._approval_storage``. - self._approval_storages_by_user: dict[str, ApprovalStorage] = {} + if ( + not self._is_workflow_agent + and isinstance(agent, RawAgent) + and not any( + _is_hosted_responses_history_sentinel(provider) + for provider in cast(Sequence[ContextProvider], agent.context_providers) + ) + ): + # The Responses provider already supplies the complete transcript on every + # call. A loading no-op provider prevents Agent from auto-injecting its + # default InMemoryHistoryProvider and replaying that transcript twice. + agent.context_providers.append( + InMemoryHistoryProvider( + source_id=_HOSTED_RESPONSES_HISTORY_SOURCE_ID, + store_inputs=False, + store_outputs=False, + ) + ) + + self._agent: SupportsAgentRun = agent + if not self._is_workflow_agent and session_store is None: + session_store = IsolationKeyScopedFileSessionStore(self._storage_root / _SESSION_DIRECTORY_NAME) + elif isinstance(session_store, FileSessionStore) and not isinstance( + session_store, + IsolationKeyScopedFileSessionStore, + ): + raise ValueError( + "ResponsesHostServer requires IsolationKeyScopedFileSessionStore for file-backed session storage." + ) + self._session_store = session_store + self._approval_storage: ApprovalStorage = InMemoryFunctionApprovalStorage() + self._approval_storage_root = self._storage_root / _FUNCTION_APPROVAL_DIRECTORY_NAME + # Per-user hosted approval stores share a process-local lock per identity. + # Local development keeps the single in-memory store. + self._approval_storages_by_isolation: dict[str, ApprovalStorage] = {} # Lazy agent lifecycle: the agent (and any MCP tools it owns) is entered on # the first request rather than at server startup, so that authentication # failures during MCP connect can be surfaced to the client as an @@ -484,7 +634,7 @@ async def _ensure_agent_ready(self) -> None: stack = AsyncExitStack() try: if isinstance(self._agent, AbstractAsyncContextManager): - await stack.enter_async_context(self._agent) + await stack.enter_async_context(cast(AbstractAsyncContextManager[Any], self._agent)) except BaseException: await stack.aclose() raise @@ -497,27 +647,14 @@ async def _cleanup_agent(self) -> None: self._agent_stack = None await stack.aclose() - def _approval_storage_for_user(self, user_id: str | None) -> ApprovalStorage: - """Return the approval storage scoped to ``user_id`` when applicable. - - For hosted multi-tenant deployments the file-based store is partitioned - by the platform per-user partition key, so one tenant can never read - another tenant's saved approval requests. Falls back to the single shared - store for local (in-memory) hosting or when no per-user partition key is - available (protocol v1 / local development). Instances are cached so - concurrent requests for the same user share one lock. - - Raises: - RuntimeError: If ``user_id`` is not a safe single path segment. - """ - if not self.config.is_hosted or not user_id: + def _approval_storage_for_isolation(self, isolation: ResolvedSessionIsolation) -> ApprovalStorage: + """Return the hosted approval store for one validated identity.""" + if not self.config.is_hosted or isolation.key is None: return self._approval_storage - storage = self._approval_storages_by_user.get(user_id) + storage = self._approval_storages_by_isolation.get(isolation.key) if storage is None: - storage = FileBasedFunctionApprovalStorage( - _approval_storage_path_for_user(self.FUNCTION_APPROVAL_STORAGE_PATH, user_id) - ) - self._approval_storages_by_user[user_id] = storage + storage = FileBasedFunctionApprovalStorage(_approval_storage_path(self._approval_storage_root, isolation)) + self._approval_storages_by_isolation[isolation.key] = storage return storage async def _handle_response( @@ -534,16 +671,23 @@ async def _handle_response( "Please upgrade your agent protocol to 2.0.0 in `agent.manifest.yaml` or `agent.yaml`, or " "downgrade the `agent-framework-foundry-hosting` package to `1.0.0a260625` or before to use 1.0.0." ) + isolation = await resolve_session_isolation( + self._session_isolation_key_resolver, + request, + context, + is_hosted=self.config.is_hosted, + ) if self._is_workflow_agent: # Workflow agents are handled differently because they require checkpoint restoration - return self._handle_inner_workflow(request, context) - return self._handle_inner_agent(request, context) + return self._handle_inner_workflow(request, context, isolation) + return self._handle_inner_agent(request, context, isolation) async def _handle_inner_agent( self, request: CreateResponse, context: ResponseContext, + isolation: ResolvedSessionIsolation = _UNSCOPED_SESSION_ISOLATION, ) -> AsyncIterable[ResponseStreamEvent | dict[str, Any]]: """Handle the creation of a response for a regular (non-workflow) agent.""" response_event_stream = ResponseEventStream(response_id=context.response_id, model=request.model) @@ -553,10 +697,44 @@ async def _handle_inner_agent( # Track the current active output item builder for streaming; # lazily created on matching content, closed when a different type arrives. tracker: _OutputItemTracker | None = None + session: AgentSession | None = None + request_session_store: SessionStore | None = None + session_store_key: str | None = None + session_save_attempted = False + + async def save_session() -> None: + nonlocal session_save_attempted + session_save_attempted = True + if request_session_store is None: + raise RuntimeError("Session storage is not configured for a regular agent.") + if session is not None and session_store_key is not None: + if isinstance(request_session_store, IsolationKeyScopedFileSessionStore): + with request_session_store.use_isolation(isolation): + await request_session_store.set(session_store_key, session) + else: + await request_session_store.set(session_store_key, session) try: - user_id = context.platform_context.user_id_key - approval_storage = self._approval_storage_for_user(user_id) + approval_storage = self._approval_storage_for_isolation(isolation) + conversation_key = _resolve_session_conversation_key(request, context) + object_id = _conversation_object_id(conversation_key) + request_session_store = self._session_store + if request_session_store is None: + raise RuntimeError("Session storage is not configured for a regular agent.") + session_store_key = ( + object_id + if isinstance(request_session_store, IsolationKeyScopedFileSessionStore) + else _custom_session_store_key(self._agent, object_id, isolation=isolation) + ) + if isinstance(request_session_store, IsolationKeyScopedFileSessionStore): + with request_session_store.use_isolation(isolation): + session = await request_session_store.get(session_store_key) + else: + session = await request_session_store.get(session_store_key) + if session is None: + session = self._agent.create_session(session_id=object_id) + _stamp_or_validate_session_isolation(session, isolation) + input_items = await context.get_input_items() input_messages = await _items_to_messages(input_items, approval_storage=approval_storage) @@ -565,7 +743,8 @@ async def _handle_inner_agent( "messages": [ *(await _output_items_to_messages(history, approval_storage=approval_storage)), *input_messages, - ] + ], + "session": session, } is_streaming_request = request.stream is not None and request.stream is True @@ -597,6 +776,7 @@ async def _handle_inner_agent( builder = response_event_stream.add_output_item(oauth_item.id) yield builder.emit_added(oauth_item) yield builder.emit_done(oauth_item) + await save_session() yield response_event_stream.emit_completed() return @@ -632,16 +812,26 @@ async def _handle_inner_agent( # Close any remaining active builder for event in tracker.close(): yield event + await save_session() yield response_event_stream.emit_completed() except Exception as ex: logger.exception("Failed to produce response for agent") + if not session_save_attempted: + try: + await save_session() + except Exception: + logger.exception("Failed to persist the Agent Framework session after an agent failure") for event in self._emit_failure(response_event_stream, tracker, ex): yield event + finally: + if session is not None and not session_save_attempted: + await save_session() async def _handle_inner_workflow( self, request: CreateResponse, context: ResponseContext, + isolation: ResolvedSessionIsolation = _UNSCOPED_SESSION_ISOLATION, ) -> AsyncIterable[ResponseStreamEvent | dict[str, Any]]: """Handle the creation of a response for a workflow agent.""" response_event_stream = ResponseEventStream(response_id=context.response_id, model=request.model) @@ -653,8 +843,7 @@ async def _handle_inner_workflow( tracker: _OutputItemTracker | None = None try: - user_id = context.platform_context.user_id_key - approval_storage = self._approval_storage_for_user(user_id) + approval_storage = self._approval_storage_for_isolation(isolation) input_items = await context.get_input_items() input_messages = await _items_to_messages(input_items, approval_storage=approval_storage) is_streaming_request = request.stream is not None and request.stream is True @@ -702,7 +891,9 @@ async def _handle_inner_workflow( restore_storage: FileCheckpointStorage | None = None if context_id is not None: restore_storage = _checkpoint_storage_for_context( - self._checkpoint_storage_path, context_id, user_id=user_id + self._checkpoint_storage_path, + context_id, + isolation=isolation, ) latest_checkpoint = await restore_storage.get_latest(workflow_name=self._agent.workflow.name) if latest_checkpoint is not None: @@ -718,7 +909,9 @@ async def _handle_inner_workflow( # directory and write_storage points at the *current* response's. write_context_id = context.conversation_id or context.response_id write_storage = _checkpoint_storage_for_context( - self._checkpoint_storage_path, write_context_id, user_id=user_id + self._checkpoint_storage_path, + write_context_id, + isolation=isolation, ) # Multi-turn pattern: when we have a prior checkpoint, restore it diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_isolation.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_isolation.py new file mode 100644 index 0000000000..8efb57afaa --- /dev/null +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_isolation.py @@ -0,0 +1,118 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import hashlib +import inspect +from collections.abc import Awaitable, Callable, Generator +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from pathlib import Path +from typing import Literal, TypeAlias + +from agent_framework import FileSessionStore, SessionStore +from azure.ai.agentserver.responses import ResponseContext +from azure.ai.agentserver.responses.models import CreateResponse + +ResponsesSessionIsolationKeyResolver: TypeAlias = Callable[ + [CreateResponse, ResponseContext], + str | Awaitable[str | None] | None, +] +"""Resolve a stable authenticated identity used to partition Responses state.""" + +_MAX_ISOLATION_KEY_LENGTH = 128 +_INVALID_PATH_CHARACTERS = frozenset('<>:"/\\|?*') + + +@dataclass(frozen=True, slots=True) +class ResolvedSessionIsolation: + """Validated per-request storage isolation.""" + + key: str | None + directory_segment: str | None + fingerprint: str | None + + +class IsolationKeyScopedFileSessionStore(FileSessionStore): + """FileSessionStore that scopes each operation to the bound user directory.""" + + def __init__( + self, + storage_path: str | Path, + *, + serialization_format: Literal["json", "msgpack"] = "json", + ) -> None: + """Initialize a scoped file store rooted at ``storage_path``.""" + super().__init__(storage_path, serialization_format=serialization_format) + self._current_directory: ContextVar[str | None] = ContextVar( + f"foundry_session_isolation_{id(self)}", + default=None, + ) + + @contextmanager + def use_isolation(self, isolation: ResolvedSessionIsolation) -> Generator[None]: + """Bind one resolved request identity for subsequent store operations.""" + token = self._current_directory.set(isolation.directory_segment) + try: + yield + finally: + self._current_directory.reset(token) + + def _session_file_path(self, session_id: str) -> Path: + """Resolve the session file beneath the currently bound user directory.""" + SessionStore.validate_session_id(session_id) + storage_root = self.storage_path.resolve() + directory_segment = self._current_directory.get() + if directory_segment is not None: + storage_root = (storage_root / directory_segment).resolve() + if not storage_root.is_relative_to(self.storage_path.resolve()): + raise ValueError(f"Session isolation path escaped storage directory: {directory_segment!r}") + storage_root.mkdir(parents=True, exist_ok=True) + return storage_root / f"{session_id}{self._file_extension}" + + +def platform_session_isolation_key_resolver( + request: CreateResponse, + context: ResponseContext, +) -> str | None: + """Resolve the trusted Foundry platform user partition key.""" + del request + return context.platform_context.user_id_key + + +def _validate_isolation_key(key: str) -> None: + """Validate a portable single-directory identity value.""" + if len(key) > _MAX_ISOLATION_KEY_LENGTH: + raise ValueError(f"Session isolation key must be at most {_MAX_ISOLATION_KEY_LENGTH} characters.") + if key in {".", ".."} or key.endswith((" ", ".")): + raise ValueError("Session isolation key is not a safe directory name.") + if any(ord(character) < 32 or character in _INVALID_PATH_CHARACTERS for character in key): + raise ValueError("Session isolation key contains characters that are unsafe in a directory name.") + + +async def resolve_session_isolation( + resolver: ResponsesSessionIsolationKeyResolver, + request: CreateResponse, + context: ResponseContext, + *, + is_hosted: bool, +) -> ResolvedSessionIsolation: + """Resolve, normalize, validate, and fingerprint one request identity.""" + value = resolver(request, context) + if inspect.isawaitable(value): + value = await value + if value is None or not value.strip(): + if is_hosted: + raise RuntimeError( + "The hosted request did not resolve a session isolation key; refusing to use unscoped session storage." + ) + return ResolvedSessionIsolation(key=None, directory_segment=None, fingerprint=None) + + key = value.strip() + _validate_isolation_key(key) + return ResolvedSessionIsolation( + key=key, + directory_segment=f"user-{key}", + fingerprint=hashlib.sha256(key.encode("utf-8")).hexdigest(), + ) diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 321ac37308..c05795b7af 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -10,27 +10,35 @@ from __future__ import annotations +import asyncio import json import uuid -from collections.abc import AsyncIterator, Awaitable, Callable, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass +from pathlib import Path from typing import Literal, overload from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest from agent_framework import ( + Agent, AgentExecutorRequest, AgentResponse, AgentResponseUpdate, AgentSession, + BaseChatClient, + ChatResponse, Content, FileCheckpointStorage, + FileSessionStore, HistoryProvider, + InMemoryHistoryProvider, Message, RawAgent, ResponseStream, ServiceSessionId, + SessionStore, SupportsAgentRun, WorkflowAgent, WorkflowBuilder, @@ -40,22 +48,38 @@ WorkflowMessage, executor, ) -from azure.ai.agentserver.responses import InMemoryResponseProvider +from azure.ai.agentserver.responses import InMemoryResponseProvider, PlatformContext, ResponseContext +from azure.ai.agentserver.responses.models import CreateResponse from mcp import McpError from mcp.types import ErrorData from typing_extensions import Any -from agent_framework_foundry_hosting import ResponsesHostServer +from agent_framework_foundry_hosting import ( + IsolationKeyScopedFileSessionStore, + ResponsesHostServer, + ResponsesSessionIsolationKeyResolver, +) from agent_framework_foundry_hosting._responses import ( _AZURE_RESPONSES_MESSAGE_ROLE_TYPE, # pyright: ignore[reportPrivateUsage] + _SESSION_ISOLATION_STATE_KEY, # pyright: ignore[reportPrivateUsage] CONSENT_ERROR_CODE, ConsentError, FileBasedFunctionApprovalStorage, # pyright: ignore[reportPrivateUsage] InMemoryFunctionApprovalStorage, # pyright: ignore[reportPrivateUsage] + _conversation_object_id, # pyright: ignore[reportPrivateUsage] + _custom_session_store_key, # pyright: ignore[reportPrivateUsage] _item_to_message, # pyright: ignore[reportPrivateUsage] _output_item_to_message, # pyright: ignore[reportPrivateUsage] + _resolve_durable_storage_root, # pyright: ignore[reportPrivateUsage] + _resolve_session_conversation_key, # pyright: ignore[reportPrivateUsage] + _response_id_partition, # pyright: ignore[reportPrivateUsage] + _stamp_or_validate_session_isolation, # pyright: ignore[reportPrivateUsage] consent_url_from_error, ) +from agent_framework_foundry_hosting._session_isolation import ( # pyright: ignore[reportPrivateUsage] + ResolvedSessionIsolation, + resolve_session_isolation, +) def _make_function_approval_request_content( @@ -88,6 +112,7 @@ def _make_agent( agent.name = "Test Agent" agent.description = "A mock agent for testing" agent.context_providers = [] + agent.create_session.side_effect = lambda *, session_id=None: AgentSession(session_id=session_id) if response is not None: @@ -112,8 +137,33 @@ def run_streaming(*args: Any, **kwargs: Any) -> Any: return agent +class _RecordingHistoryClient(BaseChatClient): + def __init__(self) -> None: + super().__init__() + self.calls: list[list[Message]] = [] + + def _inner_get_response( + self, + *, + messages: Sequence[Message], + stream: bool, + options: Mapping[str, Any], + **kwargs: Any, + ) -> Awaitable[ChatResponse]: + del options, kwargs + if stream: + raise NotImplementedError("This test client only supports non-streaming responses.") + self.calls.append(list(messages)) + + async def get_response() -> ChatResponse: + return ChatResponse(messages=[Message(role="assistant", contents=[Content.from_text("recorded")])]) + + return get_response() + + def _make_server(agent: Any, **kwargs: Any) -> ResponsesHostServer: """Create a ResponsesHostServer with an in-memory store.""" + kwargs.setdefault("session_store", SessionStore()) return ResponsesHostServer(agent, store=InMemoryResponseProvider(), **kwargs) @@ -127,6 +177,8 @@ async def _post( top_p: float | None = None, max_output_tokens: int | None = None, parallel_tool_calls: bool | None = None, + previous_response_id: str | None = None, + conversation_id: str | None = None, ) -> httpx.Response: """Send a POST /responses request through the ASGI transport.""" payload: dict[str, Any] = {"model": model, "input": input_text, "stream": stream} @@ -138,6 +190,10 @@ async def _post( payload["max_output_tokens"] = max_output_tokens if parallel_tool_calls is not None: payload["parallel_tool_calls"] = parallel_tool_calls + if previous_response_id is not None: + payload["previous_response_id"] = previous_response_id + if conversation_id is not None: + payload["conversation"] = conversation_id transport = httpx.ASGITransport(app=server) async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: @@ -186,6 +242,12 @@ def test_init_basic(self) -> None: ) server = _make_server(agent) assert server is not None + assert len(agent.context_providers) == 1 + history_sentinel = agent.context_providers[0] + assert isinstance(history_sentinel, InMemoryHistoryProvider) + assert history_sentinel.load_messages is True + assert history_sentinel.store_inputs is False + assert history_sentinel.store_outputs is False def test_init_rejects_history_provider_with_load_messages(self) -> None: @@ -213,6 +275,443 @@ async def save_messages( with pytest.raises(RuntimeError, match="history provider"): ResponsesHostServer(agent) + def test_init_rejects_unscoped_file_session_store(self, tmp_path: Path) -> None: + agent = _make_agent( + response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("hi")])]) + ) + + with pytest.raises(ValueError, match="IsolationKeyScopedFileSessionStore"): + ResponsesHostServer( + agent, + store=InMemoryResponseProvider(), + session_store=FileSessionStore(tmp_path), + ) + + async def test_hosted_request_requires_user_partition_key(self) -> None: + agent = _make_agent( + response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("hi")])]) + ) + server = _make_server(agent) + request = CreateResponse(model="m", input="hi") + context = ResponseContext( + response_id="caresp_aaaaaaaaaaaaaaaa00" + "1" * 32, + mode_flags=MagicMock(), + platform_context=PlatformContext(call_id="call-1"), + ) + + with ( + patch.object(server.config, "is_hosted", True), + pytest.raises(RuntimeError, match="did not resolve"), + ): + await server._handle_response( # pyright: ignore[reportPrivateUsage] + request, + context, + asyncio.Event(), + ) + + async def test_custom_isolation_resolver_is_called_once_per_request(self) -> None: + calls = 0 + + async def resolver(request: CreateResponse, context: ResponseContext) -> str: + nonlocal calls + calls += 1 + assert request.model == "m" + assert context.response_id == "response-1" + return "user-1" + + server = _make_server(_make_agent(), session_isolation_key_resolver=resolver) + request = CreateResponse(model="m", input="hi") + context = ResponseContext(response_id="response-1", mode_flags=MagicMock()) + expected = MagicMock() + + with patch.object(server, "_handle_inner_agent", return_value=expected) as handler: + result = await server._handle_response( # pyright: ignore[reportPrivateUsage] + request, + context, + asyncio.Event(), + ) + + assert result is expected + assert calls == 1 + isolation = handler.call_args.args[2] + assert isolation.directory_segment == "user-user-1" + + +# endregion + + +# region Session persistence + + +class TestSessionPersistenceHelpers: + def test_resolver_type_alias_is_public(self) -> None: + def resolver(request: CreateResponse, context: ResponseContext) -> str | None: + del request + return context.platform_context.user_id_key + + typed_resolver: ResponsesSessionIsolationKeyResolver = resolver + assert callable(typed_resolver) + + async def test_sync_and_async_isolation_resolvers(self) -> None: + request = CreateResponse(model="m", input="hi") + context = ResponseContext(response_id="response-1", mode_flags=MagicMock()) + + sync_isolation = await resolve_session_isolation( + lambda request, context: "user-1", + request, + context, + is_hosted=True, + ) + + async def async_resolver(request: CreateResponse, context: ResponseContext) -> str: + del request, context + return "user-2" + + async_isolation = await resolve_session_isolation( + async_resolver, + request, + context, + is_hosted=True, + ) + + assert sync_isolation.directory_segment == "user-user-1" + assert async_isolation.directory_segment == "user-user-2" + assert sync_isolation.fingerprint != async_isolation.fingerprint + + async def test_local_missing_isolation_key_is_unscoped(self) -> None: + request = CreateResponse(model="m", input="hi") + context = ResponseContext(response_id="response-1", mode_flags=MagicMock()) + + isolation = await resolve_session_isolation( + lambda request, context: None, + request, + context, + is_hosted=False, + ) + + assert isolation == ResolvedSessionIsolation(None, None, None) + + async def test_contextvar_backed_resolvers_do_not_cross_contaminate(self) -> None: + from contextvars import ContextVar + + current_user: ContextVar[str | None] = ContextVar("current_user", default=None) + request = CreateResponse(model="m", input="hi") + context = ResponseContext(response_id="response-1", mode_flags=MagicMock()) + + def resolver(request: CreateResponse, context: ResponseContext) -> str | None: + del request, context + return current_user.get() + + async def resolve_for(user_id: str) -> ResolvedSessionIsolation: + token = current_user.set(user_id) + try: + await asyncio.sleep(0) + return await resolve_session_isolation(resolver, request, context, is_hosted=True) + finally: + current_user.reset(token) + + first, second = await asyncio.gather(resolve_for("user-1"), resolve_for("user-2")) + + assert first.key == "user-1" + assert second.key == "user-2" + assert first.fingerprint != second.fingerprint + + @pytest.mark.parametrize("key", ["../escape", "user/name", "user:name", "trailing."]) + async def test_invalid_isolation_key_is_rejected(self, key: str) -> None: + request = CreateResponse(model="m", input="hi") + context = ResponseContext(response_id="response-1", mode_flags=MagicMock()) + + with pytest.raises(ValueError, match="isolation key"): + await resolve_session_isolation( + lambda request, context: key, + request, + context, + is_hosted=True, + ) + + def test_response_id_partition_supports_current_legacy_and_raw_ids(self) -> None: + current_partition = "a" * 18 + legacy_partition = "b" * 16 + + assert _response_id_partition(f"caresp_{current_partition}{'1' * 32}") == current_partition + assert _response_id_partition(f"caresp_{'2' * 32}{legacy_partition}") == legacy_partition + assert _response_id_partition("custom-response") == "custom-response" + assert _response_id_partition(None) is None + + def test_session_conversation_key_prefers_conversation_then_previous_then_response(self) -> None: + previous_partition = "a" * 18 + response_partition = "b" * 18 + previous_response_id = f"caresp_{previous_partition}{'1' * 32}" + response_id = f"caresp_{response_partition}{'2' * 32}" + request = CreateResponse(model="m", input="hi", previous_response_id=previous_response_id) + + assert ( + _resolve_session_conversation_key( + request, + ResponseContext( + response_id=response_id, + previous_response_id=previous_response_id, + conversation_id="conversation-1", + mode_flags=MagicMock(), + ), + ) + == "conversation-1" + ) + assert ( + _resolve_session_conversation_key( + request, + ResponseContext( + response_id=response_id, + previous_response_id=previous_response_id, + mode_flags=MagicMock(), + ), + ) + == previous_partition + ) + assert ( + _resolve_session_conversation_key( + CreateResponse(model="m", input="hi"), + ResponseContext(response_id=response_id, mode_flags=MagicMock()), + ) + == response_partition + ) + + def test_custom_store_key_partitions_by_agent_and_isolation(self) -> None: + agent = _make_agent() + other_agent = _make_agent() + other_agent.name = "Other Agent" + isolation_a = ResolvedSessionIsolation("user-a", "user-user-a", "a" * 64) + isolation_b = ResolvedSessionIsolation("user-b", "user-user-b", "b" * 64) + + key_a = _custom_session_store_key(agent, "conversation-1", isolation=isolation_a) + key_b = _custom_session_store_key(agent, "conversation-1", isolation=isolation_b) + key_other_agent = _custom_session_store_key(other_agent, "conversation-1", isolation=isolation_a) + + assert key_a != key_b + assert key_a != key_other_agent + assert key_a.startswith("foundry_") + assert len(key_a) == len("foundry_") + 64 + assert all(character.isascii() and (character.isalnum() or character in "-_") for character in key_a) + assert key_a == _custom_session_store_key(agent, "conversation-1", isolation=isolation_a) + + def test_conversation_object_id_preserves_safe_values_and_hashes_unsafe_values(self) -> None: + assert _conversation_object_id("conversation-1") == "conversation-1" + assert _conversation_object_id("conversation/unsafe").startswith("conversation_") + + def test_durable_storage_root_matches_hosted_and_local_layouts(self, tmp_path: Path) -> None: + home = tmp_path / "home" + current = tmp_path / "work" + + assert ( + _resolve_durable_storage_root( + is_hosted=False, + home_directory=str(home), + current_directory=str(current), + ) + == (current / ".checkpoints").resolve() + ) + assert ( + _resolve_durable_storage_root( + is_hosted=True, + home_directory=str(home), + current_directory=str(current), + ) + == (home / ".checkpoints").resolve() + ) + assert ( + _resolve_durable_storage_root( + is_hosted=True, + home_directory="/", + current_directory=str(current), + ) + == Path("/home/session/.checkpoints").resolve() + ) + + +class TestAgentSessionPersistence: + async def test_file_store_uses_per_user_child_directory_and_preserves_format(self, tmp_path: Path) -> None: + template = IsolationKeyScopedFileSessionStore(tmp_path, serialization_format="msgpack") + server = _make_server(_make_agent(), session_store=template) + isolation_a = ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64) + isolation_b = ResolvedSessionIsolation("user-B", "user-user-B", "b" * 64) + + store = server._session_store # pyright: ignore[reportPrivateUsage] + assert isinstance(store, IsolationKeyScopedFileSessionStore) + assert store.storage_path == tmp_path + assert store.serialization_format == "msgpack" + + with store.use_isolation(isolation_a): + await store.set("conversation-1", AgentSession(session_id="conversation-1")) + with store.use_isolation(isolation_b): + await store.set("conversation-1", AgentSession(session_id="conversation-1")) + + assert (tmp_path / "user-user-A" / "conversation-1.msgpack").is_file() + assert (tmp_path / "user-user-B" / "conversation-1.msgpack").is_file() + + def test_session_isolation_stamp_rejects_mismatch(self) -> None: + session = AgentSession(session_id="conversation-1") + isolation_a = ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64) + isolation_b = ResolvedSessionIsolation("user-B", "user-user-B", "b" * 64) + + _stamp_or_validate_session_isolation(session, isolation_a) + assert session.state[_SESSION_ISOLATION_STATE_KEY] == "a" * 64 + assert "user-A" not in session.state.values() + + with pytest.raises(RuntimeError, match="identity context mismatch"): + _stamp_or_validate_session_isolation(session, isolation_b) + + async def test_previous_response_chain_restores_session_state(self) -> None: + seen_counts: list[int] = [] + seen_session_ids: list[str] = [] + + async def run_with_state(*args: Any, **kwargs: Any) -> AgentResponse: + session = kwargs["session"] + assert isinstance(session, AgentSession) + count = int(session.state.get("turn_count", 0)) + 1 + session.state["turn_count"] = count + seen_counts.append(count) + seen_session_ids.append(session.session_id) + return AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text(f"turn {count}")])]) + + agent = _make_agent() + agent.run = AsyncMock(side_effect=run_with_state) + server = _make_server(agent) + + first = await _post(server) + second = await _post(server, previous_response_id=first.json()["id"]) + + assert first.status_code == 200 + assert second.status_code == 200 + assert seen_counts == [1, 2] + assert seen_session_ids[0] == seen_session_ids[1] + + async def test_responses_history_is_not_duplicated_by_default_local_history(self) -> None: + client = _RecordingHistoryClient() + agent = Agent(client=client, name="History Test Agent") + store = SessionStore() + server = ResponsesHostServer(agent, store=InMemoryResponseProvider(), session_store=store) + + first = await _post(server, input_text="first") + await _post(server, input_text="second", previous_response_id=first.json()["id"]) + + assert [[message.text for message in call] for call in client.calls] == [ + ["first"], + ["first", "recorded", "second"], + ] + assert [provider.source_id for provider in agent.context_providers] == ["_foundry_responses_history"] + + conversation_key = _response_id_partition(first.json()["id"]) + assert conversation_key is not None + stored = await store.get( + _custom_session_store_key( + agent, + _conversation_object_id(conversation_key), + isolation=ResolvedSessionIsolation(None, None, None), + ) + ) + assert stored is not None + assert InMemoryHistoryProvider.DEFAULT_SOURCE_ID not in stored.state + + async def test_file_store_restores_session_across_server_instances(self, tmp_path: Path) -> None: + seen_counts: list[int] = [] + response_store = InMemoryResponseProvider() + + def make_agent() -> MagicMock: + agent = _make_agent() + + async def run_with_state(*args: Any, **kwargs: Any) -> AgentResponse: + session = kwargs["session"] + assert isinstance(session, AgentSession) + count = int(session.state.get("turn_count", 0)) + 1 + session.state["turn_count"] = count + seen_counts.append(count) + return AgentResponse( + messages=[Message(role="assistant", contents=[Content.from_text(f"turn {count}")])] + ) + + agent.run = AsyncMock(side_effect=run_with_state) + return agent + + first_server = ResponsesHostServer( + make_agent(), + store=response_store, + session_store=IsolationKeyScopedFileSessionStore(tmp_path), + ) + first = await _post(first_server) + + second_server = ResponsesHostServer( + make_agent(), + store=response_store, + session_store=IsolationKeyScopedFileSessionStore(tmp_path), + ) + second = await _post(second_server, previous_response_id=first.json()["id"]) + + assert second.status_code == 200 + assert seen_counts == [1, 2] + + async def test_streaming_run_saves_final_session_state(self) -> None: + store = SessionStore() + agent = _make_agent() + + def run_streaming(*args: Any, **kwargs: Any) -> ResponseStream: + session = kwargs["session"] + assert isinstance(session, AgentSession) + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + session.state["stream_complete"] = True + yield AgentResponseUpdate(contents=[Content.from_text("done")], role="assistant") + + return ResponseStream(updates()) + + agent.run = MagicMock(side_effect=run_streaming) + server = _make_server(agent, session_store=store) + + response = await _post(server, stream=True) + response_id = _parse_sse_events(response.text)[-1]["data"]["response"]["id"] + conversation_key = _response_id_partition(response_id) + assert conversation_key is not None + + stored = await store.get( + _custom_session_store_key( + agent, + _conversation_object_id(conversation_key), + isolation=ResolvedSessionIsolation(None, None, None), + ) + ) + + assert stored is not None + assert stored.state["stream_complete"] is True + + async def test_failed_run_still_saves_mutated_session(self) -> None: + store = SessionStore() + agent = _make_agent() + + async def failing_run(*args: Any, **kwargs: Any) -> AgentResponse: + session = kwargs["session"] + assert isinstance(session, AgentSession) + session.state["before_failure"] = "saved" + raise RuntimeError("agent failed") + + agent.run = AsyncMock(side_effect=failing_run) + server = _make_server(agent, session_store=store) + + response = await _post(server) + body = response.json() + conversation_key = _response_id_partition(body["id"]) + assert conversation_key is not None + + stored = await store.get( + _custom_session_store_key( + agent, + _conversation_object_id(conversation_key), + isolation=ResolvedSessionIsolation(None, None, None), + ) + ) + + assert body["status"] == "failed" + assert stored is not None + assert stored.state["before_failure"] == "saved" + # endregion @@ -1735,6 +2234,7 @@ def _make_multi_response_agent( agent.name = "Test Agent" agent.description = "A mock agent for testing" agent.context_providers = [] + agent.create_session.side_effect = lambda *, session_id=None: AgentSession(session_id=session_id) call_index = [0] @@ -2834,7 +3334,8 @@ async def test_non_streaming_emits_mcp_approval_request_and_persists_to_storage( approval_request_id ) assert loaded.type == "function_approval_request" - assert loaded.function_call.name == "delete_file" # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert loaded.function_call is not None + assert loaded.function_call.name == "delete_file" async def test_streaming_emits_mcp_approval_request_and_persists_to_storage(self) -> None: request_content = _make_function_approval_request_content(request_id="apr_streaming") @@ -3210,57 +3711,53 @@ def test_url_encoded_traversal_is_treated_as_literal_segment(self, tmp_path: Any assert storage.storage_path.parent == root.resolve() assert storage.storage_path.name == "%2e%2e" - def test_user_id_scopes_storage_under_user_partition(self, tmp_path: Any) -> None: - """A per-user partition key nests the context dir under ``/``.""" + def test_isolation_scopes_storage_under_user_partition(self, tmp_path: Any) -> None: helper = self._helper() root = tmp_path / "root" root.mkdir() - storage = helper(str(root), "resp_abc123", user_id="user-A") + isolation = ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64) + storage = helper(str(root), "resp_abc123", isolation=isolation) assert storage.storage_path.is_dir() - assert storage.storage_path == (root / "user-A" / "resp_abc123").resolve() + assert storage.storage_path == (root / "user-user-A" / "resp_abc123").resolve() - @pytest.mark.parametrize("absent_user_id", [None, ""]) - def test_absent_user_id_uses_unscoped_layout(self, tmp_path: Any, absent_user_id: str | None) -> None: - """``None``/empty user id (local dev or protocol v1) falls back to the unscoped layout.""" + def test_absent_isolation_uses_unscoped_layout(self, tmp_path: Any) -> None: helper = self._helper() root = tmp_path / "root" root.mkdir() - storage = helper(str(root), "resp_abc123", user_id=absent_user_id) + storage = helper(str(root), "resp_abc123", isolation=ResolvedSessionIsolation(None, None, None)) assert storage.storage_path == (root / "resp_abc123").resolve() - def test_distinct_users_get_isolated_storage(self, tmp_path: Any) -> None: - """Two users sharing a context id must not resolve to the same directory.""" + def test_distinct_isolation_keys_get_isolated_storage(self, tmp_path: Any) -> None: helper = self._helper() root = tmp_path / "root" root.mkdir() - a = helper(str(root), "shared_context", user_id="user-A") - b = helper(str(root), "shared_context", user_id="user-B") + a = helper( + str(root), + "shared_context", + isolation=ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64), + ) + b = helper( + str(root), + "shared_context", + isolation=ResolvedSessionIsolation("user-B", "user-user-B", "b" * 64), + ) assert a.storage_path != b.storage_path - assert a.storage_path.is_relative_to((root / "user-A").resolve()) - assert b.storage_path.is_relative_to((root / "user-B").resolve()) + assert a.storage_path.is_relative_to((root / "user-user-A").resolve()) + assert b.storage_path.is_relative_to((root / "user-user-B").resolve()) - @pytest.mark.parametrize( - "bad_user_id", - [ - "../../escape", - "..", - ".", - "/tmp/escape", - "C:\\temp\\escape", - "user/../../escape", - "with\x00null", - "a/b", - ], - ) - def test_malicious_user_id_is_rejected(self, tmp_path: Any, bad_user_id: str) -> None: + def test_malicious_isolation_directory_is_rejected(self, tmp_path: Any) -> None: helper = self._helper() root = tmp_path / "root" root.mkdir() before = sorted(p.name for p in tmp_path.iterdir()) with pytest.raises(RuntimeError): - helper(str(root), "resp_abc123", user_id=bad_user_id) + helper( + str(root), + "resp_abc123", + isolation=ResolvedSessionIsolation("bad", "../../escape", "a" * 64), + ) after = sorted(p.name for p in tmp_path.iterdir()) - assert before == after, f"Unexpected filesystem artifacts created for user id {bad_user_id!r}" + assert before == after assert list(root.iterdir()) == [] @pytest.mark.parametrize( @@ -3431,56 +3928,46 @@ async def test_malicious_context_id_rejected_e2e(self, tmp_path: Any, context_fi class TestApprovalStoragePathValidation: - """Path-traversal and per-user scoping tests for function approval storage. - - Mirrors the checkpoint validation: the per-user approval directory is - derived by joining the platform-injected ``x-agent-user-id`` partition key - under the base approval directory, and the user id must be a single safe - path segment (CWE-22). - """ + """Path containment and per-user scoping tests for approval storage.""" @staticmethod def _helper() -> Callable[..., str]: from agent_framework_foundry_hosting._responses import ( # pyright: ignore[reportPrivateUsage] - _approval_storage_path_for_user, + _approval_storage_path, ) - return _approval_storage_path_for_user + return _approval_storage_path - def test_user_id_scopes_path_under_base_directory(self, tmp_path: Any) -> None: + def test_isolation_scopes_path_under_base_directory(self, tmp_path: Any) -> None: from pathlib import Path helper = self._helper() - base = tmp_path / "approvals" / "requests.json" - scoped = Path(helper(str(base), "user-A")) - assert scoped.name == "requests.json" - assert scoped.parent.name == "user-A" + base = tmp_path / "approvals" + isolation = ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64) + scoped = Path(helper(str(base), isolation)) + assert scoped.name == "approval_requests.json" + assert scoped.parent.name == "user-user-A" assert scoped.parent.parent == (tmp_path / "approvals").resolve() - def test_distinct_users_get_isolated_paths(self, tmp_path: Any) -> None: + def test_distinct_isolation_keys_get_isolated_paths(self, tmp_path: Any) -> None: helper = self._helper() - base = tmp_path / "approvals" / "requests.json" - assert helper(str(base), "user-A") != helper(str(base), "user-B") + base = tmp_path / "approvals" + assert helper( + str(base), + ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64), + ) != helper( + str(base), + ResolvedSessionIsolation("user-B", "user-user-B", "b" * 64), + ) - @pytest.mark.parametrize( - "bad_user_id", - [ - "../../escape", - "..", - ".", - "/tmp/escape", - "C:\\temp\\escape", - "user/../../escape", - "with\x00null", - "a/b", - "", - ], - ) - def test_malicious_user_id_is_rejected(self, tmp_path: Any, bad_user_id: str) -> None: + def test_malicious_isolation_directory_is_rejected(self, tmp_path: Any) -> None: helper = self._helper() - base = tmp_path / "approvals" / "requests.json" + base = tmp_path / "approvals" with pytest.raises(RuntimeError): - helper(str(base), bad_user_id) + helper( + str(base), + ResolvedSessionIsolation("bad", "../../escape", "a" * 64), + ) # region Agent lifecycle (lazy entry & OAuth consent surfacing) @@ -3738,6 +4225,7 @@ async def _raise_stream() -> AsyncIterator[AgentResponseUpdate]: agent.name = "Test Agent" agent.description = "A mock agent for testing" agent.context_providers = [] + agent.create_session.side_effect = lambda *, session_id=None: AgentSession(session_id=session_id) def run_streaming(*args: Any, **kwargs: Any) -> Any: if kwargs.get("stream"): @@ -3781,6 +4269,7 @@ async def _raise_stream() -> AsyncIterator[AgentResponseUpdate]: agent.name = "Test Agent" agent.description = "A mock agent for testing" agent.context_providers = [] + agent.create_session.side_effect = lambda *, session_id=None: AgentSession(session_id=session_id) def run_streaming(*args: Any, **kwargs: Any) -> Any: return ResponseStream(_raise_stream()) # type: ignore[arg-type] @@ -4141,7 +4630,8 @@ async def test_non_streaming_emits_mcp_approval_request_and_persists_to_storage( approval_request_id ) assert loaded.type == "function_approval_request" - assert loaded.function_call.name == "delete_file" # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert loaded.function_call is not None + assert loaded.function_call.name == "delete_file" assert mock_agent.run_count == 1 async def test_streaming_emits_mcp_approval_request_and_persists_to_storage(self) -> None: diff --git a/python/packages/hosting-responses/README.md b/python/packages/hosting-responses/README.md index f10b8b839d..cf97edd7ef 100644 --- a/python/packages/hosting-responses/README.md +++ b/python/packages/hosting-responses/README.md @@ -59,5 +59,7 @@ requests may branch from it and store their results under distinct new response ids. `conversation_id` is a mutable head instead; only one caller should advance it at a time. These helpers do not provide per-conversation locking. -The base execution-state helpers live in +`AgentState` lives in [`agent-framework-hosting`](https://pypi.org/project/agent-framework-hosting/). +The experimental in-memory and file-backed session stores live in core as +`agent_framework.SessionStore` and `agent_framework.FileSessionStore`. diff --git a/python/packages/hosting-telegram/README.md b/python/packages/hosting-telegram/README.md index b267bec015..968593998e 100644 --- a/python/packages/hosting-telegram/README.md +++ b/python/packages/hosting-telegram/README.md @@ -19,14 +19,15 @@ long-running service. Your app remains fully responsible for: a leading `/command`; your app decides what each command does. - **Sessions/storage** -- pair these helpers with [`agent-framework-hosting`](https://pypi.org/project/agent-framework-hosting/)'s - `AgentState` / `SessionStore` (or your own) to persist `AgentSession`s across turns. + `AgentState` and core's `SessionStore` or `FileSessionStore` to persist + `AgentSession`s across turns. ## Helpers - `telegram_chat_id(update)` -- the chat id an update belongs to. - `telegram_session_id(update, bot_id=...)` -- a bot-scoped `AgentState` - session id. Private chats use `telegram::`; other chats use - `telegram::`. + session id. Private chats use `telegram__user_`; other chats + use `telegram__chat_`. - `telegram_command(update)` -- a leading slash command, with `/name@bot args` normalized to `/name args`. Returns `None` if there is none. - `telegram_callback_query_id(update)` -- a callback query's id, so you can diff --git a/python/packages/hosting-telegram/agent_framework_hosting_telegram/_parsing.py b/python/packages/hosting-telegram/agent_framework_hosting_telegram/_parsing.py index 361d841867..9f91940f4f 100644 --- a/python/packages/hosting-telegram/agent_framework_hosting_telegram/_parsing.py +++ b/python/packages/hosting-telegram/agent_framework_hosting_telegram/_parsing.py @@ -92,8 +92,8 @@ def telegram_session_id(update: Mapping[str, Any], *, bot_id: int) -> str | None bot_id: The Telegram bot's numeric user id. Returns: - ``telegram::`` for a private chat, - ``telegram::`` for other chats, or ``None`` when the + ``telegram__user_`` for a private chat, + ``telegram__chat_`` for other chats, or ``None`` when the required Telegram identity is absent. """ chat_id = telegram_chat_id(update) @@ -111,7 +111,7 @@ def telegram_session_id(update: Mapping[str, Any], *, bot_id: int) -> str | None chat = cast("Mapping[str, Any]", chat_candidate) if isinstance(chat_candidate, Mapping) else None chat_type = chat.get("type") if chat is not None else None if chat_type != "private": - return f"telegram:{bot_id}:{chat_id}" + return f"telegram_{bot_id}_chat_{chat_id}" if callback_query is not None: sender_candidate = callback_query.get("from") @@ -123,7 +123,7 @@ def telegram_session_id(update: Mapping[str, Any], *, bot_id: int) -> str | None return None sender = cast("Mapping[str, Any]", sender_candidate) sender_id = sender.get("id") - return f"telegram:{bot_id}:{sender_id}" if isinstance(sender_id, int) else None + return f"telegram_{bot_id}_user_{sender_id}" if isinstance(sender_id, int) else None def telegram_callback_query_id(update: Mapping[str, Any]) -> str | None: diff --git a/python/packages/hosting-telegram/tests/hosting_telegram/test_parsing.py b/python/packages/hosting-telegram/tests/hosting_telegram/test_parsing.py index 0cb4af0041..1cf3a105de 100644 --- a/python/packages/hosting-telegram/tests/hosting_telegram/test_parsing.py +++ b/python/packages/hosting-telegram/tests/hosting_telegram/test_parsing.py @@ -57,7 +57,7 @@ def test_returns_none_for_non_int_chat_id(self) -> None: class TestTelegramSessionId: def test_private_chat_uses_bot_and_user_id(self) -> None: - assert telegram_session_id(_message_update(text="hi"), bot_id=123) == "telegram:123:777" + assert telegram_session_id(_message_update(text="hi"), bot_id=123) == "telegram_123_user_777" def test_returns_none_when_no_chat_id(self) -> None: assert telegram_session_id({"update_id": 1}, bot_id=123) is None @@ -70,7 +70,7 @@ def test_private_chat_returns_none_when_no_user_id(self) -> None: def test_group_chat_uses_bot_and_chat_id(self) -> None: update = _message_update(text="hi") update["message"]["chat"] = {"id": -555, "type": "supergroup"} - assert telegram_session_id(update, bot_id=123) == "telegram:123:-555" + assert telegram_session_id(update, bot_id=123) == "telegram_123_chat_-555" def test_private_callback_uses_callback_sender(self) -> None: update = { @@ -81,7 +81,7 @@ def test_private_callback_uses_callback_sender(self) -> None: "message": {"chat": {"id": 555, "type": "private"}}, }, } - assert telegram_session_id(update, bot_id=123) == "telegram:123:888" + assert telegram_session_id(update, bot_id=123) == "telegram_123_user_888" class TestTelegramCommand: diff --git a/python/packages/hosting/README.md b/python/packages/hosting/README.md index 0e376fd45c..b0dd74bbcf 100644 --- a/python/packages/hosting/README.md +++ b/python/packages/hosting/README.md @@ -4,12 +4,14 @@ Shared execution-state helpers for app-owned Agent Framework hosting. This package keeps Agent Framework state separate from web-framework concerns: -- `AgentState` — pairs an agent target with a `SessionStore` +- `AgentState` — pairs an agent target with a core `SessionStore` (`session_id -> AgentSession`). - `WorkflowState` — resolves a workflow target, including direct `Workflow` instances, workflow factories, `WorkflowBuilder`, and orchestration builders. -`SessionStore` provides `get`/`set`/`delete` by an app-selected id. Each +The experimental `SessionStore` and `FileSessionStore` implementations live in +`agent-framework-core` and are imported from `agent_framework`. `SessionStore` +provides `get`/`set`/`delete` by an app-selected id. Each successful `get` returns an independent copy, so a run works from a snapshot instead of mutating an older continuation point in place. The store does not know how to create a new value for an id it hasn't seen before — use @@ -19,10 +21,43 @@ checkpointing should use the existing `CheckpointStorage` abstraction directly; if an app needs per-session resume, keep a small app-owned cursor such as `session_id -> checkpoint_id`. +Session-store IDs are limited to 128 characters containing ASCII letters, +digits, `-`, and `_`. This keeps protocol-derived IDs safe to pass through +common storage backends, but it does not replace parameterized queries or +backend-specific validation in custom stores. + +`FileSessionStore` uses msgspec JSON. Custom objects placed in +`AgentSession.state` must be registered explicitly before the session is saved +or restored: + +```python +from agent_framework import register_state_type + + +class MyState: + ... + + +# Register at module import time, before any provider instance is created. +register_state_type(MyState, type_id="my_state") +``` + +Classes with `to_dict()` / `from_dict()` methods and explicitly registered +Pydantic models receive default codecs. Other classes can provide `encoder=` +and `decoder=` callbacks. Keep registration at module level so importing the +module prepares cold-start session restoration before its context provider is +instantiated. + +JSON is the default file format. Use MessagePack for a compact binary file: + +```python +store = FileSessionStore("storage/sessions", serialization_format="msgpack") +``` + Use FastAPI, Starlette, Azure Functions, Django, or another framework for route registration, auth, middleware, response construction, and background work. -> The built-in `SessionStore` is an in-memory `dict` with no eviction — every +> The core `SessionStore` is an in-memory `dict` with no eviction — every > id ever stored stays resolvable for the life of the process. That is > intentional: protocols such as OpenAI Responses' > `previous_response_id` are designed to let a caller continue from *any* @@ -35,12 +70,12 @@ registration, auth, middleware, response construction, and background work. ## Quickstart ```python +from agent_framework import FileSessionStore from agent_framework.openai import OpenAIChatClient from agent_framework_hosting import AgentState agent = OpenAIChatClient().as_agent(name="Assistant") -state = AgentState(agent) - +state = AgentState(agent, session_store=FileSessionStore("storage/sessions")) session = await state.get_or_create_session("conversation-1") result = await (await state.get_target()).run("Hello", session=session) ``` diff --git a/python/packages/hosting/agent_framework_hosting/__init__.py b/python/packages/hosting/agent_framework_hosting/__init__.py index 698ca1d303..19ef51b57c 100644 --- a/python/packages/hosting/agent_framework_hosting/__init__.py +++ b/python/packages/hosting/agent_framework_hosting/__init__.py @@ -7,7 +7,6 @@ from ._state import ( AgentRunArgs, AgentState, - SessionStore, SupportsBuild, WorkflowRunArgs, WorkflowState, @@ -21,7 +20,6 @@ __all__ = [ "AgentRunArgs", "AgentState", - "SessionStore", "SupportsBuild", "WorkflowRunArgs", "WorkflowState", diff --git a/python/packages/hosting/agent_framework_hosting/_state.py b/python/packages/hosting/agent_framework_hosting/_state.py index 148221e6b4..91638607e5 100644 --- a/python/packages/hosting/agent_framework_hosting/_state.py +++ b/python/packages/hosting/agent_framework_hosting/_state.py @@ -29,87 +29,11 @@ AgentRunInputs, AgentSession, ChatOptions, + SessionStore, SupportsAgentRun, Workflow, ) - -class SessionStore: - """Plain in-memory ``session_id -> AgentSession`` lookup. - - This store only stores and retrieves; it does not create sessions. Use - :meth:`AgentState.get_or_create_session` for that -- it resolves the - agent target and calls ``target.create_session(...)`` the first time a - given ``session_id`` is seen, then stores the result here. - - No eviction: every id ever stored stays resolvable for the life of the - process. That is intentional -- protocols such as OpenAI Responses' - ``previous_response_id`` are designed to let a caller continue from *any* - earlier point in a conversation, not just the latest turn, so every id - that has been handed out needs to stay independently resolvable. If you - back a ``SessionStore``-shaped store with real storage (Redis, a - database, ...), you are responsible for that store's own TTL/eviction - policy; this in-memory reference implementation does not model that - concern. - - The ``get`` method creates a copy of the session in order to ensure multiple - callers using the same response id can continue the session. - The behavior for that is controlled by the developer. - - So if there should be branching, then make sure to store the session with the new - session id, if the conversation should continue, then store them with the same ID. - Ensuring that multiple callers cannot simultaneously overwrite the same session is - the responsibility of the developer. - """ - - def __init__(self) -> None: - """Create an empty session store.""" - self._sessions: dict[str, AgentSession] = {} - - async def get(self, session_id: str) -> AgentSession | None: - """Return a copy of the stored session for ``session_id``, or ``None`` if absent. - - When overriding this method ensure the semantics stay the same and a copy is returned. - - Args: - session_id: Opaque app-selected session id. - - Raises: - ValueError: If ``session_id`` is empty. - """ - if not session_id: - raise ValueError("session_id must be a non-empty string") - session = self._sessions.get(session_id) - return copy.deepcopy(session) if session is not None else None - - async def set(self, session_id: str, session: AgentSession) -> None: - """Store ``session`` under ``session_id``, replacing any existing entry. - - Args: - session_id: Opaque app-selected session id. - session: The session to store. - - Raises: - ValueError: If ``session_id`` is empty. - """ - if not session_id: - raise ValueError("session_id must be a non-empty string") - self._sessions[session_id] = session - - async def delete(self, session_id: str) -> None: - """Forget the stored session for ``session_id``, if any. - - Args: - session_id: Opaque app-selected session id. - - Raises: - ValueError: If ``session_id`` is empty. - """ - if not session_id: - raise ValueError("session_id must be a non-empty string") - self._sessions.pop(session_id, None) - - AgentT = TypeVar("AgentT", bound=SupportsAgentRun) WorkflowT = TypeVar("WorkflowT", bound=Workflow) @@ -236,14 +160,14 @@ async def get_or_create_session(self, session_id: str) -> AgentSession: """Return the session for ``session_id``, creating and storing one if missing. Args: - session_id: Opaque app-selected session id. + session_id: App-selected session ID containing only ASCII letters, + digits, ``-``, and ``_``. Returns: An independent working copy of the stored or newly created ``AgentSession``. """ - if not session_id: - raise ValueError("session_id must be a non-empty string") + SessionStore.validate_session_id(session_id) session_lock = self._session_locks.setdefault(session_id, asyncio.Lock()) async with session_lock: session = await self._session_store.get(session_id) @@ -258,9 +182,11 @@ async def set_session(self, session_id: str, session: AgentSession) -> None: """Store ``session`` under ``session_id`` in this state's session store. Args: - session_id: Opaque app-selected session id. + session_id: App-selected session ID containing only ASCII letters, + digits, ``-``, and ``_``. session: Session to store. """ + SessionStore.validate_session_id(session_id) await self._session_store.set(session_id, session) diff --git a/python/packages/hosting/tests/hosting/test_state.py b/python/packages/hosting/tests/hosting/test_state.py index deb166fc96..a163bd598c 100644 --- a/python/packages/hosting/tests/hosting/test_state.py +++ b/python/packages/hosting/tests/hosting/test_state.py @@ -16,10 +16,12 @@ Content, Message, ResponseStream, + SessionStore, Workflow, ) -from agent_framework_hosting import AgentState, SessionStore, WorkflowState +import agent_framework_hosting +from agent_framework_hosting import AgentState, WorkflowState def _workflow_fixture(name: str) -> Any: @@ -100,84 +102,11 @@ async def _get_response() -> AgentResponse[Any]: return _get_response() -class TestSessionStore: - async def test_get_returns_none_for_missing_id(self) -> None: - store = SessionStore() - - assert await store.get("session-1") is None - - async def test_set_then_get_returns_session_copy(self) -> None: - store = SessionStore() - session = AgentSession(session_id="session-1") - session.state["nested"] = {"values": ["original"]} - - await store.set("session-1", session) - - stored = await store.get("session-1") - assert stored is not None - assert stored is not session - assert stored.session_id == session.session_id - assert stored.state == session.state - - stored.state["nested"]["values"].append("changed") - reread = await store.get("session-1") - assert reread is not None - assert reread.state["nested"]["values"] == ["original"] - - async def test_set_can_store_same_session_under_additional_id(self) -> None: - store = SessionStore() - session = AgentSession(session_id="resp_1") - - await store.set("resp_1", session) - await store.set("resp_2", session) - - first = await store.get("resp_1") - second = await store.get("resp_2") - assert first is not None - assert second is not None - assert first is not session - assert second is not session - assert first is not second - - async def test_set_replaces_existing_entry(self) -> None: - store = SessionStore() - first = AgentSession(session_id="session-1") - second = AgentSession(session_id="session-1") - - await store.set("session-1", first) - await store.set("session-1", second) - - stored = await store.get("session-1") - assert stored is not None - assert stored is not second - assert stored.session_id == second.session_id - - async def test_delete_forgets_session(self) -> None: - store = SessionStore() - await store.set("session-1", AgentSession(session_id="session-1")) - - await store.delete("session-1") - - assert await store.get("session-1") is None - - async def test_delete_missing_id_is_a_no_op(self) -> None: - store = SessionStore() - - await store.delete("never-stored") - - async def test_empty_session_id_raises(self) -> None: - store = SessionStore() - session = AgentSession(session_id="session-1") - - with pytest.raises(ValueError, match="session_id"): - await store.get("") - with pytest.raises(ValueError, match="session_id"): - await store.set("", session) - with pytest.raises(ValueError, match="session_id"): - await store.delete("") - - class TestAgentState: + def test_session_store_is_owned_by_core(self) -> None: + assert "SessionStore" not in agent_framework_hosting.__all__ + assert not hasattr(agent_framework_hosting, "SessionStore") + def test_default_session_store_is_fresh_in_memory_store(self) -> None: agent = _FakeAgent() state = AgentState(agent) @@ -270,6 +199,16 @@ async def test_get_or_create_session_creates_and_stores_once(self) -> None: assert second.session_id == "session-1" assert len(agent.created_sessions) == 1 + @pytest.mark.parametrize("session_id", ["two words", "tenant/user", "tenant:conversation", "' OR 1=1 --"]) + async def test_get_or_create_session_rejects_invalid_store_id(self, session_id: str) -> None: + agent = _FakeAgent() + state = AgentState(agent) + + with pytest.raises(ValueError, match="ASCII letters"): + await state.get_or_create_session(session_id) + + assert agent.created_sessions == [] + async def test_get_or_create_session_creates_once_for_concurrent_callers(self) -> None: class _YieldingSessionStore(SessionStore): async def get(self, session_id: str) -> AgentSession | None: diff --git a/python/samples/02-agents/context_providers/simple_context_provider.py b/python/samples/02-agents/context_providers/simple_context_provider.py index 10bc5e2243..30ed10e872 100644 --- a/python/samples/02-agents/context_providers/simple_context_provider.py +++ b/python/samples/02-agents/context_providers/simple_context_provider.py @@ -5,7 +5,14 @@ from contextlib import suppress from typing import Any -from agent_framework import Agent, AgentSession, ContextProvider, SessionContext, SupportsChatGetResponse +from agent_framework import ( + Agent, + AgentSession, + ContextProvider, + SessionContext, + SupportsChatGetResponse, + register_state_type, +) from agent_framework.foundry import FoundryChatClient from azure.identity import AzureCliCredential from dotenv import load_dotenv @@ -20,6 +27,13 @@ class UserInfo(BaseModel): age: int | None = None +# In order for the State to be serialized well, we need to make sure to register +# this class, and since this uses a Pydantic model, we do not need to tell the state +# how to serialize/deserialize the object. Default Python types do not need to be +# registered. +register_state_type(UserInfo, type_id="sample_user_info") + + class UserInfoMemory(ContextProvider): DEFAULT_SOURCE_ID = "user_info_memory" diff --git a/python/samples/02-agents/conversations/file_history_provider.py b/python/samples/02-agents/conversations/file_history_provider.py index cc9959b686..065669e161 100644 --- a/python/samples/02-agents/conversations/file_history_provider.py +++ b/python/samples/02-agents/conversations/file_history_provider.py @@ -20,12 +20,6 @@ from dotenv import load_dotenv from pydantic import Field -try: - import orjson # pyright: ignore[reportMissingImports] -except ImportError: - orjson = None - - # Load environment variables from .env file. load_dotenv() @@ -42,7 +36,7 @@ Key components: - `FileHistoryProvider`: Stores one message JSON object per line in a local - `.jsonl` file for each session. + `.jsonl` file for each session using msgspec JSON by default. - `lookup_weather`: A function tool that makes the persisted file show the assistant function call and tool result lines. - `json.dumps(..., indent=2)`: Pretty-prints selected records in the sample @@ -110,15 +104,7 @@ async def main() -> None: "answer with the tool result in one sentence." ), tools=[lookup_weather], - # if orjson is available, use it for faster JSON serialization in the FileHistoryProvider, - # otherwise fall back to the default json module. - context_providers=[ - FileHistoryProvider( - storage_directory, - dumps=orjson.dumps if orjson else None, - loads=orjson.loads if orjson else None, - ) - ], + context_providers=[FileHistoryProvider(storage_directory)], default_options={"store": False}, ) diff --git a/python/samples/02-agents/conversations/file_history_provider_conversation_persistence.py b/python/samples/02-agents/conversations/file_history_provider_conversation_persistence.py index 854a76d84b..6ea364a586 100644 --- a/python/samples/02-agents/conversations/file_history_provider_conversation_persistence.py +++ b/python/samples/02-agents/conversations/file_history_provider_conversation_persistence.py @@ -21,12 +21,6 @@ from dotenv import load_dotenv from pydantic import Field -try: - import orjson # pyright: ignore[reportMissingImports] -except ImportError: - orjson = None - - load_dotenv() """ @@ -42,15 +36,13 @@ Key components: - `FileHistoryProvider`: Stores one message JSON object per line in a local - `.jsonl` file for each session. + `.jsonl` file for each session using msgspec JSON by default. - `get_weather`: A function tool that makes the persisted file show the assistant function call and tool result records. - `json.dumps(..., indent=2)`: Pretty-prints a few persisted JSONL records while keeping the on-disk file compact and valid. - `load_dotenv()`: Loads `.env` values up front so the sample can stay focused on history persistence instead of manual environment variable plumbing. -- Optional `orjson`: Uses `orjson.dumps` / `orjson.loads` automatically when - available, otherwise falls back to the standard library `json` module. Security posture: - The history file is plaintext JSONL on disk, so use a trusted storage @@ -109,13 +101,7 @@ async def main() -> None: "and answer in one sentence using the tool result." ), tools=[get_weather], - context_providers=[ - FileHistoryProvider( - storage_directory, - dumps=orjson.dumps if orjson else None, - loads=orjson.loads if orjson else None, - ) - ], + context_providers=[FileHistoryProvider(storage_directory)], default_options={"store": False}, ) diff --git a/python/samples/04-hosting/af-hosting/README.md b/python/samples/04-hosting/af-hosting/README.md index 56901e0361..35747e6dc7 100644 --- a/python/samples/04-hosting/af-hosting/README.md +++ b/python/samples/04-hosting/af-hosting/README.md @@ -9,7 +9,7 @@ clients, authentication, response construction, and deployment shape. | Sample | What it shows | Packaging | |---|---|---| -| [`local_responses/`](./local_responses) | One agent + one `@tool` + native FastAPI route + Responses helper functions + `AgentState` / `SessionStore`. | **Local only.** Start here to learn the helper seam. | +| [`local_responses/`](./local_responses) | One agent + one `@tool` + native FastAPI route + Responses helper functions + hosting `AgentState` + core `SessionStore`. | **Local only.** Start here to learn the helper seam. | | [`local_responses_workflow/`](./local_responses_workflow) | A workflow target behind a native FastAPI route using Responses helper functions, `WorkflowState`, explicit `CheckpointStorage`, and an app-owned checkpoint cursor. | **Local only.** | | [`local_telegram/`](./local_telegram) | One agent + `aiogram` polling + Telegram conversion helpers + app-owned commands, media policy, and streaming edits. | **Local only.** Requires a Telegram bot token. | diff --git a/python/samples/04-hosting/af-hosting/local_responses/README.md b/python/samples/04-hosting/af-hosting/local_responses/README.md index 9b976be1c4..995eacda89 100644 --- a/python/samples/04-hosting/af-hosting/local_responses/README.md +++ b/python/samples/04-hosting/af-hosting/local_responses/README.md @@ -1,8 +1,8 @@ # local_responses — Responses helpers with native FastAPI routes The smallest end-to-end Responses hosting shape: one Foundry agent with a -`@tool`, one native FastAPI route, a small `SessionStore`, and the Responses -helper functions: +`@tool`, one native FastAPI route, core's experimental `SessionStore`, and the +Responses helper functions: - `responses_to_run(...)` - `responses_session_id(...)` @@ -42,6 +42,8 @@ What the route demonstrates: - Treats each `previous_response_id` as an immutable snapshot. Multiple callers can branch from the same response concurrently because each receives a session copy and stores its result under a newly minted response id. +- Uses core's msgspec-backed `FileSessionStore` under + `storage/sessions/snapshots`. `app:app` is a module-level FastAPI ASGI app; recommended local launch is Hypercorn. diff --git a/python/samples/04-hosting/af-hosting/local_responses/app.py b/python/samples/04-hosting/af-hosting/local_responses/app.py index 0b3503b026..63c06c6998 100644 --- a/python/samples/04-hosting/af-hosting/local_responses/app.py +++ b/python/samples/04-hosting/af-hosting/local_responses/app.py @@ -6,8 +6,8 @@ 1. ``agent-framework-hosting-responses`` converts Responses request/response payloads to and from Agent Framework run values. -2. ``agent-framework-hosting`` owns shared execution state via - ``AgentState`` and ``SessionStore``. +2. ``agent-framework-hosting`` owns ``AgentState``; core provides its + ``SessionStore``. 3. FastAPI owns the route, request parsing, policy decisions, and response object. @@ -55,7 +55,7 @@ from pathlib import Path from typing import Annotated, Any, cast -from agent_framework import Agent, FileHistoryProvider, ResponseStream, tool +from agent_framework import Agent, FileHistoryProvider, FileSessionStore, ResponseStream, tool from agent_framework_foundry import FoundryChatClient from agent_framework_hosting import AgentState from agent_framework_hosting_responses import ( @@ -105,7 +105,10 @@ def create_agent() -> Agent: app = FastAPI() -state = AgentState(create_agent) +state = AgentState( + create_agent, + session_store=FileSessionStore(SESSIONS_DIR / "snapshots"), +) ALLOWED_REQUEST_OPTIONS = frozenset({"max_tokens", "reasoning"}) diff --git a/python/samples/04-hosting/af-hosting/local_telegram/README.md b/python/samples/04-hosting/af-hosting/local_telegram/README.md index d4843691ff..239a8cecf7 100644 --- a/python/samples/04-hosting/af-hosting/local_telegram/README.md +++ b/python/samples/04-hosting/af-hosting/local_telegram/README.md @@ -71,8 +71,8 @@ process just registered. - **Session continuity:** `telegram_session_id(..., bot_id=bot.id)` follows Telegram's native identity boundaries. Private chats use - `telegram::`; groups and supergroups use - `telegram::`, creating a shared session for that group. + `telegram__user_`; groups and supergroups use + `telegram__chat_`, creating a shared session for that group. Including `bot.id` prevents two bots from accidentally sharing state. aiogram derives that numeric bot id from `TELEGRAM_BOT_TOKEN`, so these local apps do not need a separate `TELEGRAM_BOT_ID` setting. diff --git a/python/uv.lock b/python/uv.lock index 304fe7c421..cffa180cf1 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -409,6 +409,7 @@ name = "agent-framework-core" version = "1.12.1" source = { editable = "packages/core" } dependencies = [ + { name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "opentelemetry-api", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "pydantic", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "python-dotenv", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -488,6 +489,7 @@ requires-dist = [ { name = "agent-framework-redis", marker = "extra == 'all'", editable = "packages/redis" }, { name = "agent-framework-tools", marker = "extra == 'all'", editable = "packages/tools" }, { name = "mcp", marker = "extra == 'all'", specifier = ">=1.24.0,<2" }, + { name = "msgspec", specifier = ">=0.20.0,<1" }, { name = "opentelemetry-api", specifier = ">=1.39.0,<2" }, { name = "pydantic", specifier = ">=2,<3" }, { name = "python-dotenv", specifier = ">=1,<2" }, @@ -4187,6 +4189,54 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5e/75/bd9b7bb966668920f06b200e84454c8f3566b102183bc55c5473d96cb2b9/msal_extensions-1.3.1-py3-none-any.whl", hash = "sha256:96d3de4d034504e969ac5e85bae8106c8373b5c6568e4c8fa7af2eca9dbe6bca", size = 20583, upload-time = "2025-03-14T23:51:03.016Z" }, ] +[[package]] +name = "msgspec" +version = "0.21.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e3/60/f79b9b013a16fa3a58350c9295ddc6789f2e335f36ea61ed10a21b215364/msgspec-0.21.1.tar.gz", hash = "sha256:2313508e394b0d208f8f56892ca9b2799e2561329de9763b19619595a6c0f72c", size = 319193, upload-time = "2026-04-12T21:44:50.394Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ba/7f/bbc4e74cd33d316b75541149e4d35b163b63bce066530ae185a2ec3b5bfc/msgspec-0.21.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:b504b6e7f7a22a24b27232b73034421692147865162daaec9f3bf62439007c87", size = 193131, upload-time = "2026-04-12T21:43:56.094Z" }, + { url = "https://files.pythonhosted.org/packages/c1/60/504886af1aaf854112663b842d5eea9a15d9588f9bf7d0d2df736424b84d/msgspec-0.21.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:4692b7c1609155708c4418f88e92f63c13fdf08aa095c84bae82bad75b53389b", size = 186597, upload-time = "2026-04-12T21:43:57.242Z" }, + { url = "https://files.pythonhosted.org/packages/fa/54/d24ddeaa65b5278c9e67f48ce3c17a9831e8f3722f3c8322ee120aca22ef/msgspec-0.21.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d3124010b3815451494c85ff345e693cb9fe5889cfcbbef39ed8622e0e72319c", size = 215158, upload-time = "2026-04-12T21:43:58.442Z" }, + { url = "https://files.pythonhosted.org/packages/9f/75/bb79c8b89a93ae23cd33c0d802373f16feaf9633f05d8af77091350dda0a/msgspec-0.21.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6badc03b9725352219cca017bfe71c61f2fbd0fb5982b410ac17c97c213deb30", size = 219856, upload-time = "2026-04-12T21:44:00.015Z" }, + { url = "https://files.pythonhosted.org/packages/b4/9c/c5ca26b46f0ebbd3a6683695ef89396712cb9e4199fd1f0bc1dd968216b1/msgspec-0.21.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:5d2d4116ebe3035a78d9ec76e99a9d64e5fa6d44fe61a9c5de7fd1acf54bcc69", size = 220314, upload-time = "2026-04-12T21:44:01.548Z" }, + { url = "https://files.pythonhosted.org/packages/c8/31/645a351c4285dce40ed6755c3dcc0aa648e26dacb20a98018fe2cce5e87b/msgspec-0.21.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:0d1009f6715f5bff3b54d4ff5c7428ad96197e0534e1645b8e9b955890c84664", size = 223215, upload-time = "2026-04-12T21:44:02.884Z" }, + { url = "https://files.pythonhosted.org/packages/09/af/8bf15736a6dd3cb4f90c5467f6dc39197d2daaf10754490cdc0aa17b7312/msgspec-0.21.1-cp311-cp311-win_amd64.whl", hash = "sha256:c6faffe5bb644ec884052679af4dfd776d4b5ca90e4a7ec7e7e319e4e6b93a6e", size = 188554, upload-time = "2026-04-12T21:44:04.151Z" }, + { url = "https://files.pythonhosted.org/packages/ef/29/cc7db3a165b62d16e64a83f82eccb79655055cb5bc1f60459a6f9d7c82f2/msgspec-0.21.1-cp311-cp311-win_arm64.whl", hash = "sha256:ee9e3f11fa94603f7d673bf795cfa31b549c4a2c723bc39b45beb1e7f5a3fb99", size = 174517, upload-time = "2026-04-12T21:44:05.66Z" }, + { url = "https://files.pythonhosted.org/packages/6e/cf/317224852c00248c620a9bcf4b26e2e4ab8afd752f18d2a6ef73ebd423b6/msgspec-0.21.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d4248cf0b6129b7d230eacd493c17cc2d4f3989f3bb7f633a928a85b7dcfa251", size = 196188, upload-time = "2026-04-12T21:44:07.181Z" }, + { url = "https://files.pythonhosted.org/packages/6d/81/074612945c0666078f7366f40000013de9f6ba687491d450df699bceebc9/msgspec-0.21.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:5102c7e9b3acff82178449b85006d96310e690291bb1ea0142f1b24bcb8aabcb", size = 188473, upload-time = "2026-04-12T21:44:08.736Z" }, + { url = "https://files.pythonhosted.org/packages/8a/37/655101799590bcc5fddb2bd3fe0e6194e816c2d1da7c361725f5eb89a910/msgspec-0.21.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:846758412e9518252b2ac9bffd6f0e54d9ff614f5f9488df7749f81ff5c80920", size = 218871, upload-time = "2026-04-12T21:44:09.917Z" }, + { url = "https://files.pythonhosted.org/packages/b5/d1/d4cd9fe89c7d400d7a18f86ccc94daa3f0927f53558846fcb60791dce5d6/msgspec-0.21.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:21995e74b5c598c2e004110ad66ec7f1b8c20bf2bcf3b2de8fd9a3094422d3ff", size = 225025, upload-time = "2026-04-12T21:44:11.191Z" }, + { url = "https://files.pythonhosted.org/packages/24/bf/e20549e602b9edccadeeff98760345a416f9cce846a657e8b18e3396b212/msgspec-0.21.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:6129f0cca52992e898fd5344187f7c8127b63d810b2fd73e36fca73b4c6475ee", size = 222672, upload-time = "2026-04-12T21:44:12.481Z" }, + { url = "https://files.pythonhosted.org/packages/b4/68/04d7a8f0f786545cf9b8c280c57aa6befb5977af6e884b8b54191cbe44b3/msgspec-0.21.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:ef3ec2296248d1f8b9231acb051b6d471dfde8f21819e86c9adaaa9f42918521", size = 227303, upload-time = "2026-04-12T21:44:13.709Z" }, + { url = "https://files.pythonhosted.org/packages/cc/4d/619866af2840875be408047bf9e70ceafbae6ab50660de7134ed1b25eb86/msgspec-0.21.1-cp312-cp312-win_amd64.whl", hash = "sha256:d4ab834a054c6f0cbeef6df9e7e1b33d5f1bc7b86dea1d2fd7cad003873e783d", size = 190017, upload-time = "2026-04-12T21:44:14.977Z" }, + { url = "https://files.pythonhosted.org/packages/5e/2e/a8f9eca8fd00e097d7a9e99ba8a4685db994494448e3d4f0b7f6e9a3c0f7/msgspec-0.21.1-cp312-cp312-win_arm64.whl", hash = "sha256:628aaa35c74950a8c59da330d7e98917e1c7188f983745782027748ee4ca573e", size = 175345, upload-time = "2026-04-12T21:44:16.431Z" }, + { url = "https://files.pythonhosted.org/packages/7e/74/f11ede02839b19ff459f88e3145df5d711626ca84da4e23520cebf819367/msgspec-0.21.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:764173717a01743f007e9f74520ed281f24672c604514f7d76c1c3a10e8edb66", size = 196176, upload-time = "2026-04-12T21:44:17.613Z" }, + { url = "https://files.pythonhosted.org/packages/bb/40/4476c1bd341418a046c4955aff632ec769315d1e3cb94e6acf86d461f9ed/msgspec-0.21.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:344c7cd0eaed1fb81d7959f99100ef71ec9b536881a376f11b9a6c4803365697", size = 188524, upload-time = "2026-04-12T21:44:18.815Z" }, + { url = "https://files.pythonhosted.org/packages/ca/d9/9e9d7d7e5061b47540d03d640fab9b3965ba7ae49c1b2154861c8f007518/msgspec-0.21.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:48943e278b3854c2f89f955ddc6f9f430d3f0784b16e47d10604ee0463cd21f5", size = 218880, upload-time = "2026-04-12T21:44:20.028Z" }, + { url = "https://files.pythonhosted.org/packages/74/66/2bb344f34abb4b57e60c7c9c761994e0417b9718ec1460bf00c296f2a7ea/msgspec-0.21.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a9aa659ebb0101b1cbc31461212b87e341d961f0ab0772aaf068a99e001ec4aa", size = 225050, upload-time = "2026-04-12T21:44:21.577Z" }, + { url = "https://files.pythonhosted.org/packages/1a/84/7c1e412f76092277bf760cef12b7979d03314d259ab5b5cafde5d0c1722d/msgspec-0.21.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f7b27d1a8ead2b6f5b0c4f2d07b8be1ccfcc041c8a0e704781edebe3ae13c484", size = 222713, upload-time = "2026-04-12T21:44:22.83Z" }, + { url = "https://files.pythonhosted.org/packages/4e/27/0bba04b2b4ef05f3d068429410bc71d2cea925f1596a8f41152cccd5edb8/msgspec-0.21.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:38fe93e86b61328fe544cb7fd871fad5a27c8734bfda90f65e5dbe288ae50f61", size = 227259, upload-time = "2026-04-12T21:44:24.11Z" }, + { url = "https://files.pythonhosted.org/packages/b0/2d/09574b0eea02fed2c2c1383dbaae2c7f79dc16dcd6487a886000afb5d7c4/msgspec-0.21.1-cp313-cp313-win_amd64.whl", hash = "sha256:8bc666331c35fcce05a7cd2d6221adbe0f6058f8e750711413d22793c080ac6a", size = 189857, upload-time = "2026-04-12T21:44:25.359Z" }, + { url = "https://files.pythonhosted.org/packages/46/34/105b1576ad182879914f0c821f17ee1d13abb165cb060448f96fe2aff078/msgspec-0.21.1-cp313-cp313-win_arm64.whl", hash = "sha256:42bb1241e0750c1a4346f2aa84db26c5ffd99a4eb3a954927d9f149ff2f42898", size = 175403, upload-time = "2026-04-12T21:44:26.608Z" }, + { url = "https://files.pythonhosted.org/packages/5a/ad/86954e987d1d6a5c579e2c2e7832b65e0fff194179fdac4f581536086024/msgspec-0.21.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:fab48eb45fdbfbdb2c0edfec00ffc53b6b6085beefc6b50b61e01659f9f8757f", size = 196261, upload-time = "2026-04-12T21:44:27.807Z" }, + { url = "https://files.pythonhosted.org/packages/d1/a1/c5e46c3e42b866199365e35d11dddfd1fbd8bba4fdb3c52f965b1607ce94/msgspec-0.21.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:3cb779ea0c35bc807ff941d415875c1f69ca0be91a2e907ab99a171811d86a9a", size = 188729, upload-time = "2026-04-12T21:44:28.99Z" }, + { url = "https://files.pythonhosted.org/packages/85/7d/1e29a319d678d6cb962ae5bdf32a6858ebdf38f73bc654c0e9c742a0c2c8/msgspec-0.21.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:68604db36b3b4dd9bf160e436e12798a4738848144cea1aca1cb984011eb160f", size = 219866, upload-time = "2026-04-12T21:44:31.104Z" }, + { url = "https://files.pythonhosted.org/packages/25/1f/cca084ca2572810fff12ea9dbdcbe39eac048f40daf4a9077b49fcbe8cee/msgspec-0.21.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3d6b9dc50948eaf65df54d2fd0ff66e6d8c32f116037209ee861810eb9b676cb", size = 224993, upload-time = "2026-04-12T21:44:32.649Z" }, + { url = "https://files.pythonhosted.org/packages/71/94/d2120fc9d419a89a3a7c13e5b7078798c4b392a96a02a6e2b3ce43a8766c/msgspec-0.21.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:52c5e21930942302394429c5a582ce7e6b62c7f983b3760834c2ce107e0dd6df", size = 223535, upload-time = "2026-04-12T21:44:33.839Z" }, + { url = "https://files.pythonhosted.org/packages/75/17/42418b66a3ad972a89bab73dd78b79cc6282bb488a25e73c853cee7443b9/msgspec-0.21.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:abbb39d65681fa24ed394e01af3d59d869068324f900c61d06062b7fb9980f2f", size = 227222, upload-time = "2026-04-12T21:44:35.093Z" }, + { url = "https://files.pythonhosted.org/packages/c4/33/265c894268cca88ff67b144ca2b4c522fc8b9a6f1966a3640c70516e78e1/msgspec-0.21.1-cp314-cp314-win_amd64.whl", hash = "sha256:5666b1b560b97b6ec2eb3fca8a502298ebac56e13bbca1f88523538ce83d01ea", size = 193810, upload-time = "2026-04-12T21:44:36.612Z" }, + { url = "https://files.pythonhosted.org/packages/3b/8f/a6d35f25bf1fc63c492fdd88fdce01ba0875ead48c2b91f90f33653b4131/msgspec-0.21.1-cp314-cp314-win_arm64.whl", hash = "sha256:d8b8578e4c83b14ceea4cef0d0b747e31d9330fe4b03b2b2ad4063866a178f93", size = 179125, upload-time = "2026-04-12T21:44:38.198Z" }, + { url = "https://files.pythonhosted.org/packages/c6/39/74839641e64b99d87da55af0fc472854d42b46e2183b9e2a67fe1bb2a512/msgspec-0.21.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:15f523d51c00ebad412213bfe9f06f0a50ec2b93e0c19e824a2d267cabb48ea2", size = 200171, upload-time = "2026-04-12T21:44:39.414Z" }, + { url = "https://files.pythonhosted.org/packages/70/9b/ce0cca6d2d87fcd4b6ff97600790494e64f26a2c55d61507cd2755c16193/msgspec-0.21.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:4e47390360583ba3d5c6cb44cf0a9f61b0a06a899d3c2c00627cedebb2e2884b", size = 192879, upload-time = "2026-04-12T21:44:40.882Z" }, + { url = "https://files.pythonhosted.org/packages/a7/08/673a7bb05e5702dc787ddd3011195b509f9867927970da59052211929987/msgspec-0.21.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f60800e6299b798142dc40b0644da77ceac5ea0568be58228417eae14135c847", size = 226281, upload-time = "2026-04-12T21:44:42.181Z" }, + { url = "https://files.pythonhosted.org/packages/7d/45/86508cf57283e9070b3c447e3ab25b792a7a0855a3ea4e0c6d111ac34c97/msgspec-0.21.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5f8e9dfcd98419cf7568808470c4317a3fb30bef0e3715b568730a2b272a20d7", size = 229863, upload-time = "2026-04-12T21:44:43.442Z" }, + { url = "https://files.pythonhosted.org/packages/2c/62/e7c9367cd08d590559faacd711edbae36840342843e669440363f33c7d36/msgspec-0.21.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:92d89dfad13bd1ea640dc3e37e724ed380da1030b272bdf5ecafb983c3ad7c75", size = 230445, upload-time = "2026-04-12T21:44:44.806Z" }, + { url = "https://files.pythonhosted.org/packages/42/b4/c0f54632103846b658a10930025f4de41c8724b5e4805a5f3b395586cb7e/msgspec-0.21.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:0d03867786e5d7ba25d666df4b11320c27170f4aeafcb8e3a8b0a50a4fb742ca", size = 231822, upload-time = "2026-04-12T21:44:46.343Z" }, + { url = "https://files.pythonhosted.org/packages/ea/1d/0d85cc79d0ccf5508e9c846cc66552a6a16bf92abd1dbd8362617f7b35cd/msgspec-0.21.1-cp314-cp314t-win_amd64.whl", hash = "sha256:740fbf1c9d59992ca3537d6fbe9ebbf9eaf726a65fbf31448e0ecbc710697a63", size = 206650, upload-time = "2026-04-12T21:44:47.601Z" }, + { url = "https://files.pythonhosted.org/packages/90/91/56c5d560f20e6c20e9e4f55bd0e458f7f162aa689ee350346c04c48eac0b/msgspec-0.21.1-cp314-cp314t-win_arm64.whl", hash = "sha256:0d2cc73df6058d811a126ac3a8ad63a4dfa210c82f9cf5a004802eaf4712de90", size = 183149, upload-time = "2026-04-12T21:44:48.833Z" }, +] + [[package]] name = "msrest" version = "0.7.1" From 0679b7af795e509a54739f74c67aef8547df0cd1 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 27 Jul 2026 11:06:22 +0200 Subject: [PATCH 2/8] Python: Address session persistence review feedback Harden scoped file paths and corruption recovery, preserve session serialization compatibility, clarify dependency placement, and add reproducible benchmark evidence. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3e5c81ad-75e8-4e92-a883-8bbd676c6c8c --- ...0033-python-session-store-serialization.md | 68 +- python/packages/core/AGENTS.md | 6 +- .../core/agent_framework/_sessions.py | 277 +++++++-- python/packages/core/pyproject.toml | 2 +- .../packages/core/tests/core/test_sessions.py | 252 +++++++- .../_responses.py | 2 +- .../_session_isolation.py | 17 +- .../foundry_hosting/tests/test_responses.py | 70 ++- python/packages/hosting/README.md | 19 +- .../hosting/agent_framework_hosting/_state.py | 8 +- .../hosting/tests/hosting/test_state.py | 31 +- .../session_serialization_benchmark.py | 582 ++++++++++++++++++ python/uv.lock | 2 +- 13 files changed, 1196 insertions(+), 140 deletions(-) create mode 100644 python/scripts/session_serialization_benchmark.py diff --git a/docs/decisions/0033-python-session-store-serialization.md b/docs/decisions/0033-python-session-store-serialization.md index 4f06334c90..d981f873a4 100644 --- a/docs/decisions/0033-python-session-store-serialization.md +++ b/docs/decisions/0033-python-session-store-serialization.md @@ -152,6 +152,13 @@ pre-msgspec `FileHistoryProvider` design; those hooks remain only as a deprecate A benchmark using a large `AgentSession` with 2,000 `Message` objects stored through `InMemoryHistoryProvider`, nested standard dictionaries, registered custom classes, and registered Pydantic models measured the complete `AgentSession.to_dict()` / codec / `AgentSession.from_dict()` path. +The reproducible harness is +[`python/scripts/session_serialization_benchmark.py`](../../python/scripts/session_serialization_benchmark.py): + +```bash +cd python +uv run --with orjson python scripts/session_serialization_benchmark.py +``` | Codec | File size | Encode median (ms) | Decode median (ms) | Round-trip median (ms) | Disk round-trip median (ms) | | --- | ---: | ---: | ---: | ---: | ---: | @@ -169,7 +176,9 @@ MessagePack reduced file size to 92.2% of JSON (about 7.8% smaller) and produced round-trip medians. Its in-memory round-trip median was effectively tied with msgspec JSON. This supports offering it as a nice-to-have, but it is not required to justify choosing msgspec for JSON. -These results are workload- and machine-dependent, but they validate the architectural choice: +These results are workload- and machine-dependent. The small differences between optimized JSON implementations are +not the basis for the architectural choice. The benchmark instead confirms that the typed design does not impose a +material regression for this representative payload: - use msgspec JSON as the readable default; - optionally offer msgspec MessagePack when storage size or disk latency matters; @@ -189,9 +198,11 @@ implementations. `InMemorySessionStore`, protocol, or ABC is introduced. `agent-framework-hosting` consumes the core type and no longer owns or re-exports `SessionStore` (this will be a breaking change in the `hosting` package). -Session-store keys use one restricted contract suitable for the built-in file implementation: at most 128 ASCII -letters, digits, `-`, and `_`. Logical `AgentSession.session_id` values outside a store are not globally constrained by -this file-oriented key contract. +`SessionStore` accepts opaque non-empty keys so custom backends can use their native key contracts. The built-in +`FileSessionStore` restricts direct file keys to at most 128 ASCII letters, digits, `-`, and `_`. `AgentState` remains +storage-agnostic and passes keys through unchanged; each store implementation owns backend-specific validation or +normalization. Protocol-specific hosts such as Foundry may still derive their own stable storage key before calling the +store. ### Decision 2: Use msgspec codecs plus an explicit dynamic registry @@ -203,28 +214,57 @@ JSON is the required and default format. Because msgspec can reuse the same type dictionary is wrapped in one custom field; its encode/decode hooks recursively translate explicitly registered types to and from the existing tagged mappings in either format. -The dependency floor is `msgspec>=0.20.0` because core supports Python 3.14 and msgspec added Python 3.14 support in -version 0.20.0. +The dependency range is `msgspec>=0.20.0,<0.22`: version 0.20.0 added Python 3.14 support, and the upper bound limits +core to the tested 0.20/0.21 minor lines. + +Three dependency placements were considered: + +1. Make msgspec a standard core dependency. +2. Make msgspec optional in core but standard in Foundry hosting. +3. Make msgspec optional in both packages. + +Option 3 moves installation failures to application developers even though durable session persistence is required for +the primary `ResponsesHostServer` API to preserve Agent Framework state. Option 2 removes that burden from Foundry +hosting but makes core's shared `_sessions` module and public types conditionally defined or lazily imported without +removing msgspec from the default Foundry installation. Option 1 is therefore selected: msgspec is a standard core +dependency, giving both core file providers and Foundry hosting one predictable implementation path. + +Core already depends on the native `pydantic-core` extension, so native-wheel availability is not a new packaging +constraint. The msgspec project is also actively tracking upcoming Python support; its merged +[`Add 3.15-dev to CI` PR](https://github.com/msgspec/msgspec/pull/1037) exercises Python 3.15 development builds. This gives confidence that they will add support for new python version quickly. The public `AgentSession` remains a normal framework class. The msgspec Struct is an internal persistence DTO rather -than the inheritance base for runtime sessions. A targeted comparison against the dictionary-based msgspec path found -that the internal Struct improved JSON encode latency by about 9% and JSON round-trip latency by about 4%, without -changing file size. +than the inheritance base for runtime sessions. The Struct gives persistence one typed encode/decode operation, validates +the snapshot envelope, and carries an explicit payload version. The benchmark's small timing spread was not used to +choose the Struct. -`register_state_type` requires explicit registration, supports stable type IDs and optional codecs, rejects collisions, -and provides defaults for `to_dict` / `from_dict` classes and explicitly registered Pydantic models. Unknown persisted -type IDs remain raw dictionaries for compatibility. +`register_state_type` supports stable type IDs and optional codecs, rejects collisions, and provides defaults for +`to_dict` / `from_dict` classes and Pydantic models. One recursive serializer is shared by `AgentSession.to_dict()` and +the durable codecs. The established implicit Pydantic registration behavior remains temporarily for compatibility, but +now emits `DeprecationWarning` instructing applications to register the model at module import time. Same-process +round-trips continue to work; cold-start deserialization is not guaranteed without explicit registration. Unknown +persisted type IDs remain raw dictionaries. `FileHistoryProvider` also adds msgspec JSON as its default JSON Lines codec. It supports the same explicit `serialization_format="msgpack"` choice using length-prefixed append-only MessagePack records. Its existing `dumps` / `loads` extension points remain temporarily for JSON compatibility, emit `DeprecationWarning` when supplied, and do -not apply to MessagePack. New code uses the built-in msgspec codecs. +not apply to MessagePack. New code uses the built-in codecs. The default JSON reader falls back to the standard library +for legacy JSON Lines containing `NaN` or infinity, and writes those non-finite values with the standard library so +existing history semantics are preserved. ## Follow-up Work Audit the remaining file-backed stores to determine whether they benefit from the same typed msgspec treatment and optional JSON / MessagePack formats. `FileCheckpointStorage` is the first candidate because it persists large, -structured workflow state and currently uses JSON plus custom checkpoint value encoding. +structured workflow state and currently uses JSON plus custom checkpoint value encoding. Its existing +`WorkflowCheckpoint.version` field already provides a payload-shape discriminator. + +Checkpoint migration should be reader-first. A compatibility release can detect the codec from the first byte, widen +the two `glob("*.json")` readers to discover future formats, and continue writing only JSON. A later release can add +opt-in MessagePack writes while retaining JSON as the default. The payload `version` should describe the checkpoint +shape rather than the codec, which is discoverable from the bytes. MessagePack should not become the default while +mixed-version fleets may share one checkpoint directory: older readers silently ignore non-JSON files and could resume +from no checkpoint instead of surfacing an incompatibility. `MemoryContextProvider` is another candidate because its file-backed path combines `MemoryFileStore` state with transcript files and still exposes `history_dumps` / `history_loads` passthroughs to the deprecated diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 4529a19de9..e894138599 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -74,9 +74,9 @@ agent_framework/ ### Sessions (`_sessions.py`) - **`AgentSession`** - Manages conversation state and session metadata -- **`SessionStore`** - Experimental in-memory `session_id -> AgentSession` snapshot store; reads return independent copies -- **`FileSessionStore`** - Experimental msgspec file-backed session snapshot store with atomic last-writer-wins updates; JSON is the default and `serialization_format="msgpack"` enables binary MessagePack -- **`register_state_type`** - Explicitly registers custom `AgentSession.state` classes with stable type IDs and optional mapping codecs; registration must happen before persistence/restoration, and conflicting IDs are rejected +- **`SessionStore`** - Experimental in-memory opaque `session_id -> AgentSession` snapshot store; reads return independent copies +- **`FileSessionStore`** - Experimental msgspec file-backed session snapshot store with atomic last-writer-wins updates; JSON is the default, `serialization_format="msgpack"` enables binary MessagePack, file keys use a restricted portable shape, and corrupt snapshots are quarantined before an error is raised +- **`register_state_type`** - Registers custom `AgentSession.state` classes with stable type IDs and optional mapping codecs. Implicit Pydantic registration remains temporarily with `DeprecationWarning`, but module-level registration is needed to guarantee cold-start restoration. - **`ServiceSessionId`** - Mapping alias for structured service-owned continuation handles used in `AgentSession.service_session_id` - **`SessionContext`** - Context object for session-scoped data during agent runs. `extend_messages(...)` can attach ordered, deduplicated `origin_session_ids` attribution when a provider injects content from other sessions. diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index f57a2cb23d..662350770f 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -17,6 +17,7 @@ import asyncio import copy +import json import logging import math import os @@ -87,6 +88,12 @@ "LPT7", "LPT8", "LPT9", + "COM¹", + "COM²", + "COM³", + "LPT¹", + "LPT²", + "LPT³", }) @@ -100,20 +107,37 @@ def _default_json_dumps(value: Any) -> bytes: + if _contains_non_finite_float(value): + return json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode("utf-8") return _DEFAULT_JSON_ENCODER.encode(value) def _default_json_loads(value: str | bytes) -> Any: - return _DEFAULT_JSON_DECODER.decode(value) + try: + return _DEFAULT_JSON_DECODER.decode(value) + except msgspec.DecodeError: + return json.loads(value) + + +def _contains_non_finite_float(value: Any) -> bool: + """Return whether a nested JSON-compatible value contains NaN or infinity.""" + if isinstance(value, float): + return not math.isfinite(value) + if isinstance(value, Mapping): + return any(_contains_non_finite_float(item) for item in cast(Mapping[Any, Any], value).values()) + if isinstance(value, (list, tuple)): + return any(_contains_non_finite_float(item) for item in cast(Sequence[Any], value)) + return False def _is_literal_session_file_stem_safe(session_id: str) -> bool: """Return whether a session ID can be used directly as a filename stem.""" + windows_stem = session_id.split(".", maxsplit=1)[0].upper() if ( not session_id or session_id.startswith(".") or session_id.endswith((" ", ".")) - or session_id.upper() in _WINDOWS_RESERVED_FILE_STEMS + or windows_stem in _WINDOWS_RESERVED_FILE_STEMS ): return False if any(ord(character) < 32 for character in session_id): @@ -295,8 +319,19 @@ def register_state_type( _STATE_CLASS_REGISTRY[cls] = registration +def _warn_implicit_pydantic_registration(value_type: type[Any]) -> None: + """Warn that an unregistered Pydantic state type uses legacy registration.""" + warnings.warn( + f"Implicit registration of Pydantic AgentSession state type {value_type.__name__!r} is deprecated and will " + "be removed in a future version. Call register_state_type() at module import time. Cold-start deserialization " + "is not guaranteed without explicit registration.", + DeprecationWarning, + stacklevel=4, + ) + + def _serialize_value(value: Any, *, path: str) -> Any: - """Serialize one session-state value through the explicit type registry.""" + """Serialize one session-state value through the shared compatibility path.""" value_type = cast(type[Any], type(value)) registration = _STATE_CLASS_REGISTRY.get(value_type) if registration is not None: @@ -312,8 +347,14 @@ def _serialize_value(value: Any, *, path: str) -> Any: } serialized["type"] = registration.type_id return serialized - if isinstance(value, float) and not math.isfinite(value): - raise ValueError(f"Session state value at {path} must be a finite float.") + if callable(getattr(value, "to_dict", None)): + payload = value.to_dict() + if not isinstance(payload, Mapping): + raise TypeError(f"{value_type.__name__}.to_dict() must return a mapping.") + return { + str(key): _serialize_value(item, path=f"{path}.{key}") + for key, item in cast(Mapping[Any, Any], payload).items() + } if value_type in _STATE_SCALAR_TYPES: return value if isinstance(value, list): @@ -327,34 +368,69 @@ def _serialize_value(value: Any, *, path: str) -> Any: str(key): _serialize_value(item, path=f"{path}.{key}") for key, item in cast(Mapping[Any, Any], value).items() } - raise TypeError( - f"Session state value at {path} has unregistered type {type(value).__name__!r}; " - "call register_state_type() before persisting the session." - ) + from pydantic import BaseModel -def _deserialize_value(value: Any) -> Any: + if isinstance(value, BaseModel): + # Temporary compatibility fallback for unregistered Pydantic models; + # remove this branch together with implicit auto-registration. + _warn_implicit_pydantic_registration(value_type) + type_id = _resolve_state_type_id(value_type, None) + register_state_type(value_type, type_id=type_id) + return _serialize_value(value, path=path) + return value + + +def _deserialize_value(value: Any, *, path: str) -> Any: """Deserialize a single value, restoring registered types.""" if isinstance(value, list): - return [_deserialize_value(item) for item in cast(list[Any], value)] + return [_deserialize_value(item, path=f"{path}[{index}]") for index, item in enumerate(cast(list[Any], value))] if isinstance(value, Mapping): raw_mapping = {str(key): item for key, item in cast(Mapping[Any, Any], value).items()} if "type" in raw_mapping: registration = _STATE_TYPE_REGISTRY.get(str(raw_mapping["type"])) if registration is not None: - return registration.decoder(raw_mapping) - return {key: _deserialize_value(item) for key, item in raw_mapping.items()} + try: + return registration.decoder(raw_mapping) + except Exception as exc: + # Registered decoders are application extension points. Any + # ordinary failure means this payload cannot be restored by + # the active registration and should enter store recovery. + raise ValueError( + f"Failed to deserialize registered state type {registration.type_id!r} at {path}." + ) from exc + return {key: _deserialize_value(item, path=f"{path}.{key}") for key, item in raw_mapping.items()} return value def _serialize_state(state: dict[str, Any]) -> dict[str, Any]: - """Deep-serialize a state dict, converting SerializationProtocol objects to dicts.""" + """Deep-serialize a state dict using the established AgentSession contract.""" return {key: _serialize_value(value, path=f"state.{key}") for key, value in state.items()} +def _validate_durable_state_value(value: Any, *, path: str) -> None: + """Validate that serialized state can round-trip through the durable codecs.""" + if isinstance(value, float) and not math.isfinite(value): + raise ValueError(f"Session state value at {path} must be a finite float.") + if type(value) in _STATE_SCALAR_TYPES: + return + if isinstance(value, list): + for index, item in enumerate(cast(list[Any], value)): + _validate_durable_state_value(item, path=f"{path}[{index}]") + return + if isinstance(value, Mapping): + for key, item in cast(Mapping[Any, Any], value).items(): + _validate_durable_state_value(item, path=f"{path}.{key}") + return + raise TypeError( + f"Session state value at {path} has unsupported serialized type {type(value).__name__!r}; " + "call register_state_type() with a codec that returns JSON-compatible values." + ) + + def _deserialize_state(state: dict[str, Any]) -> dict[str, Any]: """Deep-deserialize a state dict, restoring SerializationProtocol objects.""" - return {k: _deserialize_value(v) for k, v in state.items()} + return {key: _deserialize_value(value, path=f"state.{key}") for key, value in state.items()} class _SessionStatePayload: @@ -366,8 +442,10 @@ def __init__(self, value: dict[str, Any]) -> None: self.value = value def serialize(self) -> dict[str, Any]: - """Serialize the wrapped state through the explicit registry.""" - return _serialize_state(self.value) + """Serialize and validate the wrapped state for durable storage.""" + serialized = _serialize_state(self.value) + _validate_durable_state_value(serialized, path="state") + return serialized class _SessionSnapshot(msgspec.Struct): @@ -377,6 +455,7 @@ class _SessionSnapshot(msgspec.Struct): session_id: str service_session_id: str | dict[str, Any] | None state: _SessionStatePayload + version: Literal["1.0"] = "1.0" def _session_snapshot_enc_hook(value: Any) -> Any: @@ -1259,9 +1338,11 @@ def session_id(self) -> str: def to_dict(self) -> dict[str, Any]: """Serialize session to a plain dict for storage/transfer. - Custom values in ``state`` must be registered with - :func:`register_state_type`. Built-in JSON scalar, list, tuple, and - mapping values are serialized recursively. + Registered custom values use their configured codecs. Unregistered + values defining ``to_dict`` retain the established dictionary behavior. + Unregistered Pydantic models are still auto-registered temporarily but + emit ``DeprecationWarning`` because cold-start restoration requires + explicit module-level registration. """ return { "type": "session", @@ -1301,14 +1382,11 @@ class SessionStore: durable storage, or distributed coordination should provide another implementation with the same async methods. - Session-store IDs are limited to 128 characters containing only ASCII - letters, digits, ``-``, and ``_``. This restricted shape is defense in - depth for storage backends; custom implementations must still use + Session IDs are opaque non-empty strings. Custom implementations are + responsible for applying any backend-specific key restrictions and must use parameterized queries and their backend's normal key-handling protections. """ - MAX_SESSION_ID_LENGTH: ClassVar[int] = 128 - def __init__(self) -> None: """Create an empty session store.""" self._sessions: dict[str, AgentSession] = {} @@ -1321,15 +1399,10 @@ def validate_session_id(session_id: str) -> None: session_id: Session-store ID to validate. Raises: - ValueError: If the ID is empty or contains characters other than - ASCII letters, digits, ``-``, and ``_``. + ValueError: If the ID is not a non-empty string. """ if not isinstance(session_id, str) or not session_id: raise ValueError("session_id must be a non-empty string") - if len(session_id) > SessionStore.MAX_SESSION_ID_LENGTH: - raise ValueError(f"session_id must be at most {SessionStore.MAX_SESSION_ID_LENGTH} characters") - if not all(character.isascii() and (character.isalnum() or character in "-_") for character in session_id): - raise ValueError("session_id must contain only ASCII letters, digits, '-' and '_'") async def get(self, session_id: str) -> AgentSession | None: """Return a copy of the stored session, or ``None`` when absent. @@ -1341,7 +1414,7 @@ async def get(self, session_id: str) -> AgentSession | None: An independent copy of the stored session, or ``None``. Raises: - ValueError: If ``session_id`` is empty or contains unsupported characters. + ValueError: If ``session_id`` is empty. """ SessionStore.validate_session_id(session_id) session = self._sessions.get(session_id) @@ -1355,7 +1428,7 @@ async def set(self, session_id: str, session: AgentSession) -> None: session: The session to store. Raises: - ValueError: If ``session_id`` is empty or contains unsupported characters. + ValueError: If ``session_id`` is empty. """ SessionStore.validate_session_id(session_id) self._sessions[session_id] = session @@ -1367,7 +1440,7 @@ async def delete(self, session_id: str) -> None: session_id: Opaque caller-selected session ID. Raises: - ValueError: If ``session_id`` is empty or contains unsupported characters. + ValueError: If ``session_id`` is empty. """ SessionStore.validate_session_id(session_id) self._sessions.pop(session_id, None) @@ -1394,13 +1467,18 @@ class FileSessionStore(SessionStore): Treat ``storage_path`` as trusted application storage, not as a secret store. Restricted store keys and resolved-path validation help prevent path traversal via ``session_id``, but they do not encrypt file contents - or coordinate concurrent updates across tasks, processes, or hosts. - Atomic replacement prevents partial writes, while concurrent writers - still use last-writer-wins semantics. Use OS-level file permissions, - trusted directories, and carefully review what session state is allowed - to be persisted. + or coordinate concurrent updates across processes or hosts. Process-local + operations are serialized per file, and atomic replacement prevents + partial writes; cross-process writers still use last-writer-wins + semantics. Use OS-level file permissions, trusted directories, and + carefully review what session state is allowed to be persisted. """ + MAX_SESSION_ID_LENGTH: ClassVar[int] = 128 + _FILE_LOCK_STRIPE_COUNT: ClassVar[int] = 64 + _FILE_OPERATION_LOCKS: ClassVar[tuple[threading.Lock, ...]] = tuple( + threading.Lock() for _ in range(_FILE_LOCK_STRIPE_COUNT) + ) _ENCODED_SESSION_PREFIX: ClassVar[str] = "~session-" def __init__( @@ -1439,28 +1517,46 @@ def __init__( async def get(self, session_id: str) -> AgentSession | None: """Load a session snapshot, or return ``None`` when it does not exist.""" file_path = self._session_file_path(session_id) + file_lock = self._session_file_lock(file_path) def _read() -> AgentSession | None: - try: - serialized = file_path.read_bytes() - except FileNotFoundError: - return None - try: - snapshot = self._decoder.decode(serialized) - session = AgentSession( - session_id=snapshot.session_id, - service_session_id=snapshot.service_session_id, - ) - session.state = snapshot.state.value - return session - except (msgspec.DecodeError, TypeError, ValueError) as exc: - raise ValueError(f"Failed to deserialize session from '{file_path}'.") from exc + with file_lock: + try: + serialized = file_path.read_bytes() + except FileNotFoundError: + return None + try: + snapshot = self._decoder.decode(serialized) + session = AgentSession( + session_id=snapshot.session_id, + service_session_id=snapshot.service_session_id, + ) + session.state = snapshot.state.value + return session + except (msgspec.DecodeError, TypeError, ValueError) as exc: + try: + quarantine_path = self._quarantine_corrupt_snapshot(file_path, serialized) + except OSError as quarantine_error: + raise ValueError( + f"Failed to deserialize session from '{file_path}', and the corrupt snapshot " + "could not be quarantined." + ) from quarantine_error + if quarantine_path is None: + raise ValueError( + f"Failed to deserialize session from '{file_path}'. " + "The snapshot changed while being read; retry." + ) from exc + raise ValueError( + f"Failed to deserialize session from '{file_path}'. The corrupt snapshot was quarantined to " + f"'{quarantine_path}'; retry to create a new session." + ) from exc return await asyncio.to_thread(_read) async def set(self, session_id: str, session: AgentSession) -> None: """Persist a session snapshot atomically.""" file_path = self._session_file_path(session_id) + file_lock = self._session_file_lock(file_path) if type(session) is not AgentSession: raise TypeError( "FileSessionStore supports AgentSession instances only; " @@ -1479,26 +1575,81 @@ async def set(self, session_id: str, session: AgentSession) -> None: serialized = self._encoder.encode(snapshot) def _write() -> None: - temp_path = file_path.with_name(f".{file_path.name}.{uuid.uuid4().hex}.tmp") - try: - temp_path.write_bytes(serialized) - os.replace(temp_path, file_path) - finally: - temp_path.unlink(missing_ok=True) + with file_lock: + temp_path = file_path.with_name(f".{file_path.name}.{uuid.uuid4().hex}.tmp") + try: + temp_path.write_bytes(serialized) + os.replace(temp_path, file_path) + finally: + temp_path.unlink(missing_ok=True) await asyncio.to_thread(_write) async def delete(self, session_id: str) -> None: """Delete a persisted session snapshot, if present.""" file_path = self._session_file_path(session_id) - await asyncio.to_thread(file_path.unlink, missing_ok=True) + file_lock = self._session_file_lock(file_path) + + def _delete() -> None: + with file_lock: + file_path.unlink(missing_ok=True) + + await asyncio.to_thread(_delete) + + @staticmethod + def validate_session_id(session_id: str) -> None: + """Validate an ID for use as a built-in file-store key. + + Args: + session_id: Session-store ID to validate. + + Raises: + ValueError: If the ID is empty, too long, or contains characters + other than ASCII letters, digits, ``-``, and ``_``. + """ + SessionStore.validate_session_id(session_id) + if len(session_id) > FileSessionStore.MAX_SESSION_ID_LENGTH: + raise ValueError(f"session_id must be at most {FileSessionStore.MAX_SESSION_ID_LENGTH} characters") + if not all(character.isascii() and (character.isalnum() or character in "-_") for character in session_id): + raise ValueError("session_id must contain only ASCII letters, digits, '-' and '_'") + + def get_session_directory(self) -> Path: + """Return the directory used for the current file-store operation. + + Subclasses may override this hook to scope operations to a child + directory. Filename encoding and containment checks remain owned by the + base implementation. + """ + return self._storage_root + + @staticmethod + def _quarantine_corrupt_snapshot(file_path: Path, serialized: bytes) -> Path | None: + """Move an unchanged corrupt snapshot aside so a retry can recover.""" + try: + current_serialized = file_path.read_bytes() + except FileNotFoundError: + return None + if current_serialized != serialized: + return None + quarantine_path = file_path.with_name(f".{file_path.name}.{uuid.uuid4().hex}.corrupt") + os.replace(file_path, quarantine_path) + return quarantine_path + + @classmethod + def _session_file_lock(cls, file_path: Path) -> threading.Lock: + """Return the process-local operation lock for a session file.""" + return cls._FILE_OPERATION_LOCKS[hash(file_path) % cls._FILE_LOCK_STRIPE_COUNT] def _session_file_path(self, session_id: str) -> Path: """Resolve the contained snapshot path for ``session_id``.""" - SessionStore.validate_session_id(session_id) + self.validate_session_id(session_id) + session_directory = self.get_session_directory().resolve() + if not session_directory.is_relative_to(self._storage_root): + raise ValueError(f"Session directory escaped storage directory: '{session_directory}'.") + session_directory.mkdir(parents=True, exist_ok=True) file_stem = _session_file_stem(session_id, encoded_prefix=self._ENCODED_SESSION_PREFIX) - file_path = (self._storage_root / f"{file_stem}{self._file_extension}").resolve() - if not file_path.is_relative_to(self._storage_root): + file_path = (session_directory / f"{file_stem}{self._file_extension}").resolve() + if not file_path.is_relative_to(session_directory): raise ValueError(f"Session path escaped storage directory: {session_id!r}") return file_path diff --git a/python/packages/core/pyproject.toml b/python/packages/core/pyproject.toml index e2ea3bdc1a..bc1e70f377 100644 --- a/python/packages/core/pyproject.toml +++ b/python/packages/core/pyproject.toml @@ -23,7 +23,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "msgspec>=0.20.0,<1", + "msgspec>=0.20.0,<0.22", "typing-extensions>=4.15.0,<5", "pydantic>=2,<3", "python-dotenv>=1,<2", diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py index 43179dfdd0..b7d17545fc 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -2,6 +2,7 @@ import asyncio import json +import math import threading import time from collections.abc import Awaitable, Callable, Mapping, Sequence @@ -681,15 +682,41 @@ class SecondState(FirstState): with pytest.raises(ValueError, match="already registered"): register_state_type(SecondState, type_id="collision_state_test") - def test_unregistered_object_fails_with_state_path(self) -> None: + def test_unregistered_to_dict_object_preserves_legacy_serialization(self) -> None: class UnregisteredState: - pass + def to_dict(self) -> dict[str, Any]: + return {"type": "unregistered_state", "value": "legacy"} session = AgentSession(session_id="unregistered") session.state["nested"] = [{"value": UnregisteredState()}] - with pytest.raises(TypeError, match=r"state\.nested\[0\]\.value.*UnregisteredState"): - session.to_dict() + assert session.to_dict()["state"]["nested"] == [{"value": {"type": "unregistered_state", "value": "legacy"}}] + + def test_unregistered_container_subclass_uses_to_dict(self) -> None: + class MappingState(dict[str, Any]): + def to_dict(self) -> dict[str, Any]: + return {"serialized": True} + + session = AgentSession(session_id="mapping") + session.state["value"] = MappingState(raw=True) + + assert session.to_dict()["state"]["value"] == {"serialized": True} + + def test_unregistered_pydantic_model_preserves_legacy_round_trip(self) -> None: + from pydantic import BaseModel + + class AutoRegisteredPydanticState(BaseModel): + value: str + + session = AgentSession(session_id="pydantic") + session.state["value"] = AutoRegisteredPydanticState(value="legacy") + + with pytest.warns(DeprecationWarning, match="Cold-start deserialization"): + serialized = session.to_dict() + restored = AgentSession.from_dict(serialized) + + assert isinstance(restored.state["value"], AutoRegisteredPydanticState) + assert restored.state["value"].value == "legacy" def test_unknown_type_tag_remains_a_raw_dict(self) -> None: session = AgentSession.from_dict({ @@ -722,12 +749,17 @@ class TextState: assert restored.state["message"].text == "hello" @pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf")]) - def test_non_finite_float_is_rejected(self, value: float) -> None: + def test_non_finite_float_preserves_legacy_dictionary_serialization(self, value: float) -> None: session = AgentSession(session_id="float") session.state["value"] = value - with pytest.raises(ValueError, match="finite float"): - session.to_dict() + serialized = session.to_dict() + + restored = serialized["state"]["value"] + if math.isnan(value): + assert math.isnan(restored) + else: + assert restored == value class TestSessionStore: @@ -777,30 +809,26 @@ async def test_delete_forgets_session(self) -> None: assert await store.get("session-1") is None - @pytest.mark.parametrize( - "session_id", - [ - "", - "two words", - "tenant/user", - "tenant:user", - "session.id", - "' OR 1=1 --", - "café", - "line\nbreak", - "a" * 129, - ], - ) - async def test_invalid_session_id_raises(self, session_id: str) -> None: + @pytest.mark.parametrize("session_id", ["two words", "tenant/user", "session.id", "café"]) + async def test_accepts_opaque_non_empty_session_id(self, session_id: str) -> None: + store = SessionStore() + session = AgentSession() + + await store.set(session_id, session) + assert await store.get(session_id) is not None + await store.delete(session_id) + assert await store.get(session_id) is None + + async def test_empty_session_id_raises(self) -> None: store = SessionStore() session = AgentSession() with pytest.raises(ValueError, match="session_id"): - await store.get(session_id) + await store.get("") with pytest.raises(ValueError, match="session_id"): - await store.set(session_id, session) + await store.set("", session) with pytest.raises(ValueError, match="session_id"): - await store.delete(session_id) + await store.delete("") class TestFileSessionStore: @@ -827,6 +855,7 @@ async def test_round_trips_session_across_store_instances(self, tmp_path: Path) assert len(files) == 1 assert files[0].parent == tmp_path assert files[0].suffix == ".json" + assert json.loads(files[0].read_bytes())["version"] == "1.0" async def test_round_trips_binary_messagepack_session(self, tmp_path: Path) -> None: store = FileSessionStore(tmp_path, serialization_format="msgpack") @@ -842,7 +871,9 @@ async def test_round_trips_binary_messagepack_session(self, tmp_path: Path) -> N files = await asyncio.to_thread(lambda: list(tmp_path.iterdir())) assert len(files) == 1 assert files[0].suffix == ".msgpack" - assert not (await asyncio.to_thread(files[0].read_bytes)).startswith(b"{") + serialized = await asyncio.to_thread(files[0].read_bytes) + assert not serialized.startswith(b"{") + assert msgspec.msgpack.decode(serialized)["version"] == "1.0" async def test_rejects_custom_agent_session_subclass(self, tmp_path: Path) -> None: class CustomAgentSession(AgentSession): @@ -853,6 +884,51 @@ class CustomAgentSession(AgentSession): with pytest.raises(TypeError, match="AgentSession instances only"): await store.set("custom-session", CustomAgentSession(session_id="custom-session")) + async def test_unregistered_to_dict_state_preserves_legacy_durable_serialization(self, tmp_path: Path) -> None: + class UnregisteredState: + def to_dict(self) -> dict[str, Any]: + return {"type": "unregistered_state", "value": "legacy"} + + session = AgentSession(session_id="unregistered") + session.state["nested"] = [{"value": UnregisteredState()}] + + store = FileSessionStore(tmp_path) + await store.set("unregistered", session) + restored = await store.get("unregistered") + + assert restored is not None + assert restored.state["nested"] == [{"value": {"type": "unregistered_state", "value": "legacy"}}] + + async def test_unregistered_pydantic_state_warns_for_durable_serialization(self, tmp_path: Path) -> None: + from pydantic import BaseModel + + class DurableAutoRegisteredPydanticState(BaseModel): + value: str + + session = AgentSession(session_id="pydantic") + session.state["value"] = DurableAutoRegisteredPydanticState(value="legacy") + + with pytest.warns(DeprecationWarning, match="Cold-start deserialization"): + await FileSessionStore(tmp_path).set("pydantic", session) + + async def test_unsupported_state_object_is_rejected_for_durable_storage(self, tmp_path: Path) -> None: + class UnsupportedState: + pass + + session = AgentSession(session_id="unsupported") + session.state["nested"] = [{"value": UnsupportedState()}] + + with pytest.raises(TypeError, match=r"state\.nested\[0\]\.value.*UnsupportedState"): + await FileSessionStore(tmp_path).set("unsupported", session) + + @pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf")]) + async def test_non_finite_state_is_rejected_for_durable_storage(self, tmp_path: Path, value: float) -> None: + session = AgentSession(session_id="float") + session.state["value"] = value + + with pytest.raises(ValueError, match="finite float"): + await FileSessionStore(tmp_path).set("float", session) + async def test_missing_and_deleted_sessions_return_none(self, tmp_path: Path) -> None: store = FileSessionStore(tmp_path) assert await store.get("session-1") is None @@ -874,15 +950,101 @@ async def test_set_replaces_atomically_without_leaving_temp_files(self, tmp_path temp_files = await asyncio.to_thread(lambda: [*tmp_path.glob("*.tmp"), *tmp_path.glob(".*.tmp")]) assert not temp_files - async def test_corrupt_session_file_raises(self, tmp_path: Path) -> None: + async def test_corrupt_session_file_is_quarantined_and_retry_returns_none(self, tmp_path: Path) -> None: store = FileSessionStore(tmp_path) await store.set("session-1", AgentSession(session_id="session-1")) session_file = await asyncio.to_thread(lambda: next(tmp_path.iterdir())) await asyncio.to_thread(session_file.write_text, "{not-json", encoding="utf-8") - with pytest.raises(ValueError, match="Failed to deserialize session"): + with pytest.raises(ValueError, match="quarantined.*retry"): + await store.get("session-1") + + assert not session_file.exists() + quarantined_files = await asyncio.to_thread(lambda: list(tmp_path.glob(".*.corrupt"))) + assert len(quarantined_files) == 1 + assert await store.get("session-1") is None + + async def test_invalid_registered_payload_is_quarantined(self, tmp_path: Path) -> None: + store = FileSessionStore(tmp_path) + session_file = store._session_file_path("session-1") + session_file.write_text( + json.dumps({ + "type": "session", + "version": "1.0", + "session_id": "session-1", + "service_session_id": None, + "state": {"value": {"type": "registered_session_state"}}, + }), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="quarantined.*retry"): await store.get("session-1") + assert not session_file.exists() + quarantined_files = await asyncio.to_thread(lambda: list(tmp_path.glob(".*.corrupt"))) + assert len(quarantined_files) == 1 + + async def test_registered_decoder_attribute_error_is_quarantined(self, tmp_path: Path) -> None: + store = FileSessionStore(tmp_path) + session_file = store._session_file_path("session-1") + session_file.write_text( + json.dumps({ + "type": "session", + "version": "1.0", + "session_id": "session-1", + "service_session_id": None, + "state": {"value": {"type": "message", "role": "user", "contents": [123]}}, + }), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="quarantined.*retry"): + await store.get("session-1") + + assert not session_file.exists() + quarantined_files = await asyncio.to_thread(lambda: list(tmp_path.glob(".*.corrupt"))) + assert len(quarantined_files) == 1 + + async def test_corrupt_quarantine_does_not_remove_concurrent_replacement( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + reader = FileSessionStore(tmp_path) + writer = FileSessionStore(tmp_path) + await reader.set("session-1", AgentSession(session_id="old")) + session_file = reader._session_file_path("session-1") + await asyncio.to_thread(session_file.write_text, "{not-json", encoding="utf-8") + quarantine_started = threading.Event() + release_quarantine = threading.Event() + original_quarantine = FileSessionStore._quarantine_corrupt_snapshot + + def blocking_quarantine(file_path: Path, serialized: bytes) -> Path | None: + quarantine_started.set() + if not release_quarantine.wait(timeout=5): + raise TimeoutError("Timed out waiting to release corrupt snapshot quarantine.") + return original_quarantine(file_path, serialized) + + monkeypatch.setattr( + FileSessionStore, + "_quarantine_corrupt_snapshot", + staticmethod(blocking_quarantine), + ) + read_task = asyncio.create_task(reader.get("session-1")) + assert await asyncio.to_thread(quarantine_started.wait, 5) + write_task = asyncio.create_task(writer.set("session-1", AgentSession(session_id="replacement"))) + await asyncio.sleep(0) + release_quarantine.set() + + with pytest.raises(ValueError, match="quarantined"): + await read_task + await write_task + + restored = await reader.get("session-1") + assert restored is not None + assert restored.session_id == "replacement" + @pytest.mark.parametrize("session_id", ["", "two words", "tenant/user", "session.id", "café"]) async def test_invalid_session_id_raises(self, tmp_path: Path, session_id: str) -> None: store = FileSessionStore(tmp_path) @@ -895,6 +1057,16 @@ async def test_invalid_session_id_raises(self, tmp_path: Path, session_id: str) with pytest.raises(ValueError, match="session_id"): await store.delete(session_id) + @pytest.mark.parametrize("session_id", ["NUL.txt", "COM¹", "LPT².log"]) + async def test_reserved_windows_filename_is_encoded(self, tmp_path: Path, session_id: str) -> None: + provider = FileHistoryProvider(tmp_path) + + await provider.save_messages(session_id, [Message(role="user", contents=["hello"])]) + + session_file = provider._session_file_path(session_id) + assert session_file.name.startswith("~session-") + assert session_file.is_file() + # --------------------------------------------------------------------------- # InMemoryHistoryProvider tests @@ -1085,6 +1257,30 @@ async def test_stores_and_loads_messages(self, tmp_path: Path) -> None: assert loaded[0].text == "hello" assert loaded[1].text == "hi there" + @pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf")]) + async def test_preserves_non_finite_json_values(self, tmp_path: Path, value: float) -> None: + provider = FileHistoryProvider(tmp_path) + message = Message(role="user", contents=["hello"], additional_properties={"value": value}) + + await provider.save_messages("non-finite", [message]) + loaded = await provider.get_messages("non-finite") + + restored = loaded[0].additional_properties["value"] + if math.isnan(value): + assert math.isnan(restored) + else: + assert restored == value + + async def test_reads_existing_stdlib_non_finite_json(self, tmp_path: Path) -> None: + provider = FileHistoryProvider(tmp_path) + message = Message(role="user", contents=["hello"], additional_properties={"value": float("inf")}) + session_file = provider._session_file_path("legacy") + await asyncio.to_thread(session_file.write_text, f"{json.dumps(message.to_dict())}\n", encoding="utf-8") + + loaded = await provider.get_messages("legacy") + + assert loaded[0].additional_properties["value"] == float("inf") + def test_creates_storage_directory(self, tmp_path: Path) -> None: nested_path = tmp_path / "nested" / "history" provider = FileHistoryProvider(nested_path) diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index d80c551e90..9d1d04d98c 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -334,7 +334,7 @@ def _resolve_session_conversation_key(request: CreateResponse, context: Response def _conversation_object_id(conversation_key: str) -> str: """Return a restricted store/file ID derived from the Responses conversation object.""" try: - SessionStore.validate_session_id(conversation_key) + FileSessionStore.validate_session_id(conversation_key) except ValueError: return f"conversation_{hashlib.sha256(conversation_key.encode('utf-8')).hexdigest()}" return conversation_key diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_isolation.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_isolation.py index 8efb57afaa..726a4de737 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_isolation.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_isolation.py @@ -11,7 +11,7 @@ from pathlib import Path from typing import Literal, TypeAlias -from agent_framework import FileSessionStore, SessionStore +from agent_framework import FileSessionStore from azure.ai.agentserver.responses import ResponseContext from azure.ai.agentserver.responses.models import CreateResponse @@ -59,17 +59,12 @@ def use_isolation(self, isolation: ResolvedSessionIsolation) -> Generator[None]: finally: self._current_directory.reset(token) - def _session_file_path(self, session_id: str) -> Path: - """Resolve the session file beneath the currently bound user directory.""" - SessionStore.validate_session_id(session_id) - storage_root = self.storage_path.resolve() + def get_session_directory(self) -> Path: + """Return the currently bound user directory.""" directory_segment = self._current_directory.get() - if directory_segment is not None: - storage_root = (storage_root / directory_segment).resolve() - if not storage_root.is_relative_to(self.storage_path.resolve()): - raise ValueError(f"Session isolation path escaped storage directory: {directory_segment!r}") - storage_root.mkdir(parents=True, exist_ok=True) - return storage_root / f"{session_id}{self._file_extension}" + if directory_segment is None: + return self.storage_path + return self.storage_path / directory_segment def platform_session_isolation_key_resolver( diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index c05795b7af..53ee26c847 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -416,7 +416,10 @@ async def resolve_for(user_id: str) -> ResolvedSessionIsolation: assert second.key == "user-2" assert first.fingerprint != second.fingerprint - @pytest.mark.parametrize("key", ["../escape", "user/name", "user:name", "trailing."]) + @pytest.mark.parametrize( + "key", + [".", "..", "../escape", "user/name", r"user\name", "user:name", "line\nbreak", "trailing."], + ) async def test_invalid_isolation_key_is_rejected(self, key: str) -> None: request = CreateResponse(model="m", input="hi") context = ResponseContext(response_id="response-1", mode_flags=MagicMock()) @@ -548,6 +551,51 @@ async def test_file_store_uses_per_user_child_directory_and_preserves_format(sel assert (tmp_path / "user-user-A" / "conversation-1.msgpack").is_file() assert (tmp_path / "user-user-B" / "conversation-1.msgpack").is_file() + async def test_scoped_file_store_reuses_base_filename_safety(self, tmp_path: Path) -> None: + store = IsolationKeyScopedFileSessionStore(tmp_path) + isolation = ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64) + + with store.use_isolation(isolation): + await store.set("CON", AgentSession(session_id="conversation-1")) + + assert not (tmp_path / "user-user-A" / "CON.json").exists() + assert len(list((tmp_path / "user-user-A").glob("~session-*.json"))) == 1 + + async def test_scoped_file_store_rejects_symlinked_session_leaf(self, tmp_path: Path) -> None: + store = IsolationKeyScopedFileSessionStore(tmp_path) + isolation = ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64) + user_directory = tmp_path / "user-user-A" + user_directory.mkdir() + outside_file = tmp_path / "outside.json" + outside_file.write_text("outside", encoding="utf-8") + try: + (user_directory / "conversation-1.json").symlink_to(outside_file) + except OSError as exc: + pytest.skip(f"Symlinks are not available: {exc}") + + with store.use_isolation(isolation), pytest.raises(ValueError, match="escaped storage directory"): + await store.get("conversation-1") + + async def test_scoped_file_store_rejects_symlinked_isolation_directory(self, tmp_path: Path) -> None: + store = IsolationKeyScopedFileSessionStore(tmp_path) + isolation = ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64) + outside_directory = tmp_path.parent / f"{tmp_path.name}-outside" + outside_directory.mkdir() + try: + (tmp_path / "user-user-A").symlink_to(outside_directory, target_is_directory=True) + except OSError as exc: + pytest.skip(f"Symlinks are not available: {exc}") + + with store.use_isolation(isolation), pytest.raises(ValueError, match="Session directory escaped"): + await store.get("conversation-1") + + async def test_scoped_file_store_rejects_traversal_directory_segment(self, tmp_path: Path) -> None: + store = IsolationKeyScopedFileSessionStore(tmp_path) + bypassed_resolver = ResolvedSessionIsolation("user-A", "../../Users", "a" * 64) + + with store.use_isolation(bypassed_resolver), pytest.raises(ValueError, match="Session directory escaped"): + await store.get("conversation-1") + def test_session_isolation_stamp_rejects_mismatch(self) -> None: session = AgentSession(session_id="conversation-1") isolation_a = ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64) @@ -649,6 +697,26 @@ async def run_with_state(*args: Any, **kwargs: Any) -> AgentResponse: assert second.status_code == 200 assert seen_counts == [1, 2] + async def test_corrupt_file_session_fails_once_and_retry_starts_clean(self, tmp_path: Path) -> None: + agent = _make_agent( + response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("recovered")])]) + ) + server = _make_server(agent, session_store=IsolationKeyScopedFileSessionStore(tmp_path)) + corrupt_file = tmp_path / "conversation-1.json" + corrupt_file.write_text("{not-json", encoding="utf-8") + + first = await _post(server, conversation_id="conversation-1") + + assert first.json()["status"] == "failed" + assert not corrupt_file.exists() + quarantined_files = await asyncio.to_thread(lambda: list(tmp_path.glob(".*.corrupt"))) + assert len(quarantined_files) == 1 + + second = await _post(server, conversation_id="conversation-1") + + assert second.json()["status"] == "completed" + agent.create_session.assert_called_once_with(session_id="conversation-1") + async def test_streaming_run_saves_final_session_state(self) -> None: store = SessionStore() agent = _make_agent() diff --git a/python/packages/hosting/README.md b/python/packages/hosting/README.md index b0dd74bbcf..05f42da4bb 100644 --- a/python/packages/hosting/README.md +++ b/python/packages/hosting/README.md @@ -21,14 +21,15 @@ checkpointing should use the existing `CheckpointStorage` abstraction directly; if an app needs per-session resume, keep a small app-owned cursor such as `session_id -> checkpoint_id`. -Session-store IDs are limited to 128 characters containing ASCII letters, -digits, `-`, and `_`. This keeps protocol-derived IDs safe to pass through -common storage backends, but it does not replace parameterized queries or -backend-specific validation in custom stores. +`SessionStore` accepts opaque non-empty IDs so custom database and Redis +implementations can use their native key contracts. `FileSessionStore` limits +its direct keys to 128 ASCII letters, digits, `-`, and `_`. `AgentState` passes +the app-selected ID to the configured store unchanged; each store implementation +owns any validation or normalization required by its backend. This does not +replace parameterized queries or backend-specific validation in custom stores. -`FileSessionStore` uses msgspec JSON. Custom objects placed in -`AgentSession.state` must be registered explicitly before the session is saved -or restored: +`FileSessionStore` uses msgspec JSON. Register custom objects placed in +`AgentSession.state` explicitly before sessions are saved or restored: ```python from agent_framework import register_state_type @@ -46,7 +47,9 @@ Classes with `to_dict()` / `from_dict()` methods and explicitly registered Pydantic models receive default codecs. Other classes can provide `encoder=` and `decoder=` callbacks. Keep registration at module level so importing the module prepares cold-start session restoration before its context provider is -instantiated. +instantiated. Unregistered Pydantic models still use legacy same-process +auto-registration for now, but emit `DeprecationWarning`; cold-start +deserialization is not guaranteed on that path. JSON is the default file format. Use MessagePack for a compact binary file: diff --git a/python/packages/hosting/agent_framework_hosting/_state.py b/python/packages/hosting/agent_framework_hosting/_state.py index 91638607e5..9c86f7515a 100644 --- a/python/packages/hosting/agent_framework_hosting/_state.py +++ b/python/packages/hosting/agent_framework_hosting/_state.py @@ -160,14 +160,12 @@ async def get_or_create_session(self, session_id: str) -> AgentSession: """Return the session for ``session_id``, creating and storing one if missing. Args: - session_id: App-selected session ID containing only ASCII letters, - digits, ``-``, and ``_``. + session_id: Opaque app-selected session ID. Returns: An independent working copy of the stored or newly created ``AgentSession``. """ - SessionStore.validate_session_id(session_id) session_lock = self._session_locks.setdefault(session_id, asyncio.Lock()) async with session_lock: session = await self._session_store.get(session_id) @@ -182,11 +180,9 @@ async def set_session(self, session_id: str, session: AgentSession) -> None: """Store ``session`` under ``session_id`` in this state's session store. Args: - session_id: App-selected session ID containing only ASCII letters, - digits, ``-``, and ``_``. + session_id: Opaque app-selected session ID. session: Session to store. """ - SessionStore.validate_session_id(session_id) await self._session_store.set(session_id, session) diff --git a/python/packages/hosting/tests/hosting/test_state.py b/python/packages/hosting/tests/hosting/test_state.py index a163bd598c..f4c991fca5 100644 --- a/python/packages/hosting/tests/hosting/test_state.py +++ b/python/packages/hosting/tests/hosting/test_state.py @@ -200,12 +200,37 @@ async def test_get_or_create_session_creates_and_stores_once(self) -> None: assert len(agent.created_sessions) == 1 @pytest.mark.parametrize("session_id", ["two words", "tenant/user", "tenant:conversation", "' OR 1=1 --"]) - async def test_get_or_create_session_rejects_invalid_store_id(self, session_id: str) -> None: + async def test_get_or_create_session_passes_opaque_id_to_store(self, session_id: str) -> None: + class _RecordingSessionStore(SessionStore): + def __init__(self) -> None: + super().__init__() + self.keys: list[str] = [] + + async def get(self, session_id: str) -> AgentSession | None: + self.keys.append(session_id) + return await super().get(session_id) + + async def set(self, session_id: str, session: AgentSession) -> None: + self.keys.append(session_id) + await super().set(session_id, session) + + agent = _FakeAgent() + store = _RecordingSessionStore() + state = AgentState(agent, session_store=store) + + session = await state.get_or_create_session(session_id) + + assert session.session_id == session_id + assert len(agent.created_sessions) == 1 + assert len(set(store.keys)) == 1 + assert store.keys[0] == session_id + + async def test_get_or_create_session_rejects_empty_id(self) -> None: agent = _FakeAgent() state = AgentState(agent) - with pytest.raises(ValueError, match="ASCII letters"): - await state.get_or_create_session(session_id) + with pytest.raises(ValueError, match="non-empty"): + await state.get_or_create_session("") assert agent.created_sessions == [] diff --git a/python/scripts/session_serialization_benchmark.py b/python/scripts/session_serialization_benchmark.py new file mode 100644 index 0000000000..1eaa9a718a --- /dev/null +++ b/python/scripts/session_serialization_benchmark.py @@ -0,0 +1,582 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Compare AgentSession serialization formats using realistic Agent Framework objects. + +Run from ``python/``: + + uv run --with orjson python scripts/session_serialization_benchmark.py + +The benchmark compares: + +- Standard-library JSON +- orjson +- Pydantic `model_dump_json` / `model_validate_json` +- msgspec JSON +- msgspec MessagePack (binary) +- An AgentSession-shaped msgspec Struct using JSON +- An AgentSession-shaped msgspec Struct using MessagePack + +The first five measurements include AgentSession ``to_dict`` / ``from_dict`` +conversion. The Struct variants instead map session fields directly and route +only the dynamic state dictionary through the same registry helpers. JSON and +MessagePack payloads are written to disk to report actual file size and cached +filesystem round-trip latency. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import math +import statistics +import tempfile +import time +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +import msgspec +import orjson +from agent_framework import ( + AgentSession, + Content, + InMemoryHistoryProvider, + Message, + register_state_type, +) +from agent_framework._sessions import ( # pyright: ignore[reportPrivateUsage] + _deserialize_state, + _serialize_state, + _validate_durable_state_value, +) +from pydantic import BaseModel + + +@dataclass(slots=True) +class BenchmarkClassState: + """Representative application-defined state class.""" + + item_id: int + label: str + scores: list[float] + attributes: dict[str, str] + + TYPE = "benchmark_class_state" + + def to_dict(self) -> dict[str, Any]: + """Serialize this state object.""" + return { + "item_id": self.item_id, + "label": self.label, + "scores": self.scores, + "attributes": self.attributes, + } + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> BenchmarkClassState: + """Restore this state object.""" + return cls( + item_id=int(value["item_id"]), + label=str(value["label"]), + scores=[float(score) for score in value["scores"]], + attributes={str(key): str(item) for key, item in value["attributes"].items()}, + ) + + +class BenchmarkProfileState(BaseModel): + """Representative explicitly registered Pydantic state.""" + + user_id: str + preferences: dict[str, str] + counters: list[int] + + +class PydanticSessionSnapshot(BaseModel): + """Typed Pydantic representation used by the Pydantic benchmark.""" + + type: Literal["session"] + session_id: str + service_session_id: str | dict[str, Any] | None = None + state: dict[str, Any] + + +class StructStatePayload: + """Opaque wrapper that forces msgspec to invoke the state registry hook.""" + + __slots__ = ("value",) + + def __init__(self, value: dict[str, Any]) -> None: + self.value = value + + +class StructAgentSession(msgspec.Struct): + """AgentSession-shaped msgspec Struct used by the direct benchmark.""" + + session_id: str + service_session_id: str | dict[str, Any] | None + state: StructStatePayload + version: Literal["1.0"] = "1.0" + + +register_state_type(BenchmarkClassState) +register_state_type(BenchmarkProfileState, type_id="benchmark_profile_state") + + +@dataclass(frozen=True, slots=True) +class Codec: + """One benchmarked serialization codec.""" + + name: str + suffix: str + encode: Callable[[AgentSession], bytes] + decode: Callable[[bytes], AgentSession] + + +@dataclass(frozen=True, slots=True) +class BenchmarkResult: + """Collected timing and size metrics for one codec.""" + + codec: str + file_size: int + encode_median_ms: float + encode_p95_ms: float + decode_median_ms: float + decode_p95_ms: float + roundtrip_median_ms: float + roundtrip_p95_ms: float + disk_roundtrip_median_ms: float + disk_roundtrip_p95_ms: float + + +def _stdlib_json_encode(value: dict[str, Any]) -> bytes: + return json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + + +def _stdlib_json_decode(value: bytes) -> Any: + return json.loads(value) + + +def _pydantic_json_encode(value: dict[str, Any]) -> bytes: + snapshot = PydanticSessionSnapshot.model_validate(value) + return snapshot.model_dump_json().encode("utf-8") + + +def _pydantic_json_decode(value: bytes) -> Any: + return PydanticSessionSnapshot.model_validate_json(value).model_dump() + + +def _encode_via_dict( + session: AgentSession, + encoder: Callable[[dict[str, Any]], bytes], +) -> bytes: + return encoder(session.to_dict()) + + +def _decode_via_dict( + payload: bytes, + decoder: Callable[[bytes], Any], + *, + codec_name: str, +) -> AgentSession: + decoded = decoder(payload) + if not isinstance(decoded, Mapping): + raise TypeError(f"{codec_name} decoded the session to {type(decoded).__name__}, not a mapping") + return AgentSession.from_dict(dict(decoded)) + + +def _struct_enc_hook(value: Any) -> Any: + if isinstance(value, StructStatePayload): + serialized = _serialize_state(value.value) + _validate_durable_state_value(serialized, path="state") + return serialized + raise NotImplementedError(f"Unsupported type: {type(value).__name__}") + + +def _struct_dec_hook(target_type: type[Any], value: Any) -> Any: + if target_type is StructStatePayload: + if not isinstance(value, Mapping): + raise TypeError("Struct state payload must decode to a mapping") + return StructStatePayload(_deserialize_state(dict(value))) + raise NotImplementedError(f"Unsupported type: {target_type.__name__}") + + +STRUCT_JSON_ENCODER = msgspec.json.Encoder(enc_hook=_struct_enc_hook) +STRUCT_JSON_DECODER = msgspec.json.Decoder(StructAgentSession, dec_hook=_struct_dec_hook) +STRUCT_MSGPACK_ENCODER = msgspec.msgpack.Encoder(enc_hook=_struct_enc_hook) +STRUCT_MSGPACK_DECODER = msgspec.msgpack.Decoder(StructAgentSession, dec_hook=_struct_dec_hook) + + +def _to_struct(session: AgentSession) -> StructAgentSession: + service_session_id = session.service_session_id + return StructAgentSession( + session_id=session.session_id, + service_session_id=dict(service_session_id) if isinstance(service_session_id, Mapping) else service_session_id, + state=StructStatePayload(session.state), + ) + + +def _from_struct(snapshot: StructAgentSession) -> AgentSession: + session = AgentSession( + session_id=snapshot.session_id, + service_session_id=snapshot.service_session_id, + ) + session.state = snapshot.state.value + return session + + +def _dict_codec( + *, + name: str, + suffix: str, + encoder: Callable[[dict[str, Any]], bytes], + decoder: Callable[[bytes], Any], +) -> Codec: + return Codec( + name=name, + suffix=suffix, + encode=lambda session: _encode_via_dict(session, encoder), + decode=lambda payload: _decode_via_dict(payload, decoder, codec_name=name), + ) + + +CODECS = ( + _dict_codec( + name="stdlib-json", + suffix=".stdlib.json", + encoder=_stdlib_json_encode, + decoder=_stdlib_json_decode, + ), + _dict_codec( + name="orjson", + suffix=".orjson.json", + encoder=orjson.dumps, + decoder=orjson.loads, + ), + _dict_codec( + name="pydantic-json", + suffix=".pydantic.json", + encoder=_pydantic_json_encode, + decoder=_pydantic_json_decode, + ), + _dict_codec( + name="msgspec-json", + suffix=".msgspec.json", + encoder=msgspec.json.encode, + decoder=msgspec.json.decode, + ), + _dict_codec( + name="msgspec-binary", + suffix=".msgspec.msgpack", + encoder=msgspec.msgpack.encode, + decoder=msgspec.msgpack.decode, + ), + Codec( + name="agent-struct-json", + suffix=".agent-struct.json", + encode=lambda session: STRUCT_JSON_ENCODER.encode(_to_struct(session)), + decode=lambda payload: _from_struct(STRUCT_JSON_DECODER.decode(payload)), + ), + Codec( + name="agent-struct-binary", + suffix=".agent-struct.msgpack", + encode=lambda session: STRUCT_MSGPACK_ENCODER.encode(_to_struct(session)), + decode=lambda payload: _from_struct(STRUCT_MSGPACK_DECODER.decode(payload)), + ), +) + + +def _build_messages(count: int, text_bytes: int) -> list[Message]: + """Build a varied conversation dominated by Message objects.""" + padding = "x" * max(0, text_bytes - 80) + messages: list[Message] = [] + for index in range(count): + role = "user" if index % 2 == 0 else "assistant" + contents = [ + Content.from_text( + text=( + f"Message {index}: benchmark conversation text with Unicode café 東京. " + f"Payload={padding}" + ) + ) + ] + if index % 20 == 5: + contents.append( + Content.from_function_call( + call_id=f"call_{index}", + name="lookup", + arguments={"query": f"item-{index}", "limit": 5}, + ) + ) + elif index % 20 == 6: + contents.append( + Content.from_function_result( + call_id=f"call_{index - 1}", + result={"items": [f"result-{index}-{item}" for item in range(5)]}, + ) + ) + messages.append( + Message( + role=role, + contents=contents, + author_name=f"participant-{index % 7}", + additional_properties={ + "sequence": index, + "trace": {"span": f"span-{index}", "sampled": index % 3 == 0}, + }, + ) + ) + return messages + + +async def build_large_session( + *, + message_count: int, + class_state_count: int, + text_bytes: int, +) -> AgentSession: + """Build a large session through InMemoryHistoryProvider.""" + session = AgentSession( + session_id="serialization-benchmark", + service_session_id={ + "conversation_id": "benchmark-conversation", + "response_id": "benchmark-response", + }, + ) + history = InMemoryHistoryProvider() + history_state = session.state.setdefault(history.source_id, {}) + await history.save_messages( + session.session_id, + _build_messages(message_count, text_bytes), + state=history_state, + ) + + session.state["plain"] = { + "flags": [True, False, None], + "numbers": list(range(class_state_count)), + "nested": { + f"key_{index}": { + "value": index, + "text": f"plain-state-{index}", + "tags": [f"tag-{item}" for item in range(5)], + } + for index in range(class_state_count) + }, + } + session.state["classes"] = [ + BenchmarkClassState( + item_id=index, + label=f"class-state-{index}", + scores=[index / 10, index / 20, index / 30], + attributes={ + "category": f"category-{index % 11}", + "partition": f"partition-{index % 17}", + }, + ) + for index in range(class_state_count) + ] + session.state["profiles"] = [ + BenchmarkProfileState( + user_id=f"user-{index}", + preferences={ + "language": "en", + "theme": "dark" if index % 2 else "light", + "timezone": f"UTC+{index % 12}", + }, + counters=[index, index * 2, index * 3], + ) + for index in range(max(1, class_state_count // 10)) + ] + return session + + +def _percentile(values: list[int], percentile: float) -> float: + ordered = sorted(values) + index = max(0, math.ceil(len(ordered) * percentile) - 1) + return ordered[index] / 1_000_000 + + +def _median_ms(values: list[int]) -> float: + return statistics.median(values) / 1_000_000 + + +def _time_ns(function: Callable[[], Any], iterations: int) -> list[int]: + timings: list[int] = [] + for _ in range(iterations): + started = time.perf_counter_ns() + function() + timings.append(time.perf_counter_ns() - started) + return timings + + +def _verify_roundtrip(original: AgentSession, restored: AgentSession) -> None: + """Verify that framework and custom objects were reconstructed.""" + original_messages = original.state[InMemoryHistoryProvider.DEFAULT_SOURCE_ID]["messages"] + restored_messages = restored.state[InMemoryHistoryProvider.DEFAULT_SOURCE_ID]["messages"] + if len(restored_messages) != len(original_messages): + raise AssertionError("Message count changed during round-trip") + if not isinstance(restored_messages[0], Message): + raise AssertionError("Message objects were not reconstructed") + if not isinstance(restored.state["classes"][0], BenchmarkClassState): + raise AssertionError("Custom class state was not reconstructed") + if not isinstance(restored.state["profiles"][0], BenchmarkProfileState): + raise AssertionError("Pydantic state was not reconstructed") + + +def benchmark_codec( + codec: Codec, + session: AgentSession, + *, + output_directory: Path, + warmups: int, + iterations: int, +) -> BenchmarkResult: + """Benchmark one codec and write its representative payload to disk.""" + payload = codec.encode(session) + restored = codec.decode(payload) + _verify_roundtrip(session, restored) + + for _ in range(warmups): + codec.encode(session) + codec.decode(payload) + codec.decode(codec.encode(session)) + + encode_timings = _time_ns(lambda: codec.encode(session), iterations) + decode_timings = _time_ns(lambda: codec.decode(payload), iterations) + roundtrip_timings = _time_ns(lambda: codec.decode(codec.encode(session)), iterations) + + output_path = output_directory / f"session{codec.suffix}" + output_path.write_bytes(payload) + + def disk_roundtrip() -> None: + current_payload = codec.encode(session) + output_path.write_bytes(current_payload) + codec.decode(output_path.read_bytes()) + + disk_roundtrip_timings = _time_ns(disk_roundtrip, iterations) + return BenchmarkResult( + codec=codec.name, + file_size=output_path.stat().st_size, + encode_median_ms=_median_ms(encode_timings), + encode_p95_ms=_percentile(encode_timings, 0.95), + decode_median_ms=_median_ms(decode_timings), + decode_p95_ms=_percentile(decode_timings, 0.95), + roundtrip_median_ms=_median_ms(roundtrip_timings), + roundtrip_p95_ms=_percentile(roundtrip_timings, 0.95), + disk_roundtrip_median_ms=_median_ms(disk_roundtrip_timings), + disk_roundtrip_p95_ms=_percentile(disk_roundtrip_timings, 0.95), + ) + + +def _format_bytes(value: int) -> str: + if value < 1024: + return f"{value} B" + if value < 1024 * 1024: + return f"{value / 1024:.1f} KiB" + return f"{value / (1024 * 1024):.2f} MiB" + + +def print_results(results: list[BenchmarkResult]) -> None: + """Print an aligned summary table and relative size ratios.""" + headers = ( + "codec", + "size", + "encode med/p95 ms", + "decode med/p95 ms", + "roundtrip med/p95 ms", + "disk med/p95 ms", + ) + rows = [ + ( + result.codec, + _format_bytes(result.file_size), + f"{result.encode_median_ms:.3f}/{result.encode_p95_ms:.3f}", + f"{result.decode_median_ms:.3f}/{result.decode_p95_ms:.3f}", + f"{result.roundtrip_median_ms:.3f}/{result.roundtrip_p95_ms:.3f}", + f"{result.disk_roundtrip_median_ms:.3f}/{result.disk_roundtrip_p95_ms:.3f}", + ) + for result in results + ] + widths = [ + max(len(headers[index]), *(len(row[index]) for row in rows)) + for index in range(len(headers)) + ] + + def render(row: tuple[str, ...]) -> str: + return " | ".join(value.ljust(widths[index]) for index, value in enumerate(row)) + + print(render(headers)) + print("-+-".join("-" * width for width in widths)) + for row in rows: + print(render(row)) + + baseline = results[0].file_size + print("\nFile size relative to stdlib JSON:") + for result in results: + print(f" {result.codec:<16} {result.file_size / baseline:>7.3f}x") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--messages", type=int, default=2_000, help="Number of Message objects in history.") + parser.add_argument("--class-state", type=int, default=500, help="Number of custom class state objects.") + parser.add_argument("--text-bytes", type=int, default=512, help="Approximate text payload per Message.") + parser.add_argument("--iterations", type=int, default=25, help="Measured iterations per operation.") + parser.add_argument("--warmups", type=int, default=5, help="Warmup iterations per codec.") + parser.add_argument( + "--output-dir", + type=Path, + help="Keep representative payload files in this directory instead of a temporary directory.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + session = asyncio.run( + build_large_session( + message_count=args.messages, + class_state_count=args.class_state, + text_bytes=args.text_bytes, + ) + ) + history_count = len(session.state[InMemoryHistoryProvider.DEFAULT_SOURCE_ID]["messages"]) + print( + f"Session: {history_count} messages, {args.class_state} class objects, " + f"{len(session.state['profiles'])} Pydantic objects" + ) + print(f"Iterations: {args.iterations} measured, {args.warmups} warmups\n") + + if args.output_dir is not None: + args.output_dir.mkdir(parents=True, exist_ok=True) + results = [ + benchmark_codec( + codec, + session, + output_directory=args.output_dir, + warmups=args.warmups, + iterations=args.iterations, + ) + for codec in CODECS + ] + print_results(results) + print(f"\nPayload files retained in: {args.output_dir.resolve()}") + return + + with tempfile.TemporaryDirectory(prefix="agent-session-serialization-") as temporary_directory: + results = [ + benchmark_codec( + codec, + session, + output_directory=Path(temporary_directory), + warmups=args.warmups, + iterations=args.iterations, + ) + for codec in CODECS + ] + print_results(results) + + +if __name__ == "__main__": + main() diff --git a/python/uv.lock b/python/uv.lock index cffa180cf1..88994e9162 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -489,7 +489,7 @@ requires-dist = [ { name = "agent-framework-redis", marker = "extra == 'all'", editable = "packages/redis" }, { name = "agent-framework-tools", marker = "extra == 'all'", editable = "packages/tools" }, { name = "mcp", marker = "extra == 'all'", specifier = ">=1.24.0,<2" }, - { name = "msgspec", specifier = ">=0.20.0,<1" }, + { name = "msgspec", specifier = ">=0.20.0,<0.22" }, { name = "opentelemetry-api", specifier = ">=1.39.0,<2" }, { name = "pydantic", specifier = ">=2,<3" }, { name = "python-dotenv", specifier = ">=1,<2" }, From fae67009b83a5bc46f5a18896745d6faba03915a Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 27 Jul 2026 15:35:55 +0200 Subject: [PATCH 3/8] Python: Preserve session snapshot compatibility Deep-copy in-memory session writes and retain existing Telegram session keys so stored conversations continue resolving. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3e5c81ad-75e8-4e92-a883-8bbd676c6c8c --- python/packages/core/agent_framework/_sessions.py | 2 +- python/packages/core/tests/core/test_sessions.py | 12 ++++++++++++ python/packages/hosting-telegram/README.md | 4 ++-- .../agent_framework_hosting_telegram/_parsing.py | 8 ++++---- .../tests/hosting_telegram/test_parsing.py | 6 +++--- .../04-hosting/af-hosting/local_telegram/README.md | 4 ++-- 6 files changed, 24 insertions(+), 12 deletions(-) diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index 662350770f..c7a816261b 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -1431,7 +1431,7 @@ async def set(self, session_id: str, session: AgentSession) -> None: ValueError: If ``session_id`` is empty. """ SessionStore.validate_session_id(session_id) - self._sessions[session_id] = session + self._sessions[session_id] = copy.deepcopy(session) async def delete(self, session_id: str) -> None: """Delete the stored session, if present. diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py index b7d17545fc..a0aaabbb3a 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -791,6 +791,18 @@ async def test_set_then_get_returns_independent_copy(self) -> None: assert reread is not None assert reread.state["nested"]["values"] == ["original"] + async def test_set_stores_independent_snapshot(self) -> None: + store = SessionStore() + session = AgentSession(session_id="session-1") + session.state["nested"] = {"values": ["original"]} + + await store.set("session-1", session) + session.state["nested"]["values"].append("changed") + + stored = await store.get("session-1") + assert stored is not None + assert stored.state["nested"]["values"] == ["original"] + async def test_set_replaces_existing_entry(self) -> None: store = SessionStore() await store.set("session-1", AgentSession(session_id="first")) diff --git a/python/packages/hosting-telegram/README.md b/python/packages/hosting-telegram/README.md index 968593998e..6c488f2d7a 100644 --- a/python/packages/hosting-telegram/README.md +++ b/python/packages/hosting-telegram/README.md @@ -26,8 +26,8 @@ long-running service. Your app remains fully responsible for: - `telegram_chat_id(update)` -- the chat id an update belongs to. - `telegram_session_id(update, bot_id=...)` -- a bot-scoped `AgentState` - session id. Private chats use `telegram__user_`; other chats - use `telegram__chat_`. + session id. Private chats use `telegram::`; other chats use + `telegram::`. - `telegram_command(update)` -- a leading slash command, with `/name@bot args` normalized to `/name args`. Returns `None` if there is none. - `telegram_callback_query_id(update)` -- a callback query's id, so you can diff --git a/python/packages/hosting-telegram/agent_framework_hosting_telegram/_parsing.py b/python/packages/hosting-telegram/agent_framework_hosting_telegram/_parsing.py index 9f91940f4f..361d841867 100644 --- a/python/packages/hosting-telegram/agent_framework_hosting_telegram/_parsing.py +++ b/python/packages/hosting-telegram/agent_framework_hosting_telegram/_parsing.py @@ -92,8 +92,8 @@ def telegram_session_id(update: Mapping[str, Any], *, bot_id: int) -> str | None bot_id: The Telegram bot's numeric user id. Returns: - ``telegram__user_`` for a private chat, - ``telegram__chat_`` for other chats, or ``None`` when the + ``telegram::`` for a private chat, + ``telegram::`` for other chats, or ``None`` when the required Telegram identity is absent. """ chat_id = telegram_chat_id(update) @@ -111,7 +111,7 @@ def telegram_session_id(update: Mapping[str, Any], *, bot_id: int) -> str | None chat = cast("Mapping[str, Any]", chat_candidate) if isinstance(chat_candidate, Mapping) else None chat_type = chat.get("type") if chat is not None else None if chat_type != "private": - return f"telegram_{bot_id}_chat_{chat_id}" + return f"telegram:{bot_id}:{chat_id}" if callback_query is not None: sender_candidate = callback_query.get("from") @@ -123,7 +123,7 @@ def telegram_session_id(update: Mapping[str, Any], *, bot_id: int) -> str | None return None sender = cast("Mapping[str, Any]", sender_candidate) sender_id = sender.get("id") - return f"telegram_{bot_id}_user_{sender_id}" if isinstance(sender_id, int) else None + return f"telegram:{bot_id}:{sender_id}" if isinstance(sender_id, int) else None def telegram_callback_query_id(update: Mapping[str, Any]) -> str | None: diff --git a/python/packages/hosting-telegram/tests/hosting_telegram/test_parsing.py b/python/packages/hosting-telegram/tests/hosting_telegram/test_parsing.py index 1cf3a105de..0cb4af0041 100644 --- a/python/packages/hosting-telegram/tests/hosting_telegram/test_parsing.py +++ b/python/packages/hosting-telegram/tests/hosting_telegram/test_parsing.py @@ -57,7 +57,7 @@ def test_returns_none_for_non_int_chat_id(self) -> None: class TestTelegramSessionId: def test_private_chat_uses_bot_and_user_id(self) -> None: - assert telegram_session_id(_message_update(text="hi"), bot_id=123) == "telegram_123_user_777" + assert telegram_session_id(_message_update(text="hi"), bot_id=123) == "telegram:123:777" def test_returns_none_when_no_chat_id(self) -> None: assert telegram_session_id({"update_id": 1}, bot_id=123) is None @@ -70,7 +70,7 @@ def test_private_chat_returns_none_when_no_user_id(self) -> None: def test_group_chat_uses_bot_and_chat_id(self) -> None: update = _message_update(text="hi") update["message"]["chat"] = {"id": -555, "type": "supergroup"} - assert telegram_session_id(update, bot_id=123) == "telegram_123_chat_-555" + assert telegram_session_id(update, bot_id=123) == "telegram:123:-555" def test_private_callback_uses_callback_sender(self) -> None: update = { @@ -81,7 +81,7 @@ def test_private_callback_uses_callback_sender(self) -> None: "message": {"chat": {"id": 555, "type": "private"}}, }, } - assert telegram_session_id(update, bot_id=123) == "telegram_123_user_888" + assert telegram_session_id(update, bot_id=123) == "telegram:123:888" class TestTelegramCommand: diff --git a/python/samples/04-hosting/af-hosting/local_telegram/README.md b/python/samples/04-hosting/af-hosting/local_telegram/README.md index 239a8cecf7..d4843691ff 100644 --- a/python/samples/04-hosting/af-hosting/local_telegram/README.md +++ b/python/samples/04-hosting/af-hosting/local_telegram/README.md @@ -71,8 +71,8 @@ process just registered. - **Session continuity:** `telegram_session_id(..., bot_id=bot.id)` follows Telegram's native identity boundaries. Private chats use - `telegram__user_`; groups and supergroups use - `telegram__chat_`, creating a shared session for that group. + `telegram::`; groups and supergroups use + `telegram::`, creating a shared session for that group. Including `bot.id` prevents two bots from accidentally sharing state. aiogram derives that numeric bot id from `TELEGRAM_BOT_TOKEN`, so these local apps do not need a separate `TELEGRAM_BOT_ID` setting. From 86de5a8eac3df1d366f4d383fac03423b45ce0cc Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 27 Jul 2026 20:59:27 +0200 Subject: [PATCH 4/8] Python: Simplify Foundry session isolation Add experimental FoundrySessionStore backed by Agent Server request context, remove resolver plumbing, and centralize v2 user isolation for sessions, checkpoints, and approvals. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3e5c81ad-75e8-4e92-a883-8bbd676c6c8c --- ...0033-python-session-store-serialization.md | 7 + python/PACKAGE_STATUS.md | 2 + python/packages/core/AGENTS.md | 2 + .../core/agent_framework/foundry/__init__.py | 1 + .../core/tests/core/test_foundry_namespace.py | 4 +- python/packages/foundry_hosting/README.md | 91 ++--- .../__init__.py | 5 +- .../_invocations.py | 13 +- .../_responses.py | 151 +++----- .../_session_isolation.py | 113 ------ .../_session_store.py | 76 ++++ .../foundry_hosting/tests/test_invocations.py | 2 +- .../foundry_hosting/tests/test_responses.py | 347 +++++++----------- 13 files changed, 292 insertions(+), 522 deletions(-) delete mode 100644 python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_isolation.py create mode 100644 python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py diff --git a/docs/decisions/0033-python-session-store-serialization.md b/docs/decisions/0033-python-session-store-serialization.md index d981f873a4..f24eb6b80c 100644 --- a/docs/decisions/0033-python-session-store-serialization.md +++ b/docs/decisions/0033-python-session-store-serialization.md @@ -204,6 +204,13 @@ storage-agnostic and passes keys through unchanged; each store implementation ow normalization. Protocol-specific hosts such as Foundry may still derive their own stable storage key before calling the store. +Foundry Hosting exposes an experimental `FoundrySessionStore` as its default +`ResponsesHostServer` store. It currently subclasses `FileSessionStore` and +derives the on-disk user partition directly from +`azure.ai.agentserver.core.get_request_context()`. The Foundry-specific type is +the host configuration seam; its implementation may later move from files to a +Foundry storage API without changing the generic core store contract. + ### Decision 2: Use msgspec codecs plus an explicit dynamic registry Chosen option: **msgspec codecs plus an explicit dynamic registry**. diff --git a/python/PACKAGE_STATUS.md b/python/PACKAGE_STATUS.md index da187a62b8..1c520a3033 100644 --- a/python/PACKAGE_STATUS.md +++ b/python/PACKAGE_STATUS.md @@ -129,6 +129,8 @@ listed below. - `agent-framework-core`: `SessionStore` and `FileSessionStore` from `agent_framework/_sessions.py` +- `agent-framework-foundry-hosting`: `FoundrySessionStore` from + `agent_framework_foundry_hosting/_session_store.py` #### `TO_PROMPT_AGENT` diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index e894138599..7f166f36b5 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -180,6 +180,8 @@ agent_framework/ ### Foundry (`foundry/`) - **`FoundryChatClient`** - Chat client for Microsoft Foundry project endpoints +- **`FoundrySessionStore`** - Experimental Foundry-hosting session store, lazily re-exported from + `agent-framework-foundry-hosting`; currently file-backed and scoped by Agent Server request context ## Key Patterns diff --git a/python/packages/core/agent_framework/foundry/__init__.py b/python/packages/core/agent_framework/foundry/__init__.py index c49b042567..7e398e9c50 100644 --- a/python/packages/core/agent_framework/foundry/__init__.py +++ b/python/packages/core/agent_framework/foundry/__init__.py @@ -32,6 +32,7 @@ "FoundryEmbeddingSettings": ("agent_framework_foundry", "agent-framework-foundry"), "FoundryEvals": ("agent_framework_foundry", "agent-framework-foundry"), "FoundryMemoryProvider": ("agent_framework_foundry", "agent-framework-foundry"), + "FoundrySessionStore": ("agent_framework_foundry_hosting", "agent-framework-foundry-hosting"), "FoundryToolbox": ("agent_framework_foundry_hosting", "agent-framework-foundry-hosting"), "FoundryLocalChatOptions": ("agent_framework_foundry_local", "agent-framework-foundry-local"), "FoundryLocalClient": ("agent_framework_foundry_local", "agent-framework-foundry-local"), diff --git a/python/packages/core/tests/core/test_foundry_namespace.py b/python/packages/core/tests/core/test_foundry_namespace.py index faadc99239..dd9f56d07e 100644 --- a/python/packages/core/tests/core/test_foundry_namespace.py +++ b/python/packages/core/tests/core/test_foundry_namespace.py @@ -2,7 +2,7 @@ import pytest from agent_framework_foundry import FoundryChatClient, FoundryMemoryProvider -from agent_framework_foundry_hosting import ResponsesHostServer +from agent_framework_foundry_hosting import FoundrySessionStore, ResponsesHostServer from agent_framework_foundry_local import FoundryLocalClient import agent_framework.azure as azure @@ -12,10 +12,12 @@ def test_foundry_namespace_exposes_cloud_and_local_symbols() -> None: assert foundry.FoundryChatClient is FoundryChatClient assert foundry.FoundryMemoryProvider is FoundryMemoryProvider + assert foundry.FoundrySessionStore is FoundrySessionStore assert foundry.ResponsesHostServer is ResponsesHostServer assert foundry.FoundryLocalClient is FoundryLocalClient assert "FoundryChatClient" in dir(foundry) assert "FoundryLocalClient" in dir(foundry) + assert "FoundrySessionStore" in dir(foundry) assert "ResponsesHostServer" in dir(foundry) diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index 93c243c540..83dfa93a2d 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -4,87 +4,38 @@ This package provides the integration of Agent Framework agents and workflows wi `ResponsesHostServer` persists the Agent Framework `AgentSession` used by regular agents in addition to the Responses provider's message history. By default it -uses core's experimental msgspec-backed `FileSessionStore` under `$HOME/.checkpoints/sessions` +uses the experimental `FoundrySessionStore` under `$HOME/.checkpoints/sessions` when hosted and `{cwd}/.checkpoints/sessions` locally. The store is partitioned -by stable agent name, the platform user key, and the Responses conversation -partition. Foundry hashes that composite scope into a restricted-alphabet -`foundry_` session ID before it reaches the store. Pass `session_store=` to -use another `SessionStore` implementation. +by the Agent Server request context's platform user key and the Responses +conversation partition. Pass `session_store=` to use another `SessionStore` +implementation. Workflow agents continue to use checkpoint storage; their checkpoints share the same `$HOME/.checkpoints` or local `.checkpoints` root. -## Custom authenticated session isolation +## Foundry session isolation -`ResponsesHostServer` accepts an explicit sync or async -`session_isolation_key_resolver(request, context)` callable. The default uses -Foundry's trusted `context.platform_context.user_id_key`. Hosted requests fail -closed when no key is resolved; local requests may remain unscoped. +`FoundrySessionStore` currently subclasses core's `FileSessionStore`. It reads +the active request through `azure.ai.agentserver.core.get_request_context()` and +hashes the platform `user_id` (the same `x-agent-user-id` value exposed as +`ResponseContext.platform_context.user_id_key`) before selecting an on-disk +directory. This makes the user partition path-safe and avoids persisting the raw +platform identity. File-backed state is partitioned under the validated identity: ```text .checkpoints/ - sessions/user-/.json - checkpoints/user-// - function-approvals/user-/approval_requests.json + sessions/user-/.json + checkpoints/user-// + function-approvals/user-/approval_requests.json ``` -The default store is `IsolationKeyScopedFileSessionStore`, a -`FileSessionStore` subclass. One store instance serves the server; each -`get`/`set`/`delete` operation resolves its path beneath the isolation directory -bound for that request. The server does not keep a store instance per user. -If `session_store=` is supplied with file-backed storage, it must already be an -`IsolationKeyScopedFileSessionStore`; an unscoped `FileSessionStore` is rejected. +Hosted requests require container protocol `2.0.0`. The v2-only request +`call_id` is checked before session, checkpoint, or approval storage is used, +and a missing platform user ID fails closed. Local requests may remain unscoped. -The resolver must return a stable, globally unique, portable path-safe identity. -The value is visible in the directory name, so return an application-owned -opaque ID rather than a sensitive claim when necessary. - -Python has no universal DI container or `IHttpContextAccessor`. Applications -that authenticate through ASGI middleware can explicitly bridge their verified -principal to the resolver with an application-owned `ContextVar`: - -```python -from contextvars import ContextVar -from typing import Any - -from agent_framework_foundry_hosting import ResponsesHostServer - - -current_subject: ContextVar[str | None] = ContextVar("current_subject", default=None) - - -class VerifiedSubjectMiddleware: - def __init__(self, app: Any) -> None: - self.app = app - - async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None: - if scope["type"] != "http": - await self.app(scope, receive, send) - return - - # An outer authentication middleware must populate this trusted user. - user = scope.get("user") - subject = getattr(user, "identity", None) if getattr(user, "is_authenticated", False) else None - token = current_subject.set(subject) - try: - await self.app(scope, receive, send) - finally: - current_subject.reset(token) - - -def claims_isolation_key(request: Any, context: Any) -> str | None: - del request, context - return current_subject.get() - - -server = ResponsesHostServer( - agent, - session_isolation_key_resolver=claims_isolation_key, -) -app = VerifiedSubjectMiddleware(server) -``` - -Authentication middleware must run before `VerifiedSubjectMiddleware`. Do not -use an unverified caller-controlled header as the isolation key. +The Foundry-specific store type intentionally hides the current filesystem +implementation from `ResponsesHostServer` setup. A future version may move +`FoundrySessionStore` to a Foundry storage API without changing the host's +default configuration. diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/__init__.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/__init__.py index 36e9027de7..56d6571670 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/__init__.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/__init__.py @@ -4,7 +4,7 @@ from ._invocations import InvocationsHostServer from ._responses import ResponsesHostServer -from ._session_isolation import IsolationKeyScopedFileSessionStore, ResponsesSessionIsolationKeyResolver +from ._session_store import FoundrySessionStore from ._toolbox import FoundryToolbox try: @@ -13,9 +13,8 @@ __version__ = "0.0.0" __all__ = [ + "FoundrySessionStore", "FoundryToolbox", "InvocationsHostServer", - "IsolationKeyScopedFileSessionStore", "ResponsesHostServer", - "ResponsesSessionIsolationKeyResolver", ] diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py index f01c7a9468..4f38b2dc06 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -1,12 +1,13 @@ # Copyright (c) Microsoft. All rights reserved. from agent_framework import AgentSession, SupportsAgentRun -from azure.ai.agentserver.core import get_request_context from azure.ai.agentserver.invocations import InvocationAgentServerHost from starlette.requests import Request from starlette.responses import Response, StreamingResponse from typing_extensions import Any, AsyncGenerator +from ._session_store import _get_foundry_request_context # pyright: ignore[reportPrivateUsage] + class InvocationsHostServer(InvocationAgentServerHost): """An invocations server host for an agent.""" @@ -49,15 +50,7 @@ def _partition_key(self) -> str: Exceptions: RuntimeError: If the context doesn't contain the expected IDs. """ - context = get_request_context() - - # Fail fast if the service is on protocol v1.0.0 - if self.config.is_hosted and context.call_id is None: - raise RuntimeError( - "The hosted environment is running on protocol 1.0.0, but the agent requires protocol 2.0.0. " - "Please upgrade your agent protocol to 2.0.0 in `agent.manifest.yaml` or `agent.yaml`, or " - "downgrade the `agent-framework-foundry-hosting` package to `1.0.0a260625` or before to use 1.0.0." - ) + context = _get_foundry_request_context(is_hosted=self.config.is_hosted) if self.config.is_hosted: if not context.session_id or not context.user_id: diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index 9d1d04d98c..f7f7bc8b39 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -120,12 +120,11 @@ from mcp import McpError from typing_extensions import Any -from ._session_isolation import ( - IsolationKeyScopedFileSessionStore, - ResolvedSessionIsolation, - ResponsesSessionIsolationKeyResolver, - platform_session_isolation_key_resolver, - resolve_session_isolation, +from ._session_store import ( + FoundrySessionStore, + _get_foundry_request_context, # pyright: ignore[reportPrivateUsage] + _request_user_directory_segment, # pyright: ignore[reportPrivateUsage] + _request_user_fingerprint, # pyright: ignore[reportPrivateUsage] ) logger = logging.getLogger(__name__) @@ -143,8 +142,6 @@ _FUNCTION_APPROVAL_DIRECTORY_NAME = "function-approvals" _FUNCTION_APPROVAL_FILE_NAME = "approval_requests.json" _HOSTED_RESPONSES_HISTORY_SOURCE_ID = "_foundry_responses_history" -_SESSION_ISOLATION_STATE_KEY = "foundry_hosting.session_isolation_fingerprint" -_UNSCOPED_SESSION_ISOLATION = ResolvedSessionIsolation(None, None, None) # region Approval Storage @@ -241,16 +238,14 @@ async def load_approval_request(self, approval_request_id: str) -> Content: return await asyncio.to_thread(self._load_sync, approval_request_id) -def _validate_path_segment(segment: str, *, kind: Literal["context id", "user id"]) -> None: +def _validate_path_segment(segment: str, *, kind: Literal["context id"]) -> None: """Validate that ``segment`` is a single safe path component (CWE-22). ``segment`` originates from caller-controlled fields (such as - ``previous_response_id``), server-generated fields (``conversation_id`` / - ``response_id``), or the platform-injected per-user partition key - (``x-agent-user-id``). In every case it must be treated as an untrusted - single path segment: path separators, drive letters, parent references and - similar would otherwise let the resulting directory escape the configured - storage root. + ``previous_response_id``) or server-generated fields (``conversation_id`` / + ``response_id``). It must be treated as an untrusted single path segment: + path separators, drive letters, parent references and similar would + otherwise let the resulting directory escape the configured storage root. We deliberately do not URL-decode the value here: the hosting layer never decodes these ids before joining them, so forms such as ``%2e%2e`` are @@ -343,34 +338,17 @@ def _conversation_object_id(conversation_key: str) -> str: def _custom_session_store_key( agent: SupportsAgentRun, object_id: str, - *, - isolation: ResolvedSessionIsolation, ) -> str: - """Create an isolation-scoped key for a non-file custom store.""" + """Create a request-user-scoped key for a non-file custom store.""" key: dict[str, str] = {"object": object_id} if agent.name: key["agent"] = agent.name - if isolation.fingerprint: - key["isolation"] = isolation.fingerprint + if user_fingerprint := _request_user_fingerprint(): + key["user"] = user_fingerprint canonical_key = json.dumps(key, ensure_ascii=False, separators=(",", ":"), sort_keys=True) return f"foundry_{hashlib.sha256(canonical_key.encode('utf-8')).hexdigest()}" -def _stamp_or_validate_session_isolation( - session: AgentSession, - isolation: ResolvedSessionIsolation, -) -> None: - """Stamp a legacy/new session or reject a cross-isolation restore.""" - existing = session.state.get(_SESSION_ISOLATION_STATE_KEY) - expected = isolation.fingerprint - if existing is None: - if expected is not None: - session.state[_SESSION_ISOLATION_STATE_KEY] = expected - return - if existing != expected: - raise RuntimeError("Hosted session identity context mismatch.") - - def _is_hosted_responses_history_sentinel(provider: ContextProvider) -> bool: """Return whether ``provider`` is the host's no-op history-loading sentinel.""" return ( @@ -385,21 +363,19 @@ def _is_hosted_responses_history_sentinel(provider: ContextProvider) -> bool: def _checkpoint_storage_for_context( root: str | Path, context_id: str, - *, - isolation: ResolvedSessionIsolation = _UNSCOPED_SESSION_ISOLATION, ) -> FileCheckpointStorage: """Build a ``FileCheckpointStorage`` for ``context_id`` rooted under ``root``. - Scoped requests use ``/user-/``. Local + Scoped requests use ``/user-/``. Local unscoped requests use ``/``. """ _validate_path_segment(context_id, kind="context id") base_path = Path(root).resolve() - if isolation.directory_segment: - user_path = (base_path / isolation.directory_segment).resolve() + if user_directory := _request_user_directory_segment(): + user_path = (base_path / user_directory).resolve() if not user_path.is_relative_to(base_path): - raise RuntimeError(f"Invalid session isolation directory: {isolation.directory_segment!r}") + raise RuntimeError(f"Invalid Foundry user directory: {user_directory!r}") base_path = user_path storage_path = (base_path / context_id).resolve() @@ -415,14 +391,13 @@ def _checkpoint_storage_for_context( def _approval_storage_path( root: str | Path, - isolation: ResolvedSessionIsolation, ) -> str: - """Return the approval file path for one resolved isolation identity.""" + """Return the approval file path for the active request user.""" base_path = Path(root).resolve() - if isolation.directory_segment: - base_path = (base_path / isolation.directory_segment).resolve() + if user_directory := _request_user_directory_segment(): + base_path = (base_path / user_directory).resolve() if not base_path.is_relative_to(Path(root).resolve()): - raise RuntimeError(f"Invalid session isolation directory: {isolation.directory_segment!r}") + raise RuntimeError(f"Invalid Foundry user directory: {user_directory!r}") return str(base_path / _FUNCTION_APPROVAL_FILE_NAME) @@ -515,7 +490,6 @@ def __init__( options: ResponsesServerOptions | None = None, store: ResponseProviderProtocol | None = None, session_store: SessionStore | None = None, - session_isolation_key_resolver: ResponsesSessionIsolationKeyResolver | None = None, **kwargs: Any, ) -> None: """Initialize a ResponsesHostServer. @@ -526,12 +500,9 @@ def __init__( options: Optional server options. store: Optional response store. session_store: Optional Agent Framework session store. Defaults to - an :class:`IsolationKeyScopedFileSessionStore` under the Foundry + a :class:`FoundrySessionStore` under the Foundry durable storage root. A caller-supplied file store must already - be an ``IsolationKeyScopedFileSessionStore``. - session_isolation_key_resolver: Optional sync/async callable that - resolves a stable authenticated identity for the request. - Defaults to the trusted Foundry platform user key. + be a ``FoundrySessionStore``. **kwargs: Additional keyword arguments. Note: @@ -563,7 +534,6 @@ def __init__( home_directory=os.getenv(_HOSTED_SESSION_DATA_DIRECTORY_ENVIRONMENT_VARIABLE), current_directory=os.getcwd(), ) - self._session_isolation_key_resolver = session_isolation_key_resolver or platform_session_isolation_key_resolver self._is_workflow_agent = False self._checkpoint_storage_path = None if isinstance(agent, WorkflowAgent): @@ -596,20 +566,15 @@ def __init__( self._agent: SupportsAgentRun = agent if not self._is_workflow_agent and session_store is None: - session_store = IsolationKeyScopedFileSessionStore(self._storage_root / _SESSION_DIRECTORY_NAME) - elif isinstance(session_store, FileSessionStore) and not isinstance( - session_store, - IsolationKeyScopedFileSessionStore, - ): - raise ValueError( - "ResponsesHostServer requires IsolationKeyScopedFileSessionStore for file-backed session storage." - ) + session_store = FoundrySessionStore(self._storage_root / _SESSION_DIRECTORY_NAME) + elif isinstance(session_store, FileSessionStore) and not isinstance(session_store, FoundrySessionStore): + raise ValueError("ResponsesHostServer requires FoundrySessionStore for file-backed session storage.") self._session_store = session_store self._approval_storage: ApprovalStorage = InMemoryFunctionApprovalStorage() self._approval_storage_root = self._storage_root / _FUNCTION_APPROVAL_DIRECTORY_NAME - # Per-user hosted approval stores share a process-local lock per identity. + # Per-user hosted approval stores share a process-local lock per user. # Local development keeps the single in-memory store. - self._approval_storages_by_isolation: dict[str, ApprovalStorage] = {} + self._approval_storages_by_user: dict[str, ApprovalStorage] = {} # Lazy agent lifecycle: the agent (and any MCP tools it owns) is entered on # the first request rather than at server startup, so that authentication # failures during MCP connect can be surfaced to the client as an @@ -647,14 +612,15 @@ async def _cleanup_agent(self) -> None: self._agent_stack = None await stack.aclose() - def _approval_storage_for_isolation(self, isolation: ResolvedSessionIsolation) -> ApprovalStorage: - """Return the hosted approval store for one validated identity.""" - if not self.config.is_hosted or isolation.key is None: + def _approval_storage_for_request(self) -> ApprovalStorage: + """Return the hosted approval store for the active request user.""" + user_fingerprint = _request_user_fingerprint() + if not self.config.is_hosted or user_fingerprint is None: return self._approval_storage - storage = self._approval_storages_by_isolation.get(isolation.key) + storage = self._approval_storages_by_user.get(user_fingerprint) if storage is None: - storage = FileBasedFunctionApprovalStorage(_approval_storage_path(self._approval_storage_root, isolation)) - self._approval_storages_by_isolation[isolation.key] = storage + storage = FileBasedFunctionApprovalStorage(_approval_storage_path(self._approval_storage_root)) + self._approval_storages_by_user[user_fingerprint] = storage return storage async def _handle_response( @@ -664,30 +630,17 @@ async def _handle_response( cancellation_signal: asyncio.Event, ) -> AsyncIterable[ResponseStreamEvent | dict[str, Any]]: """Handle the creation of a response.""" - # Fail fast if the service is on protocol v1.0.0 - if self.config.is_hosted and context.platform_context.call_id is None: - raise RuntimeError( - "The hosted environment is running on protocol 1.0.0, but the agent requires protocol 2.0.0. " - "Please upgrade your agent protocol to 2.0.0 in `agent.manifest.yaml` or `agent.yaml`, or " - "downgrade the `agent-framework-foundry-hosting` package to `1.0.0a260625` or before to use 1.0.0." - ) - isolation = await resolve_session_isolation( - self._session_isolation_key_resolver, - request, - context, - is_hosted=self.config.is_hosted, - ) + _get_foundry_request_context(is_hosted=self.config.is_hosted) if self._is_workflow_agent: # Workflow agents are handled differently because they require checkpoint restoration - return self._handle_inner_workflow(request, context, isolation) - return self._handle_inner_agent(request, context, isolation) + return self._handle_inner_workflow(request, context) + return self._handle_inner_agent(request, context) async def _handle_inner_agent( self, request: CreateResponse, context: ResponseContext, - isolation: ResolvedSessionIsolation = _UNSCOPED_SESSION_ISOLATION, ) -> AsyncIterable[ResponseStreamEvent | dict[str, Any]]: """Handle the creation of a response for a regular (non-workflow) agent.""" response_event_stream = ResponseEventStream(response_id=context.response_id, model=request.model) @@ -708,14 +661,10 @@ async def save_session() -> None: if request_session_store is None: raise RuntimeError("Session storage is not configured for a regular agent.") if session is not None and session_store_key is not None: - if isinstance(request_session_store, IsolationKeyScopedFileSessionStore): - with request_session_store.use_isolation(isolation): - await request_session_store.set(session_store_key, session) - else: - await request_session_store.set(session_store_key, session) + await request_session_store.set(session_store_key, session) try: - approval_storage = self._approval_storage_for_isolation(isolation) + approval_storage = self._approval_storage_for_request() conversation_key = _resolve_session_conversation_key(request, context) object_id = _conversation_object_id(conversation_key) request_session_store = self._session_store @@ -723,17 +672,12 @@ async def save_session() -> None: raise RuntimeError("Session storage is not configured for a regular agent.") session_store_key = ( object_id - if isinstance(request_session_store, IsolationKeyScopedFileSessionStore) - else _custom_session_store_key(self._agent, object_id, isolation=isolation) + if isinstance(request_session_store, FoundrySessionStore) + else _custom_session_store_key(self._agent, object_id) ) - if isinstance(request_session_store, IsolationKeyScopedFileSessionStore): - with request_session_store.use_isolation(isolation): - session = await request_session_store.get(session_store_key) - else: - session = await request_session_store.get(session_store_key) + session = await request_session_store.get(session_store_key) if session is None: session = self._agent.create_session(session_id=object_id) - _stamp_or_validate_session_isolation(session, isolation) input_items = await context.get_input_items() input_messages = await _items_to_messages(input_items, approval_storage=approval_storage) @@ -831,7 +775,6 @@ async def _handle_inner_workflow( self, request: CreateResponse, context: ResponseContext, - isolation: ResolvedSessionIsolation = _UNSCOPED_SESSION_ISOLATION, ) -> AsyncIterable[ResponseStreamEvent | dict[str, Any]]: """Handle the creation of a response for a workflow agent.""" response_event_stream = ResponseEventStream(response_id=context.response_id, model=request.model) @@ -843,7 +786,7 @@ async def _handle_inner_workflow( tracker: _OutputItemTracker | None = None try: - approval_storage = self._approval_storage_for_isolation(isolation) + approval_storage = self._approval_storage_for_request() input_items = await context.get_input_items() input_messages = await _items_to_messages(input_items, approval_storage=approval_storage) is_streaming_request = request.stream is not None and request.stream is True @@ -869,8 +812,8 @@ async def _handle_inner_workflow( await self._ensure_agent_ready() # Per-user checkpoint isolation for multi-tenant hosting (container - # protocol v2): the per-user partition key computed above - # (``x-agent-user-id``) scopes every checkpoint directory for this turn, + # protocol v2): the request-scoped ``x-agent-user-id`` value scopes + # every checkpoint directory for this turn, # so one tenant can never restore or observe another tenant's workflow # state -- even with a guessed or forged context id. The key is stable # per user across turns, so multi-turn continuity is preserved. Absent @@ -893,7 +836,6 @@ async def _handle_inner_workflow( restore_storage = _checkpoint_storage_for_context( self._checkpoint_storage_path, context_id, - isolation=isolation, ) latest_checkpoint = await restore_storage.get_latest(workflow_name=self._agent.workflow.name) if latest_checkpoint is not None: @@ -911,7 +853,6 @@ async def _handle_inner_workflow( write_storage = _checkpoint_storage_for_context( self._checkpoint_storage_path, write_context_id, - isolation=isolation, ) # Multi-turn pattern: when we have a prior checkpoint, restore it diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_isolation.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_isolation.py deleted file mode 100644 index 726a4de737..0000000000 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_isolation.py +++ /dev/null @@ -1,113 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -from __future__ import annotations - -import hashlib -import inspect -from collections.abc import Awaitable, Callable, Generator -from contextlib import contextmanager -from contextvars import ContextVar -from dataclasses import dataclass -from pathlib import Path -from typing import Literal, TypeAlias - -from agent_framework import FileSessionStore -from azure.ai.agentserver.responses import ResponseContext -from azure.ai.agentserver.responses.models import CreateResponse - -ResponsesSessionIsolationKeyResolver: TypeAlias = Callable[ - [CreateResponse, ResponseContext], - str | Awaitable[str | None] | None, -] -"""Resolve a stable authenticated identity used to partition Responses state.""" - -_MAX_ISOLATION_KEY_LENGTH = 128 -_INVALID_PATH_CHARACTERS = frozenset('<>:"/\\|?*') - - -@dataclass(frozen=True, slots=True) -class ResolvedSessionIsolation: - """Validated per-request storage isolation.""" - - key: str | None - directory_segment: str | None - fingerprint: str | None - - -class IsolationKeyScopedFileSessionStore(FileSessionStore): - """FileSessionStore that scopes each operation to the bound user directory.""" - - def __init__( - self, - storage_path: str | Path, - *, - serialization_format: Literal["json", "msgpack"] = "json", - ) -> None: - """Initialize a scoped file store rooted at ``storage_path``.""" - super().__init__(storage_path, serialization_format=serialization_format) - self._current_directory: ContextVar[str | None] = ContextVar( - f"foundry_session_isolation_{id(self)}", - default=None, - ) - - @contextmanager - def use_isolation(self, isolation: ResolvedSessionIsolation) -> Generator[None]: - """Bind one resolved request identity for subsequent store operations.""" - token = self._current_directory.set(isolation.directory_segment) - try: - yield - finally: - self._current_directory.reset(token) - - def get_session_directory(self) -> Path: - """Return the currently bound user directory.""" - directory_segment = self._current_directory.get() - if directory_segment is None: - return self.storage_path - return self.storage_path / directory_segment - - -def platform_session_isolation_key_resolver( - request: CreateResponse, - context: ResponseContext, -) -> str | None: - """Resolve the trusted Foundry platform user partition key.""" - del request - return context.platform_context.user_id_key - - -def _validate_isolation_key(key: str) -> None: - """Validate a portable single-directory identity value.""" - if len(key) > _MAX_ISOLATION_KEY_LENGTH: - raise ValueError(f"Session isolation key must be at most {_MAX_ISOLATION_KEY_LENGTH} characters.") - if key in {".", ".."} or key.endswith((" ", ".")): - raise ValueError("Session isolation key is not a safe directory name.") - if any(ord(character) < 32 or character in _INVALID_PATH_CHARACTERS for character in key): - raise ValueError("Session isolation key contains characters that are unsafe in a directory name.") - - -async def resolve_session_isolation( - resolver: ResponsesSessionIsolationKeyResolver, - request: CreateResponse, - context: ResponseContext, - *, - is_hosted: bool, -) -> ResolvedSessionIsolation: - """Resolve, normalize, validate, and fingerprint one request identity.""" - value = resolver(request, context) - if inspect.isawaitable(value): - value = await value - if value is None or not value.strip(): - if is_hosted: - raise RuntimeError( - "The hosted request did not resolve a session isolation key; refusing to use unscoped session storage." - ) - return ResolvedSessionIsolation(key=None, directory_segment=None, fingerprint=None) - - key = value.strip() - _validate_isolation_key(key) - return ResolvedSessionIsolation( - key=key, - directory_segment=f"user-{key}", - fingerprint=hashlib.sha256(key.encode("utf-8")).hexdigest(), - ) diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py new file mode 100644 index 0000000000..2ed0cdda70 --- /dev/null +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py @@ -0,0 +1,76 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import hashlib +from pathlib import Path +from typing import Literal + +from agent_framework import ExperimentalFeature, FileSessionStore +from agent_framework._feature_stage import experimental +from azure.ai.agentserver.core import FoundryAgentRequestContext, get_request_context + +_PROTOCOL_V2_REQUIRED_MESSAGE = ( + "The hosted environment is running on protocol 1.0.0, but the agent requires protocol 2.0.0. " + "Please upgrade your agent protocol to 2.0.0 in `agent.manifest.yaml` or `agent.yaml`, or " + "downgrade the `agent-framework-foundry-hosting` package to `1.0.0a260625` or before to use 1.0.0." +) + + +def _get_foundry_request_context( # pyright: ignore[reportUnusedFunction] + *, + is_hosted: bool, +) -> FoundryAgentRequestContext: + """Return the current request context and validate hosted v2 identity.""" + context = get_request_context() + if is_hosted and context.call_id is None: + raise RuntimeError(_PROTOCOL_V2_REQUIRED_MESSAGE) + if is_hosted and not context.user_id: + raise RuntimeError( + "The hosted environment is missing the platform user ID in the request context. " + "Please ensure that the request is coming from a valid Foundry platform service." + ) + return context + + +def _request_user_id_key() -> str | None: + """Return the platform user partition key for the active request.""" + # FoundryAgentRequestContext.user_id is populated from the same + # x-agent-user-id value exposed as ResponseContext.platform_context.user_id_key. + return get_request_context().user_id + + +def _request_user_fingerprint() -> str | None: + """Return a stable opaque fingerprint for the active request user.""" + user_id_key = _request_user_id_key() + return hashlib.sha256(user_id_key.encode("utf-8")).hexdigest() if user_id_key else None + + +def _request_user_directory_segment() -> str | None: + """Return the safe on-disk directory segment for the active request user.""" + fingerprint = _request_user_fingerprint() + return f"user-{fingerprint}" if fingerprint else None + + +@experimental(feature_id=ExperimentalFeature.SESSION_STORE) +class FoundrySessionStore(FileSessionStore): + """File-backed session store isolated by the active Foundry request user. + + This implementation currently persists through :class:`FileSessionStore`. + The Foundry-specific type leaves room to use a platform storage API later + without changing :class:`ResponsesHostServer` configuration. + """ + + def __init__( + self, + storage_path: str | Path, + *, + serialization_format: Literal["json", "msgpack"] = "json", + ) -> None: + """Initialize a Foundry-scoped file store rooted at ``storage_path``.""" + super().__init__(storage_path, serialization_format=serialization_format) + + def get_session_directory(self) -> Path: + """Return the active request user's session directory.""" + directory_segment = _request_user_directory_segment() + return self.storage_path / directory_segment if directory_segment else self.storage_path diff --git a/python/packages/foundry_hosting/tests/test_invocations.py b/python/packages/foundry_hosting/tests/test_invocations.py index 05bcdc1d5b..0fad1c72a0 100644 --- a/python/packages/foundry_hosting/tests/test_invocations.py +++ b/python/packages/foundry_hosting/tests/test_invocations.py @@ -183,7 +183,7 @@ def test_hosted_missing_user_id_raises(self) -> None: server.config.is_hosted = True with ( _request_context(call_id="call-1", session_id="sess-1"), - pytest.raises(RuntimeError, match="missing session_id or user_id"), + pytest.raises(RuntimeError, match="missing the platform user ID"), ): server._partition_key() # pyright: ignore[reportPrivateUsage] diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 53ee26c847..bb497ce119 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -11,9 +11,11 @@ from __future__ import annotations import asyncio +import hashlib import json import uuid -from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path from typing import Literal, overload @@ -30,6 +32,7 @@ BaseChatClient, ChatResponse, Content, + ExperimentalFeature, FileCheckpointStorage, FileSessionStore, HistoryProvider, @@ -48,20 +51,23 @@ WorkflowMessage, executor, ) -from azure.ai.agentserver.responses import InMemoryResponseProvider, PlatformContext, ResponseContext +from azure.ai.agentserver.core import ( + FoundryAgentRequestContext, + reset_request_context, + set_request_context, +) +from azure.ai.agentserver.responses import InMemoryResponseProvider, ResponseContext from azure.ai.agentserver.responses.models import CreateResponse from mcp import McpError from mcp.types import ErrorData from typing_extensions import Any from agent_framework_foundry_hosting import ( - IsolationKeyScopedFileSessionStore, + FoundrySessionStore, ResponsesHostServer, - ResponsesSessionIsolationKeyResolver, ) from agent_framework_foundry_hosting._responses import ( _AZURE_RESPONSES_MESSAGE_ROLE_TYPE, # pyright: ignore[reportPrivateUsage] - _SESSION_ISOLATION_STATE_KEY, # pyright: ignore[reportPrivateUsage] CONSENT_ERROR_CODE, ConsentError, FileBasedFunctionApprovalStorage, # pyright: ignore[reportPrivateUsage] @@ -73,13 +79,8 @@ _resolve_durable_storage_root, # pyright: ignore[reportPrivateUsage] _resolve_session_conversation_key, # pyright: ignore[reportPrivateUsage] _response_id_partition, # pyright: ignore[reportPrivateUsage] - _stamp_or_validate_session_isolation, # pyright: ignore[reportPrivateUsage] consent_url_from_error, ) -from agent_framework_foundry_hosting._session_isolation import ( # pyright: ignore[reportPrivateUsage] - ResolvedSessionIsolation, - resolve_session_isolation, -) def _make_function_approval_request_content( @@ -97,6 +98,21 @@ def _make_function_approval_request_content( return Content.from_function_approval_request(request_id, function_call) +@contextmanager +def _request_context( + *, + call_id: str | None = None, + user_id: str | None = None, + session_id: str | None = None, +) -> Iterator[None]: + """Install a Foundry request context for the duration of the block.""" + token = set_request_context(FoundryAgentRequestContext(call_id=call_id, user_id=user_id, session_id=session_id)) + try: + yield + finally: + reset_request_context(token) + + # region Helpers @@ -249,6 +265,18 @@ def test_init_basic(self) -> None: assert history_sentinel.store_inputs is False assert history_sentinel.store_outputs is False + def test_init_uses_foundry_session_store_by_default( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.chdir(tmp_path) + agent = _make_agent( + response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("hi")])]) + ) + server = ResponsesHostServer(agent, store=InMemoryResponseProvider()) + assert isinstance(server._session_store, FoundrySessionStore) # pyright: ignore[reportPrivateUsage] + def test_init_rejects_history_provider_with_load_messages(self) -> None: class _LoadMessagesHistoryProvider(HistoryProvider): @@ -280,7 +308,7 @@ def test_init_rejects_unscoped_file_session_store(self, tmp_path: Path) -> None: response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("hi")])]) ) - with pytest.raises(ValueError, match="IsolationKeyScopedFileSessionStore"): + with pytest.raises(ValueError, match="FoundrySessionStore"): ResponsesHostServer( agent, store=InMemoryResponseProvider(), @@ -296,12 +324,12 @@ async def test_hosted_request_requires_user_partition_key(self) -> None: context = ResponseContext( response_id="caresp_aaaaaaaaaaaaaaaa00" + "1" * 32, mode_flags=MagicMock(), - platform_context=PlatformContext(call_id="call-1"), ) with ( patch.object(server.config, "is_hosted", True), - pytest.raises(RuntimeError, match="did not resolve"), + _request_context(call_id="call-1"), + pytest.raises(RuntimeError, match="platform user ID"), ): await server._handle_response( # pyright: ignore[reportPrivateUsage] request, @@ -309,33 +337,22 @@ async def test_hosted_request_requires_user_partition_key(self) -> None: asyncio.Event(), ) - async def test_custom_isolation_resolver_is_called_once_per_request(self) -> None: - calls = 0 - - async def resolver(request: CreateResponse, context: ResponseContext) -> str: - nonlocal calls - calls += 1 - assert request.model == "m" - assert context.response_id == "response-1" - return "user-1" - - server = _make_server(_make_agent(), session_isolation_key_resolver=resolver) + async def test_hosted_request_requires_protocol_v2(self) -> None: + server = _make_server(_make_agent()) request = CreateResponse(model="m", input="hi") context = ResponseContext(response_id="response-1", mode_flags=MagicMock()) - expected = MagicMock() - with patch.object(server, "_handle_inner_agent", return_value=expected) as handler: - result = await server._handle_response( # pyright: ignore[reportPrivateUsage] + with ( + patch.object(server.config, "is_hosted", True), + _request_context(user_id="user-1"), + pytest.raises(RuntimeError, match="protocol 2.0.0"), + ): + await server._handle_response( # pyright: ignore[reportPrivateUsage] request, context, asyncio.Event(), ) - assert result is expected - assert calls == 1 - isolation = handler.call_args.args[2] - assert isolation.directory_segment == "user-user-1" - # endregion @@ -344,93 +361,11 @@ async def resolver(request: CreateResponse, context: ResponseContext) -> str: class TestSessionPersistenceHelpers: - def test_resolver_type_alias_is_public(self) -> None: - def resolver(request: CreateResponse, context: ResponseContext) -> str | None: - del request - return context.platform_context.user_id_key - - typed_resolver: ResponsesSessionIsolationKeyResolver = resolver - assert callable(typed_resolver) - - async def test_sync_and_async_isolation_resolvers(self) -> None: - request = CreateResponse(model="m", input="hi") - context = ResponseContext(response_id="response-1", mode_flags=MagicMock()) - - sync_isolation = await resolve_session_isolation( - lambda request, context: "user-1", - request, - context, - is_hosted=True, - ) - - async def async_resolver(request: CreateResponse, context: ResponseContext) -> str: - del request, context - return "user-2" - - async_isolation = await resolve_session_isolation( - async_resolver, - request, - context, - is_hosted=True, - ) - - assert sync_isolation.directory_segment == "user-user-1" - assert async_isolation.directory_segment == "user-user-2" - assert sync_isolation.fingerprint != async_isolation.fingerprint - - async def test_local_missing_isolation_key_is_unscoped(self) -> None: - request = CreateResponse(model="m", input="hi") - context = ResponseContext(response_id="response-1", mode_flags=MagicMock()) - - isolation = await resolve_session_isolation( - lambda request, context: None, - request, - context, - is_hosted=False, - ) - - assert isolation == ResolvedSessionIsolation(None, None, None) - - async def test_contextvar_backed_resolvers_do_not_cross_contaminate(self) -> None: - from contextvars import ContextVar - - current_user: ContextVar[str | None] = ContextVar("current_user", default=None) - request = CreateResponse(model="m", input="hi") - context = ResponseContext(response_id="response-1", mode_flags=MagicMock()) - - def resolver(request: CreateResponse, context: ResponseContext) -> str | None: - del request, context - return current_user.get() - - async def resolve_for(user_id: str) -> ResolvedSessionIsolation: - token = current_user.set(user_id) - try: - await asyncio.sleep(0) - return await resolve_session_isolation(resolver, request, context, is_hosted=True) - finally: - current_user.reset(token) - - first, second = await asyncio.gather(resolve_for("user-1"), resolve_for("user-2")) - - assert first.key == "user-1" - assert second.key == "user-2" - assert first.fingerprint != second.fingerprint - - @pytest.mark.parametrize( - "key", - [".", "..", "../escape", "user/name", r"user\name", "user:name", "line\nbreak", "trailing."], - ) - async def test_invalid_isolation_key_is_rejected(self, key: str) -> None: - request = CreateResponse(model="m", input="hi") - context = ResponseContext(response_id="response-1", mode_flags=MagicMock()) - - with pytest.raises(ValueError, match="isolation key"): - await resolve_session_isolation( - lambda request, context: key, - request, - context, - is_hosted=True, - ) + def test_foundry_session_store_is_public_and_experimental(self) -> None: + assert FoundrySessionStore.__feature_stage__ == "experimental" # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert FoundrySessionStore.__feature_id__ == ExperimentalFeature.SESSION_STORE.value # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert FoundrySessionStore.__doc__ is not None + assert ".. warning:: Experimental" in FoundrySessionStore.__doc__ def test_response_id_partition_supports_current_legacy_and_raw_ids(self) -> None: current_partition = "a" * 18 @@ -479,23 +414,24 @@ def test_session_conversation_key_prefers_conversation_then_previous_then_respon == response_partition ) - def test_custom_store_key_partitions_by_agent_and_isolation(self) -> None: + def test_custom_store_key_partitions_by_agent_and_request_user(self) -> None: agent = _make_agent() other_agent = _make_agent() other_agent.name = "Other Agent" - isolation_a = ResolvedSessionIsolation("user-a", "user-user-a", "a" * 64) - isolation_b = ResolvedSessionIsolation("user-b", "user-user-b", "b" * 64) - key_a = _custom_session_store_key(agent, "conversation-1", isolation=isolation_a) - key_b = _custom_session_store_key(agent, "conversation-1", isolation=isolation_b) - key_other_agent = _custom_session_store_key(other_agent, "conversation-1", isolation=isolation_a) + with _request_context(user_id="user-a"): + key_a = _custom_session_store_key(agent, "conversation-1") + key_other_agent = _custom_session_store_key(other_agent, "conversation-1") + repeated_key_a = _custom_session_store_key(agent, "conversation-1") + with _request_context(user_id="user-b"): + key_b = _custom_session_store_key(agent, "conversation-1") assert key_a != key_b assert key_a != key_other_agent assert key_a.startswith("foundry_") assert len(key_a) == len("foundry_") + 64 assert all(character.isascii() and (character.isalnum() or character in "-_") for character in key_a) - assert key_a == _custom_session_store_key(agent, "conversation-1", isolation=isolation_a) + assert key_a == repeated_key_a def test_conversation_object_id_preserves_safe_values_and_hashes_unsafe_values(self) -> None: assert _conversation_object_id("conversation-1") == "conversation-1" @@ -533,38 +469,37 @@ def test_durable_storage_root_matches_hosted_and_local_layouts(self, tmp_path: P class TestAgentSessionPersistence: async def test_file_store_uses_per_user_child_directory_and_preserves_format(self, tmp_path: Path) -> None: - template = IsolationKeyScopedFileSessionStore(tmp_path, serialization_format="msgpack") + template = FoundrySessionStore(tmp_path, serialization_format="msgpack") server = _make_server(_make_agent(), session_store=template) - isolation_a = ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64) - isolation_b = ResolvedSessionIsolation("user-B", "user-user-B", "b" * 64) store = server._session_store # pyright: ignore[reportPrivateUsage] - assert isinstance(store, IsolationKeyScopedFileSessionStore) + assert isinstance(store, FoundrySessionStore) assert store.storage_path == tmp_path assert store.serialization_format == "msgpack" - with store.use_isolation(isolation_a): + with _request_context(user_id="user-A"): await store.set("conversation-1", AgentSession(session_id="conversation-1")) - with store.use_isolation(isolation_b): + with _request_context(user_id="user-B"): await store.set("conversation-1", AgentSession(session_id="conversation-1")) - assert (tmp_path / "user-user-A" / "conversation-1.msgpack").is_file() - assert (tmp_path / "user-user-B" / "conversation-1.msgpack").is_file() + user_a_directory = f"user-{hashlib.sha256(b'user-A').hexdigest()}" + user_b_directory = f"user-{hashlib.sha256(b'user-B').hexdigest()}" + assert (tmp_path / user_a_directory / "conversation-1.msgpack").is_file() + assert (tmp_path / user_b_directory / "conversation-1.msgpack").is_file() async def test_scoped_file_store_reuses_base_filename_safety(self, tmp_path: Path) -> None: - store = IsolationKeyScopedFileSessionStore(tmp_path) - isolation = ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64) + store = FoundrySessionStore(tmp_path) + user_directory = tmp_path / f"user-{hashlib.sha256(b'user-A').hexdigest()}" - with store.use_isolation(isolation): + with _request_context(user_id="user-A"): await store.set("CON", AgentSession(session_id="conversation-1")) - assert not (tmp_path / "user-user-A" / "CON.json").exists() - assert len(list((tmp_path / "user-user-A").glob("~session-*.json"))) == 1 + assert not (user_directory / "CON.json").exists() + assert len(list(user_directory.glob("~session-*.json"))) == 1 async def test_scoped_file_store_rejects_symlinked_session_leaf(self, tmp_path: Path) -> None: - store = IsolationKeyScopedFileSessionStore(tmp_path) - isolation = ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64) - user_directory = tmp_path / "user-user-A" + store = FoundrySessionStore(tmp_path) + user_directory = tmp_path / f"user-{hashlib.sha256(b'user-A').hexdigest()}" user_directory.mkdir() outside_file = tmp_path / "outside.json" outside_file.write_text("outside", encoding="utf-8") @@ -573,40 +508,27 @@ async def test_scoped_file_store_rejects_symlinked_session_leaf(self, tmp_path: except OSError as exc: pytest.skip(f"Symlinks are not available: {exc}") - with store.use_isolation(isolation), pytest.raises(ValueError, match="escaped storage directory"): + with _request_context(user_id="user-A"), pytest.raises(ValueError, match="escaped storage directory"): await store.get("conversation-1") async def test_scoped_file_store_rejects_symlinked_isolation_directory(self, tmp_path: Path) -> None: - store = IsolationKeyScopedFileSessionStore(tmp_path) - isolation = ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64) + store = FoundrySessionStore(tmp_path) + user_directory = f"user-{hashlib.sha256(b'user-A').hexdigest()}" outside_directory = tmp_path.parent / f"{tmp_path.name}-outside" outside_directory.mkdir() try: - (tmp_path / "user-user-A").symlink_to(outside_directory, target_is_directory=True) + (tmp_path / user_directory).symlink_to(outside_directory, target_is_directory=True) except OSError as exc: pytest.skip(f"Symlinks are not available: {exc}") - with store.use_isolation(isolation), pytest.raises(ValueError, match="Session directory escaped"): + with _request_context(user_id="user-A"), pytest.raises(ValueError, match="Session directory escaped"): await store.get("conversation-1") - async def test_scoped_file_store_rejects_traversal_directory_segment(self, tmp_path: Path) -> None: - store = IsolationKeyScopedFileSessionStore(tmp_path) - bypassed_resolver = ResolvedSessionIsolation("user-A", "../../Users", "a" * 64) - - with store.use_isolation(bypassed_resolver), pytest.raises(ValueError, match="Session directory escaped"): - await store.get("conversation-1") - - def test_session_isolation_stamp_rejects_mismatch(self) -> None: - session = AgentSession(session_id="conversation-1") - isolation_a = ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64) - isolation_b = ResolvedSessionIsolation("user-B", "user-user-B", "b" * 64) - - _stamp_or_validate_session_isolation(session, isolation_a) - assert session.state[_SESSION_ISOLATION_STATE_KEY] == "a" * 64 - assert "user-A" not in session.state.values() - - with pytest.raises(RuntimeError, match="identity context mismatch"): - _stamp_or_validate_session_isolation(session, isolation_b) + async def test_local_file_store_without_user_uses_root_directory(self, tmp_path: Path) -> None: + store = FoundrySessionStore(tmp_path) + with _request_context(): + await store.set("conversation-1", AgentSession(session_id="conversation-1")) + assert (tmp_path / "conversation-1.json").is_file() async def test_previous_response_chain_restores_session_state(self) -> None: seen_counts: list[int] = [] @@ -654,7 +576,6 @@ async def test_responses_history_is_not_duplicated_by_default_local_history(self _custom_session_store_key( agent, _conversation_object_id(conversation_key), - isolation=ResolvedSessionIsolation(None, None, None), ) ) assert stored is not None @@ -683,14 +604,14 @@ async def run_with_state(*args: Any, **kwargs: Any) -> AgentResponse: first_server = ResponsesHostServer( make_agent(), store=response_store, - session_store=IsolationKeyScopedFileSessionStore(tmp_path), + session_store=FoundrySessionStore(tmp_path), ) first = await _post(first_server) second_server = ResponsesHostServer( make_agent(), store=response_store, - session_store=IsolationKeyScopedFileSessionStore(tmp_path), + session_store=FoundrySessionStore(tmp_path), ) second = await _post(second_server, previous_response_id=first.json()["id"]) @@ -701,7 +622,7 @@ async def test_corrupt_file_session_fails_once_and_retry_starts_clean(self, tmp_ agent = _make_agent( response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("recovered")])]) ) - server = _make_server(agent, session_store=IsolationKeyScopedFileSessionStore(tmp_path)) + server = _make_server(agent, session_store=FoundrySessionStore(tmp_path)) corrupt_file = tmp_path / "conversation-1.json" corrupt_file.write_text("{not-json", encoding="utf-8") @@ -743,7 +664,6 @@ async def updates() -> AsyncIterator[AgentResponseUpdate]: _custom_session_store_key( agent, _conversation_object_id(conversation_key), - isolation=ResolvedSessionIsolation(None, None, None), ) ) @@ -772,7 +692,6 @@ async def failing_run(*args: Any, **kwargs: Any) -> AgentResponse: _custom_session_store_key( agent, _conversation_object_id(conversation_key), - isolation=ResolvedSessionIsolation(None, None, None), ) ) @@ -3779,54 +3698,46 @@ def test_url_encoded_traversal_is_treated_as_literal_segment(self, tmp_path: Any assert storage.storage_path.parent == root.resolve() assert storage.storage_path.name == "%2e%2e" - def test_isolation_scopes_storage_under_user_partition(self, tmp_path: Any) -> None: + def test_request_user_scopes_storage_under_user_partition(self, tmp_path: Any) -> None: helper = self._helper() root = tmp_path / "root" root.mkdir() - isolation = ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64) - storage = helper(str(root), "resp_abc123", isolation=isolation) + with _request_context(user_id="user-A"): + storage = helper(str(root), "resp_abc123") + user_directory = f"user-{hashlib.sha256(b'user-A').hexdigest()}" assert storage.storage_path.is_dir() - assert storage.storage_path == (root / "user-user-A" / "resp_abc123").resolve() + assert storage.storage_path == (root / user_directory / "resp_abc123").resolve() - def test_absent_isolation_uses_unscoped_layout(self, tmp_path: Any) -> None: + def test_absent_request_user_uses_unscoped_layout(self, tmp_path: Any) -> None: helper = self._helper() root = tmp_path / "root" root.mkdir() - storage = helper(str(root), "resp_abc123", isolation=ResolvedSessionIsolation(None, None, None)) + with _request_context(): + storage = helper(str(root), "resp_abc123") assert storage.storage_path == (root / "resp_abc123").resolve() - def test_distinct_isolation_keys_get_isolated_storage(self, tmp_path: Any) -> None: + def test_distinct_request_users_get_isolated_storage(self, tmp_path: Any) -> None: helper = self._helper() root = tmp_path / "root" root.mkdir() - a = helper( - str(root), - "shared_context", - isolation=ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64), - ) - b = helper( - str(root), - "shared_context", - isolation=ResolvedSessionIsolation("user-B", "user-user-B", "b" * 64), - ) + with _request_context(user_id="user-A"): + a = helper(str(root), "shared_context") + with _request_context(user_id="user-B"): + b = helper(str(root), "shared_context") + user_a_directory = root / f"user-{hashlib.sha256(b'user-A').hexdigest()}" + user_b_directory = root / f"user-{hashlib.sha256(b'user-B').hexdigest()}" assert a.storage_path != b.storage_path - assert a.storage_path.is_relative_to((root / "user-user-A").resolve()) - assert b.storage_path.is_relative_to((root / "user-user-B").resolve()) + assert a.storage_path.is_relative_to(user_a_directory.resolve()) + assert b.storage_path.is_relative_to(user_b_directory.resolve()) - def test_malicious_isolation_directory_is_rejected(self, tmp_path: Any) -> None: + def test_unsafe_request_user_is_hashed_before_path_join(self, tmp_path: Any) -> None: helper = self._helper() root = tmp_path / "root" root.mkdir() - before = sorted(p.name for p in tmp_path.iterdir()) - with pytest.raises(RuntimeError): - helper( - str(root), - "resp_abc123", - isolation=ResolvedSessionIsolation("bad", "../../escape", "a" * 64), - ) - after = sorted(p.name for p in tmp_path.iterdir()) - assert before == after - assert list(root.iterdir()) == [] + with _request_context(user_id="../../escape"): + storage = helper(str(root), "resp_abc123") + expected_directory = root / f"user-{hashlib.sha256(b'../../escape').hexdigest()}" + assert storage.storage_path == (expected_directory / "resp_abc123").resolve() @pytest.mark.parametrize( "context_field,bad_id", @@ -4006,36 +3917,34 @@ def _helper() -> Callable[..., str]: return _approval_storage_path - def test_isolation_scopes_path_under_base_directory(self, tmp_path: Any) -> None: + def test_request_user_scopes_path_under_base_directory(self, tmp_path: Any) -> None: from pathlib import Path helper = self._helper() base = tmp_path / "approvals" - isolation = ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64) - scoped = Path(helper(str(base), isolation)) + with _request_context(user_id="user-A"): + scoped = Path(helper(str(base))) + user_directory = f"user-{hashlib.sha256(b'user-A').hexdigest()}" assert scoped.name == "approval_requests.json" - assert scoped.parent.name == "user-user-A" + assert scoped.parent.name == user_directory assert scoped.parent.parent == (tmp_path / "approvals").resolve() - def test_distinct_isolation_keys_get_isolated_paths(self, tmp_path: Any) -> None: + def test_distinct_request_users_get_isolated_paths(self, tmp_path: Any) -> None: helper = self._helper() base = tmp_path / "approvals" - assert helper( - str(base), - ResolvedSessionIsolation("user-A", "user-user-A", "a" * 64), - ) != helper( - str(base), - ResolvedSessionIsolation("user-B", "user-user-B", "b" * 64), - ) + with _request_context(user_id="user-A"): + path_a = helper(str(base)) + with _request_context(user_id="user-B"): + path_b = helper(str(base)) + assert path_a != path_b - def test_malicious_isolation_directory_is_rejected(self, tmp_path: Any) -> None: + def test_unsafe_request_user_is_hashed_before_path_join(self, tmp_path: Any) -> None: helper = self._helper() base = tmp_path / "approvals" - with pytest.raises(RuntimeError): - helper( - str(base), - ResolvedSessionIsolation("bad", "../../escape", "a" * 64), - ) + with _request_context(user_id="../../escape"): + path = Path(helper(str(base))) + expected_directory = base / f"user-{hashlib.sha256(b'../../escape').hexdigest()}" + assert path == (expected_directory / "approval_requests.json").resolve() # region Agent lifecycle (lazy entry & OAuth consent surfacing) From 769aa6710f1e51858ed85aed2dd2351cf3ce4006 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 27 Jul 2026 22:47:02 +0200 Subject: [PATCH 5/8] Python: Reduce Foundry session helper layering Inline the single-use request user accessor while keeping separate context validation, fingerprint, and directory helpers for their distinct callers. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3e5c81ad-75e8-4e92-a883-8bbd676c6c8c --- .../agent_framework_foundry_hosting/_session_store.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py index 2ed0cdda70..e4cc971383 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py @@ -33,16 +33,11 @@ def _get_foundry_request_context( # pyright: ignore[reportUnusedFunction] return context -def _request_user_id_key() -> str | None: - """Return the platform user partition key for the active request.""" - # FoundryAgentRequestContext.user_id is populated from the same - # x-agent-user-id value exposed as ResponseContext.platform_context.user_id_key. - return get_request_context().user_id - - def _request_user_fingerprint() -> str | None: """Return a stable opaque fingerprint for the active request user.""" - user_id_key = _request_user_id_key() + # FoundryAgentRequestContext.user_id is populated from the same + # x-agent-user-id value exposed as ResponseContext.platform_context.user_id_key. + user_id_key = get_request_context().user_id return hashlib.sha256(user_id_key.encode("utf-8")).hexdigest() if user_id_key else None From ba214306dfe44a24f36c506d6c427fb13eaf5df4 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 27 Jul 2026 22:49:10 +0200 Subject: [PATCH 6/8] Python: Clarify Foundry request context validation Separate fail-fast request validation from context retrieval so Responses no longer appears to discard a returned context. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3e5c81ad-75e8-4e92-a883-8bbd676c6c8c --- .../agent_framework_foundry_hosting/_invocations.py | 6 ++++-- .../agent_framework_foundry_hosting/_responses.py | 5 +++-- .../agent_framework_foundry_hosting/_session_store.py | 9 ++++----- 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py index 4f38b2dc06..0e1eb0894f 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -1,12 +1,13 @@ # Copyright (c) Microsoft. All rights reserved. from agent_framework import AgentSession, SupportsAgentRun +from azure.ai.agentserver.core import get_request_context from azure.ai.agentserver.invocations import InvocationAgentServerHost from starlette.requests import Request from starlette.responses import Response, StreamingResponse from typing_extensions import Any, AsyncGenerator -from ._session_store import _get_foundry_request_context # pyright: ignore[reportPrivateUsage] +from ._session_store import _validate_foundry_request_context # pyright: ignore[reportPrivateUsage] class InvocationsHostServer(InvocationAgentServerHost): @@ -50,7 +51,8 @@ def _partition_key(self) -> str: Exceptions: RuntimeError: If the context doesn't contain the expected IDs. """ - context = _get_foundry_request_context(is_hosted=self.config.is_hosted) + context = get_request_context() + _validate_foundry_request_context(context, is_hosted=self.config.is_hosted) if self.config.is_hosted: if not context.session_id or not context.user_id: diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index f7f7bc8b39..7e3a6777ab 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -32,6 +32,7 @@ WorkflowAgent, ) from agent_framework.exceptions import AgentFrameworkException +from azure.ai.agentserver.core import get_request_context from azure.ai.agentserver.responses import ( ResponseContext, ResponseEventStream, @@ -122,9 +123,9 @@ from ._session_store import ( FoundrySessionStore, - _get_foundry_request_context, # pyright: ignore[reportPrivateUsage] _request_user_directory_segment, # pyright: ignore[reportPrivateUsage] _request_user_fingerprint, # pyright: ignore[reportPrivateUsage] + _validate_foundry_request_context, # pyright: ignore[reportPrivateUsage] ) logger = logging.getLogger(__name__) @@ -630,7 +631,7 @@ async def _handle_response( cancellation_signal: asyncio.Event, ) -> AsyncIterable[ResponseStreamEvent | dict[str, Any]]: """Handle the creation of a response.""" - _get_foundry_request_context(is_hosted=self.config.is_hosted) + _validate_foundry_request_context(get_request_context(), is_hosted=self.config.is_hosted) if self._is_workflow_agent: # Workflow agents are handled differently because they require checkpoint restoration diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py index e4cc971383..fd050b45fd 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py @@ -17,12 +17,12 @@ ) -def _get_foundry_request_context( # pyright: ignore[reportUnusedFunction] +def _validate_foundry_request_context( # pyright: ignore[reportUnusedFunction] + context: FoundryAgentRequestContext, *, is_hosted: bool, -) -> FoundryAgentRequestContext: - """Return the current request context and validate hosted v2 identity.""" - context = get_request_context() +) -> None: + """Validate that a hosted request contains protocol-v2 user identity.""" if is_hosted and context.call_id is None: raise RuntimeError(_PROTOCOL_V2_REQUIRED_MESSAGE) if is_hosted and not context.user_id: @@ -30,7 +30,6 @@ def _get_foundry_request_context( # pyright: ignore[reportUnusedFunction] "The hosted environment is missing the platform user ID in the request context. " "Please ensure that the request is coming from a valid Foundry platform service." ) - return context def _request_user_fingerprint() -> str | None: From e6a139f0a3b4f6399fbe0edab2c8e9b4146b6ddd Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 27 Jul 2026 22:51:46 +0200 Subject: [PATCH 7/8] Python: Share Foundry request context helpers Move protocol validation and user-scope derivation into a dedicated request-context module, leaving the session-store module focused on storage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3e5c81ad-75e8-4e92-a883-8bbd676c6c8c --- .../_invocations.py | 2 +- .../_request_context.py | 42 +++++++++++++++++++ .../_responses.py | 4 +- .../_session_store.py | 37 +--------------- 4 files changed, 46 insertions(+), 39 deletions(-) create mode 100644 python/packages/foundry_hosting/agent_framework_foundry_hosting/_request_context.py diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py index 0e1eb0894f..24f5bc8fe4 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -7,7 +7,7 @@ from starlette.responses import Response, StreamingResponse from typing_extensions import Any, AsyncGenerator -from ._session_store import _validate_foundry_request_context # pyright: ignore[reportPrivateUsage] +from ._request_context import _validate_foundry_request_context # pyright: ignore[reportPrivateUsage] class InvocationsHostServer(InvocationAgentServerHost): diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_request_context.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_request_context.py new file mode 100644 index 0000000000..74522d4a5c --- /dev/null +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_request_context.py @@ -0,0 +1,42 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import hashlib + +from azure.ai.agentserver.core import FoundryAgentRequestContext, get_request_context + +_PROTOCOL_V2_REQUIRED_MESSAGE = ( + "The hosted environment is running on protocol 1.0.0, but the agent requires protocol 2.0.0. " + "Please upgrade your agent protocol to 2.0.0 in `agent.manifest.yaml` or `agent.yaml`, or " + "downgrade the `agent-framework-foundry-hosting` package to `1.0.0a260625` or before to use 1.0.0." +) + + +def _validate_foundry_request_context( # pyright: ignore[reportUnusedFunction] + context: FoundryAgentRequestContext, + *, + is_hosted: bool, +) -> None: + """Validate that a hosted request contains protocol-v2 user identity.""" + if is_hosted and context.call_id is None: + raise RuntimeError(_PROTOCOL_V2_REQUIRED_MESSAGE) + if is_hosted and not context.user_id: + raise RuntimeError( + "The hosted environment is missing the platform user ID in the request context. " + "Please ensure that the request is coming from a valid Foundry platform service." + ) + + +def _request_user_fingerprint() -> str | None: + """Return a stable opaque fingerprint for the active request user.""" + # FoundryAgentRequestContext.user_id is populated from the same + # x-agent-user-id value exposed as ResponseContext.platform_context.user_id_key. + user_id_key = get_request_context().user_id + return hashlib.sha256(user_id_key.encode("utf-8")).hexdigest() if user_id_key else None + + +def _request_user_directory_segment() -> str | None: # pyright: ignore[reportUnusedFunction] + """Return the safe on-disk directory segment for the active request user.""" + fingerprint = _request_user_fingerprint() + return f"user-{fingerprint}" if fingerprint else None diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index 7e3a6777ab..4c1464f22c 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -121,12 +121,12 @@ from mcp import McpError from typing_extensions import Any -from ._session_store import ( - FoundrySessionStore, +from ._request_context import ( _request_user_directory_segment, # pyright: ignore[reportPrivateUsage] _request_user_fingerprint, # pyright: ignore[reportPrivateUsage] _validate_foundry_request_context, # pyright: ignore[reportPrivateUsage] ) +from ._session_store import FoundrySessionStore logger = logging.getLogger(__name__) diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py index fd050b45fd..f58296f0c1 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py @@ -2,48 +2,13 @@ from __future__ import annotations -import hashlib from pathlib import Path from typing import Literal from agent_framework import ExperimentalFeature, FileSessionStore from agent_framework._feature_stage import experimental -from azure.ai.agentserver.core import FoundryAgentRequestContext, get_request_context -_PROTOCOL_V2_REQUIRED_MESSAGE = ( - "The hosted environment is running on protocol 1.0.0, but the agent requires protocol 2.0.0. " - "Please upgrade your agent protocol to 2.0.0 in `agent.manifest.yaml` or `agent.yaml`, or " - "downgrade the `agent-framework-foundry-hosting` package to `1.0.0a260625` or before to use 1.0.0." -) - - -def _validate_foundry_request_context( # pyright: ignore[reportUnusedFunction] - context: FoundryAgentRequestContext, - *, - is_hosted: bool, -) -> None: - """Validate that a hosted request contains protocol-v2 user identity.""" - if is_hosted and context.call_id is None: - raise RuntimeError(_PROTOCOL_V2_REQUIRED_MESSAGE) - if is_hosted and not context.user_id: - raise RuntimeError( - "The hosted environment is missing the platform user ID in the request context. " - "Please ensure that the request is coming from a valid Foundry platform service." - ) - - -def _request_user_fingerprint() -> str | None: - """Return a stable opaque fingerprint for the active request user.""" - # FoundryAgentRequestContext.user_id is populated from the same - # x-agent-user-id value exposed as ResponseContext.platform_context.user_id_key. - user_id_key = get_request_context().user_id - return hashlib.sha256(user_id_key.encode("utf-8")).hexdigest() if user_id_key else None - - -def _request_user_directory_segment() -> str | None: - """Return the safe on-disk directory segment for the active request user.""" - fingerprint = _request_user_fingerprint() - return f"user-{fingerprint}" if fingerprint else None +from ._request_context import _request_user_directory_segment # pyright: ignore[reportPrivateUsage] @experimental(feature_id=ExperimentalFeature.SESSION_STORE) From 253aee44405a552b2b046c50e9f866f2194a3c0f Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 27 Jul 2026 23:23:15 +0200 Subject: [PATCH 8/8] Restore Foundry checkpoint storage paths Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3e5c81ad-75e8-4e92-a883-8bbd676c6c8c --- python/packages/foundry_hosting/README.md | 18 ++- .../_responses.py | 104 ++++++++++------ .../foundry_hosting/tests/test_responses.py | 113 +++++++++++------- 3 files changed, 149 insertions(+), 86 deletions(-) diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index 83dfa93a2d..be0a3cc71f 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -10,8 +10,7 @@ by the Agent Server request context's platform user key and the Responses conversation partition. Pass `session_store=` to use another `SessionStore` implementation. -Workflow agents continue to use checkpoint storage; their checkpoints share the -same `$HOME/.checkpoints` or local `.checkpoints` root. +Workflow agents continue to use their existing checkpoint storage layout. ## Foundry session isolation @@ -22,15 +21,24 @@ hashes the platform `user_id` (the same `x-agent-user-id` value exposed as directory. This makes the user partition path-safe and avoids persisting the raw platform identity. -File-backed state is partitioned under the validated identity: +Regular-agent session snapshots use a hashed platform user directory: ```text .checkpoints/ sessions/user-/.json - checkpoints/user-// - function-approvals/user-/approval_requests.json ``` +Workflow checkpoints and function approvals preserve the existing Foundry +Hosting layout. Hosted paths insert the validated raw platform user ID: + +```text +/.checkpoints/// +/.function_approvals//approval_requests.json +``` + +Local workflow checkpoints use `{cwd}/.checkpoints//`, and local +function approvals remain in memory. + Hosted requests require container protocol `2.0.0`. The v2-only request `call_id` is checked before session, checkpoint, or approval storage is used, and a missing platform user ID fails closed. Local requests may remain unscoped. diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index 4c1464f22c..c19ce303f8 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -122,7 +122,6 @@ from typing_extensions import Any from ._request_context import ( - _request_user_directory_segment, # pyright: ignore[reportPrivateUsage] _request_user_fingerprint, # pyright: ignore[reportPrivateUsage] _validate_foundry_request_context, # pyright: ignore[reportPrivateUsage] ) @@ -139,9 +138,6 @@ _DEFAULT_HOSTED_SESSION_DATA_DIRECTORY = "/home/session" _CHECKPOINT_DIRECTORY_NAME = ".checkpoints" _SESSION_DIRECTORY_NAME = "sessions" -_WORKFLOW_CHECKPOINT_DIRECTORY_NAME = "checkpoints" -_FUNCTION_APPROVAL_DIRECTORY_NAME = "function-approvals" -_FUNCTION_APPROVAL_FILE_NAME = "approval_requests.json" _HOSTED_RESPONSES_HISTORY_SOURCE_ID = "_foundry_responses_history" @@ -239,14 +235,16 @@ async def load_approval_request(self, approval_request_id: str) -> Content: return await asyncio.to_thread(self._load_sync, approval_request_id) -def _validate_path_segment(segment: str, *, kind: Literal["context id"]) -> None: +def _validate_path_segment(segment: str, *, kind: Literal["context id", "user id"]) -> None: """Validate that ``segment`` is a single safe path component (CWE-22). ``segment`` originates from caller-controlled fields (such as - ``previous_response_id``) or server-generated fields (``conversation_id`` / - ``response_id``). It must be treated as an untrusted single path segment: - path separators, drive letters, parent references and similar would - otherwise let the resulting directory escape the configured storage root. + ``previous_response_id``), server-generated fields (``conversation_id`` / + ``response_id``), or the platform-injected per-user partition key + (``x-agent-user-id``). In every case it must be treated as an untrusted + single path segment: path separators, drive letters, parent references and + similar would otherwise let the resulting directory escape the configured + storage root. We deliberately do not URL-decode the value here: the hosting layer never decodes these ids before joining them, so forms such as ``%2e%2e`` are @@ -290,7 +288,7 @@ def _resolve_durable_storage_root( home_directory: str | None, current_directory: str, ) -> Path: - """Resolve the shared root for Foundry session snapshots and workflow checkpoints.""" + """Resolve the root for regular-agent Foundry session snapshots.""" if is_hosted: if home_directory is not None and _is_usable_hosted_home_directory(home_directory): base_directory = Path(home_directory).expanduser() @@ -362,21 +360,34 @@ def _is_hosted_responses_history_sentinel(provider: ContextProvider) -> bool: def _checkpoint_storage_for_context( - root: str | Path, + root: str, context_id: str, + *, + user_id: str | None = None, ) -> FileCheckpointStorage: """Build a ``FileCheckpointStorage`` for ``context_id`` rooted under ``root``. - Scoped requests use ``/user-/``. Local - unscoped requests use ``/``. + When the platform supplies a per-user partition key (``user_id``, from the + ``x-agent-user-id`` header on container protocol v2), the per-conversation + checkpoint directory is nested under it: ``//``. + This isolates each tenant's workflow state so one user can never restore or + observe another user's checkpoint, even with a guessed or forged + ``context_id``. An absent (``None``) or empty ``user_id`` -- local + development or protocol v1 -- falls back to the unscoped + ``/`` layout. + + Both ``context_id`` and ``user_id`` are validated as single safe path + segments, and each resolved directory is verified to stay under its parent + before any directory is created on disk (CWE-22). """ _validate_path_segment(context_id, kind="context id") base_path = Path(root).resolve() - if user_directory := _request_user_directory_segment(): - user_path = (base_path / user_directory).resolve() + if user_id: + _validate_path_segment(user_id, kind="user id") + user_path = (base_path / user_id).resolve() if not user_path.is_relative_to(base_path): - raise RuntimeError(f"Invalid Foundry user directory: {user_directory!r}") + raise RuntimeError(f"Invalid user id: {user_id!r}") base_path = user_path storage_path = (base_path / context_id).resolve() @@ -390,16 +401,23 @@ def _checkpoint_storage_for_context( ) -def _approval_storage_path( - root: str | Path, -) -> str: - """Return the approval file path for the active request user.""" - base_path = Path(root).resolve() - if user_directory := _request_user_directory_segment(): - base_path = (base_path / user_directory).resolve() - if not base_path.is_relative_to(Path(root).resolve()): - raise RuntimeError(f"Invalid Foundry user directory: {user_directory!r}") - return str(base_path / _FUNCTION_APPROVAL_FILE_NAME) +def _approval_storage_path_for_user(base_path: str, user_id: str) -> str: + """Return the per-user approval storage file path under the base directory. + + Inserts the validated ``user_id`` as a directory segment between the base + directory and the file name (``//``), mirroring the + per-user checkpoint partitioning so one tenant can never read another + tenant's saved approval requests. The user id is validated as a single safe + path segment and the resulting directory is verified to stay under the base + directory before use (CWE-22). + """ + _validate_path_segment(user_id, kind="user id") + directory, filename = os.path.split(base_path) + base_dir = Path(directory or ".").resolve() + user_dir = (base_dir / user_id).resolve() + if not user_dir.is_relative_to(base_dir): + raise RuntimeError(f"Invalid user id: {user_id!r}") + return str(user_dir / filename) # endregion Approval Storage @@ -483,6 +501,10 @@ def consent_url_from_error(exc: BaseException) -> list[ConsentError] | None: class ResponsesHostServer(ResponsesAgentServerHost): """A responses server host for an agent.""" + # TODO(@taochen): Allow a different checkpoint storage that stores checkpoints externally + CHECKPOINT_STORAGE_PATH = "/.checkpoints" + FUNCTION_APPROVAL_STORAGE_PATH = "/.function_approvals/approval_requests.json" + def __init__( self, agent: SupportsAgentRun, @@ -543,7 +565,11 @@ def __init__( "There should not be a checkpoint storage already present in the workflow agent. " "The hosting infrastructure will manage checkpoints instead." ) - self._checkpoint_storage_path = str(self._storage_root / _WORKFLOW_CHECKPOINT_DIRECTORY_NAME) + self._checkpoint_storage_path = ( + self.CHECKPOINT_STORAGE_PATH + if self.config.is_hosted + else os.path.join(os.getcwd(), self.CHECKPOINT_STORAGE_PATH.lstrip("/")) + ) self._is_workflow_agent = True if ( @@ -571,10 +597,11 @@ def __init__( elif isinstance(session_store, FileSessionStore) and not isinstance(session_store, FoundrySessionStore): raise ValueError("ResponsesHostServer requires FoundrySessionStore for file-backed session storage.") self._session_store = session_store - self._approval_storage: ApprovalStorage = InMemoryFunctionApprovalStorage() - self._approval_storage_root = self._storage_root / _FUNCTION_APPROVAL_DIRECTORY_NAME - # Per-user hosted approval stores share a process-local lock per user. - # Local development keeps the single in-memory store. + self._approval_storage: ApprovalStorage = ( + FileBasedFunctionApprovalStorage(self.FUNCTION_APPROVAL_STORAGE_PATH) + if self.config.is_hosted + else InMemoryFunctionApprovalStorage() + ) self._approval_storages_by_user: dict[str, ApprovalStorage] = {} # Lazy agent lifecycle: the agent (and any MCP tools it owns) is entered on # the first request rather than at server startup, so that authentication @@ -615,13 +642,15 @@ async def _cleanup_agent(self) -> None: def _approval_storage_for_request(self) -> ApprovalStorage: """Return the hosted approval store for the active request user.""" - user_fingerprint = _request_user_fingerprint() - if not self.config.is_hosted or user_fingerprint is None: + user_id = get_request_context().user_id + if not self.config.is_hosted or not user_id: return self._approval_storage - storage = self._approval_storages_by_user.get(user_fingerprint) + storage = self._approval_storages_by_user.get(user_id) if storage is None: - storage = FileBasedFunctionApprovalStorage(_approval_storage_path(self._approval_storage_root)) - self._approval_storages_by_user[user_fingerprint] = storage + storage = FileBasedFunctionApprovalStorage( + _approval_storage_path_for_user(self.FUNCTION_APPROVAL_STORAGE_PATH, user_id) + ) + self._approval_storages_by_user[user_id] = storage return storage async def _handle_response( @@ -831,12 +860,14 @@ async def _handle_inner_workflow( # on every turn we restore the latest checkpoint and feed the new # input back into the start executor as a continuation rather than # a fresh run. + user_id = get_request_context().user_id latest_checkpoint_id: str | None = None restore_storage: FileCheckpointStorage | None = None if context_id is not None: restore_storage = _checkpoint_storage_for_context( self._checkpoint_storage_path, context_id, + user_id=user_id, ) latest_checkpoint = await restore_storage.get_latest(workflow_name=self._agent.workflow.name) if latest_checkpoint is not None: @@ -854,6 +885,7 @@ async def _handle_inner_workflow( write_storage = _checkpoint_storage_for_context( self._checkpoint_storage_path, write_context_id, + user_id=user_id, ) # Multi-turn pattern: when we have a prior checkpoint, restore it diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index bb497ce119..297e8e413b 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -3698,46 +3698,58 @@ def test_url_encoded_traversal_is_treated_as_literal_segment(self, tmp_path: Any assert storage.storage_path.parent == root.resolve() assert storage.storage_path.name == "%2e%2e" - def test_request_user_scopes_storage_under_user_partition(self, tmp_path: Any) -> None: + def test_user_id_scopes_storage_under_user_partition(self, tmp_path: Any) -> None: + """A per-user partition key nests the context dir under ``/``.""" helper = self._helper() root = tmp_path / "root" root.mkdir() - with _request_context(user_id="user-A"): - storage = helper(str(root), "resp_abc123") - user_directory = f"user-{hashlib.sha256(b'user-A').hexdigest()}" + storage = helper(str(root), "resp_abc123", user_id="user-A") assert storage.storage_path.is_dir() - assert storage.storage_path == (root / user_directory / "resp_abc123").resolve() + assert storage.storage_path == (root / "user-A" / "resp_abc123").resolve() - def test_absent_request_user_uses_unscoped_layout(self, tmp_path: Any) -> None: + @pytest.mark.parametrize("absent_user_id", [None, ""]) + def test_absent_user_id_uses_unscoped_layout(self, tmp_path: Any, absent_user_id: str | None) -> None: + """``None``/empty user id (local dev or protocol v1) falls back to the unscoped layout.""" helper = self._helper() root = tmp_path / "root" root.mkdir() - with _request_context(): - storage = helper(str(root), "resp_abc123") + storage = helper(str(root), "resp_abc123", user_id=absent_user_id) assert storage.storage_path == (root / "resp_abc123").resolve() - def test_distinct_request_users_get_isolated_storage(self, tmp_path: Any) -> None: + def test_distinct_users_get_isolated_storage(self, tmp_path: Any) -> None: + """Two users sharing a context id must not resolve to the same directory.""" helper = self._helper() root = tmp_path / "root" root.mkdir() - with _request_context(user_id="user-A"): - a = helper(str(root), "shared_context") - with _request_context(user_id="user-B"): - b = helper(str(root), "shared_context") - user_a_directory = root / f"user-{hashlib.sha256(b'user-A').hexdigest()}" - user_b_directory = root / f"user-{hashlib.sha256(b'user-B').hexdigest()}" + a = helper(str(root), "shared_context", user_id="user-A") + b = helper(str(root), "shared_context", user_id="user-B") assert a.storage_path != b.storage_path - assert a.storage_path.is_relative_to(user_a_directory.resolve()) - assert b.storage_path.is_relative_to(user_b_directory.resolve()) + assert a.storage_path.is_relative_to((root / "user-A").resolve()) + assert b.storage_path.is_relative_to((root / "user-B").resolve()) - def test_unsafe_request_user_is_hashed_before_path_join(self, tmp_path: Any) -> None: + @pytest.mark.parametrize( + "bad_user_id", + [ + "../../escape", + "..", + ".", + "/tmp/escape", + "C:\\temp\\escape", + "user/../../escape", + "with\x00null", + "a/b", + ], + ) + def test_malicious_user_id_is_rejected(self, tmp_path: Any, bad_user_id: str) -> None: helper = self._helper() root = tmp_path / "root" root.mkdir() - with _request_context(user_id="../../escape"): - storage = helper(str(root), "resp_abc123") - expected_directory = root / f"user-{hashlib.sha256(b'../../escape').hexdigest()}" - assert storage.storage_path == (expected_directory / "resp_abc123").resolve() + before = sorted(p.name for p in tmp_path.iterdir()) + with pytest.raises(RuntimeError): + helper(str(root), "resp_abc123", user_id=bad_user_id) + after = sorted(p.name for p in tmp_path.iterdir()) + assert before == after, f"Unexpected filesystem artifacts created for user id {bad_user_id!r}" + assert list(root.iterdir()) == [] @pytest.mark.parametrize( "context_field,bad_id", @@ -3907,44 +3919,55 @@ async def test_malicious_context_id_rejected_e2e(self, tmp_path: Any, context_fi class TestApprovalStoragePathValidation: - """Path containment and per-user scoping tests for approval storage.""" + """Path-traversal and per-user scoping tests for function approval storage. + + Mirrors the checkpoint validation: the per-user approval directory is + derived by joining the platform-injected ``x-agent-user-id`` partition key + under the base approval directory, and the user id must be a single safe + path segment (CWE-22). + """ @staticmethod def _helper() -> Callable[..., str]: from agent_framework_foundry_hosting._responses import ( # pyright: ignore[reportPrivateUsage] - _approval_storage_path, + _approval_storage_path_for_user, ) - return _approval_storage_path + return _approval_storage_path_for_user - def test_request_user_scopes_path_under_base_directory(self, tmp_path: Any) -> None: + def test_user_id_scopes_path_under_base_directory(self, tmp_path: Any) -> None: from pathlib import Path helper = self._helper() - base = tmp_path / "approvals" - with _request_context(user_id="user-A"): - scoped = Path(helper(str(base))) - user_directory = f"user-{hashlib.sha256(b'user-A').hexdigest()}" - assert scoped.name == "approval_requests.json" - assert scoped.parent.name == user_directory + base = tmp_path / "approvals" / "requests.json" + scoped = Path(helper(str(base), "user-A")) + assert scoped.name == "requests.json" + assert scoped.parent.name == "user-A" assert scoped.parent.parent == (tmp_path / "approvals").resolve() - def test_distinct_request_users_get_isolated_paths(self, tmp_path: Any) -> None: + def test_distinct_users_get_isolated_paths(self, tmp_path: Any) -> None: helper = self._helper() - base = tmp_path / "approvals" - with _request_context(user_id="user-A"): - path_a = helper(str(base)) - with _request_context(user_id="user-B"): - path_b = helper(str(base)) - assert path_a != path_b + base = tmp_path / "approvals" / "requests.json" + assert helper(str(base), "user-A") != helper(str(base), "user-B") - def test_unsafe_request_user_is_hashed_before_path_join(self, tmp_path: Any) -> None: + @pytest.mark.parametrize( + "bad_user_id", + [ + "../../escape", + "..", + ".", + "/tmp/escape", + "C:\\temp\\escape", + "user/../../escape", + "with\x00null", + "a/b", + ], + ) + def test_malicious_user_id_is_rejected(self, tmp_path: Any, bad_user_id: str) -> None: helper = self._helper() - base = tmp_path / "approvals" - with _request_context(user_id="../../escape"): - path = Path(helper(str(base))) - expected_directory = base / f"user-{hashlib.sha256(b'../../escape').hexdigest()}" - assert path == (expected_directory / "approval_requests.json").resolve() + base = tmp_path / "approvals" / "requests.json" + with pytest.raises(RuntimeError): + helper(str(base), bad_user_id) # region Agent lifecycle (lazy entry & OAuth consent surfacing)