From 5dfe089ac94c36228120faf1cd75e9affbb7cfa7 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 24 Jul 2026 16:30:17 +0200 Subject: [PATCH 01/14] 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 00000000000..4f06334c909 --- /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 43ab1b73ad4..91931ac8b2e 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. @@ -144,6 +146,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 b99f0ec7dad..c423aadcc8c 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 9e8033c8858..763b9ccdaec 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 7ebfaa59f44..4529a19de9c 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 ae09f7610b2..62a385aa8dc 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 d4547535e85..cf9a7ce004a 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 39f9c74f4bb..7f6c26ffc08 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 04d80cf0491..cc30512ab39 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 9a6dab9a7c1..f57a2cb23db 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 ac535dc693e..e2ea3bdc1aa 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 ba05281d1fe..bf8be941ea9 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 f5b2733e577..a23f9e58d93 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 c4f4061d42c..43179dfdd0f 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 c0794fd4b05..93c243c5409 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 9233f5a7631..36e9027de72 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 114ff4f0055..d80c551e906 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 00000000000..8efb57afaa9 --- /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 321ac373080..c05795b7af0 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 f10b8b839df..cf97edd7ef1 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 b267bec0151..968593998e2 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 361d8418676..9f91940f4f9 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 0cb4af0041b..1cf3a105de4 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 0e376fd45c6..b0dd74bbcf1 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 698ca1d3037..19ef51b57c0 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 148221e6b44..91638607e5f 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 deb166fc966..a163bd598ce 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 10bc5e2243a..30ed10e8720 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 cc9959b686a..065669e1612 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 854a76d84b1..6ea364a5865 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 56901e03610..35747e6dc77 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 9b976be1c4b..995eacda897 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 0b3503b0262..63c06c69986 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 d4843691ffc..239a8cecf7a 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 7a8effc0636..f01dc589825 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 674ca0c9ed2dc8532798ce04fbf100bf165ce700 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 27 Jul 2026 11:06:22 +0200 Subject: [PATCH 02/14] 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 4f06334c909..d981f873a48 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 4529a19de9c..e8941385990 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 f57a2cb23db..662350770f9 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 e2ea3bdc1aa..bc1e70f3777 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 43179dfdd0f..b7d17545fcb 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 d80c551e906..9d1d04d98c6 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 8efb57afaa9..726a4de737d 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 c05795b7af0..53ee26c8471 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 b0dd74bbcf1..05f42da4bb0 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 91638607e5f..9c86f7515ac 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 a163bd598ce..f4c991fca50 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 00000000000..1eaa9a718a2 --- /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 f01dc589825..57a0ca49135 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 67b67eafb25bfa4aeb6599fdde2b8fdb8bd63a9e Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 27 Jul 2026 15:35:55 +0200 Subject: [PATCH 03/14] 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 662350770f9..c7a816261b1 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 b7d17545fcb..a0aaabbb3a0 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 968593998e2..6c488f2d7a5 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 9f91940f4f9..361d8418676 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 1cf3a105de4..0cb4af0041b 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 239a8cecf7a..d4843691ffc 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 459e476291acff59df32b7f35ac84625d6e0c729 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 27 Jul 2026 20:59:27 +0200 Subject: [PATCH 04/14] 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 d981f873a48..f24eb6b80cd 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 763b9ccdaec..12ca2fa766d 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 e8941385990..7f166f36b5d 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 c49b042567f..7e398e9c508 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 faadc99239a..dd9f56d07eb 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 93c243c5409..83dfa93a2dd 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 36e9027de72..56d65716703 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 f01c7a94684..4f38b2dc06e 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 9d1d04d98c6..f7f7bc8b39c 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 726a4de737d..00000000000 --- 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 00000000000..2ed0cdda702 --- /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 05bcdc1d5ba..0fad1c72a01 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 53ee26c8471..bb497ce1195 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 2dc827f683f532a353dd7b3307339da2022be78c Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 27 Jul 2026 22:47:02 +0200 Subject: [PATCH 05/14] 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 2ed0cdda702..e4cc971383b 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 1b1ca4d2f29718b6fd036cabc320d79464d7545c Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 27 Jul 2026 22:49:10 +0200 Subject: [PATCH 06/14] 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 4f38b2dc06e..0e1eb0894fa 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 f7f7bc8b39c..7e3a6777ab2 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 e4cc971383b..fd050b45fd5 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 cbe78a5e4552e6dc987b07a51a0094ba50cea4f9 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 27 Jul 2026 22:51:46 +0200 Subject: [PATCH 07/14] 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 0e1eb0894fa..24f5bc8fe4f 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 00000000000..74522d4a5c7 --- /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 7e3a6777ab2..4c1464f22c9 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 fd050b45fd5..f58296f0c1c 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 00268b786c55e3472e249c71c1ae22952488649e Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 27 Jul 2026 23:23:15 +0200 Subject: [PATCH 08/14] 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 83dfa93a2dd..be0a3cc71fe 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 4c1464f22c9..c19ce303f82 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 bb497ce1195..297e8e413b2 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) From 59fe012ad50b36bcca256ae3e6e7f36fcb20822b Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 27 Jul 2026 23:41:32 +0200 Subject: [PATCH 09/14] Simplify Foundry session storage paths Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3e5c81ad-75e8-4e92-a883-8bbd676c6c8c --- ...0033-python-session-store-serialization.md | 16 +- python/packages/foundry_hosting/README.md | 26 +-- .../_request_context.py | 44 +++-- .../_responses.py | 123 ++++--------- .../_session_store.py | 5 +- .../foundry_hosting/tests/test_responses.py | 167 ++++++++++++------ 6 files changed, 208 insertions(+), 173 deletions(-) diff --git a/docs/decisions/0033-python-session-store-serialization.md b/docs/decisions/0033-python-session-store-serialization.md index f24eb6b80cd..85f2fce957f 100644 --- a/docs/decisions/0033-python-session-store-serialization.md +++ b/docs/decisions/0033-python-session-store-serialization.md @@ -204,12 +204,16 @@ 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. +Foundry Hosting exposes an experimental `FoundrySessionStore`, which is the +default `ResponsesHostServer` store when hosted; local hosting defaults to the +in-memory `SessionStore`. `FoundrySessionStore` currently subclasses +`FileSessionStore`, stores snapshots under +`/.sessions//.json`, and derives the validated platform +IDs from `azure.ai.agentserver.core.get_request_context()`. Callers can +explicitly override either default through `session_store=`. 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 diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index be0a3cc71fe..32b4f3f2566 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -4,11 +4,11 @@ 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 the experimental `FoundrySessionStore` under `$HOME/.checkpoints/sessions` -when hosted and `{cwd}/.checkpoints/sessions` locally. The store is partitioned -by the Agent Server request context's platform user key and the Responses -conversation partition. Pass `session_store=` to use another `SessionStore` -implementation. +uses the experimental `FoundrySessionStore` under `/.sessions` when hosted and +an in-memory `SessionStore` locally. Hosted snapshots are partitioned by the +Agent Server request context's platform user ID, and their filenames come from +its platform session ID. Pass `session_store=` to explicitly override either +default. Workflow agents continue to use their existing checkpoint storage layout. @@ -16,16 +16,14 @@ Workflow agents continue to use their existing checkpoint storage layout. `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. +validates the platform `user_id` (the same `x-agent-user-id` value exposed as +`ResponseContext.platform_context.user_id_key`) before selecting its on-disk +directory. -Regular-agent session snapshots use a hashed platform user directory: +Regular-agent session snapshots use the platform user and session IDs: ```text -.checkpoints/ - sessions/user-/.json +/.sessions//.json ``` Workflow checkpoints and function approvals preserve the existing Foundry @@ -41,7 +39,9 @@ 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. +and a missing platform user ID fails closed. Regular agents also require the +platform session ID used for their snapshot filename. Local requests may remain +unscoped. The Foundry-specific store type intentionally hides the current filesystem implementation from `ResponsesHostServer` setup. A future version may move 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 index 74522d4a5c7..fc510c8cae0 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_request_context.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_request_context.py @@ -2,7 +2,8 @@ from __future__ import annotations -import hashlib +import os +from typing import Literal from azure.ai.agentserver.core import FoundryAgentRequestContext, get_request_context @@ -13,6 +14,30 @@ ) +def _validate_path_segment( + segment: str, + *, + kind: Literal["context id", "user id"], +) -> None: + """Validate that ``segment`` is a single safe path component (CWE-22). + + Request context values are untrusted when used as path segments. Reject + separators, drive letters, parent references, and similar values rather + than attempting to sanitize them and risk collisions. + """ + if not isinstance(segment, str) or not segment: + raise RuntimeError(f"Invalid {kind}: must be a non-empty string.") + if ( + "/" in segment + or "\\" in segment + or "\x00" in segment + or segment.strip(".") == "" + or os.path.isabs(segment) + or os.path.splitdrive(segment)[0] + ): + raise RuntimeError(f"Invalid {kind}: {segment!r}") + + def _validate_foundry_request_context( # pyright: ignore[reportUnusedFunction] context: FoundryAgentRequestContext, *, @@ -28,15 +53,10 @@ def _validate_foundry_request_context( # pyright: ignore[reportUnusedFunction] ) -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 + """Return the validated platform user ID for an on-disk directory.""" + user_id = get_request_context().user_id + if not user_id: + return None + _validate_path_segment(user_id, kind="user id") + return 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 c19ce303f82..d28af39b5df 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -122,8 +122,8 @@ from typing_extensions import Any from ._request_context import ( - _request_user_fingerprint, # pyright: ignore[reportPrivateUsage] _validate_foundry_request_context, # pyright: ignore[reportPrivateUsage] + _validate_path_segment, # pyright: ignore[reportPrivateUsage] ) from ._session_store import FoundrySessionStore @@ -134,10 +134,6 @@ _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" _HOSTED_RESPONSES_HISTORY_SOURCE_ID = "_foundry_responses_history" @@ -235,70 +231,6 @@ 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: - """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. - - 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 - accepted as literal directory names. Do NOT add decoding here without - re-validating after the decode -- decode-then-join is exactly the pattern - that reintroduces traversal. We also do not attempt to "sanitize" by - stripping characters because that can introduce collisions between distinct - ids. - """ - if not isinstance(segment, str) or not segment: - raise RuntimeError(f"Invalid {kind}: must be a non-empty string.") - # Reject any value that is not a single safe path component. This covers - # POSIX/Windows separators, NUL bytes, drive letters, and all-dot segments - # (``.``, ``..``, ``...``, ...). - if ( - "/" in segment - or "\\" in segment - or "\x00" in segment - # All-dot segments (``.``, ``..``, ``...``, ...) reduce to "" after stripping dots. - or segment.strip(".") == "" - or os.path.isabs(segment) - or os.path.splitdrive(segment)[0] - ): - raise RuntimeError(f"Invalid {kind}: {segment!r}") - - -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 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() - 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(): @@ -342,8 +274,8 @@ def _custom_session_store_key( key: dict[str, str] = {"object": object_id} if agent.name: key["agent"] = agent.name - if user_fingerprint := _request_user_fingerprint(): - key["user"] = user_fingerprint + if user_id := get_request_context().user_id: + key["user"] = user_id canonical_key = json.dumps(key, ensure_ascii=False, separators=(",", ":"), sort_keys=True) return f"foundry_{hashlib.sha256(canonical_key.encode('utf-8')).hexdigest()}" @@ -504,6 +436,7 @@ class ResponsesHostServer(ResponsesAgentServerHost): # TODO(@taochen): Allow a different checkpoint storage that stores checkpoints externally CHECKPOINT_STORAGE_PATH = "/.checkpoints" FUNCTION_APPROVAL_STORAGE_PATH = "/.function_approvals/approval_requests.json" + SESSION_STORAGE_PATH = "/.sessions" def __init__( self, @@ -522,10 +455,9 @@ 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 - a :class:`FoundrySessionStore` under the Foundry - durable storage root. A caller-supplied file store must already - be a ``FoundrySessionStore``. + session_store: Optional Agent Framework session store override. + Defaults to a :class:`FoundrySessionStore` under ``/.sessions`` + when hosted and an in-memory :class:`SessionStore` locally. **kwargs: Additional keyword arguments. Note: @@ -552,11 +484,6 @@ 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._is_workflow_agent = False self._checkpoint_storage_path = None if isinstance(agent, WorkflowAgent): @@ -593,9 +520,7 @@ def __init__( self._agent: SupportsAgentRun = agent if not self._is_workflow_agent and session_store is None: - 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.") + session_store = FoundrySessionStore(self.SESSION_STORAGE_PATH) if self.config.is_hosted else SessionStore() self._session_store = session_store self._approval_storage: ApprovalStorage = ( FileBasedFunctionApprovalStorage(self.FUNCTION_APPROVAL_STORAGE_PATH) @@ -660,7 +585,18 @@ async def _handle_response( cancellation_signal: asyncio.Event, ) -> AsyncIterable[ResponseStreamEvent | dict[str, Any]]: """Handle the creation of a response.""" - _validate_foundry_request_context(get_request_context(), is_hosted=self.config.is_hosted) + request_context = get_request_context() + _validate_foundry_request_context(request_context, is_hosted=self.config.is_hosted) + if ( + self.config.is_hosted + and not self._is_workflow_agent + and isinstance(self._session_store, FoundrySessionStore) + and not request_context.session_id + ): + raise RuntimeError( + "The hosted environment is missing the platform session ID in the request context. " + "Please ensure that the request is associated with a valid Foundry session." + ) if self._is_workflow_agent: # Workflow agents are handled differently because they require checkpoint restoration @@ -700,14 +636,21 @@ async def save_session() -> None: 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, FoundrySessionStore) - else _custom_session_store_key(self._agent, object_id) - ) + if isinstance(request_session_store, FoundrySessionStore) and self.config.is_hosted: + session_id = get_request_context().session_id + if session_id is None: # Guarded by _handle_response. + raise RuntimeError("The hosted request context is missing its platform session ID.") + session_store_key = session_id + else: + session_id = object_id + session_store_key = ( + session_id + if isinstance(request_session_store, FoundrySessionStore) + else _custom_session_store_key(self._agent, object_id) + ) session = await request_session_store.get(session_store_key) if session is None: - session = self._agent.create_session(session_id=object_id) + session = self._agent.create_session(session_id=session_id) input_items = await context.get_input_items() input_messages = await _items_to_messages(input_items, approval_storage=approval_storage) 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 f58296f0c1c..f32e61d2c80 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 @@ -13,9 +13,10 @@ @experimental(feature_id=ExperimentalFeature.SESSION_STORE) class FoundrySessionStore(FileSessionStore): - """File-backed session store isolated by the active Foundry request user. + """File-backed session store partitioned by the active Foundry request user. This implementation currently persists through :class:`FileSessionStore`. + Each request's validated platform user ID is used as a child directory. The Foundry-specific type leaves room to use a platform storage API later without changing :class:`ResponsesHostServer` configuration. """ @@ -30,6 +31,6 @@ def __init__( super().__init__(storage_path, serialization_format=serialization_format) def get_session_directory(self) -> Path: - """Return the active request user's session directory.""" + """Return the active request user's validated 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_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 297e8e413b2..1ae0782c193 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -11,7 +11,6 @@ from __future__ import annotations import asyncio -import hashlib import json import uuid from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence @@ -76,7 +75,6 @@ _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] consent_url_from_error, @@ -265,17 +263,32 @@ 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( + def test_init_uses_in_memory_session_store_by_default_locally(self) -> None: + agent = _make_agent( + response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("hi")])]) + ) + server = ResponsesHostServer(agent, store=InMemoryResponseProvider()) + session_store = server._session_store # pyright: ignore[reportPrivateUsage] + assert type(session_store) is SessionStore + + def test_init_uses_foundry_session_store_by_default_when_hosted( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.chdir(tmp_path) + monkeypatch.setenv("FOUNDRY_HOSTING_ENVIRONMENT", "1") + session_storage_path = tmp_path / ".sessions" 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] + + with patch.object(ResponsesHostServer, "SESSION_STORAGE_PATH", str(session_storage_path)): + server = ResponsesHostServer(agent, store=InMemoryResponseProvider()) + + session_store = server._session_store # pyright: ignore[reportPrivateUsage] + assert isinstance(session_store, FoundrySessionStore) + assert session_store.storage_path == session_storage_path + assert server.SESSION_STORAGE_PATH == "/.sessions" def test_init_rejects_history_provider_with_load_messages(self) -> None: @@ -303,17 +316,25 @@ 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: + def test_init_accepts_explicit_file_session_store_override( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setenv("FOUNDRY_HOSTING_ENVIRONMENT", "1") agent = _make_agent( response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("hi")])]) ) + session_store = FileSessionStore(tmp_path) - with pytest.raises(ValueError, match="FoundrySessionStore"): - ResponsesHostServer( - agent, - store=InMemoryResponseProvider(), - session_store=FileSessionStore(tmp_path), - ) + server = ResponsesHostServer( + agent, + store=InMemoryResponseProvider(), + session_store=session_store, + ) + + assert server._session_store is session_store # pyright: ignore[reportPrivateUsage] + assert server.config.is_hosted async def test_hosted_request_requires_user_partition_key(self) -> None: agent = _make_agent( @@ -353,6 +374,36 @@ async def test_hosted_request_requires_protocol_v2(self) -> None: asyncio.Event(), ) + async def test_hosted_foundry_store_requires_platform_session_id(self, tmp_path: Path) -> None: + server = _make_server(_make_agent(), session_store=FoundrySessionStore(tmp_path)) + request = CreateResponse(model="m", input="hi") + context = ResponseContext(response_id="response-1", mode_flags=MagicMock()) + + with ( + patch.object(server.config, "is_hosted", True), + _request_context(call_id="call-1", user_id="user-1"), + pytest.raises(RuntimeError, match="platform session ID"), + ): + await server._handle_response( # pyright: ignore[reportPrivateUsage] + request, + context, + asyncio.Event(), + ) + + async def test_hosted_explicit_store_does_not_require_platform_session_id(self) -> None: + server = _make_server(_make_agent(), session_store=SessionStore()) + request = CreateResponse(model="m", input="hi") + context = ResponseContext(response_id="response-1", mode_flags=MagicMock()) + + with patch.object(server.config, "is_hosted", True), _request_context(call_id="call-1", user_id="user-1"): + handler = await server._handle_response( # pyright: ignore[reportPrivateUsage] + request, + context, + asyncio.Event(), + ) + + assert handler is not None + # endregion @@ -437,38 +488,9 @@ def test_conversation_object_id_preserves_safe_values_and_hashes_unsafe_values(s 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: + async def test_file_store_uses_raw_user_directory_and_preserves_format(self, tmp_path: Path) -> None: template = FoundrySessionStore(tmp_path, serialization_format="msgpack") server = _make_server(_make_agent(), session_store=template) @@ -482,14 +504,12 @@ async def test_file_store_uses_per_user_child_directory_and_preserves_format(sel with _request_context(user_id="user-B"): await store.set("conversation-1", AgentSession(session_id="conversation-1")) - 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() + assert (tmp_path / "user-A" / "conversation-1.msgpack").is_file() + assert (tmp_path / "user-B" / "conversation-1.msgpack").is_file() async def test_scoped_file_store_reuses_base_filename_safety(self, tmp_path: Path) -> None: store = FoundrySessionStore(tmp_path) - user_directory = tmp_path / f"user-{hashlib.sha256(b'user-A').hexdigest()}" + user_directory = tmp_path / "user-A" with _request_context(user_id="user-A"): await store.set("CON", AgentSession(session_id="conversation-1")) @@ -499,7 +519,7 @@ async def test_scoped_file_store_reuses_base_filename_safety(self, tmp_path: Pat async def test_scoped_file_store_rejects_symlinked_session_leaf(self, tmp_path: Path) -> None: store = FoundrySessionStore(tmp_path) - user_directory = tmp_path / f"user-{hashlib.sha256(b'user-A').hexdigest()}" + user_directory = tmp_path / "user-A" user_directory.mkdir() outside_file = tmp_path / "outside.json" outside_file.write_text("outside", encoding="utf-8") @@ -513,7 +533,7 @@ async def test_scoped_file_store_rejects_symlinked_session_leaf(self, tmp_path: async def test_scoped_file_store_rejects_symlinked_isolation_directory(self, tmp_path: Path) -> None: store = FoundrySessionStore(tmp_path) - user_directory = f"user-{hashlib.sha256(b'user-A').hexdigest()}" + user_directory = "user-A" outside_directory = tmp_path.parent / f"{tmp_path.name}-outside" outside_directory.mkdir() try: @@ -524,12 +544,59 @@ async def test_scoped_file_store_rejects_symlinked_isolation_directory(self, tmp with _request_context(user_id="user-A"), pytest.raises(ValueError, match="Session directory escaped"): await store.get("conversation-1") + @pytest.mark.parametrize("user_id", ["../../escape", "/tmp/escape", "user/subdirectory", "user\\subdirectory"]) + async def test_file_store_rejects_unsafe_user_directory(self, tmp_path: Path, user_id: str) -> None: + store = FoundrySessionStore(tmp_path) + + with _request_context(user_id=user_id), pytest.raises(RuntimeError, match="Invalid user id"): + await store.set("conversation-1", AgentSession(session_id="conversation-1")) + 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_hosted_file_store_persists_by_platform_user_and_session_ids(self, tmp_path: Path) -> None: + seen_counts: list[int] = [] + + 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 = _make_agent() + agent.run = AsyncMock(side_effect=run_with_state) + server = _make_server(agent, session_store=FoundrySessionStore(tmp_path)) + + with ( + patch.object(server.config, "is_hosted", True), + patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])), + patch.object(ResponseContext, "get_history", new=AsyncMock(return_value=[])), + _request_context(call_id="call-1", user_id="user-1", session_id="session-1"), + ): + first_stream = await server._handle_response( # pyright: ignore[reportPrivateUsage] + CreateResponse(model="m", input="hi"), + ResponseContext(response_id="response-1", conversation_id="conversation-1", mode_flags=MagicMock()), + asyncio.Event(), + ) + first_events = [event async for event in first_stream] + second_stream = await server._handle_response( # pyright: ignore[reportPrivateUsage] + CreateResponse(model="m", input="again"), + ResponseContext(response_id="response-2", conversation_id="conversation-2", mode_flags=MagicMock()), + asyncio.Event(), + ) + second_events = [event async for event in second_stream] + + assert any(getattr(event, "type", None) == "response.completed" for event in first_events) + assert any(getattr(event, "type", None) == "response.completed" for event in second_events) + assert seen_counts == [1, 2] + assert (tmp_path / "user-1" / "session-1.json").is_file() + agent.create_session.assert_called_once_with(session_id="session-1") + async def test_previous_response_chain_restores_session_state(self) -> None: seen_counts: list[int] = [] seen_session_ids: list[str] = [] From 73bea32b9b141846e184efbc19e2471f69db2728 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 27 Jul 2026 23:45:22 +0200 Subject: [PATCH 10/14] Persist Foundry sessions under hosted home Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3e5c81ad-75e8-4e92-a883-8bbd676c6c8c --- docs/decisions/0033-python-session-store-serialization.md | 3 ++- python/packages/foundry_hosting/README.md | 3 +++ .../agent_framework_foundry_hosting/_responses.py | 6 +++++- python/packages/foundry_hosting/tests/test_responses.py | 7 +++---- 4 files changed, 13 insertions(+), 6 deletions(-) diff --git a/docs/decisions/0033-python-session-store-serialization.md b/docs/decisions/0033-python-session-store-serialization.md index 85f2fce957f..a3301d0eadd 100644 --- a/docs/decisions/0033-python-session-store-serialization.md +++ b/docs/decisions/0033-python-session-store-serialization.md @@ -213,7 +213,8 @@ IDs from `azure.ai.agentserver.core.get_request_context()`. Callers can explicitly override either default through `session_store=`. 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. +core store contract. The session file API maps `/` to the hosted `$HOME` +directory, so this API path is persisted on disk under `$HOME/.sessions`. ### Decision 2: Use msgspec codecs plus an explicit dynamic registry diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index 32b4f3f2566..39d98adaa24 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -10,6 +10,9 @@ Agent Server request context's platform user ID, and their filenames come from its platform session ID. Pass `session_store=` to explicitly override either default. +Foundry's session file API exposes the hosted `$HOME` directory as `/`, so the +API path `/.sessions` is stored on disk at `$HOME/.sessions`. + Workflow agents continue to use their existing checkpoint storage layout. ## Foundry session isolation 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 d28af39b5df..552f5dc12d6 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -520,7 +520,11 @@ def __init__( self._agent: SupportsAgentRun = agent if not self._is_workflow_agent and session_store is None: - session_store = FoundrySessionStore(self.SESSION_STORAGE_PATH) if self.config.is_hosted else SessionStore() + session_store = ( + FoundrySessionStore(Path.home() / self.SESSION_STORAGE_PATH.lstrip("/")) + if self.config.is_hosted + else SessionStore() + ) self._session_store = session_store self._approval_storage: ApprovalStorage = ( FileBasedFunctionApprovalStorage(self.FUNCTION_APPROVAL_STORAGE_PATH) diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 1ae0782c193..63b79d56403 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -277,17 +277,16 @@ def test_init_uses_foundry_session_store_by_default_when_hosted( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("FOUNDRY_HOSTING_ENVIRONMENT", "1") - session_storage_path = tmp_path / ".sessions" + monkeypatch.setenv("HOME", str(tmp_path)) agent = _make_agent( response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("hi")])]) ) - with patch.object(ResponsesHostServer, "SESSION_STORAGE_PATH", str(session_storage_path)): - server = ResponsesHostServer(agent, store=InMemoryResponseProvider()) + server = ResponsesHostServer(agent, store=InMemoryResponseProvider()) session_store = server._session_store # pyright: ignore[reportPrivateUsage] assert isinstance(session_store, FoundrySessionStore) - assert session_store.storage_path == session_storage_path + assert session_store.storage_path == tmp_path / ".sessions" assert server.SESSION_STORAGE_PATH == "/.sessions" def test_init_rejects_history_provider_with_load_messages(self) -> None: From a3e514d9e6a939866fc6ed1661e498f44169a286 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 27 Jul 2026 23:50:44 +0200 Subject: [PATCH 11/14] Make hosted path test platform independent Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3e5c81ad-75e8-4e92-a883-8bbd676c6c8c --- python/packages/foundry_hosting/tests/test_responses.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 63b79d56403..65bc28a6ea7 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -277,12 +277,12 @@ def test_init_uses_foundry_session_store_by_default_when_hosted( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("FOUNDRY_HOSTING_ENVIRONMENT", "1") - monkeypatch.setenv("HOME", str(tmp_path)) agent = _make_agent( response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("hi")])]) ) - server = ResponsesHostServer(agent, store=InMemoryResponseProvider()) + with patch.object(Path, "home", return_value=tmp_path): + server = ResponsesHostServer(agent, store=InMemoryResponseProvider()) session_store = server._session_store # pyright: ignore[reportPrivateUsage] assert isinstance(session_store, FoundrySessionStore) From 2735cfc02b83f44eeefeca80aabfe9c7f40b1d05 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 28 Jul 2026 09:42:12 +0200 Subject: [PATCH 12/14] Address session persistence review feedback Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3e5c81ad-75e8-4e92-a883-8bbd676c6c8c --- ...0033-python-session-store-serialization.md | 43 +++--- python/packages/core/AGENTS.md | 4 +- .../core/agent_framework/_sessions.py | 131 +++++++++++------- .../core/agent_framework/foundry/__init__.pyi | 8 +- .../packages/core/tests/core/test_sessions.py | 96 +++++++++++-- python/packages/foundry_hosting/README.md | 8 +- .../_invocations.py | 4 +- .../_request_context.py | 8 +- .../_responses.py | 37 +++-- .../_session_store.py | 17 ++- .../foundry_hosting/tests/test_responses.py | 82 ++++++++++- 11 files changed, 331 insertions(+), 107 deletions(-) diff --git a/docs/decisions/0033-python-session-store-serialization.md b/docs/decisions/0033-python-session-store-serialization.md index a3301d0eadd..a0a465a087e 100644 --- a/docs/decisions/0033-python-session-store-serialization.md +++ b/docs/decisions/0033-python-session-store-serialization.md @@ -51,7 +51,7 @@ The framework therefore needs to decide: - 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. +- Preserve dynamic registration of nested state types by the provider modules that own them. - Fail before persistence when an object cannot be restored after a cold start. - Keep the existing serialized `{"type": "", ...}` representation compatible. @@ -89,8 +89,9 @@ overriding the same async methods. ## 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. +`AgentSession`, including nested framework and application-defined state. Serialization belongs to each durable store +implementation rather than the `SessionStore` API: the default in-memory store does not serialize, and custom stores +remain free to choose another protocol. 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. @@ -198,8 +199,9 @@ 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). -`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 +`SessionStore` accepts opaque non-empty keys so custom backends can use their native key contracts. `FileSessionStore` +accepts opaque keys up to 128 characters and encodes values that are not portable filename stems; this supports +provider IDs such as `telegram::` without permitting path traversal. `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. @@ -209,12 +211,16 @@ default `ResponsesHostServer` store when hosted; local hosting defaults to the in-memory `SessionStore`. `FoundrySessionStore` currently subclasses `FileSessionStore`, stores snapshots under `/.sessions//.json`, and derives the validated platform -IDs from `azure.ai.agentserver.core.get_request_context()`. Callers can -explicitly override either default through `session_store=`. 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. The session file API maps `/` to the hosted `$HOME` -directory, so this API path is persisted on disk under `$HOME/.sessions`. +IDs from `azure.ai.agentserver.core.get_request_context()`. A Foundry session +controls hosted compute and filesystem lifetime, while a MAF `AgentSession` +contains framework context state; using the same identifier correlates them but +does not make their semantics equivalent. Callers can override the default +through `session_store=` when snapshots must be stored outside Foundry, such as +in a database or blob store. 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. The session file API maps `/` +to the hosted `$HOME` directory, so this API path is persisted on disk under +`$HOME/.sessions`. ### Decision 2: Use msgspec codecs plus an explicit dynamic registry @@ -251,11 +257,16 @@ the snapshot envelope, and carries an explicit payload version. The benchmark's choose the Struct. `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. +`to_dict` / `from_dict` classes and Pydantic models. Type IDs share one process-wide registry, so provider packages +should use stable package-qualified identifiers and register their own state types at module import time; consumers do +not need to know those implementation details. 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`. Same-process round-trips continue to work; cold-start deserialization is not guaranteed +without explicit provider registration. Unknown persisted type IDs remain raw dictionaries. + +File snapshots are quarantined only when their bytes cannot be parsed as the selected JSON or MessagePack format. +Schema errors, unsupported snapshot versions, and registered state-decoder failures leave the original file in place so +an application fix, rollback, or compatible reader can recover it. `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` / diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 7f166f36b5d..78d6db51176 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -75,8 +75,8 @@ agent_framework/ - **`AgentSession`** - Manages conversation state and session metadata - **`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. +- **`FileSessionStore`** - Experimental msgspec file-backed session snapshot store with atomic last-writer-wins updates; JSON is the default, `serialization_format="msgpack"` enables binary MessagePack, opaque keys are encoded to portable filenames, and only syntactically malformed snapshots are quarantined (schema, version, and state-decoder failures preserve the original file) +- **`register_state_type`** - Registers custom `AgentSession.state` classes with stable, process-wide type IDs and optional mapping codecs. Provider modules own registration for their custom state types and should use package-qualified IDs. 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 c7a816261b1..fdd6a39f57c 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 hashlib import json import logging import math @@ -104,6 +105,8 @@ _JSON_FILE_EXTENSION = ".json" _JSON_LINES_FILE_EXTENSION = ".jsonl" _MSGPACK_FILE_EXTENSION = ".msgpack" +_SESSION_SNAPSHOT_VERSION = "1.0" +_MAX_ENCODED_SESSION_FILE_STEM_LENGTH = 180 def _default_json_dumps(value: Any) -> bytes: @@ -120,7 +123,12 @@ def _default_json_loads(value: str | bytes) -> Any: def _contains_non_finite_float(value: Any) -> bool: - """Return whether a nested JSON-compatible value contains NaN or infinity.""" + """Return whether legacy JSON encoding is needed for NaN or infinity. + + msgspec normalizes non-finite floats to ``null``, while the previous + FileHistoryProvider JSON encoder emitted Python's ``NaN`` and ``Infinity`` + tokens. Detecting them keeps existing history files byte-compatible. + """ if isinstance(value, float): return not math.isfinite(value) if isinstance(value, Mapping): @@ -131,7 +139,12 @@ def _contains_non_finite_float(value: Any) -> bool: def _is_literal_session_file_stem_safe(session_id: str) -> bool: - """Return whether a session ID can be used directly as a filename stem.""" + """Return whether an opaque session ID is a portable filename stem. + + FileSessionStore and FileHistoryProvider accept opaque IDs, including IDs + with separators and platform-reserved names such as ``CON``. Unsafe values + are encoded by :func:`_session_file_stem` rather than rejected. + """ windows_stem = session_id.split(".", maxsplit=1)[0].upper() if ( not session_id @@ -142,7 +155,7 @@ def _is_literal_session_file_stem_safe(session_id: str) -> bool: 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 all(character.isascii() and (character.isalnum() or character in "._-") for character in session_id) def _session_file_stem(session_id: str, *, encoded_prefix: str) -> str: @@ -150,7 +163,11 @@ def _session_file_stem(session_id: str, *, encoded_prefix: str) -> str: 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}" + encoded_stem = f"{encoded_prefix}{encoded_session_id}" + if len(encoded_stem) <= _MAX_ENCODED_SESSION_FILE_STEM_LENGTH: + return encoded_stem + digest = hashlib.sha256(session_id.encode("utf-8")).hexdigest() + return f"{encoded_prefix}sha256-{digest}" def _deduplicate_origin_session_ids(origin_session_ids: Iterable[str]) -> list[str]: @@ -189,7 +206,7 @@ class _StateTypeRegistration: def _resolve_state_type_id(cls: type[Any], type_id: str | None) -> str: - """Resolve the stable serialized ID for a registered session-state type.""" + """Resolve the stable, process-wide ID for a registered state type.""" if type_id is not None: resolved = type_id elif callable(identifier := getattr(cls, "_get_type_identifier", None)): @@ -264,10 +281,14 @@ def register_state_type( """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. + process restart. Type identifiers share one process-wide registry and must + therefore be globally unique; provider packages should pass a stable, + package-qualified ``type_id``. For compatibility with existing framework + types, the identifier otherwise falls back to ``_get_type_identifier()``, + then ``TYPE``, and finally the lowercased class name. Conflicting + registrations fail immediately rather than silently selecting one type. + 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 @@ -378,6 +399,13 @@ def _serialize_value(value: Any, *, path: str) -> Any: type_id = _resolve_state_type_id(value_type, None) register_state_type(value_type, type_id=type_id) return _serialize_value(value, path=path) + warnings.warn( + f"AgentSession state value at {path} has unsupported type {value_type.__name__!r}; " + "AgentSession.to_dict() is leaving it unchanged for compatibility, but durable session stores will reject it. " + "Call register_state_type() with a restorable codec.", + RuntimeWarning, + stacklevel=4, + ) return value @@ -454,8 +482,8 @@ class _SessionSnapshot(msgspec.Struct): type: Literal["session"] session_id: str service_session_id: str | dict[str, Any] | None - state: _SessionStatePayload - version: Literal["1.0"] = "1.0" + state: object + version: str = _SESSION_SNAPSHOT_VERSION def _session_snapshot_enc_hook(value: Any) -> Any: @@ -465,19 +493,10 @@ def _session_snapshot_enc_hook(value: Any) -> Any: 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_DECODER = msgspec.json.Decoder(_SessionSnapshot) _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) +_SESSION_SNAPSHOT_MSGPACK_DECODER = msgspec.msgpack.Decoder(_SessionSnapshot) # Register known types @@ -735,12 +754,12 @@ class ContextProvider: 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 + custom class or Pydantic model, the provider module 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. + after its definition. Importing the provider then registers the type before + a persisted session is loaded, without requiring the application to know + which internal state types the provider uses. Framework-owned state types + such as :class:`Message` are registered by Agent Framework. Attributes: source_id: Unique identifier for this provider instance (required). @@ -814,11 +833,12 @@ class HistoryProvider(ContextProvider): 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 + provider module 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. + provider itself is instantiated, without requiring consumer registration. + 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). @@ -1459,19 +1479,21 @@ class FileSessionStore(SessionStore): 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. + explicit :func:`register_state_type` registry during encoding and restored + after the typed envelope and snapshot version are validated. 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 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. + store. Opaque store keys are encoded as portable filename stems, and + resolved-path validation prevents path traversal via ``session_id``. + These protections do not encrypt file contents 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 @@ -1527,13 +1549,9 @@ def _read() -> AgentSession | None: 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: + except msgspec.ValidationError as exc: + raise ValueError(f"Session snapshot '{file_path}' has an invalid schema.") from exc + except msgspec.DecodeError as exc: try: quarantine_path = self._quarantine_corrupt_snapshot(file_path, serialized) except OSError as quarantine_error: @@ -1550,6 +1568,24 @@ def _read() -> AgentSession | None: 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 + if snapshot.version != _SESSION_SNAPSHOT_VERSION: + raise ValueError( + f"Unsupported session snapshot version {snapshot.version!r} in '{file_path}'; " + f"expected {_SESSION_SNAPSHOT_VERSION!r}." + ) + raw_state: object = snapshot.state + if not isinstance(raw_state, dict): + raise ValueError(f"Session snapshot state in '{file_path}' must be a mapping.") + try: + state = _deserialize_state(cast(dict[str, Any], raw_state)) + except (TypeError, ValueError) as exc: + raise ValueError(f"Failed to restore session state from '{file_path}'.") from exc + session = AgentSession( + session_id=snapshot.session_id, + service_session_id=snapshot.service_session_id, + ) + session.state = state + return session return await asyncio.to_thread(_read) @@ -1604,14 +1640,11 @@ def validate_session_id(session_id: str) -> None: 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 ``_``. + ValueError: If the ID is empty or too long. """ 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. diff --git a/python/packages/core/agent_framework/foundry/__init__.pyi b/python/packages/core/agent_framework/foundry/__init__.pyi index 8413e538643..2e19fed8d65 100644 --- a/python/packages/core/agent_framework/foundry/__init__.pyi +++ b/python/packages/core/agent_framework/foundry/__init__.pyi @@ -29,7 +29,12 @@ from agent_framework_foundry import ( evaluate_traces, to_prompt_agent, ) -from agent_framework_foundry_hosting import FoundryToolbox, InvocationsHostServer, ResponsesHostServer +from agent_framework_foundry_hosting import ( + FoundrySessionStore, + FoundryToolbox, + InvocationsHostServer, + ResponsesHostServer, +) from agent_framework_foundry_local import ( FoundryLocalChatOptions, FoundryLocalClient, @@ -54,6 +59,7 @@ __all__ = [ "FoundryLocalClient", "FoundryLocalSettings", "FoundryMemoryProvider", + "FoundrySessionStore", "FoundryToolbox", "GeneratedEvaluatorRef", "InvocationsHostServer", diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py index a0aaabbb3a0..a9438496285 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -930,7 +930,10 @@ class UnsupportedState: session = AgentSession(session_id="unsupported") session.state["nested"] = [{"value": UnsupportedState()}] - with pytest.raises(TypeError, match=r"state\.nested\[0\]\.value.*UnsupportedState"): + with ( + pytest.warns(RuntimeWarning, match="leaving it unchanged"), + 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")]) @@ -976,7 +979,7 @@ async def test_corrupt_session_file_is_quarantined_and_retry_returns_none(self, 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: + async def test_invalid_registered_payload_preserves_snapshot(self, tmp_path: Path) -> None: store = FileSessionStore(tmp_path) session_file = store._session_file_path("session-1") session_file.write_text( @@ -990,14 +993,13 @@ async def test_invalid_registered_payload_is_quarantined(self, tmp_path: Path) - encoding="utf-8", ) - with pytest.raises(ValueError, match="quarantined.*retry"): + with pytest.raises(ValueError, match="Failed to restore session state"): 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 session_file.exists() + assert not await asyncio.to_thread(lambda: list(tmp_path.glob(".*.corrupt"))) - async def test_registered_decoder_attribute_error_is_quarantined(self, tmp_path: Path) -> None: + async def test_registered_decoder_attribute_error_preserves_snapshot(self, tmp_path: Path) -> None: store = FileSessionStore(tmp_path) session_file = store._session_file_path("session-1") session_file.write_text( @@ -1011,12 +1013,51 @@ async def test_registered_decoder_attribute_error_is_quarantined(self, tmp_path: encoding="utf-8", ) - with pytest.raises(ValueError, match="quarantined.*retry"): + with pytest.raises(ValueError, match="Failed to restore session state"): 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 session_file.exists() + assert not await asyncio.to_thread(lambda: list(tmp_path.glob(".*.corrupt"))) + + async def test_unsupported_snapshot_version_preserves_snapshot(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.1", + "session_id": "session-1", + "service_session_id": None, + "state": {}, + }), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="Unsupported session snapshot version '1.1'"): + await store.get("session-1") + + assert session_file.exists() + assert not await asyncio.to_thread(lambda: list(tmp_path.glob(".*.corrupt"))) + + async def test_invalid_snapshot_schema_preserves_snapshot(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": 123, + "service_session_id": None, + "state": {}, + }), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="invalid schema"): + await store.get("session-1") + + assert session_file.exists() + assert not await asyncio.to_thread(lambda: list(tmp_path.glob(".*.corrupt"))) async def test_corrupt_quarantine_does_not_remove_concurrent_replacement( self, @@ -1057,7 +1098,7 @@ def blocking_quarantine(file_path: Path, serialized: bytes) -> Path | None: assert restored is not None assert restored.session_id == "replacement" - @pytest.mark.parametrize("session_id", ["", "two words", "tenant/user", "session.id", "café"]) + @pytest.mark.parametrize("session_id", ["", "x" * (FileSessionStore.MAX_SESSION_ID_LENGTH + 1)]) async def test_invalid_session_id_raises(self, tmp_path: Path, session_id: str) -> None: store = FileSessionStore(tmp_path) session = AgentSession() @@ -1069,6 +1110,37 @@ 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", "is_encoded"), + [ + ("two words", True), + ("tenant/user", True), + ("session.id", False), + ("café", True), + ("telegram:123:456", True), + ("CON", True), + ("x" * 128, False), + ], + ) + async def test_opaque_session_id_round_trips_through_safe_filename( + self, + tmp_path: Path, + session_id: str, + is_encoded: bool, + ) -> None: + store = FileSessionStore(tmp_path) + session = AgentSession(session_id="framework-session") + + await store.set(session_id, session) + restored = await store.get(session_id) + + assert restored is not None + assert restored.session_id == "framework-session" + session_file = store._session_file_path(session_id) + assert session_file.is_file() + assert session_file.parent == tmp_path + assert session_file.name.startswith("~session-") is is_encoded + @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) diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index 39d98adaa24..7286d8f4800 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -8,7 +8,8 @@ uses the experimental `FoundrySessionStore` under `/.sessions` when hosted and an in-memory `SessionStore` locally. Hosted snapshots are partitioned by the Agent Server request context's platform user ID, and their filenames come from its platform session ID. Pass `session_store=` to explicitly override either -default. +default when MAF session snapshots must be stored outside Foundry, such as in a +database or blob store. Foundry's session file API exposes the hosted `$HOME` directory as `/`, so the API path `/.sessions` is stored on disk at `$HOME/.sessions`. @@ -29,6 +30,11 @@ Regular-agent session snapshots use the platform user and session IDs: /.sessions//.json ``` +A Foundry session controls hosted compute and filesystem lifetime. A MAF +`AgentSession` contains framework context state. Responses hosting uses the same +identifier to correlate and locate the snapshot, but the two session concepts +do not otherwise share lifecycle semantics. + Workflow checkpoints and function approvals preserve the existing Foundry Hosting layout. Hosted paths insert the validated raw platform user ID: 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 24f5bc8fe4f..64d96e1e3f6 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 ._request_context import _validate_foundry_request_context # pyright: ignore[reportPrivateUsage] +from ._request_context import validate_foundry_request_context class InvocationsHostServer(InvocationAgentServerHost): @@ -52,7 +52,7 @@ def _partition_key(self) -> str: RuntimeError: If the context doesn't contain the expected IDs. """ context = get_request_context() - _validate_foundry_request_context(context, is_hosted=self.config.is_hosted) + 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/_request_context.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_request_context.py index fc510c8cae0..a76b204c29d 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_request_context.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_request_context.py @@ -14,7 +14,7 @@ ) -def _validate_path_segment( +def validate_path_segment( segment: str, *, kind: Literal["context id", "user id"], @@ -38,7 +38,7 @@ def _validate_path_segment( raise RuntimeError(f"Invalid {kind}: {segment!r}") -def _validate_foundry_request_context( # pyright: ignore[reportUnusedFunction] +def validate_foundry_request_context( context: FoundryAgentRequestContext, *, is_hosted: bool, @@ -53,10 +53,10 @@ def _validate_foundry_request_context( # pyright: ignore[reportUnusedFunction] ) -def _request_user_directory_segment() -> str | None: # pyright: ignore[reportUnusedFunction] +def request_user_directory_segment() -> str | None: """Return the validated platform user ID for an on-disk directory.""" user_id = get_request_context().user_id if not user_id: return None - _validate_path_segment(user_id, kind="user id") + validate_path_segment(user_id, kind="user id") return 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 552f5dc12d6..f16b1ec1a7e 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -122,8 +122,8 @@ from typing_extensions import Any from ._request_context import ( - _validate_foundry_request_context, # pyright: ignore[reportPrivateUsage] - _validate_path_segment, # pyright: ignore[reportPrivateUsage] + validate_foundry_request_context, + validate_path_segment, ) from ._session_store import FoundrySessionStore @@ -312,11 +312,11 @@ def _checkpoint_storage_for_context( 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") + validate_path_segment(context_id, kind="context id") base_path = Path(root).resolve() if user_id: - _validate_path_segment(user_id, kind="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 user id: {user_id!r}") @@ -343,7 +343,7 @@ def _approval_storage_path_for_user(base_path: str, user_id: str) -> str: 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") + 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() @@ -458,6 +458,9 @@ def __init__( session_store: Optional Agent Framework session store override. Defaults to a :class:`FoundrySessionStore` under ``/.sessions`` when hosted and an in-memory :class:`SessionStore` locally. + Provide another implementation when MAF session snapshots must + be persisted outside Foundry, such as in a database or blob + store. **kwargs: Additional keyword arguments. Note: @@ -508,8 +511,9 @@ def __init__( ) ): # 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. + # call. Agent.run would otherwise mutate the same user-owned agent by + # auto-injecting its default InMemoryHistoryProvider. Install a loading + # no-op provider up front so that transcript is not replayed twice. agent.context_providers.append( InMemoryHistoryProvider( source_id=_HOSTED_RESPONSES_HISTORY_SOURCE_ID, @@ -590,7 +594,7 @@ async def _handle_response( ) -> AsyncIterable[ResponseStreamEvent | dict[str, Any]]: """Handle the creation of a response.""" request_context = get_request_context() - _validate_foundry_request_context(request_context, is_hosted=self.config.is_hosted) + validate_foundry_request_context(request_context, is_hosted=self.config.is_hosted) if ( self.config.is_hosted and not self._is_workflow_agent @@ -612,7 +616,12 @@ async def _handle_inner_agent( request: CreateResponse, context: ResponseContext, ) -> AsyncIterable[ResponseStreamEvent | dict[str, Any]]: - """Handle the creation of a response for a regular (non-workflow) agent.""" + """Handle a regular agent with MAF state scoped by a Foundry session. + + Foundry sessions govern hosted compute and filesystem lifetime. MAF + AgentSession objects hold framework context state. The shared identifier + correlates the two without making their lifecycle semantics equivalent. + """ response_event_stream = ResponseEventStream(response_id=context.response_id, model=request.model) yield response_event_stream.emit_created() yield response_event_stream.emit_in_progress() @@ -623,6 +632,9 @@ async def _handle_inner_agent( session: AgentSession | None = None request_session_store: SessionStore | None = None session_store_key: str | None = None + # Successful runs save before response.completed so a persistence failure + # can become response.failed. The finally path covers cancellation or an + # abandoned stream and must not save a second time. session_save_attempted = False async def save_session() -> None: @@ -746,7 +758,12 @@ async def save_session() -> None: yield event finally: if session is not None and not session_save_attempted: - await save_session() + try: + await save_session() + except Exception: + logger.exception( + "Failed to persist the Agent Framework session while unwinding an interrupted request" + ) async def _handle_inner_workflow( self, 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 f32e61d2c80..a4d9783b349 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 @@ -8,16 +8,21 @@ from agent_framework import ExperimentalFeature, FileSessionStore from agent_framework._feature_stage import experimental -from ._request_context import _request_user_directory_segment # pyright: ignore[reportPrivateUsage] +from ._request_context import request_user_directory_segment @experimental(feature_id=ExperimentalFeature.SESSION_STORE) class FoundrySessionStore(FileSessionStore): - """File-backed session store partitioned by the active Foundry request user. + """Persist MAF AgentSession snapshots within a Foundry hosted session. - This implementation currently persists through :class:`FileSessionStore`. - Each request's validated platform user ID is used as a child directory. - The Foundry-specific type leaves room to use a platform storage API later + A Foundry hosted session controls platform compute and filesystem lifetime; + a MAF :class:`AgentSession` contains framework context state. They remain + distinct concepts even though Responses hosting uses the Foundry session ID + as the MAF session identifier and snapshot filename for correlation. + + This implementation currently persists through :class:`FileSessionStore`, + with each validated platform user ID as a child directory. The + Foundry-specific type leaves room to use a platform storage API later without changing :class:`ResponsesHostServer` configuration. """ @@ -32,5 +37,5 @@ def __init__( def get_session_directory(self) -> Path: """Return the active request user's validated session directory.""" - directory_segment = _request_user_directory_segment() + 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_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 65bc28a6ea7..99d53905275 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -13,11 +13,11 @@ import asyncio import json import uuid -from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path -from typing import Literal, overload +from typing import Literal, cast, overload from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -175,6 +175,17 @@ async def get_response() -> ChatResponse: return get_response() +class _FailingSessionStore(SessionStore): + def __init__(self) -> None: + super().__init__() + self.set_attempts = 0 + + async def set(self, session_id: str, session: AgentSession) -> None: + del session_id, session + self.set_attempts += 1 + raise OSError("session storage is full") + + def _make_server(agent: Any, **kwargs: Any) -> ResponsesHostServer: """Create a ResponsesHostServer with an in-memory store.""" kwargs.setdefault("session_store", SessionStore()) @@ -483,9 +494,9 @@ def test_custom_store_key_partitions_by_agent_and_request_user(self) -> None: assert all(character.isascii() and (character.isalnum() or character in "-_") for character in key_a) assert key_a == repeated_key_a - def test_conversation_object_id_preserves_safe_values_and_hashes_unsafe_values(self) -> None: + def test_conversation_object_id_preserves_supported_opaque_values(self) -> None: assert _conversation_object_id("conversation-1") == "conversation-1" - assert _conversation_object_id("conversation/unsafe").startswith("conversation_") + assert _conversation_object_id("conversation/opaque") == "conversation/opaque" class TestAgentSessionPersistence: @@ -765,6 +776,69 @@ async def failing_run(*args: Any, **kwargs: Any) -> AgentResponse: assert stored is not None assert stored.state["before_failure"] == "saved" + async def test_streaming_save_failure_emits_failed_response(self) -> None: + store = _FailingSessionStore() + 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) + event_types = _sse_event_types(_parse_sse_events(response.text)) + + assert event_types[-1] == "response.failed" + assert "response.completed" not in event_types + assert store.set_attempts == 1 + + async def test_cancellation_is_preserved_when_best_effort_save_fails( + self, + caplog: pytest.LogCaptureFixture, + ) -> None: + store = _FailingSessionStore() + 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["started"] = True + yield AgentResponseUpdate(contents=[Content.from_text("started")], role="assistant") + + return ResponseStream(updates()) + + agent.run = MagicMock(side_effect=run_streaming) + server = _make_server(agent, session_store=store) + request = CreateResponse(model="m", input="hi", stream=True) + context = ResponseContext(response_id="response-1", mode_flags=MagicMock()) + + with ( + patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])), + patch.object(ResponseContext, "get_history", new=AsyncMock(return_value=[])), + ): + handler = cast( + AsyncGenerator[Any, None], + server._handle_inner_agent(request, context), # pyright: ignore[reportPrivateUsage] + ) + await anext(handler) + await anext(handler) + await anext(handler) + with pytest.raises(asyncio.CancelledError): + await handler.athrow(asyncio.CancelledError()) + + assert store.set_attempts == 1 + assert "while unwinding an interrupted request" in caplog.text + # endregion From 30be30b086c8f2d4e9e0fc56a4d89ed78ab92675 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 28 Jul 2026 10:47:34 +0200 Subject: [PATCH 13/14] Isolate Foundry session path handling Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3e5c81ad-75e8-4e92-a883-8bbd676c6c8c --- .../core/agent_framework/_sessions.py | 26 ++++++--------- .../packages/core/tests/core/test_sessions.py | 13 ++++++++ .../_session_store.py | 20 ++++++++++-- .../foundry_hosting/tests/test_responses.py | 32 +++++++++++++++++++ 4 files changed, 71 insertions(+), 20 deletions(-) diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index fdd6a39f57c..f4a48cc7522 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -1646,15 +1646,6 @@ def validate_session_id(session_id: str) -> None: if len(session_id) > FileSessionStore.MAX_SESSION_ID_LENGTH: raise ValueError(f"session_id must be at most {FileSessionStore.MAX_SESSION_ID_LENGTH} characters") - 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.""" @@ -1675,17 +1666,18 @@ def _session_file_lock(cls, file_path: Path) -> threading.Lock: def _session_file_path(self, session_id: str) -> Path: """Resolve the contained snapshot path for ``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 = (session_directory / f"{file_stem}{self._file_extension}").resolve() - if not file_path.is_relative_to(session_directory): + candidate_path = self._storage_root / self._session_file_name(session_id) + file_path = candidate_path.resolve() + if file_path != candidate_path or not file_path.is_relative_to(self._storage_root): raise ValueError(f"Session path escaped storage directory: {session_id!r}") return file_path + def _session_file_name(self, session_id: str) -> str: + """Return the portable snapshot filename for ``session_id``.""" + self.validate_session_id(session_id) + file_stem = _session_file_stem(session_id, encoded_prefix=self._ENCODED_SESSION_PREFIX) + return f"{file_stem}{self._file_extension}" + class InMemoryHistoryProvider(HistoryProvider): """Built-in history provider that stores messages in session.state. diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py index a9438496285..d71555ec10c 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -1141,6 +1141,19 @@ async def test_opaque_session_id_round_trips_through_safe_filename( assert session_file.parent == tmp_path assert session_file.name.startswith("~session-") is is_encoded + async def test_rejects_symlinked_session_file_within_storage_root(self, tmp_path: Path) -> None: + store = FileSessionStore(tmp_path) + target_file = tmp_path / "other-session.json" + session_file = tmp_path / "session-1.json" + await asyncio.to_thread(target_file.write_text, "{}", encoding="utf-8") + try: + await asyncio.to_thread(session_file.symlink_to, target_file) + except OSError as exc: + pytest.skip(f"Symlinks are not available: {exc}") + + with pytest.raises(ValueError, match="escaped storage directory"): + await store.get("session-1") + @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) 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 a4d9783b349..77ed9de28b9 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 @@ -35,7 +35,21 @@ def __init__( """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 validated session directory.""" + def _session_file_path(self, session_id: str) -> Path: + """Resolve a snapshot path within the active Foundry user's directory.""" directory_segment = request_user_directory_segment() - return self.storage_path / directory_segment if directory_segment else self.storage_path + candidate_directory = self._storage_root / directory_segment if directory_segment else self._storage_root + session_directory = candidate_directory.resolve() + if session_directory != candidate_directory or 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) + + candidate_path = session_directory / self._session_file_name(session_id) + file_path = candidate_path.resolve() + if ( + file_path != candidate_path + or not file_path.is_relative_to(session_directory) + or not file_path.is_relative_to(self._storage_root) + ): + raise ValueError(f"Session path escaped storage directory: {session_id!r}") + return file_path diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 99d53905275..6d6d9747f00 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -541,6 +541,23 @@ async def test_scoped_file_store_rejects_symlinked_session_leaf(self, tmp_path: 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_session_leaf_within_user_directory( + self, + tmp_path: Path, + ) -> None: + store = FoundrySessionStore(tmp_path) + user_directory = tmp_path / "user-A" + user_directory.mkdir() + target_file = user_directory / "conversation-2.json" + target_file.write_text("outside", encoding="utf-8") + try: + (user_directory / "conversation-1.json").symlink_to(target_file) + except OSError as exc: + pytest.skip(f"Symlinks are not available: {exc}") + + 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 = FoundrySessionStore(tmp_path) user_directory = "user-A" @@ -554,6 +571,21 @@ async def test_scoped_file_store_rejects_symlinked_isolation_directory(self, tmp 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_symlinked_isolation_directory_within_storage_root( + self, + tmp_path: Path, + ) -> None: + store = FoundrySessionStore(tmp_path) + other_user_directory = tmp_path / "user-B" + other_user_directory.mkdir() + try: + (tmp_path / "user-A").symlink_to(other_user_directory, target_is_directory=True) + except OSError as exc: + pytest.skip(f"Symlinks are not available: {exc}") + + with _request_context(user_id="user-A"), pytest.raises(ValueError, match="Session directory escaped"): + await store.get("conversation-1") + @pytest.mark.parametrize("user_id", ["../../escape", "/tmp/escape", "user/subdirectory", "user\\subdirectory"]) async def test_file_store_rejects_unsafe_user_directory(self, tmp_path: Path, user_id: str) -> None: store = FoundrySessionStore(tmp_path) From 5a47cca2e0baab2b87a6175606d0b86e4440ef8a Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 28 Jul 2026 10:55:49 +0200 Subject: [PATCH 14/14] Clarify Foundry session path terminology Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3e5c81ad-75e8-4e92-a883-8bbd676c6c8c --- .../_request_context.py | 11 +----- .../_session_store.py | 34 +++++++++++-------- .../foundry_hosting/tests/test_responses.py | 8 ++--- 3 files changed, 24 insertions(+), 29 deletions(-) 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 index a76b204c29d..b2db9b7f224 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_request_context.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_request_context.py @@ -5,7 +5,7 @@ import os from typing import Literal -from azure.ai.agentserver.core import FoundryAgentRequestContext, get_request_context +from azure.ai.agentserver.core import FoundryAgentRequestContext _PROTOCOL_V2_REQUIRED_MESSAGE = ( "The hosted environment is running on protocol 1.0.0, but the agent requires protocol 2.0.0. " @@ -51,12 +51,3 @@ def validate_foundry_request_context( "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_directory_segment() -> str | None: - """Return the validated platform user ID for an on-disk directory.""" - user_id = get_request_context().user_id - if not user_id: - return None - validate_path_segment(user_id, kind="user id") - return user_id 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 77ed9de28b9..f584fe4672a 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 @@ -7,8 +7,9 @@ from agent_framework import ExperimentalFeature, FileSessionStore from agent_framework._feature_stage import experimental +from azure.ai.agentserver.core import get_request_context -from ._request_context import request_user_directory_segment +from ._request_context import validate_path_segment @experimental(feature_id=ExperimentalFeature.SESSION_STORE) @@ -37,19 +38,22 @@ def __init__( def _session_file_path(self, session_id: str) -> Path: """Resolve a snapshot path within the active Foundry user's directory.""" - directory_segment = request_user_directory_segment() - candidate_directory = self._storage_root / directory_segment if directory_segment else self._storage_root - session_directory = candidate_directory.resolve() - if session_directory != candidate_directory or 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) - - candidate_path = session_directory / self._session_file_name(session_id) - file_path = candidate_path.resolve() + user_directory = self._storage_root + if user_id := get_request_context().user_id: + validate_path_segment(user_id, kind="user id") + candidate_user_directory = self._storage_root / user_id + user_directory = candidate_user_directory.resolve() + if user_directory != candidate_user_directory or not user_directory.is_relative_to(self._storage_root): + raise ValueError(f"User directory escaped storage directory: '{user_directory}'.") + user_directory.mkdir(parents=True, exist_ok=True) + + session_file_name = self._session_file_name(session_id) + candidate_session_file_path = user_directory / session_file_name + session_file_path = candidate_session_file_path.resolve() if ( - file_path != candidate_path - or not file_path.is_relative_to(session_directory) - or not file_path.is_relative_to(self._storage_root) + session_file_path != candidate_session_file_path + or not session_file_path.is_relative_to(user_directory) + or not session_file_path.is_relative_to(self._storage_root) ): - raise ValueError(f"Session path escaped storage directory: {session_id!r}") - return file_path + raise ValueError(f"Session file path escaped user directory: {session_id!r}") + return session_file_path diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 6d6d9747f00..e66686559ba 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -538,7 +538,7 @@ 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 _request_context(user_id="user-A"), pytest.raises(ValueError, match="escaped storage directory"): + with _request_context(user_id="user-A"), pytest.raises(ValueError, match="escaped user directory"): await store.get("conversation-1") async def test_scoped_file_store_rejects_symlinked_session_leaf_within_user_directory( @@ -555,7 +555,7 @@ async def test_scoped_file_store_rejects_symlinked_session_leaf_within_user_dire except OSError as exc: pytest.skip(f"Symlinks are not available: {exc}") - with _request_context(user_id="user-A"), pytest.raises(ValueError, match="escaped storage directory"): + with _request_context(user_id="user-A"), pytest.raises(ValueError, match="escaped user directory"): await store.get("conversation-1") async def test_scoped_file_store_rejects_symlinked_isolation_directory(self, tmp_path: Path) -> None: @@ -568,7 +568,7 @@ async def test_scoped_file_store_rejects_symlinked_isolation_directory(self, tmp except OSError as exc: pytest.skip(f"Symlinks are not available: {exc}") - with _request_context(user_id="user-A"), pytest.raises(ValueError, match="Session directory escaped"): + with _request_context(user_id="user-A"), pytest.raises(ValueError, match="User directory escaped"): await store.get("conversation-1") async def test_scoped_file_store_rejects_symlinked_isolation_directory_within_storage_root( @@ -583,7 +583,7 @@ async def test_scoped_file_store_rejects_symlinked_isolation_directory_within_st except OSError as exc: pytest.skip(f"Symlinks are not available: {exc}") - with _request_context(user_id="user-A"), pytest.raises(ValueError, match="Session directory escaped"): + with _request_context(user_id="user-A"), pytest.raises(ValueError, match="User directory escaped"): await store.get("conversation-1") @pytest.mark.parametrize("user_id", ["../../escape", "/tmp/escape", "user/subdirectory", "user\\subdirectory"])