Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 25 additions & 7 deletions python/packages/foundry_hosting/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,21 +4,39 @@ This package provides the integration of Agent Framework agents and workflows wi

## State store

### Local persistence

Outside the Foundry hosting environment, state is persisted as JSON files under
`~/.agentserver/state_stores` by default. Set `AGENTSERVER_STATE_ROOT` to use a
different root directory; the files will be written to its `state_stores`
subdirectory instead.

Each logical store is saved as one JSON file whose name is a URL-safe Base64
encoding of the store name. For example:

- Agent sessions: `YWdlbnRfc2Vzc2lvbnM.json`
- Function approvals: `ZnVuY3Rpb25fYXBwcm92YWxz.json`
- Workflow checkpoints: one file per context, encoded from `checkpoints/<context_id>`

> Read more about the Foundry durable state store in the [developer guide](https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/agentserver/azure-ai-agentserver-core/docs/state-store-guide.md).

### Agent Sessions

`ResponsesHostServer` persists the Agent Framework `AgentSession` durably. By default it
uses the `FoundryAgentSessionStore` when hosted and an in-memory `SessionStore` locally.
When hosted, the stored sessions will be isolated by the platform user ID and scoped
under `agent_sessions`.
uses the `FoundryAgentSessionStore`, backed by Foundry storage when hosted and file-based
storage locally. Stored sessions are scoped under `agent_sessions`.

See the [custom storage provider sample](../../samples/04-hosting/foundry-hosted-agents/responses/custom_storage/)
for an example that uses an in-memory session store locally and Azure Cosmos DB when hosted.

### Workflow checkpoints

`ResponsesHostServer` persists workflow checkpoints durably. By default, it uses the
`FoundryCheckpointStore` when hosted and an in-memory `InMemoryCheckpointStorage` locally.
When hosted, the stored checkpoints will be isolated by the platform user ID and scoped
under `checkpoints`.
`FoundryCheckpointStore`, backed by Foundry storage when hosted and file-based storage
locally. Stored checkpoints are scoped under `checkpoints`.

### Function approvals

`ResponsesHostServer` persists function approvals durably. By default, it uses the
`FoundryFunctionApprovalStore` when hosted and an in-memory `InMemoryFunctionApprovalStore` locally. When hosted, the stored approvals will be isolated by the platform user ID and scoped under `function_approvals`.
`FoundryFunctionApprovalStore`, backed by Foundry storage when hosted and file-based
storage locally. Stored approvals are scoped under `function_approvals`.
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,6 @@

from __future__ import annotations

import os
from typing import Literal

from azure.ai.agentserver.core import FoundryAgentRequestContext

_PROTOCOL_V2_REQUIRED_MESSAGE = (
Expand All @@ -14,30 +11,6 @@
)


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(
context: FoundryAgentRequestContext,
*,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,6 @@
from ._feature_usage import FeatureIndex
from ._request_context import (
validate_foundry_request_context,
validate_path_segment,
)
from ._state_store import (
AgentSessionStoreProvider,
Expand Down Expand Up @@ -372,29 +371,10 @@ async def _handle_inner_agent(
try:
approval_storage = self._function_approval_storage_provider.get_store(config=self.config)
session_storage = self._session_storage_provider.get_store(config=self.config)
# Agent sessions are either tied to the conversation_id (for multi-turn conversation mode)
# or the previous_response_id (for response chaining). If neither is present, a new session
# is created for this request and stored under the current response_id. The current response_id
# will become the previous_response_id for the next request in a response chain, allowing the
# session to be retrieved.
if (previous_response_id := request.get("previous_response_id")) is not None:
session = await session_storage.get(previous_response_id)
if session is None:
raise RuntimeError(
f"Cannot find an existing agent session for previous_response_id={previous_response_id}. "
"Ensure that the previous response was created successfully and that the ID is correct."
)
elif (conversation_id := context.conversation_id) is not None:
session = await session_storage.get(conversation_id)
if session is None:
# Note that we cannot determine if the session was deleted or never existed,
# so we log a warning and create a new session.
logger.info(
"Cannot find an existing agent session for id=%s. Creating a new session.",
conversation_id,
)
session = self._agent.create_session()
else:

context_id = context.conversation_chain_id

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This key no longer restores regular-agent state for previous_response_id requests in the actual endpoint path: the existing generated-ID chain test observes [1, 1, 1] instead of restoring the prior session, and a missing previous response silently creates a fresh session. That breaks the core multi-turn workflow and discards conversation state without notifying the caller. Please ensure the first response and its continuations resolve to the same persisted key, while preserving an error for a genuinely unknown previous response.

session = await session_storage.get(context_id)
Comment on lines +375 to +376
if session is None:
session = self._agent.create_session()
except Exception as ex:
logger.error("Failed to prepare state storage: %s", ex, exc_info=(type(ex), ex, ex.__traceback__))
Expand Down Expand Up @@ -456,7 +436,7 @@ async def _handle_inner_agent(
if self._uses_hosted_responses_history:
session.state.pop(_HOSTED_RESPONSES_HISTORY_SOURCE_ID, None)
try:
await session_storage.set(context.conversation_id or context.response_id, session)
await session_storage.set(context_id, session)
except Exception as save_error:
save_failure = save_error
if request_interrupted:
Expand Down Expand Up @@ -518,42 +498,19 @@ async def _handle_inner_workflow(
# any future async resources owned by the workflow are entered here.
await self._ensure_agent_ready()

context_id = context.conversation_chain_id

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new SDK safely hashes malformed context values into opaque IDs, but the three existing checkpoint-context tests still require these inputs to produce response.failed; all three now complete successfully and fail the suite. Please update the tests to assert the new safe-key behavior, or retain explicit rejection if that remains the package contract.

checkpoint_storage = self._checkpoint_storage_provider.get_store(config=self.config, context_id=context_id)

Comment on lines +501 to +503
# Determine the latest checkpoint (if any) so we can resume the
# workflow's prior state for this turn. The directory is keyed by
# the inbound context id (conversation_id when set, otherwise
# previous_response_id). Multi-turn declarative workflows need the
# workflow's internal state (e.g. Conversation.messages,
# the platform derived context_id. Multi-turn declarative workflows
# need the workflow's internal state (e.g. Conversation.messages,
# intermediate Local.* variables) to survive across user turns;
# the only place that state lives is the workflow checkpoint, so
# 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.
latest_checkpoint_id: str | None = None
restore_storage: CheckpointStorage | None = None
if context_id is not None:
validate_path_segment(context_id, kind="context id")
restore_storage = self._checkpoint_storage_provider.get_store(
config=self.config,
context_id=context_id,
)
latest_checkpoint = await restore_storage.get_latest(workflow_name=self._agent.workflow.name)
if latest_checkpoint is not None:
latest_checkpoint_id = latest_checkpoint.checkpoint_id

# Storage that will receive checkpoints written during this turn.
# When the caller chains with previous_response_id, the next turn
# will reference the current response_id as its previous_response_id,
# so new checkpoints must land under the current response_id (or the
# conversation_id when set). When conversation_id is set, this
# matches restore_storage; when only previous_response_id was
# supplied, restore_storage points at the *prior* response's
# directory and write_storage points at the *current* response's.
write_context_id = context.conversation_id or context.response_id
validate_path_segment(write_context_id, kind="context id")
write_storage = self._checkpoint_storage_provider.get_store(
config=self.config,
context_id=write_context_id,
)
latest_checkpoint = await checkpoint_storage.get_latest(workflow_name=self._agent.workflow.name)

# Multi-turn pattern: when we have a prior checkpoint, restore it
# first (drive the workflow back to idle with prior state intact),
Expand All @@ -571,11 +528,11 @@ async def _handle_inner_workflow(
# ``run(input_messages, ...)`` call may contain ``function_call_output``
# items (carried as FunctionResult/FunctionApprovalResponse content)
# that fulfill them via :meth:`WorkflowAgent._process_pending_requests`.
if latest_checkpoint_id is not None:
if latest_checkpoint is not None:
async for _ in self._agent.run(
stream=True,
checkpoint_id=latest_checkpoint_id,
checkpoint_storage=restore_storage,
checkpoint_id=latest_checkpoint.checkpoint_id,
checkpoint_storage=checkpoint_storage,
):
pass

Expand All @@ -585,7 +542,7 @@ async def _handle_inner_workflow(
async for update in self._agent.run(
input_messages,
stream=True,
checkpoint_storage=write_storage,
checkpoint_storage=checkpoint_storage,
):
for content in update.contents:
for event in tracker.handle(content):
Expand All @@ -600,27 +557,12 @@ async def _handle_inner_workflow(
# Close any remaining active builder
for event in tracker.close():
yield event

await self._delete_not_latest_checkpoints(write_storage, self._agent.workflow.name)
yield response_event_stream.emit_completed()
except Exception as ex:
logger.exception("Failed to produce response for workflow agent")
for event in self._emit_failure(response_event_stream, tracker, ex):
yield event

@staticmethod
async def _delete_not_latest_checkpoints(checkpoint_storage: CheckpointStorage, workflow_name: str) -> None:
"""Delete all checkpoints except the latest one.

We only need the last checkpoint for each invocation.
"""
latest_checkpoint = await checkpoint_storage.get_latest(workflow_name=workflow_name)
if latest_checkpoint is not None:
all_checkpoints = await checkpoint_storage.list_checkpoints(workflow_name=workflow_name)
for checkpoint in all_checkpoints:
if checkpoint.checkpoint_id != latest_checkpoint.checkpoint_id:
await checkpoint_storage.delete(checkpoint.checkpoint_id)

@staticmethod
def _emit_failure(
response_event_stream: ResponseEventStream,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@
CheckpointID,
CheckpointStorage,
Content,
InMemoryCheckpointStorage,
SessionStore,
WorkflowCheckpoint,
WorkflowCheckpointException,
Expand Down Expand Up @@ -169,13 +168,11 @@ class CheckpointStoreProvider(ContextScopedStoreProvider[CheckpointStorage]):
collection boundary used to list checkpoints, restore the latest checkpoint,
and clean up older checkpoints without affecting another workflow context.

This will default to using the `FoundryCheckpointStore` when hosted in Foundry,
and an in-memory store otherwise.
This defaults to using the `FoundryCheckpointStore` in all environments.
"""

def __init__(self) -> None:
self._foundry_storages: dict[str, CheckpointStorage] = {}
self._in_memory_storages: dict[str, CheckpointStorage] = {}
self._storages: dict[str, CheckpointStorage] = {}

def get_store(
self,
Expand All @@ -184,14 +181,12 @@ def get_store(
context_id: str,
) -> CheckpointStorage:
"""Get checkpoint store for the requested hosting environment."""
stores = self._foundry_storages if config.is_hosted else self._in_memory_storages

if not context_id:
raise ValueError("context_id must be provided to get a checkpoint store.")

if context_id not in stores:
stores[context_id] = FoundryCheckpointStore(context_id) if config.is_hosted else InMemoryCheckpointStorage()
return stores[context_id]
if context_id not in self._storages:
self._storages[context_id] = FoundryCheckpointStore(context_id)
return self._storages[context_id]


# endregion Checkpoint persistence
Expand Down Expand Up @@ -242,43 +237,20 @@ async def load_approval_request(self, approval_request_id: str) -> Content:
return Content.from_dict(item.value)


class InMemoryFunctionApprovalStore:
"""An in-memory store for function approval requests."""

def __init__(self) -> None:
self._store: dict[str, Content] = {}

async def save_approval_request(self, approval_request_id: str, request: Content) -> None:
if approval_request_id in self._store:
raise ValueError(f"Approval request with ID '{approval_request_id}' already exists.")
self._store[approval_request_id] = request

async def load_approval_request(self, approval_request_id: str) -> Content:
if approval_request_id not in self._store:
raise KeyError(f"Approval request with ID '{approval_request_id}' does not exist.")
return self._store[approval_request_id]


class FunctionApprovalStoreProvider(StoreProvider[FunctionApprovalStore]):
"""Provide function approval store for the active hosting environment.

This will default to using the `FoundryFunctionApprovalStore` when hosted in Foundry,
and an in-memory store otherwise.
This defaults to using the `FoundryFunctionApprovalStore` in all environments.
"""

def __init__(self) -> None:
self._foundry_storage: FunctionApprovalStore | None = None
self._in_memory_storage: FunctionApprovalStore | None = None
self._storage: FunctionApprovalStore | None = None

def get_store(self, *, config: AgentConfig) -> FunctionApprovalStore:
"""Get function approval store for the requested hosting environment."""
if config.is_hosted:
if self._foundry_storage is None:
self._foundry_storage = FoundryFunctionApprovalStore()
return self._foundry_storage
if self._in_memory_storage is None:
self._in_memory_storage = InMemoryFunctionApprovalStore()
return self._in_memory_storage
if self._storage is None:
self._storage = FoundryFunctionApprovalStore()
return self._storage


# endregion Function approval persistence
Expand Down Expand Up @@ -316,23 +288,17 @@ async def delete(self, session_id: str) -> None:
class AgentSessionStoreProvider(StoreProvider[SessionStore]):
"""Provide agent session store for the active hosting environment.

This will default to using the `FoundryAgentSessionStore` when hosted in Foundry,
and an in-memory store otherwise.
This defaults to using the `FoundryAgentSessionStore` in all environments.
"""

def __init__(self) -> None:
self._foundry_storage: SessionStore | None = None
self._in_memory_storage: SessionStore | None = None
self._storage: SessionStore | None = None

def get_store(self, *, config: AgentConfig) -> SessionStore:
"""Get agent session store for the requested hosting environment."""
if config.is_hosted:
if self._foundry_storage is None:
self._foundry_storage = FoundryAgentSessionStore()
return self._foundry_storage
if self._in_memory_storage is None:
self._in_memory_storage = SessionStore()
return self._in_memory_storage
if self._storage is None:
self._storage = FoundryAgentSessionStore()
return self._storage


# endregion Agent session persistence
6 changes: 3 additions & 3 deletions python/packages/foundry_hosting/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,9 @@ classifiers = [
]
dependencies = [
"agent-framework-core>=1.13.0,<2",
"azure-ai-agentserver-core>=2.0.0b11,<3",
"azure-ai-agentserver-responses>=2.0.0b1,<3",
"azure-ai-agentserver-invocations>=1.0.0b8,<2",
"azure-ai-agentserver-core>=2.1.0b1,<3",
"azure-ai-agentserver-responses>=2.0.0,<3",
"azure-ai-agentserver-invocations>=1.0.0,<2",
"httpx>=0.28,<1",
"mcp>=1.24.0,<2",
]
Expand Down
Loading
Loading