agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Simplify Python hosting core (#6492)

Browse Source 
Remove linking, multicast, durable delivery, and host push machinery from the v1 hosting core. Keep those scenarios in a proposed follow-up ADR and update channel packages, samples, docs, tests, and workspace metadata around the smaller host/channel contract.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
2026-06-12 08:34:08 +02:00
committed by GitHub
Unverified
parent e5a6e35843
commit 36ce0950e4
50 changed files with 1290 additions and 11651 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/hosting/README.md
+44 -133
View File
@@ -2,94 +2,43 @@
Multi-channel hosting for Microsoft Agent Framework agents.
`agent-framework-hosting` lets you serve a single agent (or workflow)
target through one or more **channels** — pluggable adapters that
expose the target over different transports. The result is a single
Starlette ASGI application you can host anywhere (local Hypercorn,
Azure Container Apps, Foundry Hosted Agents, …).
`agent-framework-hosting` lets you serve a single agent or workflow target
through one or more **channels**. The host owns one Starlette ASGI app,
route/lifecycle composition, and per-`isolation_key` session resolution.
Each channel owns its protocol parsing and response rendering.
The base package contains only the channel-neutral plumbing:
The base package contains only channel-neutral plumbing:
- `AgentFrameworkHost` — the Starlette host
- `Channel` / `ChannelPush` — the channel protocols
- `ChannelRequest` / `ChannelSession` / `ChannelIdentity` / `ResponseTarget`
— the request envelope and routing primitives
- `ChannelContext` / `ChannelContribution` / `ChannelCommand` — the
channel-side hooks for invoking the target and contributing routes,
commands, and lifecycle callbacks
- `ChannelRunHook` / `ChannelStreamTransformHook` — the per-request
customization seams
- `DurableTaskRunner` + `InProcessTaskRunner` — the seam used to
dispatch non-originating push fan-out; the in-process runner is the
default. Plug in a durable adapter (e.g.
`agent-framework-hosting-durabletask`) for `runtime_mode="ephemeral"`
deployments.
- `AgentFrameworkHost` — the Starlette host.
- `Channel` — the channel protocol.
- `ChannelRequest` / `ChannelSession` / `ChannelIdentity` — the request
envelope and optional channel metadata.
- `ChannelContext` / `ChannelContribution` / `ChannelCommand` — channel-side
hooks for invoking the target and contributing routes, commands, and
lifecycle callbacks.
- `ChannelRunHook` / `ChannelResponseHook` / `ChannelStreamUpdateHook` —
host-invoked customization seams.
Concrete channels live in their own packages so you only install what
you use:
`ChannelStreamUpdateHook` applies to streamed updates only. It is not a
substitute for final-response redaction.
Concrete channels live in their own packages so you only install what you use:
| Package | Transport |
|---|---|
| `agent-framework-hosting-responses` | OpenAI Responses API |
| `agent-framework-hosting-invocations` | Foundry-native invocation envelope |
| `agent-framework-hosting-telegram` | Telegram Bot API |
| `agent-framework-hosting-activity-protocol` | Bot Framework Activity Protocol (Teams, Direct Line, Web Chat, …) |
| `agent-framework-hosting-teams` | Microsoft Teams (Teams SDK) |
| `agent-framework-hosting-entra` | Entra (OAuth) identity-link sidecar |
## Architecture
```mermaid
graph LR
Caller[External caller /<br/>messaging app]
subgraph Host[AgentFrameworkHost]
direction TB
ASGI[Starlette app]
Router[Channel router]
Parse{parse →<br/>command or<br/>message?}
Auth[host.authorize]
Resolver[IdentityResolver]
Delivery[_deliver_response]
Push[_handle_push_task]
end
Channels[Channels<br/>Responses · Invocations ·<br/>Telegram · Activity ·<br/>IdentityLinker]
CmdHandler[CommandHandler<br/>via ChannelCommandContext]
Target[(Agent or Workflow)]
Runner[DurableTaskRunner]
StateStore[(HostStateStore)]
Caller --> ASGI
ASGI --> Router
Router --> Parse
Parse -- /command --> CmdHandler
Parse -- message --> Auth
CmdHandler -- ctx.run --> Auth
CmdHandler -- local reply --> Channels
Auth --> Resolver
Resolver --> StateStore
Auth --> Target
Target --> Delivery
Delivery -- originating sync --> Channels
Delivery -- non-originating --> Runner
Runner --> Push
Push --> Channels
Channels --> ASGI
```
For a richer set of flow diagrams — identity linking, multi-channel
fan-out, server-side relays, background runs, durable-runner codec
envelopes, echo idempotency, workflow targets — see the
[Python hosting spec](https://github.com/microsoft/agent-framework/blob/main/docs/specs/002-python-hosting-channels.md).
| `agent-framework-hosting-activity-protocol` | Bot Framework Activity Protocol |
| `agent-framework-hosting-discord` | Discord HTTP Interactions |
## Install
```bash
pip install agent-framework-hosting agent-framework-hosting-responses
# or with uvicorn pre-installed for the demo `host.serve(...)` helper
# or with Hypercorn pre-installed for the demo `host.serve(...)` helper
pip install "agent-framework-hosting[serve]" agent-framework-hosting-responses
# add the [disk] extra to opt in to on-disk persistence (see below)
# add the [disk] extra to persist reset-session aliases
pip install "agent-framework-hosting[disk]"
```
@@ -109,84 +58,46 @@ host = AgentFrameworkHost(target=agent, channels=channels)
host.serve(port=8000)
```
See the [hosting samples](https://github.com/microsoft/agent-framework/tree/main/python/samples/04-hosting/af-hosting)
for richer multi-channel apps (Telegram + Teams + Responses fan-out,
identity linking, `ResponseTarget` routing, etc.).
## Session state and workflow checkpoints
## Optional disk persistence (`state_dir`)
By default the host keeps live `AgentSession` objects and reset-session aliases
in memory. Channels opt into continuity by setting
`ChannelRequest.session = ChannelSession(isolation_key=...)`; requests with the
same isolation key reuse the same host-created session.
By default the host keeps everything in memory: the durable-task runner's
pending push queue, the per-isolation-key session aliases, the active-channel
map, and the per-channel `ChannelIdentity` map. That is the right shape for
**ephemeral** runtimes (Foundry Hosted Agents et al.) where the host is
restarted per request and persistence lives behind a service like the Foundry
response store, and for short-lived local dev.
For **long-running** deployments (an always-on container, a local dev server
you restart often, a single-VM bot) opt in to disk persistence by passing
`state_dir` to `AgentFrameworkHost`. The runner queue and the session
bookkeeping use [`diskcache`](https://grantjenks.com/docs/diskcache/)
(installed via the `[disk]` extra) protected by an OS-level advisory file
lock so two hosts pointed at the same directory can't double-execute
scheduled pushes. Workflow checkpoints (when the target is a `Workflow`)
use the framework's `FileCheckpointStorage` — no extra dependency. The
identity-link store path is offered to linkers that implement
`SupportsLinkStorePath`; linkers that manage persistence themselves should
be configured directly.
For long-running deployments that need `reset_session(...)` aliases to survive
restart, pass `state_dir`:
```python
from agent_framework_hosting import AgentFrameworkHost
# Single path → host auto-derives `runner/`, `sessions/`, `links/`, and
# (for workflow targets) `checkpoints/` subpaths.
host = AgentFrameworkHost(
target=agent,
channels=channels,
state_dir="./.host-state",
)
```
# Or route components to different roots — use the HostStatePaths TypedDict
# (or a plain dict with the same keys) for editor autocomplete on the keys.
# Omit a key to opt that component out of persistence.
This creates `./.host-state/sessions/` and stores only lightweight alias
bookkeeping. Live `AgentSession` objects are still rehydrated lazily by the
configured history provider on the next turn.
For workflow targets, `checkpoint_location=...` is the clearest way to enable
checkpoint persistence. As a convenience, `state_dir="./.host-state"` also
derives `./.host-state/checkpoints/` for workflow targets. Use the mapping form
when you want only one component:
```python
from agent_framework_hosting import HostStatePaths
host = AgentFrameworkHost(
target=workflow,
channels=channels,
state_dir=HostStatePaths(
runner="/var/lib/myapp/tasks",
sessions="/var/lib/myapp/state",
sessions="/var/lib/myapp/sessions",
checkpoints="/var/lib/myapp/checkpoints",
links="/var/lib/myapp/links",
),
)
```
What survives a restart:
- **Pending durable-task records** — scheduled but not-yet-completed push
deliveries replay on the next host startup via `runner.resume()`. Records
that crashed mid-attempt resume with their already-consumed retry budget.
- **`_session_aliases`** — per-isolation-key session-id rewrites (via the
reset-session command).
- **`_active`** — the most recently active channel for each isolation key
(consumed by `ResponseTarget.active`).
- **`_identities`** — channel-native `ChannelIdentity` rows used by
`ResponseTarget.channels([...])` / `.all_linked` fan-out.
- **Workflow checkpoints** — when the target is a `Workflow`, the host wraps
the `checkpoints` path in a per-isolation-key `FileCheckpointStorage`
(equivalent to passing `checkpoint_location=...` directly; the explicit
parameter takes precedence and emits a warning when both are set).
- **Identity-link store** — when the configured linker implements
`SupportsLinkStorePath`, the host passes the `links` path to it so pending
challenges, linked identities, and verified claims can survive restarts.
What doesn't:
- Live `AgentSession` objects (rehydrated lazily by the history provider on the
next turn).
- The `ContinuationToken` store (separate concern, plug in your own).
Unpicklable push payloads raise `PushPayloadNotPicklable` *eagerly* from
`schedule()` so issues surface at the call site, not on the next restart.
Cross-channel identity linking, multicast delivery, background runs,
continuation tokens, and durable delivery runners are follow-up enhancements,
not part of this v1 host contract.
python/packages/hosting/agent_framework_hosting/__init__.py
+3 -80
View File
@@ -13,30 +13,7 @@ they need.
import importlib.metadata
from ._authorization import (
AllOfAllowlists,
AllowAll,
Allowed,
AllowlistDecision,
AnyOfAllowlists,
AuthorizationContext,
AuthorizationOutcome,
AuthPolicy,
CallableAllowlist,
ChannelConfigurationError,
ClaimValue,
Denied,
IdentityAllowlist,
IdentityLinker,
LinkChallenge,
LinkedClaimAllowlist,
LinkedIdentity,
LinkRequired,
LinkResolution,
NativeIdAllowlist,
SupportsLinkStorePath,
)
from ._host import AgentFrameworkHost, ChannelContext, RuntimeMode, logger
from ._host import AgentFrameworkHost, ChannelContext, logger
from ._isolation import (
ISOLATION_HEADER_CHAT,
ISOLATION_HEADER_USER,
@@ -45,35 +22,19 @@ from ._isolation import (
reset_current_isolation_keys,
set_current_isolation_keys,
)
from ._runner import InProcessTaskRunner
from ._types import (
Channel,
ChannelCommand,
ChannelCommandContext,
ChannelContribution,
ChannelIdentity,
ChannelPush,
ChannelPushCodec,
ChannelRequest,
ChannelResponseContext,
ChannelResponseHook,
ChannelRunHook,
ChannelSession,
ChannelStreamTransformHook,
DurableTaskPayloadMode,
DurableTaskRunner,
ChannelStreamUpdateHook,
HostedRunResult,
HostStatePaths,
PushPayloadNotPicklable,
PushPayloadNotSerializable,
ResponseTarget,
ResponseTargetKind,
RetryPolicy,
TaskHandle,
TaskStatus,
apply_channel_response_hook,
apply_response_hook,
apply_run_hook,
)
try:
@@ -85,59 +46,21 @@ __all__ = [
"ISOLATION_HEADER_CHAT",
"ISOLATION_HEADER_USER",
"AgentFrameworkHost",
"AllOfAllowlists",
"AllowAll",
"Allowed",
"AllowlistDecision",
"AnyOfAllowlists",
"AuthPolicy",
"AuthorizationContext",
"AuthorizationOutcome",
"CallableAllowlist",
"Channel",
"ChannelCommand",
"ChannelCommandContext",
"ChannelConfigurationError",
"ChannelContext",
"ChannelContribution",
"ChannelIdentity",
"ChannelPush",
"ChannelPushCodec",
"ChannelRequest",
"ChannelResponseContext",
"ChannelResponseHook",
"ChannelRunHook",
"ChannelSession",
"ChannelStreamTransformHook",
"ClaimValue",
"Denied",
"DurableTaskPayloadMode",
"DurableTaskRunner",
"ChannelStreamUpdateHook",
"HostStatePaths",
"HostedRunResult",
"IdentityAllowlist",
"IdentityLinker",
"InProcessTaskRunner",
"IsolationKeys",
"LinkChallenge",
"LinkRequired",
"LinkResolution",
"LinkedClaimAllowlist",
"LinkedIdentity",
"NativeIdAllowlist",
"PushPayloadNotPicklable",
"PushPayloadNotSerializable",
"ResponseTarget",
"ResponseTargetKind",
"RetryPolicy",
"RuntimeMode",
"SupportsLinkStorePath",
"TaskHandle",
"TaskStatus",
"__version__",
"apply_channel_response_hook",
"apply_response_hook",
"apply_run_hook",
"get_current_isolation_keys",
"logger",
"reset_current_isolation_keys",
python/packages/hosting/agent_framework_hosting/_authorization.py
-485
View File
@@ -1,485 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Authorization seam — :class:`IdentityAllowlist`, :class:`IdentityLinker`, and outcomes.
Channels that emit a :class:`ChannelIdentity` compose authorization from
two **orthogonal** parameters set per channel:
- ``require_link: bool`` — "identity must be linked to an IdP claim". The
host delegates this to the configured :class:`IdentityLinker`; pairing
``require_link=True`` with no linker is rejected at construction
(silent-deny-everyone is the worst possible default).
- ``allowlist: IdentityAllowlist | Literal["inherit"] | None`` — "identity
is on the accept list". The host evaluates the allowlist on every
inbound message via :func:`AgentFrameworkHost.authorize`.
The two axes compose into the three named profiles **open** (no gate),
**forced-link** (any authenticated identity), and **allowlist** (only
listed identities, keyed either on the channel-native id pre-link or on
a verified IdP claim post-link). See
``docs/specs/002-python-hosting-channels.md`` §
"Authorization profiles and the IdentityAllowlist seam".
This module ships the channel-neutral core pieces. Provider-specific
linking channels (for example Entra OAuth helpers) can implement
:class:`IdentityLinker` without the core package taking a dependency on
their transport or identity-provider SDKs.
"""
from __future__ import annotations
import os
from collections.abc import Awaitable, Callable, Collection, Mapping, Sequence
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Literal, Protocol, TypeAlias, runtime_checkable
from ._types import ChannelIdentity
class AllowlistDecision(str, Enum):
"""Tri-state allowlist evaluation outcome.
``ABSTAIN`` is **not** a denial — it means "this allowlist has no
information yet" (typically a claim-based allowlist evaluated at
``pre_link``). The host's :meth:`AgentFrameworkHost.authorize`
pipeline is what turns an all-``ABSTAIN`` outcome into the next
step (allow when open, escalate to a link ceremony when the config
calls for one). Boolean composition cannot distinguish "claim
allowlist denies you" from "claim allowlist hasn't seen any claims
yet" — a critical distinction for the **Mixed** profile.
"""
ALLOW = "allow"
DENY = "deny"
ABSTAIN = "abstain"
ClaimValue: TypeAlias = str | Sequence[str]
"""Verified claim value shape understood by :class:`LinkedClaimAllowlist`."""
def _empty_claim_mapping() -> Mapping[str, ClaimValue]:
return {}
def _empty_any_mapping() -> Mapping[str, Any]:
return {}
@dataclass(frozen=True)
class AuthorizationContext:
"""Inputs to a single :meth:`IdentityAllowlist.evaluate` call."""
identity: ChannelIdentity
phase: Literal["pre_link", "post_link"]
isolation_key: str | None = None
verified_claims: Mapping[str, ClaimValue] = field(default_factory=_empty_claim_mapping)
claim_source: Literal["linker", "channel", "none"] = "none"
@runtime_checkable
class IdentityAllowlist(Protocol):
"""Per-channel accept/deny gate evaluated by the host.
``requires_linked_claims`` declares that this allowlist's
:meth:`evaluate` cannot ``ALLOW`` until verified claims are
available — the host's construction-time validator rejects
configurations that would silently deny everyone (e.g. a
:class:`LinkedClaimAllowlist` on a channel that neither has
``require_link=True`` nor natively emits verified claims).
"""
requires_linked_claims: bool
async def evaluate(self, context: AuthorizationContext) -> AllowlistDecision: ...
class AllowAll:
"""Explicit "open" sentinel.
Useful for tests, sample code, and for **overriding** a host-level
``default_allowlist`` on a specific channel that should be public
inside an otherwise locked-down host.
"""
requires_linked_claims: bool = False
async def evaluate(self, context: AuthorizationContext) -> AllowlistDecision:
return AllowlistDecision.ALLOW
class NativeIdAllowlist:
"""Accept only listed channel-native ids.
Telegram ``chat_id``, WhatsApp number, Slack user id, etc. The
list can be a plain collection or an async loader so allowlist
sources can be config files, secret stores, or feature flags.
Pre-link and post-link behaviour is identical — native-id
allowlists do not depend on link state.
When ``channel`` is set, the allowlist participates in
:class:`AnyOfAllowlists` composition by returning ``ABSTAIN`` for
requests from other channels — this lets per-channel native lists
coexist under a single combinator without one channel's ``DENY``
masking another channel's ``ALLOW``.
Keyword Args:
native_ids: A static collection of ids, or an async loader.
channel: When set, only requests whose
``ChannelIdentity.channel`` matches participate; others
``ABSTAIN``.
"""
requires_linked_claims: bool = False
def __init__(
self,
native_ids: Collection[str] | Callable[[], Awaitable[Collection[str]]],
*,
channel: str | None = None,
) -> None:
self._native_ids: Collection[str] | None
self._loader: Callable[[], Awaitable[Collection[str]]] | None
if callable(native_ids):
self._native_ids = None
self._loader = native_ids
else:
self._native_ids = frozenset(native_ids)
self._loader = None
self.channel = channel
async def _resolve(self) -> Collection[str]:
if self._native_ids is not None:
return self._native_ids
loader = self._loader
if loader is None: # pragma: no cover - defensive
raise RuntimeError("NativeIdAllowlist: loader missing after cache miss")
loaded = await loader()
# Cache the resolved set so subsequent calls avoid re-loading.
self._native_ids = frozenset(loaded)
self._loader = None
return self._native_ids
async def evaluate(self, context: AuthorizationContext) -> AllowlistDecision:
if self.channel is not None and context.identity.channel != self.channel:
return AllowlistDecision.ABSTAIN
ids = await self._resolve()
if context.identity.native_id in ids:
return AllowlistDecision.ALLOW
return AllowlistDecision.DENY
class LinkedClaimAllowlist:
"""Accept only identities whose verified IdP claim is on the list.
``evaluate`` returns ``ABSTAIN`` at ``pre_link`` (no claims yet)
and ``ALLOW``/``DENY`` at ``post_link``. Claim values may be plain
strings or a sequence of strings (for multi-valued claims such as
group ids); any intersection with ``values`` allows the identity.
Keyword Args:
claim: The verified-claim key to inspect (e.g. ``"oid"``,
``"tid"``, ``"groups"``).
values: Accepted values.
"""
requires_linked_claims: bool = True
def __init__(self, claim: str, values: Collection[str]) -> None:
self.claim = claim
self.values = frozenset(values)
async def evaluate(self, context: AuthorizationContext) -> AllowlistDecision:
if context.phase == "pre_link":
return AllowlistDecision.ABSTAIN
value = context.verified_claims.get(self.claim)
if value is None:
return AllowlistDecision.DENY
if isinstance(value, str):
return AllowlistDecision.ALLOW if value in self.values else AllowlistDecision.DENY
return AllowlistDecision.ALLOW if any(item in self.values for item in value) else AllowlistDecision.DENY
class AnyOfAllowlists:
"""Combinator: any child ``ALLOW`` wins; ``DENY`` only if all children ``DENY``.
Use this for the **Mixed** profile (native id OR linked claim).
Returns ``ABSTAIN`` when no child decides.
"""
def __init__(self, *allowlists: IdentityAllowlist) -> None:
self._children = allowlists
self.requires_linked_claims = any(getattr(a, "requires_linked_claims", False) for a in allowlists)
async def evaluate(self, context: AuthorizationContext) -> AllowlistDecision:
any_abstain = False
all_deny = True
for child in self._children:
decision = await child.evaluate(context)
if decision is AllowlistDecision.ALLOW:
return AllowlistDecision.ALLOW
if decision is AllowlistDecision.ABSTAIN:
any_abstain = True
all_deny = False
# DENY contributes to all_deny without short-circuit.
if all_deny and self._children:
return AllowlistDecision.DENY
if any_abstain:
return AllowlistDecision.ABSTAIN
# No children — treat as ABSTAIN to avoid surprise DENY.
return AllowlistDecision.ABSTAIN
class AllOfAllowlists:
"""Combinator: any child ``DENY`` wins; ``ALLOW`` only if all children ``ALLOW``.
Use this to require multiple conditions (e.g. tenancy
**and** group membership). Returns ``ABSTAIN`` when no child
denies but at least one ``ABSTAIN``s.
"""
def __init__(self, *allowlists: IdentityAllowlist) -> None:
self._children = allowlists
self.requires_linked_claims = any(getattr(a, "requires_linked_claims", False) for a in allowlists)
async def evaluate(self, context: AuthorizationContext) -> AllowlistDecision:
any_abstain = False
for child in self._children:
decision = await child.evaluate(context)
if decision is AllowlistDecision.DENY:
return AllowlistDecision.DENY
if decision is AllowlistDecision.ABSTAIN:
any_abstain = True
if not self._children:
return AllowlistDecision.ABSTAIN
if any_abstain:
return AllowlistDecision.ABSTAIN
return AllowlistDecision.ALLOW
class CallableAllowlist:
"""Escape hatch: wrap an arbitrary async function as an allowlist.
Recommended only after exhausting the structured variants —
composition is harder to reason about with opaque callables.
"""
def __init__(
self,
fn: Callable[[AuthorizationContext], Awaitable[AllowlistDecision]],
*,
requires_linked_claims: bool = False,
) -> None:
self._fn = fn
self.requires_linked_claims = requires_linked_claims
async def evaluate(self, context: AuthorizationContext) -> AllowlistDecision:
return await self._fn(context)
# --------------------------------------------------------------------------- #
# Outcome types #
# --------------------------------------------------------------------------- #
@dataclass(frozen=True)
class LinkChallenge:
"""Challenge a channel can render to complete an identity link.
Attributes:
challenge_id: Opaque linker-owned id for correlating the challenge
with the later completion callback.
url: Optional URL (OAuth authorization URL, device-flow URL, etc.)
the user should open.
expires_at: Optional challenge expiry time.
message: Optional safe text a channel may render with the challenge.
attributes: Linker-specific structured metadata. Channels should
only use keys documented by the concrete linker they integrate.
"""
challenge_id: str
url: str | None = None
expires_at: datetime | None = None
message: str | None = None
attributes: Mapping[str, Any] = field(default_factory=_empty_any_mapping)
@dataclass(frozen=True)
class LinkedIdentity:
"""Resolved IdP-backed identity returned by :class:`IdentityLinker`.
Attributes:
isolation_key: Stable key the host should use for the linked user.
verified_claims: Claims verified by the linker or by a channel that
natively authenticates the user.
claim_source: Where the claims came from.
"""
isolation_key: str
verified_claims: Mapping[str, ClaimValue] = field(default_factory=_empty_claim_mapping)
claim_source: Literal["linker", "channel"] = "linker"
LinkResolution: TypeAlias = LinkedIdentity | LinkChallenge
"""Result returned by :meth:`IdentityLinker.resolve`."""
class IdentityLinker(Protocol):
"""Resolve a channel-native identity or return a challenge to link it.
Concrete linker packages own the storage, OAuth/device-code routes, and
provider-specific claim mapping. The core host only consumes the single
resolution call so authorization can be a one-round-trip decision.
"""
async def resolve(self, identity: ChannelIdentity) -> LinkResolution:
"""Return a linked identity or the challenge needed to create one."""
...
@runtime_checkable
class SupportsLinkStorePath(Protocol):
"""Optional protocol for linkers that accept host-provided persistence.
When ``AgentFrameworkHost(state_dir=...)`` derives a ``links`` path, the
host calls this hook on identity linkers that implement it. Linkers that
manage their own persistence can ignore this protocol and should be
configured directly by the application.
"""
def configure_link_store_path(self, path: str | os.PathLike[str]) -> None:
"""Configure where the linker should persist its link store."""
...
@dataclass(frozen=True)
class Allowed:
"""The identity is authorized; ``isolation_key`` is its stable key."""
isolation_key: str
verified_claims: Mapping[str, ClaimValue] = field(default_factory=_empty_claim_mapping)
claim_source: Literal["linker", "channel", "none"] = "none"
@dataclass(frozen=True)
class LinkRequired:
"""The identity must complete the link ceremony before proceeding.
Channels render ``challenge`` through their native UX (the same
path the ``link`` command uses).
"""
challenge: LinkChallenge
@dataclass(frozen=True)
class Denied:
"""The identity is rejected.
Attributes:
reason_code: Stable, machine-readable token (e.g.
``"allowlist_denied_pre_link"``). Never echoed to end
users.
user_message: Safe to render publicly (group-chat-safe);
``None`` falls back to a bland default ("You don't have
access to this bot.").
log_details: Structured payload for audit/observability;
**never** shown to users.
"""
reason_code: str
user_message: str | None = None
log_details: Mapping[str, Any] = field(default_factory=_empty_any_mapping)
AuthorizationOutcome = Allowed | LinkRequired | Denied
"""Result of :func:`AgentFrameworkHost.authorize`. Channels render
each variant through their native UX."""
class AuthPolicy:
"""Factory helpers for common authorization policies.
These helpers are thin wrappers over the concrete allowlist types; they
exist so application code can describe authorization intent without
importing each building block separately.
"""
@staticmethod
def open() -> AllowAll:
"""Allow every identity."""
return AllowAll()
@staticmethod
def native_ids(
native_ids: Collection[str] | Callable[[], Awaitable[Collection[str]]],
*,
channel: str | None = None,
) -> NativeIdAllowlist:
"""Allow listed channel-native ids."""
return NativeIdAllowlist(native_ids, channel=channel)
@staticmethod
def linked_claim(claim: str, values: Collection[str]) -> LinkedClaimAllowlist:
"""Allow identities whose verified claim matches one of ``values``."""
return LinkedClaimAllowlist(claim, values)
@staticmethod
def any_of(*allowlists: IdentityAllowlist) -> AnyOfAllowlists:
"""Allow when any child allowlist allows."""
return AnyOfAllowlists(*allowlists)
@staticmethod
def all_of(*allowlists: IdentityAllowlist) -> AllOfAllowlists:
"""Allow only when every child allowlist allows."""
return AllOfAllowlists(*allowlists)
@staticmethod
def custom(
fn: Callable[[AuthorizationContext], Awaitable[AllowlistDecision]],
*,
requires_linked_claims: bool = False,
) -> CallableAllowlist:
"""Wrap a custom async allowlist function."""
return CallableAllowlist(fn, requires_linked_claims=requires_linked_claims)
# --------------------------------------------------------------------------- #
# Configuration error #
# --------------------------------------------------------------------------- #
class ChannelConfigurationError(ValueError):
"""Raised at host construction for authorization config that would deny all users.
The host validator runs three rules (see spec §"Configuration
validation"); any failure is reported here rather than letting
the misconfigured host start up and reject every request.
"""
__all__ = [
"AllOfAllowlists",
"AllowAll",
"Allowed",
"AllowlistDecision",
"AnyOfAllowlists",
"AuthPolicy",
"AuthorizationContext",
"AuthorizationOutcome",
"CallableAllowlist",
"ChannelConfigurationError",
"ClaimValue",
"Denied",
"IdentityAllowlist",
"IdentityLinker",
"LinkChallenge",
"LinkRequired",
"LinkResolution",
"LinkedClaimAllowlist",
"LinkedIdentity",
"NativeIdAllowlist",
"SupportsLinkStorePath",
]
python/packages/hosting/agent_framework_hosting/_host.py
+237 -1236
View File
File diff suppressed because it is too large Load Diff
python/packages/hosting/agent_framework_hosting/_persistence.py
+20 -87
View File
@@ -2,25 +2,10 @@
"""Shared persistence primitives for the hosting package.
The hosting core ships with an opt-in disk-persistence layer for the
in-process task runner and the host's session-related state. The
on-disk format is provided by the ``diskcache`` package (a small,
pure-Python, sqlite-backed dependency installed via the ``[disk]``
optional extra).
This module centralises:
- :func:`load_diskcache` — lazy import that raises a helpful error when
the optional extra is missing.
- :func:`acquire_state_dir_lock` — single-owner file lock that fails
fast when a second process points at the same directory.
- :func:`normalize_state_dir` — turn the host-level ``state_dir``
parameter (``str`` / ``PathLike`` / :class:`HostStatePaths` /
``Mapping``) into a normalised ``dict[component_name -> Path | None]``.
Everything in this module is internal — public callers should go
through :class:`AgentFrameworkHost` or
:class:`InProcessTaskRunner` directly.
The simplified hosting core keeps disk persistence only for session aliases
created by :meth:`AgentFrameworkHost.reset_session` and for workflow
checkpoint path derivation. The on-disk session-alias store uses the optional
``diskcache`` package installed via the ``[disk]`` extra.
"""
from __future__ import annotations
@@ -35,29 +20,19 @@ from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from ._types import HostStatePaths
# Known component keys recognised by the host's ``state_dir`` normaliser.
# Adding a new component is a non-breaking change: extend this tuple and
# add the matching key to :class:`HostStatePaths` in ``_types.py``.
_KNOWN_COMPONENTS: tuple[str, ...] = ("runner", "sessions", "checkpoints", "links")
_KNOWN_COMPONENTS: tuple[str, ...] = ("sessions", "checkpoints")
def load_diskcache() -> Any:
"""Lazy-import :mod:`diskcache` with a helpful error when missing.
The ``diskcache`` package is an optional dependency installed via
the ``agent-framework-hosting[disk]`` extra. Users that never set
``state_dir`` never trigger the import. This wrapper produces a
single, consistent error message when the import is needed but the
extra was not installed.
"""
"""Lazy-import :mod:`diskcache` with a helpful error when missing."""
try:
import diskcache # type: ignore[import-untyped]
except ImportError as exc: # pragma: no cover - exercised via tests by monkeypatching
raise ImportError(
"agent-framework-hosting was asked to persist state to disk "
"(state_dir is set) but the optional `diskcache` dependency "
"agent-framework-hosting was asked to persist session aliases to disk "
"(state_dir['sessions'] is set) but the optional `diskcache` dependency "
"is not installed. Install the disk extra: "
"`pip install 'agent-framework-hosting[disk]'`."
"`pip install 'agent-framework-hosting[disk]`."
) from exc
return diskcache
@@ -65,24 +40,11 @@ def load_diskcache() -> Any:
def acquire_state_dir_lock(component_dir: Path) -> Any:
"""Acquire an exclusive single-owner lock on a component's state dir.
Two processes pointing at the same state directory would both scan
pending records on startup and could execute the same task twice;
we therefore enforce single-owner semantics with an OS-level
advisory lock. The lock file lives at ``<component_dir>/.lock`` and
is held for the lifetime of the returned file handle. Closing the
handle (or process exit) releases it.
On Unix this uses :func:`fcntl.flock`. On Windows it uses
:func:`msvcrt.locking`. The lock is *advisory* — the OS will not
enforce it against processes that ignore it, but no
well-behaved component of this package will.
Raises ``RuntimeError`` if another process already holds the lock.
Raises:
RuntimeError: If another process already holds the lock.
"""
component_dir.mkdir(parents=True, exist_ok=True)
lock_path = component_dir / ".lock"
# Open in append mode so we don't truncate an existing lock file
# (some monitoring tools may inspect it).
fh = open(lock_path, "a+", encoding="utf-8") # noqa: SIM115 - kept open for lifetime
try:
if sys.platform == "win32":
@@ -94,8 +56,7 @@ def acquire_state_dir_lock(component_dir: Path) -> Any:
fh.close()
raise RuntimeError(
f"Another process already holds the hosting state lock at {lock_path}. "
"Two hosts (or two runners) pointing at the same state directory would "
"double-execute scheduled tasks; point each host at its own state_dir."
"Point each host at its own state_dir."
) from exc
else:
import fcntl
@@ -106,8 +67,7 @@ def acquire_state_dir_lock(component_dir: Path) -> Any:
fh.close()
raise RuntimeError(
f"Another process already holds the hosting state lock at {lock_path}. "
"Two hosts (or two runners) pointing at the same state directory would "
"double-execute scheduled tasks; point each host at its own state_dir."
"Point each host at its own state_dir."
) from exc
except RuntimeError:
raise
@@ -118,15 +78,10 @@ def acquire_state_dir_lock(component_dir: Path) -> Any:
def release_state_dir_lock(handle: Any) -> None:
"""Release a lock previously acquired by :func:`acquire_state_dir_lock`.
Closing the file handle is sufficient to drop the lock on both
platforms, but we make the intent explicit so the caller doesn't
have to know which mechanism (``fcntl`` vs ``msvcrt``) is in use.
"""
"""Release a lock previously acquired by :func:`acquire_state_dir_lock`."""
if handle is None:
return
with contextlib.suppress(Exception): # close errors are not actionable
with contextlib.suppress(Exception):
handle.close()
@@ -135,40 +90,27 @@ def normalize_state_dir(
) -> dict[str, Path | None]:
"""Resolve the host-level ``state_dir`` parameter into a per-component map.
Accepts any of:
- ``None`` → all components return ``None`` (fully in-memory; today's behavior).
- ``str`` / :class:`os.PathLike` → all components share a parent
directory and get an auto-allocated subfolder (``runner/``,
``sessions/``, ``checkpoints/``, ``links/``).
- :class:`HostStatePaths` typed dict / plain ``Mapping`` → per-key
override. Components missing from the mapping fall back to ``None``
(in-memory only). Unknown keys raise ``ValueError`` to surface
typos early.
Returns a ``dict[component_name -> Path | None]`` covering every
component in :data:`_KNOWN_COMPONENTS`.
Accepts ``None``, a single root path, or a mapping with ``sessions`` and
``checkpoints`` keys. Unknown keys raise ``ValueError`` so obsolete
``runner`` / ``links`` configuration is rejected instead of silently
doing nothing.
"""
result: dict[str, Path | None] = {name: None for name in _KNOWN_COMPONENTS}
if state_dir is None:
return result
# Strings and PathLikes use the default subfolder layout.
if isinstance(state_dir, (str, os.PathLike)):
root = Path(os.fspath(state_dir))
for name in _KNOWN_COMPONENTS:
result[name] = root / name
return result
# Mappings (incl. TypedDict at runtime) get per-component overrides.
if isinstance(state_dir, Mapping):
unknown = [k for k in state_dir if k not in _KNOWN_COMPONENTS]
if unknown:
raise ValueError(
f"state_dir mapping contains unknown component key(s): {unknown!r}. "
f"Known components are: {list(_KNOWN_COMPONENTS)!r}. "
"If you are trying to use a future component, upgrade "
"agent-framework-hosting to a version that supports it."
f"Known components are: {list(_KNOWN_COMPONENTS)!r}."
)
for name in _KNOWN_COMPONENTS:
raw_value: Any = state_dir.get(name)
@@ -184,12 +126,3 @@ def normalize_state_dir(
raise TypeError(
f"state_dir must be a str, PathLike, HostStatePaths mapping, or None — got {type(state_dir).__name__}"
)
__all__ = [
"_KNOWN_COMPONENTS",
"acquire_state_dir_lock",
"load_diskcache",
"normalize_state_dir",
"release_state_dir_lock",
]
python/packages/hosting/agent_framework_hosting/_runner.py
-751
View File
@@ -1,751 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""In-process implementation of :class:`DurableTaskRunner`.
This is the default runner the host wires in when the operator does not
supply one. It runs tasks via :func:`asyncio.create_task` with a bounded
retry loop following the supplied :class:`RetryPolicy`.
Two modes:
* **In-memory** (``state_dir=None``, default) — pending tasks live as
``asyncio.Task`` references in process memory. In-flight tasks are
lost on process death. Cheap, zero dependencies, suitable for unit
tests and for long-running deployments where "the process dies,
queued pushes are lost" is an acceptable failure mode.
* **Disk-persistent** (``state_dir=<path>``) — pending tasks are
pickled into a :mod:`diskcache`-backed sqlite store before the
``asyncio.Task`` is created. On the next startup the host calls
:meth:`InProcessTaskRunner.resume` which re-schedules every
surviving ``"pending"`` record with its persisted attempt count.
Graceful shutdown cancellations leave records in ``"pending"`` so
they replay on the next boot. Suitable for ``runtime_mode="long_running"``
deployments that survive container moves / OOMs.
For ``runtime_mode="ephemeral"`` deployments (Foundry Hosted Agent,
Azure Functions, Lambda) plug in a durable adapter package
(``agent-framework-hosting-durabletask`` for the gRPC TaskHub backend,
a future Foundry adapter, …) — they all implement the same
:class:`DurableTaskRunner` Protocol.
See ``docs/specs/002-python-hosting-channels.md`` § "Durable task runner".
"""
from __future__ import annotations
import asyncio
import contextlib
import logging
import os
import pickle # noqa: S403 # nosec B403 - used only to validate user payloads round-trip
import time
import uuid
from collections.abc import Awaitable, Callable, Mapping
from pathlib import Path
from typing import Any, cast
from ._persistence import (
acquire_state_dir_lock,
load_diskcache,
release_state_dir_lock,
)
from ._types import (
DurableTaskPayloadMode,
DurableTaskRunner,
PushPayloadNotPicklable,
RetryPolicy,
TaskHandle,
TaskStatus,
)
logger = logging.getLogger(__name__)
# Keys used inside the per-task on-disk record. Kept as module constants
# so the schema is documented once and refactors are mechanical.
_REC_HANDLER_NAME = "handler_name"
_REC_PAYLOAD = "payload"
_REC_RETRY_POLICY = "retry_policy"
_REC_ATTEMPTS = "attempts_completed"
_REC_STATUS = "status"
_REC_CREATED_AT = "created_at"
_REC_TERMINAL_AT = "terminal_at"
_REC_NAME = "name"
# Deque key inside the cache holding terminal task ids in insertion order.
# Used for FIFO eviction of terminal records once the bounded cap is hit.
_TERMINAL_ORDER_KEY = "__terminal_order__"
class _PersistedPayloadDict(dict[str, Any]):
"""Drop-in :class:`dict` that mirrors mutations back to disk.
Used by :class:`InProcessTaskRunner` when ``state_dir`` is set so
handler-side cursors (``echo_done``) survive process restarts. The
handler interacts with this object exactly as it would with a plain
dict; the override on :meth:`__setitem__` is the only difference.
Held weakly by the runner so handlers that capture the dict in
long-lived closures don't keep the runner alive past its natural
lifetime.
"""
# Type annotation for the persist callback; the actual attribute is
# assigned via the __slots__-aware ``object.__setattr__`` dance
# below so PyPy doesn't reject the assignment on a ``dict`` subclass.
_persist_cb: Callable[[Mapping[str, Any]], None]
__slots__ = ("_persist_cb",)
def __init__(
self,
data: Mapping[str, Any],
persist_cb: Callable[[Mapping[str, Any]], None],
) -> None:
super().__init__(data)
# Use object.__setattr__ to bypass the __slots__ checker on
# dict subclasses (CPython is liberal here but PyPy is strict).
object.__setattr__(self, "_persist_cb", persist_cb)
def __setitem__(self, key: str, value: Any) -> None:
super().__setitem__(key, value)
# Re-serialise after each mutation. The cache stores opaque
# pickled values, so partial-field updates aren't possible —
# we send the whole payload mapping every time. Mutations on
# the runner's hot path are rare (just the ``echo_done``
# cursor today) so this is fine.
self._persist_cb(dict(self))
class InProcessTaskRunner(DurableTaskRunner):
"""In-memory or disk-persistent :class:`DurableTaskRunner`.
Schedules each task as an :func:`asyncio.create_task` coroutine and
retries on exception up to ``RetryPolicy.max_attempts`` times with
exponential backoff. Terminal status (``succeeded`` / ``failed`` /
``cancelled``) is reported via :meth:`get`.
Re-registration of the same handler name after :meth:`schedule` has
been called is rejected to avoid silent re-orderings of in-flight
work; the host registers all handlers at startup, before serving
traffic.
Keyword Args:
default_retry_policy: Per-runner default :class:`RetryPolicy`;
overridable per-task at :meth:`schedule` call sites.
terminal_cache_size: Maximum number of terminal task records to
retain. Older entries are FIFO-evicted so a long-running
host can't accumulate unbounded status entries.
shutdown_grace_seconds: Window :meth:`shutdown` waits for
in-flight tasks to drain before cancelling stragglers.
state_dir: When set, the runner persists pending and terminal
task records under this directory (a :mod:`diskcache`
sqlite store at ``<state_dir>/cache.db`` and a single-owner
lock at ``<state_dir>/.lock``). Persisted pending records
survive process restarts and are replayed by :meth:`resume`.
When ``None`` (default) the runner is purely in-memory and
in-flight tasks are lost on process death. Requires the
optional ``diskcache`` dependency — install with
``pip install 'agent-framework-hosting[disk]'``.
"""
# Declared at class level so the ``DurableTaskRunner`` Protocol's
# ``payload_mode`` attribute resolves on instances without needing
# to assign it in ``__init__``.
payload_mode: DurableTaskPayloadMode = DurableTaskPayloadMode.OBJECT
def __init__(
self,
*,
default_retry_policy: RetryPolicy | None = None,
terminal_cache_size: int = 1024,
shutdown_grace_seconds: float = 5.0,
state_dir: str | os.PathLike[str] | None = None,
) -> None:
self._handlers: dict[str, Callable[[Mapping[str, Any]], Awaitable[None]]] = {}
self._default_retry_policy = default_retry_policy or RetryPolicy()
self._terminal_cache_size = terminal_cache_size
# How long ``shutdown()`` waits for in-flight tasks to finish on
# their own before cancelling them. Channels may legitimately
# schedule a final push during their own shutdown callback
# (goodbye message, telemetry flush), so the runner gives them
# this window to complete before cancellation kicks in.
self._shutdown_grace_seconds = shutdown_grace_seconds
# Operational state. ``_pending`` holds asyncio tasks that are
# scheduled or running. ``_terminal`` is an in-memory mirror of
# the most recent terminal statuses (kept in-memory regardless of
# ``state_dir`` so ``get`` is fast and works before/without the
# cache being opened).
self._pending: dict[str, asyncio.Task[None]] = {}
self._terminal: dict[str, TaskStatus] = {}
self._terminal_order: list[str] = []
# Set to True on the first ``schedule``/``resume`` call so subsequent
# ``register`` calls fail loudly rather than silently swapping a
# handler out from under in-flight work.
self._started = False
# Set to True when ``shutdown()`` starts so the retry loop's
# ``CancelledError`` handler distinguishes "the runner is going
# down, leave my record in 'pending' for resume()" from "this
# task was explicitly cancelled, mark it 'cancelled'".
self._shutting_down = False
# Disk persistence — opt-in via ``state_dir``. ``None`` keeps
# the runner pure-memory (the default behaviour).
self._state_dir: Path | None = Path(os.fspath(state_dir)) if state_dir is not None else None
self._cache: Any = None
self._terminal_deque: Any = None
self._lock_handle: Any = None
if self._state_dir is not None:
self._open_cache()
# ------------------------------------------------------------------ #
# Cache lifecycle
# ------------------------------------------------------------------ #
def _open_cache(self) -> None:
"""Open the disk cache and acquire the single-owner lock.
Called from ``__init__`` when ``state_dir`` is set. Splitting it
out keeps the constructor body readable and gives tests a clean
seam for monkeypatching.
"""
if self._state_dir is None: # pragma: no cover - guarded by caller
raise RuntimeError("_open_cache called without state_dir")
diskcache = load_diskcache()
# Acquire the directory lock *before* opening the cache so two
# runners pointed at the same dir don't both try to initialise
# sqlite. The lock handle stays open for the runner's lifetime.
self._lock_handle = acquire_state_dir_lock(self._state_dir)
try:
self._cache = diskcache.Cache(str(self._state_dir))
# Re-hydrate the in-memory terminal mirror so ``get`` works
# for task ids that completed in a prior process. Doing this
# here (rather than lazily) means the mirror is consistent
# the moment construction returns.
order: Any = self._cache.get(_TERMINAL_ORDER_KEY, default=[])
if not isinstance(order, list):
# Defensive: a corrupted ordering list shouldn't take
# the host down. Reset and continue — at worst we lose
# ordering for FIFO eviction, not correctness.
logger.warning(
"InProcessTaskRunner: terminal-order entry in %s is not a list; resetting", self._state_dir
)
order = []
self._cache.set(_TERMINAL_ORDER_KEY, order)
self._terminal_order = [str(x) for x in cast(list[Any], order)]
for task_id in self._terminal_order:
rec_obj: Any
try:
rec_obj = self._cache.get(task_id)
except Exception: # pragma: no cover - exercised via corrupt-entry test
rec_obj = None
if not isinstance(rec_obj, dict):
continue
rec = cast(dict[str, Any], rec_obj)
status = rec.get(_REC_STATUS)
if status in {"succeeded", "failed", "cancelled"}:
self._terminal[task_id] = status
except Exception:
release_state_dir_lock(self._lock_handle)
self._lock_handle = None
raise
# ------------------------------------------------------------------ #
# DurableTaskRunner Protocol
# ------------------------------------------------------------------ #
def register(
self,
name: str,
handler: Callable[[Mapping[str, Any]], Awaitable[None]],
) -> None:
if self._started:
raise RuntimeError(
f"InProcessTaskRunner.register({name!r}) called after the "
"runner started scheduling tasks — register all handlers at "
"host startup, before serving traffic, to avoid silently "
"reordering in-flight work."
)
if name in self._handlers:
logger.warning("InProcessTaskRunner: replacing handler registered under %r", name)
self._handlers[name] = handler
async def schedule(
self,
name: str,
payload: Mapping[str, Any],
*,
retry_policy: RetryPolicy | None = None,
) -> TaskHandle:
if name not in self._handlers:
raise KeyError(
f"InProcessTaskRunner.schedule({name!r}): no handler "
"registered under this name. Call register(name, handler) "
"at host startup before scheduling."
)
self._started = True
policy = retry_policy or self._default_retry_policy
task_id = uuid.uuid4().hex
handle = TaskHandle(task_id=task_id, name=name)
# Persist the record (when state_dir is set) BEFORE we spawn the
# asyncio task — if the persistence write fails we surface it as
# a synchronous error from ``schedule`` rather than silently
# downgrading to in-memory.
if self._cache is not None:
record = self._build_record(name, dict(payload), policy)
self._validate_picklable(record)
self._cache.set(task_id, record)
# When persisted, wrap the payload so handler-side mutations
# (e.g. ``payload["echo_done"] = True``) flow back to disk.
runtime_payload: Mapping[str, Any]
if self._cache is not None:
captured_task_id = task_id
def _persist_cb(new_payload: Mapping[str, Any]) -> None:
self._update_record_payload(captured_task_id, new_payload)
runtime_payload = _PersistedPayloadDict(payload, _persist_cb)
else:
runtime_payload = payload
handler = self._handlers[name]
task = asyncio.create_task(
self._run_with_retry(handle, handler, runtime_payload, policy),
name=f"hosting.task[{name}]:{task_id}",
)
self._pending[task_id] = task
def _on_done(_t: asyncio.Task[None], tid: str = task_id) -> None:
self._pending.pop(tid, None)
task.add_done_callback(_on_done)
return handle
async def get(self, handle: TaskHandle) -> TaskStatus | None:
if handle.task_id in self._pending:
task = self._pending[handle.task_id]
if task.cancelled():
return "cancelled"
return "running"
# In-memory terminal mirror covers both pure-memory and
# disk-persistent runs (we re-hydrate on cache open).
if handle.task_id in self._terminal:
return self._terminal[handle.task_id]
# Disk fallback for very-aged task ids that left the in-memory
# mirror but still have a record on disk (extremely unlikely
# given that we re-hydrate all terminals at open, but defensive).
if self._cache is not None:
rec_obj: Any = self._cache.get(handle.task_id)
if isinstance(rec_obj, dict):
rec = cast(dict[str, Any], rec_obj)
status = rec.get(_REC_STATUS)
# Records on disk only live in one of four states:
# ``pending`` (queued or in-flight — resume picks these
# up) or one of the terminals. There is no transient
# ``running`` status; the in-flight asyncio task is
# observable via ``_pending`` only inside its own
# process.
if status in {"succeeded", "failed", "cancelled", "pending"}:
return cast(TaskStatus, status)
return None
# ------------------------------------------------------------------ #
# Resume — replay persisted pending records on startup
# ------------------------------------------------------------------ #
async def resume(self) -> int:
"""Re-schedule pending tasks persisted by a previous process.
Walks the cache for records in ``"pending"`` status, looks up
their handler in :attr:`_handlers`, and re-creates an
:class:`asyncio.Task` for each — preserving the persisted
attempt count so retry budgets resume mid-way through their
backoff schedule.
Records whose handler is no longer registered are marked
``"failed"`` with a clear reason in the log; they will not be
retried again. Records that fail to deserialise (corrupted
sqlite row, schema drift, …) are quarantined: their entry is
removed from the cache and the task id is logged. Both classes
of error are non-fatal — the host should boot even when a
small number of legacy records can't be replayed.
Returns the number of records successfully re-scheduled.
Called automatically from :class:`AgentFrameworkHost`'s lifespan
startup hook when the runner is host-owned. Callers driving the
runner directly (tests, bespoke ASGI setups) MUST call this
once after registering handlers and before serving traffic.
"""
if self._cache is None:
return 0
# Mark started so subsequent register() calls fail loudly — we
# don't want handler swaps after replay begins.
self._started = True
replayed = 0
# iterkeys returns a live view; we copy to a list because we may
# delete entries inside the loop (quarantine / drop-on-missing-handler).
task_ids: list[str] = [str(k) for k in self._cache.iterkeys() if k != _TERMINAL_ORDER_KEY]
for task_id in task_ids:
rec_obj: Any
try:
rec_obj = self._cache.get(task_id)
except Exception:
logger.exception("InProcessTaskRunner.resume: failed to read record %s; quarantining", task_id)
with contextlib.suppress(KeyError):
del self._cache[task_id]
continue
if not isinstance(rec_obj, dict) or _REC_STATUS not in rec_obj:
logger.warning("InProcessTaskRunner.resume: record %s is not a task dict; quarantining", task_id)
with contextlib.suppress(KeyError):
del self._cache[task_id]
continue
rec = cast(dict[str, Any], rec_obj)
status = rec[_REC_STATUS]
if status != "pending":
continue
handler_name = rec.get(_REC_HANDLER_NAME)
if not isinstance(handler_name, str) or handler_name not in self._handlers:
logger.warning(
"InProcessTaskRunner.resume: no handler registered for record %s (handler=%r); marking failed",
task_id,
handler_name,
)
self._mark_terminal(task_id, "failed")
continue
handler = self._handlers[handler_name]
policy_value = rec.get(_REC_RETRY_POLICY) or self._default_retry_policy
if not isinstance(policy_value, RetryPolicy):
# Legacy / corrupt entry — fall back to the default rather
# than failing the whole resume.
policy_value = self._default_retry_policy
policy: RetryPolicy = policy_value
payload_value: Any = rec.get(_REC_PAYLOAD) or {}
payload: dict[str, Any]
if isinstance(payload_value, dict):
payload = cast(dict[str, Any], payload_value)
elif hasattr(payload_value, "keys"):
payload = dict(cast(Mapping[str, Any], payload_value))
else:
payload = {}
name_value = rec.get(_REC_NAME, handler_name)
handle = TaskHandle(task_id=task_id, name=str(name_value))
attempts_value = rec.get(_REC_ATTEMPTS, 0)
attempts_completed = int(attempts_value or 0)
def _make_resume_persist_cb(tid: str) -> Callable[[Mapping[str, Any]], None]:
def _cb(new_payload: Mapping[str, Any]) -> None:
self._update_record_payload(tid, new_payload)
return _cb
runtime_payload = _PersistedPayloadDict(payload, _make_resume_persist_cb(task_id))
task = asyncio.create_task(
self._run_with_retry(handle, handler, runtime_payload, policy, _resume_from_attempt=attempts_completed),
name=f"hosting.task[{handle.name}]:{task_id}(resumed)",
)
self._pending[task_id] = task
def _on_done(_t: asyncio.Task[None], tid: str = task_id) -> None:
self._pending.pop(tid, None)
task.add_done_callback(_on_done)
replayed += 1
if replayed:
logger.info(
"InProcessTaskRunner.resume: re-scheduled %d pending task(s) from %s", replayed, self._state_dir
)
return replayed
# ------------------------------------------------------------------ #
# Lifecycle helper (the host calls this from ``on_shutdown``)
# ------------------------------------------------------------------ #
async def shutdown(self, *, timeout: float | None = None) -> None:
"""Wait briefly for pending tasks to drain, then cancel anything still running.
Called by the host on ``on_shutdown`` so a graceful shutdown does
not orphan in-flight push retries. Channels may legitimately
schedule a final push from their own shutdown callback (e.g. a
goodbye message); the runner therefore *waits* up to
``timeout`` seconds (default: the runner's
``shutdown_grace_seconds`` configured at construction) for the
in-flight set to finish on its own before cancelling stragglers.
Tasks that don't honour cancellation within the same window are
abandoned — the runner makes no synchronous durability claim,
so cleanup is best-effort.
When ``state_dir`` is set, tasks that didn't drain are left in
``"pending"`` status on disk so the next process replays them
via :meth:`resume`. The disk cache is closed and the
single-owner lock is released regardless of drain outcome.
"""
self._shutting_down = True
try:
if self._pending:
grace = timeout if timeout is not None else self._shutdown_grace_seconds
tasks = list(self._pending.values())
# Phase 1 — wait for natural completion within the grace window.
if grace > 0:
await asyncio.wait(tasks, timeout=grace)
# Phase 2 — cancel anything still pending, then wait briefly for
# cancellation to propagate.
still_pending = [t for t in tasks if not t.done()]
if still_pending:
logger.info(
"InProcessTaskRunner.shutdown: %d task(s) still running after %.2fs grace; cancelling",
len(still_pending),
grace,
)
for task in still_pending:
task.cancel()
cancellation_window = max(grace, 1.0)
try:
await asyncio.wait_for(
asyncio.gather(*still_pending, return_exceptions=True),
timeout=cancellation_window,
)
except (TimeoutError, asyncio.TimeoutError):
logger.warning(
"InProcessTaskRunner.shutdown: %d task(s) did not exit within %.2fs "
"of cancellation; abandoning",
sum(not t.done() for t in still_pending),
cancellation_window,
)
finally:
# Release disk resources after the in-flight set has been
# given a chance to drain — tasks that mutate the payload
# mid-shutdown will fail to persist after this point, which
# is the correct behaviour (the next process will replay
# from whatever the last fully-committed state was).
if self._cache is not None:
try:
self._cache.close()
except Exception: # pragma: no cover - close errors aren't actionable
logger.exception("InProcessTaskRunner.shutdown: failed to close cache cleanly")
self._cache = None
if self._lock_handle is not None:
release_state_dir_lock(self._lock_handle)
self._lock_handle = None
# ------------------------------------------------------------------ #
# Internals — retry loop
# ------------------------------------------------------------------ #
async def _run_with_retry(
self,
handle: TaskHandle,
handler: Callable[[Mapping[str, Any]], Awaitable[None]],
payload: Mapping[str, Any],
policy: RetryPolicy,
*,
_resume_from_attempt: int = 0,
) -> None:
delay = policy.initial_backoff_seconds
attempt = _resume_from_attempt
try:
while True:
attempt += 1
# Persist the attempt counter BEFORE we invoke the
# handler so a crash mid-handler doesn't lose the fact
# that we tried — replay sees the bumped counter and
# respects the original retry budget. Trade-off: a
# crash before the external call is made still consumes
# one attempt (at-most-once semantics around the bump);
# we document this as best-effort across crashes.
self._update_record_attempts(handle.task_id, attempt)
try:
await handler(payload)
except asyncio.CancelledError:
# On a graceful shutdown of a disk-persistent runner
# we deliberately *don't* mark the record terminal —
# ``resume()`` will pick it up on the next boot and
# replay it with the persisted attempt counter. For
# in-memory runners (no cache) there's nothing to
# resume from, so we still mark ``cancelled`` so
# callers holding the handle can observe the
# outcome.
if not (self._shutting_down and self._cache is not None):
self._mark_terminal(handle.task_id, "cancelled")
raise
except Exception as exc:
if attempt >= policy.max_attempts:
logger.exception(
"InProcessTaskRunner: task %s (%s) failed after %d attempts",
handle.name,
handle.task_id,
attempt,
)
self._mark_terminal(handle.task_id, "failed")
return
logger.warning(
"InProcessTaskRunner: task %s (%s) attempt %d/%d failed (%s); retrying in %.2fs",
handle.name,
handle.task_id,
attempt,
policy.max_attempts,
exc,
delay,
)
try:
await asyncio.sleep(delay)
except asyncio.CancelledError:
if not (self._shutting_down and self._cache is not None):
self._mark_terminal(handle.task_id, "cancelled")
raise
delay = min(delay * policy.backoff_multiplier, policy.max_backoff_seconds)
else:
self._mark_terminal(handle.task_id, "succeeded")
return
except asyncio.CancelledError:
# Propagate so the outer ``asyncio.Task`` records cancellation
# in its own state for any observer that holds the raw task.
return
# ------------------------------------------------------------------ #
# Internals — record / disk helpers
# ------------------------------------------------------------------ #
def _build_record(
self,
name: str,
payload: Mapping[str, Any],
policy: RetryPolicy,
) -> dict[str, Any]:
"""Construct the on-disk record dict for a freshly-scheduled task."""
return {
_REC_HANDLER_NAME: name,
_REC_NAME: name,
_REC_PAYLOAD: dict(payload),
_REC_RETRY_POLICY: policy,
_REC_ATTEMPTS: 0,
_REC_STATUS: "pending",
_REC_CREATED_AT: time.time(),
}
def _validate_picklable(self, record: Mapping[str, Any]) -> None:
"""Pickle-probe a record at schedule time so misconfig is loud.
We only do this when the cache is open (i.e. persistence is on).
The probe runs ``pickle.dumps`` on the record and raises a
framework-typed :class:`PushPayloadNotPicklable` if it fails.
Loud failure here is better than silent data loss after the
next restart.
"""
try:
pickle.dumps(record) # nosec B301 - dumps only, no untrusted load
except Exception as exc:
raise PushPayloadNotPicklable(
"InProcessTaskRunner: scheduled task payload is not picklable; "
"disk persistence (state_dir) requires payloads to round-trip "
"through pickle. Common causes: a user-supplied response that "
"embeds a live network client, asyncio.Lock, or generator. "
f"Underlying pickle error: {exc!r}"
) from exc
def _update_record_attempts(self, task_id: str, attempt: int) -> None:
"""Bump the attempt counter on the persisted record (if any).
Status stays ``"pending"`` while the task is in-flight — there
is no transient ``"running"`` status. This keeps the resume
contract simple: anything ``"pending"`` on disk is a candidate
for replay, whether it was never picked up or crashed mid-attempt.
"""
if self._cache is None:
return
rec = self._cache.get(task_id)
if not isinstance(rec, dict):
# Record was evicted / quarantined since schedule; nothing
# to persist. The asyncio task continues — it just won't
# be resumable on next boot.
return
rec[_REC_ATTEMPTS] = attempt
try:
self._cache.set(task_id, rec)
except Exception: # pragma: no cover - cache write failures aren't actionable
logger.exception("InProcessTaskRunner: failed to persist attempt counter for %s", task_id)
def _update_record_payload(self, task_id: str, new_payload: Mapping[str, Any]) -> None:
"""Persist a handler-side payload mutation back to disk.
Called from :class:`_PersistedPayloadDict.__setitem__`. The whole
payload mapping is re-written (the cache stores opaque pickled
values, so partial-field updates aren't possible). Handler-side
mutations on the runner's hot path are rare (today: only the
``echo_done`` cursor) so the extra write is acceptable.
"""
if self._cache is None:
return
rec = self._cache.get(task_id)
if not isinstance(rec, dict):
return
rec[_REC_PAYLOAD] = dict(new_payload)
try:
self._cache.set(task_id, rec)
except Exception: # pragma: no cover - cache write failures aren't actionable
logger.exception("InProcessTaskRunner: failed to persist payload mutation for %s", task_id)
def _mark_terminal(self, task_id: str, status: TaskStatus) -> None:
"""Move a task to a terminal status, updating both memory and disk.
Records are first updated on disk (so a crash between the disk
write and the in-memory write doesn't lose the terminal status),
then mirrored to the in-memory cache, then FIFO-bounded.
"""
# Disk side first.
if self._cache is not None:
rec = self._cache.get(task_id)
if isinstance(rec, dict):
rec[_REC_STATUS] = status
rec[_REC_TERMINAL_AT] = time.time()
# Truncate heavy fields (payload, retry_policy) — once
# the task is terminal we never need them again, and
# keeping them around bloats disk on long-lived hosts.
rec[_REC_PAYLOAD] = None
rec[_REC_RETRY_POLICY] = None
try:
self._cache.set(task_id, rec)
except Exception: # pragma: no cover
logger.exception("InProcessTaskRunner: failed to persist terminal status for %s", task_id)
# In-memory side.
if task_id not in self._terminal:
self._terminal_order.append(task_id)
self._terminal[task_id] = status
# FIFO-evict from BOTH layers once we exceed the cap.
while len(self._terminal_order) > self._terminal_cache_size:
evicted = self._terminal_order.pop(0)
self._terminal.pop(evicted, None)
if self._cache is not None:
try:
del self._cache[evicted]
except KeyError:
pass
except Exception: # pragma: no cover
logger.exception("InProcessTaskRunner: failed to evict %s from disk cache", evicted)
# Persist the new ordering list so a restart sees the same FIFO
# ordering for further eviction decisions.
if self._cache is not None:
try:
self._cache.set(_TERMINAL_ORDER_KEY, list(self._terminal_order))
except Exception: # pragma: no cover
logger.exception("InProcessTaskRunner: failed to persist terminal-order list")
__all__ = ["InProcessTaskRunner"]
python/packages/hosting/agent_framework_hosting/_state_store.py
+26 -282
View File
@@ -1,49 +1,11 @@
# Copyright (c) Microsoft. All rights reserved.
"""Disk-backed wrappers for the host's in-memory state dicts.
"""Disk-backed wrapper for the host's session-alias map.
The host keeps three in-process dictionaries that need to survive a
process restart when the operator opts in to disk persistence:
- ``_session_aliases`` (``isolation_key -> active session_id``): rotated
by :meth:`AgentFrameworkHost.reset_session`; without persistence a
restart silently re-uses the pre-rotation session_id and the user sees
history they were supposed to have walked away from.
- ``_active`` (``isolation_key -> last-seen channel name``): drives
:class:`ResponseTarget` ``.active`` fan-out; losing it on restart makes
:class:`ResponseTarget.active` raise ``"no active channel"`` for every
user the host has previously talked to.
- ``_identities``
(``isolation_key -> {channel_name -> ChannelIdentity}``): the per-user
channel registry that powers :class:`ResponseTarget` ``.channel(name)``,
``.channels([...])`` and ``.all_linked``; losing it on restart turns
every linked-identity push target into a not-found.
Both wrappers are :class:`dict` subclasses so the rest of the host code
doesn't need to know whether persistence is on or off; the only
difference is that mutations are mirrored back to a
:mod:`diskcache`-backed sqlite store. Reads stay fast because the
in-memory copy is the source of truth — disk is purely a backing
store for write-through and re-hydration.
Layout under ``<state_dir>/sessions/`` (the ``sessions`` component
chosen because all three dicts share the same per-user-life cycle):
<state_dir>/sessions/
.lock # single-owner lock (advisory)
cache.db, … # diskcache sqlite files
keyed by:
"aliases:<isolation_key>" -> str (session_id)
"active:<isolation_key>" -> str (channel name)
"identities:<isolation_key>" -> dict[channel_name, ChannelIdentity]
Pickle is what diskcache uses by default; the wrappers do not impose
their own serialisation. :class:`ChannelIdentity` is a frozen dataclass
of plain scalars and so round-trips cleanly.
Everything in this module is internal. Public consumers should use
:class:`AgentFrameworkHost(state_dir=...)` and let the host wire the
wrappers up.
``AgentFrameworkHost.reset_session(isolation_key)`` rotates future requests for
that isolation key onto a new session id. Persisting the alias map lets that
rotation survive a host restart without introducing cross-channel identity or
delivery state into the core host.
"""
from __future__ import annotations
@@ -52,7 +14,7 @@ import logging
import os
from collections.abc import Mapping
from pathlib import Path
from typing import Any, TypeVar, cast
from typing import Any, TypeVar
from ._persistence import (
acquire_state_dir_lock,
@@ -62,27 +24,12 @@ from ._persistence import (
logger = logging.getLogger(__name__)
_V = TypeVar("_V")
# Key prefixes inside the shared sessions cache. Three logical maps live
# in one diskcache so they share a single sqlite handle and a single
# directory lock — opening multiple diskcaches against the same
# directory is supported but doubles file-handle pressure and the
# per-open lock acquisition cost.
_ALIASES_PREFIX = "aliases:"
_ACTIVE_PREFIX = "active:"
_IDENTITIES_PREFIX = "identities:"
class SessionsStateStore:
"""One disk cache + lock shared by every host-side persisted dict.
The host constructs one of these per ``state_dir["sessions"]`` value
and threads it into each :class:`_PersistedDict` it creates. Closing
the store releases the lock and the cache handle.
"""
"""One disk cache + lock for host-side session aliases."""
def __init__(self, sessions_dir: str | os.PathLike[str]) -> None:
self._sessions_dir: Path = Path(os.fspath(sessions_dir))
@@ -97,22 +44,11 @@ class SessionsStateStore:
@property
def cache(self) -> Any:
"""Return the underlying :mod:`diskcache` Cache.
Intended for the wrapper classes in this module only. Callers
outside the module should go through the typed wrappers — direct
cache access bypasses the key-prefix discipline that keeps the
three maps from colliding.
"""
"""Return the underlying :mod:`diskcache` Cache."""
return self._cache
def close(self) -> None:
"""Close the cache and release the directory lock.
Safe to call multiple times. The host invokes this from its
lifespan shutdown hook so a second host can re-open the same
``state_dir`` cleanly after the first exits.
"""
"""Close the cache and release the directory lock."""
if self._cache is not None:
try:
self._cache.close()
@@ -125,15 +61,7 @@ class SessionsStateStore:
class _PersistedDict(dict[str, _V]):
"""Drop-in :class:`dict` whose mutations mirror to a diskcache prefix.
Used for the host's flat ``str -> V`` dicts (``_session_aliases``
and ``_active``). The in-memory copy is the source of truth for
reads; writes update memory first and then mirror to disk so a
crash between the two leaves the in-memory state correct (which is
what subsequent reads will see anyway) and only loses the last
not-yet-flushed value on next restart.
"""
"""Drop-in :class:`dict` whose mutations mirror to a diskcache prefix."""
def __init__(
self,
@@ -144,26 +72,20 @@ class _PersistedDict(dict[str, _V]):
super().__init__()
self._store = store
self._prefix = key_prefix
# Rehydrate from disk into memory exactly once at construction.
# Doing this here (rather than lazily) keeps the in-memory dict
# behaviour consistent with the non-persisted code path —
# ``len(host._session_aliases)`` reflects all known users from
# the moment the host is constructed.
cache: Any = store.cache
for raw_key in cache.iterkeys():
if not isinstance(raw_key, str) or not raw_key.startswith(key_prefix):
continue
value: Any
try:
value = cache.get(raw_key)
value: Any = cache.get(raw_key)
except Exception:
logger.exception("SessionsStateStore: failed to rehydrate %s; skipping", raw_key)
continue
logical_key = raw_key[len(key_prefix) :]
super().__setitem__(logical_key, value)
if initial:
for k, v in initial.items():
self[k] = v
for key, value in initial.items():
self[key] = value
def __setitem__(self, key: str, value: _V) -> None:
super().__setitem__(key, value)
@@ -182,9 +104,7 @@ class _PersistedDict(dict[str, _V]):
logger.exception("SessionsStateStore: failed to evict %s%s", self._prefix, key)
def pop(self, key: str, *args: Any) -> _V:
# ``dict.pop`` doesn't go through ``__delitem__``, so we mirror
# the disk side here explicitly. Forward the default sentinel
# only when present so we match ``dict.pop`` semantics exactly.
"""Mirror ``dict.pop`` to disk."""
value: _V = super().pop(key, *args)
try:
del self._store.cache[self._prefix + key]
@@ -195,16 +115,17 @@ class _PersistedDict(dict[str, _V]):
return value
def clear(self) -> None:
"""Mirror ``dict.clear`` to disk."""
keys = list(self.keys())
super().clear()
cache = self._store.cache
for k in keys:
for key in keys:
try:
del cache[self._prefix + k]
del cache[self._prefix + key]
except KeyError:
pass
except Exception: # pragma: no cover
logger.exception("SessionsStateStore: failed to evict %s%s during clear", self._prefix, k)
logger.exception("SessionsStateStore: failed to evict %s%s during clear", self._prefix, key)
def update( # type: ignore[override]
self,
@@ -212,191 +133,14 @@ class _PersistedDict(dict[str, _V]):
/,
**kwargs: _V,
) -> None:
# Defer to __setitem__ so every entry is mirrored to disk; the
# default ``dict.update`` writes into the underlying storage
# directly and would skip our persistence hook.
"""Mirror ``dict.update`` to disk one item at a time."""
if other is not None:
for k in other:
self[k] = other[k]
for k, v in kwargs.items():
self[k] = v
for key in other:
self[key] = other[key]
for key, value in kwargs.items():
self[key] = value
class _PersistedNestedDict(dict[str, dict[str, _V]]):
"""Disk-backed wrapper for the per-isolation-key identity map.
The host's ``_identities`` is a nested dict
``isolation_key -> {channel_name -> ChannelIdentity}``. The whole
inner dict for a given isolation_key is small (one entry per channel
the user has appeared on), so we persist the inner dict as a single
cache value rather than per-channel — fewer cache hits, simpler
schema, no need for a separate sub-prefix.
To make mutations of the inner dict mirror to disk, ``__getitem__``
returns a ``_NestedInnerProxy`` that mutates the parent's cache slot
on each ``__setitem__`` / ``__delitem__``. The wrapper is purely
additive — callers that pass a plain dict in via ``__setitem__`` get
the same write-through behaviour for free.
"""
def __init__(
self,
store: SessionsStateStore,
key_prefix: str = _IDENTITIES_PREFIX,
) -> None:
super().__init__()
self._store = store
self._prefix = key_prefix
cache: Any = store.cache
for raw_key in cache.iterkeys():
if not isinstance(raw_key, str) or not raw_key.startswith(key_prefix):
continue
value: Any
try:
value = cache.get(raw_key)
except Exception:
logger.exception("SessionsStateStore: failed to rehydrate %s; skipping", raw_key)
continue
if not isinstance(value, dict):
continue
inner_value = cast(dict[str, _V], value)
logical_key = raw_key[len(key_prefix) :]
# Wrap so caller-side mutations on the inner dict mirror back.
inner: _NestedInnerProxy[_V] = _NestedInnerProxy(self, logical_key, inner_value)
super().__setitem__(logical_key, inner)
def __setitem__(self, key: str, value: dict[str, _V]) -> None:
# Wrap whatever the caller passes in so subsequent ``inner[ch] = ...``
# mutations are mirrored to disk. We always wrap (even
# ``_NestedInnerProxy`` inputs) so the proxy's ``_outer`` link
# points at us rather than at any previous outer dict.
wrapped = _NestedInnerProxy(self, key, dict(value))
super().__setitem__(key, wrapped)
self.persist_inner(key, dict(value))
def __delitem__(self, key: str) -> None:
super().__delitem__(key)
try:
del self._store.cache[self._prefix + key]
except KeyError:
pass
except Exception: # pragma: no cover
logger.exception("SessionsStateStore: failed to evict %s%s", self._prefix, key)
def setdefault(self, key: str, default: dict[str, _V] | None = None) -> dict[str, _V]: # type: ignore[override]
if key in self:
return self[key]
if default is None:
default = {}
self[key] = default
return self[key]
def persist_inner(self, isolation_key: str, snapshot: Mapping[str, _V]) -> None:
"""Write the full inner dict for ``isolation_key`` back to disk.
Called from :class:`_NestedInnerProxy` on every mutation and by
:meth:`__setitem__` when a new outer key is added. A single
write per change keeps the schema simple — there is no
partial-row update — and is fine for the access pattern
(mutations on the host's hot path are rare: identity registry
writes are once-per-channel-per-user).
"""
try:
self._store.cache.set(self._prefix + isolation_key, snapshot)
except Exception: # pragma: no cover - cache write failures aren't actionable
logger.exception(
"SessionsStateStore: failed to persist identities for %s%s",
self._prefix,
isolation_key,
)
class _NestedInnerProxy(dict[str, _V]):
"""Inner-dict proxy that mirrors mutations back to its outer.
Returned by :class:`_PersistedNestedDict.__getitem__` (via the
rehydration / ``__setitem__`` wrap). When the channel-registry code
does ``self._identities[ik][channel_name] = identity``, the
``__setitem__`` on this proxy fires and re-writes the whole inner
dict to disk via the parent's ``persist_inner``. Behavioural
identity with ``dict`` is preserved otherwise (``len``, iteration,
``__contains__``, …).
"""
_outer: _PersistedNestedDict[_V]
_key: str
__slots__ = ("_key", "_outer")
def __init__(
self,
outer: _PersistedNestedDict[_V],
key: str,
data: Mapping[str, _V],
) -> None:
super().__init__(data)
# ``__slots__`` on a ``dict`` subclass requires the back-door —
# CPython is lenient, PyPy is strict.
object.__setattr__(self, "_outer", outer)
object.__setattr__(self, "_key", key)
def __setitem__(self, key: str, value: _V) -> None:
super().__setitem__(key, value)
self._outer.persist_inner(self._key, dict(self))
def __delitem__(self, key: str) -> None:
super().__delitem__(key)
self._outer.persist_inner(self._key, dict(self))
def pop(self, key: str, *args: Any) -> _V:
value: _V = super().pop(key, *args)
self._outer.persist_inner(self._key, dict(self))
return value
def clear(self) -> None:
super().clear()
self._outer.persist_inner(self._key, dict(self))
def update( # type: ignore[override]
self,
other: Mapping[str, _V] | None = None,
/,
**kwargs: _V,
) -> None:
if other is not None:
for k in other:
super().__setitem__(k, other[k])
for k, v in kwargs.items():
super().__setitem__(k, v)
self._outer.persist_inner(self._key, dict(self))
def build_session_dicts(
store: SessionsStateStore,
) -> tuple[
_PersistedDict[str],
_PersistedDict[str],
_PersistedNestedDict[Any],
]:
"""Construct the three host-side persisted dicts against a single store.
Returns ``(session_aliases, active, identities)`` in the order the
host assigns them, so the call site reads
``self._session_aliases, self._active, self._identities = build_session_dicts(store)``.
"""
aliases: _PersistedDict[str] = _PersistedDict(store, _ALIASES_PREFIX)
active: _PersistedDict[str] = _PersistedDict(store, _ACTIVE_PREFIX)
identities: _PersistedNestedDict[Any] = _PersistedNestedDict(store)
return aliases, active, identities
# Re-export keys for tests / power users that want to inspect the cache.
__all__ = [
"_ACTIVE_PREFIX",
"_ALIASES_PREFIX",
"_IDENTITIES_PREFIX",
"SessionsStateStore",
"_PersistedDict",
"_PersistedNestedDict",
"build_session_dicts",
]
def build_session_aliases(store: SessionsStateStore) -> dict[str, str]:
"""Return the disk-backed session-alias map for ``store``."""
return _PersistedDict[str](store, _ALIASES_PREFIX)
python/packages/hosting/agent_framework_hosting/_types.py
+22 -772
View File
@@ -11,12 +11,7 @@
These types form the boundary between the host and individual channels.
A channel parses its native payload, builds a :class:`ChannelRequest`, and
hands it to :class:`ChannelContext.run` (or ``run_stream``) on the host.
The host normalizes the request into a single agent invocation and either
returns the result to the originating channel or fans out via
:class:`ResponseTarget` to other channels that implement
:class:`ChannelPush`.
See ``docs/specs/002-python-hosting-channels.md`` for the full design.
The channel owns rendering the result back onto its originating protocol.
"""
from __future__ import annotations
@@ -24,16 +19,11 @@ from __future__ import annotations
import os
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass, field
from enum import Enum
from typing import TYPE_CHECKING, Any, Generic, Literal, Protocol, TypedDict, TypeVar, cast, runtime_checkable
from typing import TYPE_CHECKING, Any, Generic, Protocol, TypedDict, TypeVar, runtime_checkable
from agent_framework import (
AgentResponse,
AgentResponseUpdate,
AgentRunInputs,
ResponseStream,
SupportsAgentRun,
Workflow,
)
from starlette.routing import BaseRoute
@@ -41,11 +31,6 @@ if TYPE_CHECKING:
from ._host import ChannelContext
# --------------------------------------------------------------------------- #
# Channel-neutral request envelope
# --------------------------------------------------------------------------- #
class ChannelSession:
"""Channel-supplied session hint.
@@ -59,15 +44,12 @@ class ChannelSession:
class ChannelIdentity:
"""Channel-native identity the host sees on each request.
"""Channel-native identity metadata observed on a request.
Consumed by the host's identity registry. The host uses it for two things:
1. Recording the active channel for an ``isolation_key`` so
``ResponseTarget.active`` resolves correctly.
2. Telling :class:`ChannelPush` ``push`` recipients **where** in their
native namespace to deliver — Telegram uses ``native_id`` as the
chat id, Teams as the conversation/AAD id, etc.
The simplified hosting core records this only on the persisted input
message's ``additional_properties["hosting"]`` block and forwards it
through run/response hooks. Cross-channel linking and recipient lookup are
follow-up concerns, not part of the v1 host contract.
"""
def __init__(
@@ -81,176 +63,6 @@ class ChannelIdentity:
self.attributes: Mapping[str, Any] = attributes if attributes is not None else dict()
class ResponseTargetKind(str, Enum):
"""Discriminator for :class:`ResponseTarget` variants."""
ORIGINATING = "originating"
ACTIVE = "active"
CHANNELS = "channels"
ALL_LINKED = "all_linked"
IDENTITIES = "identities"
NONE = "none"
class ResponseTarget:
"""Per-request directive controlling **where** the host delivers the agent reply.
Independent of ``session_mode``. Construct via the classmethod helpers or
use the module-level singletons rather than touching ``kind`` directly.
Variants:
- ``ResponseTarget.originating`` (default) — synchronous response on the
originating channel only.
- ``ResponseTarget.active`` — push to the channel most recently observed
for the resolved ``isolation_key``.
- ``ResponseTarget.channel("teams")`` / ``.channels([...])`` — push to
one or more named destinations. Each entry is either a bare channel
name (host resolves the native id from its identity registry) or a
``"channel:native_id"`` token (used verbatim). The pseudo-name
``"originating"`` includes the originating channel in the fan-out.
- ``ResponseTarget.identity(ChannelIdentity)`` /
``.identities([ChannelIdentity, ...])`` — push to one or more
**fully-specified identities**. Preferred over the ``"channel:native_id"``
string variant when the destination needs ``identity.attributes``
preserved (Teams conversation/thread metadata, Slack channel+thread,
Bot Framework service-url, etc.).
- ``ResponseTarget.all_linked`` — push to every channel where the
resolved ``isolation_key`` has been observed.
- ``ResponseTarget.none`` — background-only; in the prototype this just
suppresses the originating reply (no ``ContinuationToken`` yet).
Instances are intended to be treated as immutable; the singletons are
shared across the process.
"""
def __init__(
self,
kind: ResponseTargetKind = ResponseTargetKind.ORIGINATING,
targets: tuple[str, ...] = (),
identities: tuple[ChannelIdentity, ...] = (),
*,
echo_input: bool = False,
) -> None:
self.kind = kind
self.targets = targets
# Stored under a non-clashing name so the ``identities``
# *classmethod* (the public builder) can coexist with the
# value accessor (the ``identities`` property below). At
# runtime instance attributes shadow class attributes anyway,
# but type checkers see the classmethod and reject reassignment.
self._target_identities: tuple[ChannelIdentity, ...] = tuple(identities)
# When True, the host first pushes the originating user message
# to every non-originating destination (so end-user apps observing
# those channels can keep their UI in sync) before pushing the
# agent response. Defaults to False — opt-in only, because not
# every channel knows how to render ``role="user"`` content
# gracefully on its own surface.
self.echo_input = echo_input
@property
def target_identities(self) -> tuple[ChannelIdentity, ...]:
"""Destination identities for ``kind == IDENTITIES`` targets.
Public name distinct from the :meth:`identities` classmethod
builder. Empty for non-``IDENTITIES`` kinds.
"""
return self._target_identities
# -- builders ---------------------------------------------------------- #
@classmethod
def channel(cls, name: str, *, echo_input: bool = False) -> ResponseTarget:
"""Target a single named destination channel."""
return cls(kind=ResponseTargetKind.CHANNELS, targets=(name,), echo_input=echo_input)
@classmethod
def channels(cls, names: Sequence[str], *, echo_input: bool = False) -> ResponseTarget:
"""Target an explicit list of destination channels."""
return cls(kind=ResponseTargetKind.CHANNELS, targets=tuple(names), echo_input=echo_input)
@classmethod
def identity(cls, identity: ChannelIdentity, *, echo_input: bool = False) -> ResponseTarget:
"""Target a single fully-specified :class:`ChannelIdentity`.
Preferred over the ``"channel:native_id"`` string token in
:meth:`channels` when ``identity.attributes`` carries metadata the
destination channel needs (Teams conversation/thread ids and
service-url, Slack channel + thread, Bot Framework activity-locator
fields, etc.). The host pushes to the named identity verbatim
without consulting its own identity registry.
"""
return cls(kind=ResponseTargetKind.IDENTITIES, identities=(identity,), echo_input=echo_input)
@classmethod
def identities(cls, identities: Sequence[ChannelIdentity], *, echo_input: bool = False) -> ResponseTarget:
"""Target an explicit list of fully-specified :class:`ChannelIdentity` objects.
See :meth:`identity` for the single-destination variant.
"""
return cls(kind=ResponseTargetKind.IDENTITIES, identities=tuple(identities), echo_input=echo_input)
# -- value semantics --------------------------------------------------- #
# ``ResponseTarget`` is treated as immutable, so two instances with the
# same ``kind`` + ``targets`` + ``identities`` + ``echo_input`` are
# interchangeable. Tests and channel parsers compare instances with
# ``==`` and use them as dict keys.
def __eq__(self, other: object) -> bool:
if not isinstance(other, ResponseTarget):
return NotImplemented
return (
self.kind is other.kind
and self.targets == other.targets
and _identities_equal(self._target_identities, other._target_identities)
and self.echo_input == other.echo_input
)
def __hash__(self) -> int:
# ``ChannelIdentity`` is not itself hashable (mutable attributes
# mapping); fold the identifying triple so two ``identities``
# tuples with the same channel/native_id/attributes content hash
# the same.
identities_key = tuple(
(i.channel, i.native_id, tuple(sorted(i.attributes.items()))) for i in self._target_identities
)
return hash((self.kind, self.targets, identities_key, self.echo_input))
def __repr__(self) -> str:
suffix = ", echo_input=True" if self.echo_input else ""
if self.kind is ResponseTargetKind.CHANNELS:
return f"ResponseTarget.channels({list(self.targets)!r}{suffix})"
if self.kind is ResponseTargetKind.IDENTITIES:
return f"ResponseTarget.identities({list(self._target_identities)!r}{suffix})"
return f"ResponseTarget.{self.kind.value}{suffix}"
def _identities_equal(left: tuple[ChannelIdentity, ...], right: tuple[ChannelIdentity, ...]) -> bool:
"""Structural-equality helper for ``ResponseTarget.identities`` comparisons.
``ChannelIdentity`` is a plain class without ``__eq__``, so ``tuple`` /
``list`` comparisons fall back to identity equality which is too strict
for value-typed ``ResponseTarget`` callers (two equivalent identity
tuples produced independently would otherwise compare unequal).
"""
if len(left) != len(right):
return False
for a, b in zip(left, right, strict=True):
if a.channel != b.channel or a.native_id != b.native_id:
return False
if dict(a.attributes) != dict(b.attributes):
return False
return True
# Module-level singletons so callers can write ``ResponseTarget.originating``
# (matching the spec's classmethod-style notation) without juggling Python's
# no-zero-arg-classmethod-property limitation.
ResponseTarget.originating = ResponseTarget(kind=ResponseTargetKind.ORIGINATING) # type: ignore[attr-defined]
ResponseTarget.active = ResponseTarget(kind=ResponseTargetKind.ACTIVE) # type: ignore[attr-defined]
ResponseTarget.all_linked = ResponseTarget(kind=ResponseTargetKind.ALL_LINKED) # type: ignore[attr-defined]
ResponseTarget.none = ResponseTarget(kind=ResponseTargetKind.NONE) # type: ignore[attr-defined]
@dataclass
class ChannelRequest:
"""Uniform invocation envelope every channel produces from its native payload.
@@ -260,16 +72,15 @@ class ChannelRequest:
"""
channel: str
operation: str # e.g. "message.create", "command.invoke"
operation: str
input: AgentRunInputs
session: ChannelSession | None = None
options: Mapping[str, Any] | None = None
session_mode: str = "auto" # "auto" | "required" | "disabled"
session_mode: str = "auto"
metadata: Mapping[str, Any] = field(default_factory=lambda: {})
attributes: Mapping[str, Any] = field(default_factory=lambda: {})
stream: bool = False
identity: ChannelIdentity | None = None
response_target: ResponseTarget = field(default_factory=lambda: ResponseTarget.originating) # type: ignore[attr-defined]
class ChannelCommand:
@@ -335,42 +146,11 @@ TResult = TypeVar("TResult")
class HostedRunResult(Generic[TResult]):
r"""Channel-neutral envelope around the target's full-fidelity result.
"""Channel-neutral envelope around the target's full-fidelity result.
Carries the underlying execution payload **unchanged** so channels
(and developer-supplied ``response_hook``\\s) can read everything the
target produced — full multi-modal contents, structured ``value``,
``usage_details``, ``response_id``, workflow per-executor outputs,
final ``WorkflowRunState``, etc.
``result`` is generic in ``TResult`` so callers retain static typing:
* Agent targets always produce
``HostedRunResult[AgentResponse]`` — channels read
``result.messages``, ``result.value``, ``result.usage_details``, …
directly.
* Workflow targets produce ``HostedRunResult[WorkflowRunResult]``
today (``Workflow`` is not itself generic, so the static narrowing
is only as tight as ``Workflow.run``'s return). Channels iterate
``result.get_outputs()`` and inspect ``result.get_final_state()``
to render workflow-specific UX. When a host author drives the
workflow themselves and knows the final-output type, they may
narrow to ``HostedRunResult[MyOutput]`` in their own
``response_hook`` signatures.
* The echo-input phase synthesises an ``HostedRunResult[AgentResponse]``
wrapping the originating user turn so the same per-destination
delivery machinery applies.
The optional ``session`` slot carries the resolved
:class:`~agent_framework.AgentSession` the host bound to this
invocation (``None`` for workflow targets, which do not own session
state in the agent sense). Channels that want to surface session
metadata (e.g. echo the resolved isolation key into a response
header) read it here.
Treat instances as immutable: the host clones per-destination before
invoking a per-channel ``response_hook`` so one channel's transform
cannot perturb the payload another destination observes.
The host does not flatten or pre-shape the target output. Channels and
response hooks read the underlying result type directly and serialize the
subset their wire format can carry.
"""
def __init__(
@@ -388,575 +168,45 @@ class HostedRunResult(Generic[TResult]):
result: TResult | _Unset = _UNSET,
session: Any | _Unset | None = _UNSET,
) -> HostedRunResult[TResult]:
"""Return a shallow copy with the supplied fields overridden.
Used by the host's delivery layer to clone the envelope before
applying a per-destination ``response_hook``, so one channel's
transform cannot mutate the payload another destination sees.
The clone is shallow — channels that need to mutate
``result.messages`` (or any other nested mutable container) are
responsible for deep-cloning that container themselves.
"""
"""Return a shallow copy with the supplied fields overridden."""
new: HostedRunResult[TResult] = HostedRunResult.__new__(HostedRunResult) # pyright: ignore[reportUnknownVariableType]
new.result = self.result if isinstance(result, _Unset) else result
new.session = self.session if isinstance(session, _Unset) else session
return new
class DurableTaskPayloadMode(str, Enum):
"""How a :class:`DurableTaskRunner` consumes scheduled-task payloads.
Used by the host's startup validator to pair a runner's persistence
expectations with the channels' push-codec capabilities. Adapter packages
pick the right value for their backing store.
* ``OBJECT`` — the runner accepts live Python objects in the payload.
No serialization is required; the host's
:class:`InProcessTaskRunner` is the canonical example. Suitable for
``runtime_mode="long_running"`` deployments where the runner shares
address space with the producer.
* ``JSON`` — the runner persists the payload (database, durable queue,
Foundry scheduled-task store, …) and replays it after a process
restart. Payloads MUST be JSON-serializable, which constrains what
the host can put on the wire. The host validates at construction
that every push-capable channel exposes a
:class:`ChannelPushCodec` (so :class:`HostedRunResult` payloads can
be reduced to a JSON envelope before scheduling).
"""
OBJECT = "object"
JSON = "json"
# A push-codec implementation reduces the ``(result, request, identity)``
# triple a destination channel will receive into a JSON-safe envelope that
# a durable :class:`DurableTaskRunner` can persist, and reconstructs the
# rendering inputs on the consumer side. The host *invokes* the codec
# during scheduling; the destination channel implements it (the channel
# knows what shape of payload it can render).
#
# Channels with no push codec are usable only with object-mode runners
# (the default :class:`InProcessTaskRunner`) — the host validates this at
# construction so the mismatch surfaces eagerly rather than on first push.
class ChannelPushCodec(Protocol):
"""Optional capability: serialise the push envelope for a durable task runner.
Implementations live on the destination channel (alongside ``push``)
as a duck-typed ``push_codec`` attribute. The host's
:meth:`_deliver_response` invokes :meth:`encode` once per scheduled
push (in JSON-mode runner deployments) to produce a JSON-safe
envelope for the runner; the handler calls :meth:`decode`
immediately before invoking :meth:`ChannelPush.push`. Object-mode
runners (the default in-process runner) bypass the codec entirely
and pass live references through verbatim.
Encoded envelopes MUST be JSON-serialisable
(``dict``/``list``/``str``/``int``/``float``/``bool``/``None``).
Channels that cannot satisfy this for some inputs (e.g. arbitrary
workflow result objects without a stable schema) SHOULD raise a
typed :class:`PushPayloadNotSerializable` from :meth:`encode`
rather than return a best-effort representation; the host surfaces
that as a schedule-time error and the destination is treated as
skipped (other destinations still get their chance).
"""
async def encode(
self,
*,
result: HostedRunResult[Any],
request: ChannelRequest,
identity: ChannelIdentity,
echo_result: HostedRunResult[Any] | None,
) -> Mapping[str, Any]:
"""Project the in-memory push triple into a JSON-safe envelope."""
...
async def decode(
self,
envelope: Mapping[str, Any],
) -> tuple[HostedRunResult[Any], ChannelRequest, ChannelIdentity, HostedRunResult[Any] | None]:
"""Reconstruct ``(result, request, identity, echo_result)`` from an envelope."""
...
class PushPayloadNotSerializable(RuntimeError):
"""Raised by a :class:`ChannelPushCodec` when the payload cannot be serialised.
Channels raise this from :meth:`ChannelPushCodec.encode` when the
inbound :class:`HostedRunResult` carries content the codec has no
JSON projection for (e.g. an arbitrary workflow result with no
declared schema). The host surfaces the error eagerly at schedule
time rather than letting the runner discover it after persisting
a half-formed envelope.
"""
class PushPayloadNotPicklable(RuntimeError):
"""Raised when a disk-persistent runner cannot pickle a scheduled task payload.
The in-process runner falls back to pickle when ``state_dir`` is set
so a long-running host can resume in-flight pushes across restarts.
Most :class:`HostedRunResult` payloads (frozen dataclasses wrapping
:class:`AgentResponse` or workflow output) pickle without issue, but
a user-supplied workflow result or response hook may embed an
unpickleable object (live network client, ``asyncio.Lock``, generator).
The runner raises this at schedule time so the misconfig is loud
rather than silently downgrading to no-persistence.
"""
class HostStatePaths(TypedDict, total=False):
"""Per-component disk paths for host-managed state.
Pass an instance of this typed dict to
:class:`~agent_framework_hosting._host.AgentFrameworkHost`'s
``state_dir`` parameter when you want to place individual components
on different volumes — for example, a fast local SSD for the runner
task queue and a network-attached durable volume for session state
that needs to survive container moves.
All keys are optional (``total=False``): unset components fall back
to in-memory storage (or, for ``checkpoints``, to no checkpoint
persistence). Pass a single ``str``/``PathLike`` to ``state_dir``
instead to get the default subfolder layout
(``state_dir/runner/``, ``state_dir/sessions/``,
``state_dir/checkpoints/``, ``state_dir/links/``).
Future components (continuations, ledger) will be added as additional
keys in subsequent releases.
Only session aliases and workflow checkpoints remain in the simplified
host. Linking stores, active-channel maps, identity registries, and runner
queues are follow-up concerns.
"""
runner: str | os.PathLike[str]
"""Where :class:`~agent_framework_hosting._runner.InProcessTaskRunner`
persists its pending-task queue and bounded terminal-status cache.
Required for in-flight push retries to survive process restarts."""
sessions: str | os.PathLike[str]
"""Where the host persists session aliases (from
:meth:`AgentFrameworkHost.reset_session`), the per-isolation-key
identity registry, and the last-active-channel map. Required for
``ResponseTarget.active``/``.channel``/``.all_linked`` to find
destinations after a restart, and for ``reset_session`` rotations
to survive a restart."""
"""Where the host persists session aliases created by ``reset_session``."""
checkpoints: str | os.PathLike[str]
"""Where the host persists workflow checkpoints for ``Workflow``
targets. Equivalent to passing ``checkpoint_location=<this path>``
directly: the host wraps it in a per-isolation-key
:class:`~agent_framework.FileCheckpointStorage`. Ignored when the
target is a ``SupportsAgentRun`` agent (a warning is emitted if you
set it explicitly via the mapping form). Pass the legacy
``checkpoint_location`` parameter instead when you need to supply a
:class:`~agent_framework.CheckpointStorage` instance — it takes
precedence over this key."""
links: str | os.PathLike[str]
"""Where identity-linker implementations persist their link store:
pending link challenges/grants, channel-native identity to linked
isolation-key mappings, and verified-claim metadata. The core host
does not impose a storage format; concrete :class:`IdentityLinker`
implementations that support host-provided persistence receive this
path via ``configure_link_store_path``. If a linker manages its own
persistence, omit this key or configure that linker directly."""
"""Where the host persists workflow checkpoints for ``Workflow`` targets."""
# A transform hook runs over each AgentResponseUpdate as the channel consumes
# the stream. It can return a replacement update, ``None`` to drop the update,
# or be async. Channels apply it during iteration so that channel-specific
# concerns (e.g. masking, redaction, formatting for the wire) live close to
# the channel rather than on the agent.
ChannelStreamTransformHook = Callable[
ChannelStreamUpdateHook = Callable[
[AgentResponseUpdate],
"AgentResponseUpdate | Awaitable[AgentResponseUpdate | None] | None",
]
# --------------------------------------------------------------------------- #
# Channel run hook
# --------------------------------------------------------------------------- #
# Run hooks accept the channel-built ``ChannelRequest`` and return a
# (possibly modified) replacement. Channels invoke the hook with both the
# request and the channel-side context as keyword arguments — the call
# convention is ``await hook(request, target=..., protocol_request=...)``.
#
# The ergonomic minimum for a hook implementation is therefore a function
# accepting ``request`` positionally plus ``**kwargs`` and returning a
# (possibly mutated) :class:`ChannelRequest`. Hooks that need the agent
# target or the raw channel-native payload pull them off the keyword
# arguments by name (``target`` / ``protocol_request``).
#
# ``protocol_request`` is the raw, channel-native payload the channel
# parsed (the JSON body for Responses, the Telegram ``Update`` dict, the
# Bot Framework ``Activity`` for Teams). Use it when the hook needs a
# field the channel did not lift onto ``ChannelRequest`` (e.g. OpenAI's
# ``safety_identifier``, Teams' ``from.aadObjectId``, …).
ChannelRunHook = Callable[..., "Awaitable[ChannelRequest] | ChannelRequest"]
async def apply_run_hook(
hook: ChannelRunHook,
request: ChannelRequest,
*,
target: SupportsAgentRun | Workflow,
protocol_request: Any | None,
) -> ChannelRequest:
"""Channel-side helper to invoke a :data:`ChannelRunHook` with the standard kwargs.
Channels call this rather than calling the hook directly so the
invocation convention (``request`` positional, ``target`` /
``protocol_request`` keyword) is enforced in one place.
"""
result = hook(request, target=target, protocol_request=protocol_request)
if isinstance(result, Awaitable):
return await result
return result
# --------------------------------------------------------------------------- #
# Channel response hook
# --------------------------------------------------------------------------- #
class ChannelResponseContext:
"""Per-destination context handed to a :data:`ChannelResponseHook`.
Response hooks run on the *output* side of the host pipeline, after
the agent / workflow has produced a :class:`HostedRunResult` but
before the destination channel serialises it to its wire format.
Hooks may need to make decisions based on *where* the payload is
headed — e.g. flatten multi-modal output to text for a text-only
destination, or pick which content variant to deliver to a card-
capable channel. The context captures that information without
forcing hooks to parse stringly destination tokens.
"""
def __init__(
self,
request: ChannelRequest,
channel_name: str,
destination_identity: ChannelIdentity | None,
originating: bool,
is_echo: bool = False,
) -> None:
self.request = request
self.channel_name = channel_name
# ``None`` when the originating channel is rendering its own reply
# (no push identity needed for "respond on the wire you came in
# on") or when the destination is named without a known native id.
self.destination_identity = destination_identity
# True when this hook invocation is for the originating channel's
# synchronous reply. False for non-originating push targets.
self.originating = originating
# True when the payload being shaped is the user-message echo
# rather than the agent response (only happens when
# ``ResponseTarget.echo_input`` is set).
self.is_echo = is_echo
# Response hooks accept the :class:`HostedRunResult` the host has assembled
# and return a (possibly modified) replacement. Channels invoke the hook
# with both the payload and the per-destination
# :class:`ChannelResponseContext` as keyword arguments — the call
# convention is ``await hook(result, context=...)``.
#
# The ergonomic minimum for a hook implementation is a function accepting
# ``result`` positionally plus ``**kwargs`` and returning a (possibly
# rewritten) :class:`HostedRunResult`. Hooks that need to branch on the
# destination read it off the ``context`` keyword argument.
#
# ``HostedRunResult`` is generic in the underlying ``result`` type; the
# hook callable signature stays ``Any``-typed so a single
# ``response_hook`` attribute on a channel can serve both agent
# (``HostedRunResult[AgentResponse]``) and workflow
# (``HostedRunResult[WorkflowRunResult]``) payloads — channels narrow
# at hook entry if they need static checking.
ChannelResponseHook = Callable[..., "Awaitable[HostedRunResult[Any]] | HostedRunResult[Any]"]
async def apply_response_hook(
hook: ChannelResponseHook,
result: HostedRunResult[Any],
*,
context: ChannelResponseContext,
) -> HostedRunResult[Any]:
"""Channel-side helper to invoke a :data:`ChannelResponseHook` with the standard kwargs.
Channels (and the host's delivery layer) call this rather than calling
the hook directly so the invocation convention (``result`` positional,
``context`` keyword) is enforced in one place.
"""
out = hook(result, context=context)
if isinstance(out, Awaitable):
return await out
return out
# --------------------------------------------------------------------------- #
# Channel protocols
# --------------------------------------------------------------------------- #
@runtime_checkable
class Channel(Protocol):
"""A pluggable adapter that exposes one transport on the host.
Channels publish their routes, commands, and lifecycle callbacks via
:meth:`contribute`. The host mounts them under the channel's ``path``
(or at the app root when ``path == ""``) and gives the channel a
:class:`ChannelContext` so it can call back into the host to invoke
the agent target and deliver responses.
"""
"""A pluggable adapter that exposes one transport on the host."""
name: str
path: str # default endpoint path (e.g. "/responses"); use "" to mount contributed routes at the app root
path: str
def contribute(self, context: ChannelContext) -> ChannelContribution: ...
@runtime_checkable
class ChannelPush(Protocol):
r"""Optional capability: a channel that can deliver outbound messages without a prior request.
Per SPEC-002 (req #13), channels that can do proactive delivery
(Telegram bot proactive message, Teams proactive bot message,
webhook callbacks, SSE broadcasts) implement ``push`` on top of the
base :class:`Channel` protocol. Channels without push can only be
addressed as the ``originating`` :class:`ResponseTarget`.
Distinguishing user echoes from agent replies
---------------------------------------------
When the originating :class:`ResponseTarget` opts in to
``echo_input=True``, the host pushes the user's input message to
each non-originating destination **before** the agent reply. Both
pushes go through the same ``push(identity, payload)`` entry point;
the channel distinguishes them by inspecting the role on the
payload's underlying :class:`~agent_framework.Message`\\(s):
* ``payload.result.messages[i].role == "user"`` → the echo phase
(originating user's turn mirrored onto this destination so the
channel's UX can stay coherent with the user's actual prompt).
Channels that cannot impersonate the user (most chat bots can
only send AS the bot) typically render echoes as a quoted /
prefixed block, drop them, or skip them via a
``response_hook`` — see below.
* ``payload.result.messages[i].role == "assistant"`` → the agent's
reply.
Channels that want to branch on phase WITHOUT inspecting roles can
instead expose a ``response_hook`` attribute on the channel
instance: the host calls the hook with a
:class:`ChannelResponseContext` whose ``is_echo`` flag carries the
same phase information explicitly, and the hook returns a
(possibly rewritten) :class:`HostedRunResult` that the host then
hands to ``push``. The hook seam is duck-typed and intentionally
NOT part of this Protocol so adding hook support to an existing
channel never breaks its public contract.
"""
name: str
async def push(self, identity: ChannelIdentity, payload: HostedRunResult[Any]) -> None: ...
async def apply_channel_response_hook(
channel: Channel | ChannelPush,
result: HostedRunResult[Any],
*,
request: ChannelRequest,
originating: bool,
destination_identity: ChannelIdentity | None = None,
is_echo: bool = False,
clone: bool = False,
) -> HostedRunResult[Any]:
"""Apply a channel's optional response hook with the standard context.
Channels and the host call this helper when they need to shape a
:class:`HostedRunResult` for one destination. The helper centralizes the
response-hook convention: hooks are discovered from a duck-typed
``response_hook`` attribute, called through :func:`apply_response_hook`,
and receive a :class:`ChannelResponseContext` that identifies the channel,
destination identity, originating-vs-push phase, and echo phase.
Args:
channel: Channel whose ``response_hook`` attribute may shape the payload.
result: Hosted run result to pass to the hook.
request: Originating channel request.
originating: Whether this is the originating channel's synchronous reply.
destination_identity: Destination identity for non-originating pushes, or
``None`` for originating replies.
is_echo: Whether the payload is an echo of the user input.
clone: Whether to shallow-clone ``result`` before applying the hook.
Returns:
The original, cloned, or hook-shaped hosted run result.
"""
shaped = result.replace() if clone else result
hook = cast(ChannelResponseHook | None, getattr(channel, "response_hook", None))
if not callable(hook):
return shaped
context = ChannelResponseContext(
request=request,
channel_name=channel.name,
destination_identity=destination_identity,
originating=originating,
is_echo=is_echo,
)
return await apply_response_hook(hook, shaped, context=context)
# --------------------------------------------------------------------------- #
# Durable task runner — pluggable seam for non-originating push fan-out and
# (in v1 fast-follow) background runs. See spec §"Durable task runner".
# --------------------------------------------------------------------------- #
@dataclass(frozen=True)
class RetryPolicy:
"""Retry contract a :class:`DurableTaskRunner` honours per scheduled task.
Defaults are deliberately conservative — five attempts on a 1s/2x/60s
exponential backoff — so a transient channel outage (Telegram returning
502, Activity Protocol token refresh) is rerouted to retry without the
operator wiring anything. Adapter backends (TaskHub, Foundry durable
tasks) MAY translate this into their native retry primitive; the
in-process runner implements it directly via ``asyncio.sleep``.
"""
max_attempts: int = 5
initial_backoff_seconds: float = 1.0
backoff_multiplier: float = 2.0
max_backoff_seconds: float = 60.0
@dataclass(frozen=True)
class TaskHandle:
"""Opaque, runner-issued handle for a scheduled task.
Callers receive one of these from :meth:`DurableTaskRunner.schedule` and
pass it back to :meth:`DurableTaskRunner.get` to poll status. ``task_id``
is opaque — its shape is implementation-defined (UUID for the in-process
runner, instance id for TaskHub, scheduled-task arn for Foundry). The
``name`` mirrors the handler name supplied to :meth:`schedule` so the
caller does not have to track it separately.
"""
task_id: str
name: str
TaskStatus = Literal["scheduled", "running", "succeeded", "failed", "cancelled"]
@runtime_checkable
class DurableTaskRunner(Protocol):
"""Pluggable seam the host uses to schedule out-of-band work.
The host registers a single internal handler — ``"hosting.push"`` — at
startup; each non-originating push destination becomes a
``runner.schedule("hosting.push", payload)`` call. The handler resolves
the destination channel, runs its ``response_hook`` (if any), and calls
:meth:`ChannelPush.push`. Failures inside the handler are caught by the
runner, retried per the supplied :class:`RetryPolicy`, and ultimately
marked terminal-failed when ``max_attempts`` is exhausted.
Two implementations ship in the framework: an in-process default
(``InProcessTaskRunner``, asyncio + bounded retry, no cross-restart
persistence) suitable for ``runtime_mode="long_running"`` deployments,
plus adapter packages (``agent-framework-hosting-durabletask``, a future
Foundry adapter) for ``runtime_mode="ephemeral"`` deployments that need
cross-restart durability.
Adapters MUST publish their ``payload_mode`` so the host's startup
validator can pair runner persistence expectations with channel
push-codec capabilities. Object-mode runners accept live Python
references in the payload (the in-process default does this for
speed); JSON-mode runners persist payloads across process restarts
and therefore require every push-capable channel to expose a
:class:`ChannelPushCodec`.
"""
# Adapter classes set this explicitly; the host inspects it at
# construction time. Default is conservative ("object") so a runner
# that omits the attribute is treated as in-process-only and does
# not silently impose a JSON requirement on channels.
payload_mode: DurableTaskPayloadMode
def register(
self,
name: str,
handler: Callable[[Mapping[str, Any]], Awaitable[None]],
) -> None:
"""Register a named handler the runner will invoke when a task fires.
Re-registering under the same name replaces the previous handler.
Implementations SHOULD raise :class:`RuntimeError` if called after
the runner has been started, to avoid silent reorderings of in-flight
work; the in-process runner enforces this.
"""
...
async def schedule(
self,
name: str,
payload: Mapping[str, Any],
*,
retry_policy: RetryPolicy | None = None,
) -> TaskHandle:
"""Schedule a previously-registered handler invocation.
``name`` MUST match a name previously passed to :meth:`register`. The
``payload`` is forwarded verbatim to the handler; implementations
MUST treat it as opaque (no introspection, no normalization).
``retry_policy`` overrides the runner's default for this task only;
``None`` means "use the runner-wide default".
Returns a :class:`TaskHandle` the caller may use with :meth:`get` to
poll status. Returning the handle MUST NOT wait for the task to run
— scheduling is fire-and-forget from the caller's perspective.
"""
...
async def get(self, handle: TaskHandle) -> TaskStatus | None:
"""Return the current status of a scheduled task.
Returns ``None`` if the runner no longer has any record of the task
(e.g. it was scheduled in a prior process and the runner has no
persistent backing). Otherwise one of the :data:`TaskStatus` values.
"""
...
__all__ = [
"AgentResponse",
"AgentResponseUpdate",
"Channel",
"ChannelCommand",
"ChannelCommandContext",
"ChannelContribution",
"ChannelIdentity",
"ChannelPush",
"ChannelPushCodec",
"ChannelRequest",
"ChannelResponseContext",
"ChannelResponseHook",
"ChannelRunHook",
"ChannelSession",
"ChannelStreamTransformHook",
"DurableTaskPayloadMode",
"DurableTaskRunner",
"HostStatePaths",
"HostedRunResult",
"PushPayloadNotPicklable",
"PushPayloadNotSerializable",
"ResponseStream",
"ResponseTarget",
"ResponseTargetKind",
"RetryPolicy",
"TaskHandle",
"TaskStatus",
"apply_channel_response_hook",
"apply_response_hook",
"apply_run_hook",
]
python/packages/hosting/tests/test_authorization.py
-580
View File
@@ -1,580 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for the authorization and identity-linking seam."""
from __future__ import annotations
from collections.abc import Collection
from typing import Any
import pytest
from agent_framework_hosting import (
AgentFrameworkHost,
AllOfAllowlists,
AllowAll,
Allowed,
AllowlistDecision,
AnyOfAllowlists,
AuthorizationContext,
AuthPolicy,
CallableAllowlist,
ChannelConfigurationError,
ChannelContext,
ChannelContribution,
ChannelIdentity,
Denied,
LinkChallenge,
LinkedClaimAllowlist,
LinkedIdentity,
LinkRequired,
NativeIdAllowlist,
)
# --------------------------------------------------------------------------- #
# Fakes #
# --------------------------------------------------------------------------- #
class _ChannelStub:
name: str = "stub"
path: str = "/stub"
require_link: bool = False
allowlist: Any = "inherit"
emits_verified_claims: bool = False
def __init__(
self,
*,
name: str = "stub",
require_link: bool = False,
allowlist: Any = "inherit",
emits_verified_claims: bool = False,
) -> None:
self.name = name
self.path = f"/{name}"
self.require_link = require_link
self.allowlist = allowlist
self.emits_verified_claims = emits_verified_claims
def contribute(self, context: ChannelContext) -> ChannelContribution:
return ChannelContribution(routes=[])
class _AgentStub:
"""Bare minimum target — the validators run during ``__init__``,
not on first request, so the target is never actually invoked."""
async def run(self, *args: Any, **kwargs: Any) -> Any: # pragma: no cover
raise NotImplementedError
class _StaticLinker:
"""Test linker returning either a linked identity or a challenge."""
def __init__(self, result: LinkedIdentity | LinkChallenge) -> None:
self.result = result
self.calls: list[ChannelIdentity] = []
async def resolve(self, identity: ChannelIdentity) -> LinkedIdentity | LinkChallenge:
self.calls.append(identity)
return self.result
def _ctx_pre_link(channel: str = "telegram", native_id: str = "42") -> AuthorizationContext:
return AuthorizationContext(
identity=ChannelIdentity(channel=channel, native_id=native_id),
phase="pre_link",
)
def _ctx_post_link(claims: dict[str, str] | None = None) -> AuthorizationContext:
return AuthorizationContext(
identity=ChannelIdentity(channel="telegram", native_id="42"),
phase="post_link",
isolation_key="alice",
verified_claims=claims or {},
claim_source="linker",
)
# --------------------------------------------------------------------------- #
# Built-in allowlists #
# --------------------------------------------------------------------------- #
class TestAllowAll:
async def test_allows_both_phases(self) -> None:
a = AllowAll()
assert await a.evaluate(_ctx_pre_link()) is AllowlistDecision.ALLOW
assert await a.evaluate(_ctx_post_link()) is AllowlistDecision.ALLOW
def test_does_not_require_linked_claims(self) -> None:
assert AllowAll().requires_linked_claims is False
class TestNativeIdAllowlist:
async def test_allows_listed_id(self) -> None:
a = NativeIdAllowlist({"42", "99"})
assert await a.evaluate(_ctx_pre_link(native_id="42")) is AllowlistDecision.ALLOW
async def test_denies_unlisted_id(self) -> None:
a = NativeIdAllowlist({"42"})
assert await a.evaluate(_ctx_pre_link(native_id="99")) is AllowlistDecision.DENY
async def test_channel_filter_abstains_for_other_channels(self) -> None:
# The native-id list is scoped to "telegram" — a request from
# another channel should ABSTAIN so a combinator can give a
# parallel allowlist a chance to ALLOW.
a = NativeIdAllowlist({"42"}, channel="telegram")
assert await a.evaluate(_ctx_pre_link(channel="slack", native_id="42")) is AllowlistDecision.ABSTAIN
async def test_channel_filter_evaluates_matching_channel(self) -> None:
a = NativeIdAllowlist({"42"}, channel="telegram")
assert await a.evaluate(_ctx_pre_link(channel="telegram", native_id="42")) is AllowlistDecision.ALLOW
assert await a.evaluate(_ctx_pre_link(channel="telegram", native_id="99")) is AllowlistDecision.DENY
async def test_async_loader_caches_after_first_call(self) -> None:
# The loader should run once; subsequent ``evaluate`` calls hit
# the cache so a slow / costly source isn't re-queried per
# message.
calls = {"n": 0}
async def loader() -> Collection[str]:
calls["n"] += 1
return {"42"}
a = NativeIdAllowlist(loader)
assert await a.evaluate(_ctx_pre_link(native_id="42")) is AllowlistDecision.ALLOW
assert await a.evaluate(_ctx_pre_link(native_id="42")) is AllowlistDecision.ALLOW
assert calls["n"] == 1
class TestLinkedClaimAllowlist:
"""Claim allowlists abstain pre-link and decide once claims are available."""
def test_declares_requires_linked_claims(self) -> None:
a = LinkedClaimAllowlist("oid", ["abc"])
assert a.requires_linked_claims is True
async def test_pre_link_abstains(self) -> None:
a = LinkedClaimAllowlist("oid", ["abc"])
assert await a.evaluate(_ctx_pre_link()) is AllowlistDecision.ABSTAIN
async def test_post_link_allows_matching_claim(self) -> None:
a = LinkedClaimAllowlist("oid", ["abc"])
assert await a.evaluate(_ctx_post_link({"oid": "abc"})) is AllowlistDecision.ALLOW
async def test_post_link_allows_matching_multi_value_claim(self) -> None:
a = LinkedClaimAllowlist("groups", ["admins"])
ctx = AuthorizationContext(
identity=ChannelIdentity(channel="telegram", native_id="42"),
phase="post_link",
isolation_key="alice",
verified_claims={"groups": ("users", "admins")},
claim_source="linker",
)
assert await a.evaluate(ctx) is AllowlistDecision.ALLOW
async def test_post_link_denies_missing_or_nonmatching_claim(self) -> None:
a = LinkedClaimAllowlist("oid", ["abc"])
assert await a.evaluate(_ctx_post_link({"oid": "def"})) is AllowlistDecision.DENY
assert await a.evaluate(_ctx_post_link({"tid": "abc"})) is AllowlistDecision.DENY
class TestAnyOfAllowlists:
async def test_any_allow_wins(self) -> None:
a = AnyOfAllowlists(NativeIdAllowlist({"42"}), NativeIdAllowlist({"99"}))
# native_id=42 → first ALLOWs, short-circuit.
assert await a.evaluate(_ctx_pre_link(native_id="42")) is AllowlistDecision.ALLOW
async def test_all_deny_yields_deny(self) -> None:
# Both lists deny native_id=7.
a = AnyOfAllowlists(NativeIdAllowlist({"42"}), NativeIdAllowlist({"99"}))
assert await a.evaluate(_ctx_pre_link(native_id="7")) is AllowlistDecision.DENY
async def test_abstain_when_no_decision(self) -> None:
# Channel-scoped lists both ABSTAIN on a "slack" request.
a = AnyOfAllowlists(
NativeIdAllowlist({"42"}, channel="telegram"),
NativeIdAllowlist({"99"}, channel="teams"),
)
assert await a.evaluate(_ctx_pre_link(channel="slack", native_id="42")) is AllowlistDecision.ABSTAIN
async def test_empty_is_abstain(self) -> None:
# No children → ABSTAIN (not DENY) to avoid silent deny-all.
a = AnyOfAllowlists()
assert await a.evaluate(_ctx_pre_link()) is AllowlistDecision.ABSTAIN
def test_propagates_requires_linked_claims(self) -> None:
a = AnyOfAllowlists(NativeIdAllowlist({"42"}), LinkedClaimAllowlist("oid", []))
assert a.requires_linked_claims is True
class TestAllOfAllowlists:
async def test_any_deny_short_circuits(self) -> None:
a = AllOfAllowlists(NativeIdAllowlist({"42"}), NativeIdAllowlist({"99"}))
assert await a.evaluate(_ctx_pre_link(native_id="42")) is AllowlistDecision.DENY
async def test_all_allow_yields_allow(self) -> None:
a = AllOfAllowlists(NativeIdAllowlist({"42"}), NativeIdAllowlist({"42", "99"}))
assert await a.evaluate(_ctx_pre_link(native_id="42")) is AllowlistDecision.ALLOW
async def test_abstain_when_no_deny_but_no_unanimous_allow(self) -> None:
a = AllOfAllowlists(
NativeIdAllowlist({"42"}, channel="telegram"),
NativeIdAllowlist({"42"}, channel="teams"),
)
# ABSTAIN from teams (different channel), ALLOW from telegram → ABSTAIN.
assert await a.evaluate(_ctx_pre_link(channel="telegram", native_id="42")) is AllowlistDecision.ABSTAIN
async def test_empty_is_abstain(self) -> None:
a = AllOfAllowlists()
assert await a.evaluate(_ctx_pre_link()) is AllowlistDecision.ABSTAIN
class TestCallableAllowlist:
async def test_wraps_async_fn(self) -> None:
async def fn(ctx: AuthorizationContext) -> AllowlistDecision:
if ctx.identity.native_id == "42":
return AllowlistDecision.ALLOW
return AllowlistDecision.DENY
a = CallableAllowlist(fn)
assert await a.evaluate(_ctx_pre_link(native_id="42")) is AllowlistDecision.ALLOW
assert await a.evaluate(_ctx_pre_link(native_id="99")) is AllowlistDecision.DENY
def test_requires_linked_claims_passthrough(self) -> None:
async def fn(_: AuthorizationContext) -> AllowlistDecision: # pragma: no cover
return AllowlistDecision.ALLOW
a = CallableAllowlist(fn, requires_linked_claims=True)
assert a.requires_linked_claims is True
class TestAuthPolicy:
async def test_factory_helpers_return_working_allowlists(self) -> None:
assert await AuthPolicy.open().evaluate(_ctx_pre_link()) is AllowlistDecision.ALLOW
assert await AuthPolicy.native_ids({"42"}).evaluate(_ctx_pre_link()) is AllowlistDecision.ALLOW
assert await AuthPolicy.linked_claim("oid", {"abc"}).evaluate(_ctx_post_link({"oid": "abc"})) is (
AllowlistDecision.ALLOW
)
async def test_custom_factory(self) -> None:
async def fn(_: AuthorizationContext) -> AllowlistDecision:
return AllowlistDecision.ALLOW
policy = AuthPolicy.custom(fn, requires_linked_claims=True)
assert policy.requires_linked_claims is True
assert await policy.evaluate(_ctx_pre_link()) is AllowlistDecision.ALLOW
# --------------------------------------------------------------------------- #
# Host configuration validator #
# --------------------------------------------------------------------------- #
class TestChannelAuthorizationValidator:
"""The host's startup validator catches three classes of misconfig
so they fail at construction rather than silently denying every
user at runtime."""
def test_require_link_without_linker_raises(self) -> None:
# ``require_link=True`` with no linker would silently reject
# every request — caught at construction.
with pytest.raises(ChannelConfigurationError, match="identity_linker"):
AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub(require_link=True)],
)
def test_require_link_with_linker_passes(self) -> None:
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub(require_link=True)],
identity_linker=_StaticLinker(LinkedIdentity("alice", {"oid": "abc"})),
)
assert host.runtime_mode == "long_running"
def test_linked_claim_allowlist_without_claim_source_raises(self) -> None:
# The channel has no ``require_link=True`` AND doesn't emit
# claims natively → the allowlist would always DENY / ABSTAIN.
with pytest.raises(ChannelConfigurationError, match="verified IdP claims"):
AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub(allowlist=LinkedClaimAllowlist("oid", []))],
)
def test_linked_claim_allowlist_with_native_claim_source_passes(self) -> None:
# When the channel declares ``emits_verified_claims=True``
# (e.g. Activity Protocol with AAD bearer) the validator
# accepts the LinkedClaimAllowlist without needing a linker.
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[
_ChannelStub(
allowlist=LinkedClaimAllowlist("oid", ["abc"]),
emits_verified_claims=True,
)
],
)
assert host.default_allowlist is None
def test_linked_claim_allowlist_with_require_link_and_linker_passes(self) -> None:
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub(require_link=True, allowlist=LinkedClaimAllowlist("oid", ["abc"]))],
identity_linker=_StaticLinker(LinkedIdentity("alice", {"oid": "abc"})),
)
assert host.runtime_mode == "long_running"
def test_native_id_allowlist_unknown_channel_raises(self) -> None:
with pytest.raises(ChannelConfigurationError, match="unknown channel 'mystery'"):
AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub(allowlist=NativeIdAllowlist({"42"}, channel="mystery"))],
)
def test_native_id_allowlist_known_channel_passes(self) -> None:
# A channel-scoped native list pointing at a peer channel is
# the supported way to compose per-channel allowlists.
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[
_ChannelStub(name="telegram", allowlist=NativeIdAllowlist({"42"}, channel="telegram")),
_ChannelStub(name="slack"),
],
)
assert host.runtime_mode == "long_running"
def test_default_allowlist_applies_to_inheriting_channel(self) -> None:
# ``allowlist="inherit"`` (the default) picks up the host-level
# ``default_allowlist``. This is the "lock down a whole bot in
# one place" ergonomic.
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub(name="telegram")],
default_allowlist=NativeIdAllowlist({"42"}),
)
# The default flowed through; channel sees the host's allowlist.
assert host.default_allowlist is not None
def test_explicit_none_carve_out_overrides_default(self) -> None:
# ``allowlist=None`` on a channel explicitly opts out of the
# host default — useful for a public endpoint inside an
# otherwise locked-down host.
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub(name="public", allowlist=None)],
default_allowlist=NativeIdAllowlist({"42"}),
)
# Construction succeeded; the validator did not raise.
assert host.default_allowlist is not None
def test_combinator_with_unknown_nested_channel_raises(self) -> None:
# The validator walks ``AnyOfAllowlists`` / ``AllOfAllowlists``
# so a typo'd channel name nested under a combinator is still
# caught at construction.
with pytest.raises(ChannelConfigurationError, match="unknown channel 'typo'"):
AgentFrameworkHost(
target=_AgentStub(),
channels=[
_ChannelStub(
allowlist=AnyOfAllowlists(
NativeIdAllowlist({"42"}, channel="stub"),
NativeIdAllowlist({"99"}, channel="typo"),
)
)
],
)
# --------------------------------------------------------------------------- #
# host.authorize pipeline #
# --------------------------------------------------------------------------- #
class TestHostAuthorize:
"""Host authorization pipeline across open, native-id, and linked-claim profiles."""
def _host(self) -> AgentFrameworkHost:
return AgentFrameworkHost(target=_AgentStub(), channels=[_ChannelStub()])
async def test_open_profile_returns_allowed_with_auto_isolation_key(self) -> None:
host = self._host()
outcome = await host.authorize(ChannelIdentity(channel="telegram", native_id="42"))
assert isinstance(outcome, Allowed)
assert outcome.isolation_key == "telegram:42"
async def test_native_allowlist_allows_listed_id(self) -> None:
host = self._host()
outcome = await host.authorize(
ChannelIdentity(channel="telegram", native_id="42"),
allowlist=NativeIdAllowlist({"42"}),
)
assert isinstance(outcome, Allowed)
assert outcome.isolation_key == "telegram:42"
async def test_native_allowlist_denies_unlisted_id(self) -> None:
host = self._host()
outcome = await host.authorize(
ChannelIdentity(channel="telegram", native_id="99"),
allowlist=NativeIdAllowlist({"42"}),
)
assert isinstance(outcome, Denied)
assert outcome.reason_code == "allowlist_denied_pre_link"
assert outcome.user_message is not None
# The bland default leaks neither tenant nor list size.
assert "telegram" not in (outcome.user_message or "")
async def test_abstain_with_claim_requirement_yields_link_required_message(self) -> None:
# Without a linker and without channel-emitted claims, a claim-required
# allowlist cannot make progress and the host returns a safe denial.
async def abstain(_: AuthorizationContext) -> AllowlistDecision:
return AllowlistDecision.ABSTAIN
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub(emits_verified_claims=True)],
)
outcome = await host.authorize(
ChannelIdentity(channel="telegram", native_id="42"),
allowlist=CallableAllowlist(abstain, requires_linked_claims=True),
)
assert isinstance(outcome, Denied)
assert outcome.reason_code == "allowlist_requires_link"
async def test_abstain_without_claim_requirement_falls_through_to_allowed(self) -> None:
async def abstain(_: AuthorizationContext) -> AllowlistDecision:
return AllowlistDecision.ABSTAIN
host = self._host()
outcome = await host.authorize(
ChannelIdentity(channel="telegram", native_id="42"),
allowlist=CallableAllowlist(abstain),
)
assert isinstance(outcome, Allowed)
async def test_auto_issue_returns_existing_key_when_known(self) -> None:
# When an identity has already been observed, the auto-issued
# key matches the existing one rather than coining a fresh
# token. This is the linker-free equivalent of identity resolution.
host = self._host()
host._identities["alice"] = {"telegram": ChannelIdentity(channel="telegram", native_id="42")}
outcome = await host.authorize(ChannelIdentity(channel="telegram", native_id="42"))
assert isinstance(outcome, Allowed)
assert outcome.isolation_key == "alice"
async def test_verified_claims_propagate_to_context(self) -> None:
# Channels that natively carry verified claims (e.g. Activity
# Protocol bearer with AAD oid) pass them through to
# ``authorize`` — the allowlist sees them on the
# ``AuthorizationContext``.
seen: list[AuthorizationContext] = []
async def capture(ctx: AuthorizationContext) -> AllowlistDecision:
seen.append(ctx)
return AllowlistDecision.ALLOW
host = self._host()
await host.authorize(
ChannelIdentity(channel="telegram", native_id="42"),
allowlist=CallableAllowlist(capture),
verified_claims={"oid": "abc"},
)
assert len(seen) == 1
assert seen[0].claim_source == "channel"
assert dict(seen[0].verified_claims) == {"oid": "abc"}
async def test_require_link_returns_challenge_when_unlinked(self) -> None:
challenge = LinkChallenge("c1", url="https://login.example/c1")
linker = _StaticLinker(challenge)
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub(require_link=True)],
identity_linker=linker,
)
outcome = await host.authorize(
ChannelIdentity(channel="telegram", native_id="42"),
require_link=True,
)
assert isinstance(outcome, LinkRequired)
assert outcome.challenge is challenge
assert [call.native_id for call in linker.calls] == ["42"]
async def test_require_link_returns_linked_identity_when_resolved(self) -> None:
linked = LinkedIdentity("entra:abc", {"oid": "abc"})
linker = _StaticLinker(linked)
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub(require_link=True)],
identity_linker=linker,
)
outcome = await host.authorize(
ChannelIdentity(channel="telegram", native_id="42"),
require_link=True,
)
assert isinstance(outcome, Allowed)
assert outcome.isolation_key == "entra:abc"
assert dict(outcome.verified_claims) == {"oid": "abc"}
assert outcome.claim_source == "linker"
# authorize() is decision-only; identity registry writes remain on
# the request execution path.
assert host._identities == {}
async def test_linked_claim_allowlist_with_linker_allows_matching_claim(self) -> None:
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub(require_link=True, allowlist=LinkedClaimAllowlist("oid", ["abc"]))],
identity_linker=_StaticLinker(LinkedIdentity("entra:abc", {"oid": "abc"})),
)
outcome = await host.authorize(
ChannelIdentity(channel="telegram", native_id="42"),
require_link=True,
allowlist=LinkedClaimAllowlist("oid", ["abc"]),
)
assert isinstance(outcome, Allowed)
assert outcome.isolation_key == "entra:abc"
async def test_linked_claim_allowlist_with_linker_denies_nonmatching_claim(self) -> None:
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub(require_link=True, allowlist=LinkedClaimAllowlist("oid", ["abc"]))],
identity_linker=_StaticLinker(LinkedIdentity("entra:def", {"oid": "def"})),
)
outcome = await host.authorize(
ChannelIdentity(channel="telegram", native_id="42"),
require_link=True,
allowlist=LinkedClaimAllowlist("oid", ["abc"]),
)
assert isinstance(outcome, Denied)
assert outcome.reason_code == "allowlist_denied_post_link"
async def test_linked_claim_allowlist_with_linker_returns_challenge_when_unlinked(self) -> None:
challenge = LinkChallenge("c1")
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub(require_link=True, allowlist=LinkedClaimAllowlist("oid", ["abc"]))],
identity_linker=_StaticLinker(challenge),
)
outcome = await host.authorize(
ChannelIdentity(channel="telegram", native_id="42"),
require_link=True,
allowlist=LinkedClaimAllowlist("oid", ["abc"]),
)
assert isinstance(outcome, LinkRequired)
assert outcome.challenge is challenge
async def test_linked_claim_allowlist_uses_channel_verified_claims_without_linker(self) -> None:
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub(emits_verified_claims=True, allowlist=LinkedClaimAllowlist("oid", ["abc"]))],
)
outcome = await host.authorize(
ChannelIdentity(channel="activity", native_id="aad-user"),
allowlist=LinkedClaimAllowlist("oid", ["abc"]),
verified_claims={"oid": "abc"},
)
assert isinstance(outcome, Allowed)
assert outcome.isolation_key == "activity:aad-user"
assert outcome.claim_source == "channel"
python/packages/hosting/tests/test_host.py
+77 -629
View File
@@ -4,7 +4,7 @@
from __future__ import annotations
from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence
from collections.abc import AsyncIterator, Sequence
from dataclasses import dataclass, field
from typing import Any
@@ -21,16 +21,9 @@ from agent_framework_hosting import (
ChannelContext,
ChannelContribution,
ChannelIdentity,
ChannelPush,
ChannelRequest,
ChannelSession,
DurableTaskPayloadMode,
DurableTaskRunner,
HostedRunResult,
ResponseTarget,
RetryPolicy,
TaskHandle,
TaskStatus,
)
@@ -75,27 +68,33 @@ class _FakeAgent:
self.created_sessions.append(s)
return s
async def run(self, messages: Any = None, *, stream: bool = False, session: Any = None, **kwargs: Any) -> Any:
def run(self, messages: Any = None, *, stream: bool = False, session: Any = None, **kwargs: Any) -> Any:
self.calls.append({"messages": messages, "stream": stream, "session": session, "kwargs": kwargs})
if stream: # pragma: no cover - not used by these tests
if stream:
updates = [AgentResponseUpdate(contents=[Content.from_text(text=self._reply)], role="assistant")]
async def _gen() -> AsyncIterator[Any]:
yield self._reply
async def _gen() -> AsyncIterator[AgentResponseUpdate]:
for update in updates:
yield update
return _gen()
return _FakeAgentResponse(text=self._reply)
async def _finalize(items: Sequence[AgentResponseUpdate]) -> AgentResponse: # noqa: RUF029
return AgentResponse.from_updates(items)
return ResponseStream[AgentResponseUpdate, AgentResponse](_gen(), finalizer=_finalize)
async def _coro() -> _FakeAgentResponse:
return _FakeAgentResponse(text=self._reply)
return _coro()
class _RecordingChannel:
"""Minimal :class:`Channel` + :class:`ChannelPush` for routing tests."""
"""Minimal :class:`Channel` for host tests."""
def __init__(self, name: str = "fake", path: str = "/fake", supports_push: bool = True) -> None:
def __init__(self, name: str = "fake", path: str = "/fake") -> None:
self.name = name
self.path = path
self.context: ChannelContext | None = None
self.pushes: list[tuple[ChannelIdentity, HostedRunResult[Any]]] = []
self._push_raises: Exception | None = None
self._supports_push = supports_push
# Provide a single trivial route so contribute() exercises the endpoint path.
self._routes: Sequence[BaseRoute] = (Route("/ping", _ping),)
@@ -103,88 +102,6 @@ class _RecordingChannel:
self.context = context
return ChannelContribution(routes=self._routes)
async def push(self, identity: ChannelIdentity, payload: HostedRunResult[Any]) -> None:
if self._push_raises is not None:
raise self._push_raises
self.pushes.append((identity, payload))
class _NoPushChannel:
"""A channel that does NOT implement :class:`ChannelPush`."""
def __init__(self, name: str = "nopush", path: str = "/nopush") -> None:
self.name = name
self.path = path
def contribute(self, context: ChannelContext) -> ChannelContribution:
return ChannelContribution()
class _SyncTaskRunner(DurableTaskRunner):
"""A :class:`DurableTaskRunner` that runs handlers inline.
Tests of the delivery routing want deterministic, synchronous
behaviour. The real :class:`InProcessTaskRunner` schedules via
``asyncio.create_task`` so push side effects only land *after*
the test has yielded control — awkward for assertions that read
a channel's recorded pushes immediately after
:meth:`ChannelContext.deliver_response` returns.
Two knobs control failure handling:
- ``schedule_raises``: when set, every call to :meth:`schedule`
raises this exception. Mimics a host-side outage (the durable
backend is unreachable).
- ``swallow_handler_errors`` (default ``True``): when the
handler raises, the error is recorded in
:attr:`handler_errors` but :meth:`schedule` still returns
successfully — matching the real durable contract that
"scheduled" is a separate signal from "delivered". Set to
``False`` to surface handler exceptions through
:meth:`schedule` for the few tests that want to assert on
handler-raised failures inline.
"""
def __init__(self, *, swallow_handler_errors: bool = True) -> None:
self._handlers: dict[str, Callable[[Mapping[str, Any]], Awaitable[None]]] = {}
self.scheduled: list[tuple[str, Mapping[str, Any]]] = []
self.handler_errors: list[BaseException] = []
self.schedule_raises: BaseException | None = None
self.swallow_handler_errors = swallow_handler_errors
# Default object-mode matches the real ``InProcessTaskRunner`` —
# tests that want to exercise the JSON-mode path override this on
# the instance.
payload_mode = DurableTaskPayloadMode.OBJECT
def register(
self,
name: str,
handler: Callable[[Mapping[str, Any]], Awaitable[None]],
) -> None:
self._handlers[name] = handler
async def schedule(
self,
name: str,
payload: Mapping[str, Any],
*,
retry_policy: RetryPolicy | None = None,
) -> TaskHandle:
if self.schedule_raises is not None:
raise self.schedule_raises
self.scheduled.append((name, payload))
try:
await self._handlers[name](payload)
except Exception as exc:
self.handler_errors.append(exc)
if not self.swallow_handler_errors:
raise
return TaskHandle(task_id=f"sync-{len(self.scheduled)}", name=name)
async def get(self, handle: TaskHandle) -> TaskStatus | None: # pragma: no cover - unused
return "succeeded"
def _assistant_response(text: str) -> AgentResponse:
"""Build a one-message ``AgentResponse`` to use as a ``HostedRunResult.result``."""
@@ -227,7 +144,6 @@ class TestHostWiring:
def test_channel_is_recognized(self) -> None:
ch = _RecordingChannel()
assert isinstance(ch, Channel)
assert isinstance(ch, ChannelPush)
def test_app_mounts_channel_routes_under_path(self) -> None:
agent = _FakeAgent()
@@ -313,10 +229,6 @@ class TestHostInvoke:
"native_id": "user:1",
"attributes": {},
}
assert msg.additional_properties["hosting"]["response_target"] == {
"kind": "originating",
"targets": [],
}
async def test_invoke_caches_session_per_isolation_key(self) -> None:
agent = _FakeAgent()
@@ -398,6 +310,56 @@ class TestHostInvoke:
assert agent.calls[0]["kwargs"]["options"] == {"temperature": 0.4}
class TestHostOwnedHooks:
async def test_context_run_applies_run_hook_before_invocation(self) -> None:
agent = _FakeAgent()
ch = _RecordingChannel()
host = AgentFrameworkHost(target=agent, channels=[ch])
_ = host.app
assert ch.context is not None
captured: dict[str, Any] = {}
async def hook(request: ChannelRequest, **kwargs: Any) -> ChannelRequest:
captured["target"] = kwargs["target"]
captured["protocol_request"] = kwargs["protocol_request"]
return ChannelRequest(
channel=request.channel,
operation=request.operation,
input="rewritten",
session=request.session,
)
req = ChannelRequest(channel=ch.name, operation="op", input="original", session=ChannelSession("alice"))
await ch.context.run(req, run_hook=hook, protocol_request={"raw": True})
assert captured["target"] is agent
assert captured["protocol_request"] == {"raw": True}
assert agent.calls[0]["messages"].text == "rewritten"
async def test_context_run_stream_applies_run_hook_before_opening_stream(self) -> None:
agent = _FakeAgent()
ch = _RecordingChannel()
host = AgentFrameworkHost(target=agent, channels=[ch])
_ = host.app
assert ch.context is not None
def hook(request: ChannelRequest, **_: Any) -> ChannelRequest:
return ChannelRequest(channel=request.channel, operation=request.operation, input="streamed")
stream = await ch.context.run_stream(
ChannelRequest(channel=ch.name, operation="op", input="original"),
run_hook=hook,
stream_update_hook=lambda update: AgentResponseUpdate(
contents=[Content.from_text(text=update.text.upper())],
role="assistant",
),
)
chunks = [update.text async for update in stream]
assert chunks == ["OK"]
assert agent.calls[0]["messages"].text == "streamed"
# --------------------------------------------------------------------------- #
# Workflow target #
# --------------------------------------------------------------------------- #
@@ -436,7 +398,7 @@ class TestHostWorkflowTarget:
assert ch.context is not None
req = ChannelRequest(channel="fake", operation="message.create", input="hi")
stream = ch.context.run_stream(req)
stream = await ch.context.run_stream(req)
updates: list[AgentResponseUpdate] = []
async for update in stream:
@@ -464,7 +426,7 @@ class TestHostWorkflowTarget:
assert ch.context is not None
req = ChannelRequest(channel="fake", operation="message.create", input="x")
stream = ch.context.run_stream(req)
stream = await ch.context.run_stream(req)
chunks: list[str] = []
async for update in stream:
@@ -561,7 +523,7 @@ class TestHostWorkflowCheckpointing:
input="hi",
session=ChannelSession(isolation_key="bob"),
)
stream = ch.context.run_stream(req)
stream = await ch.context.run_stream(req)
async for _ in stream:
pass
await stream.get_final_response()
@@ -702,520 +664,6 @@ class TestHostWorkflowCheckpointingPathTraversal:
assert list(tmp_path.iterdir()) == []
# --------------------------------------------------------------------------- #
# Delivery routing #
# --------------------------------------------------------------------------- #
def _make_host_with_two_channels(
*,
runner: DurableTaskRunner | None = None,
) -> tuple[AgentFrameworkHost, _RecordingChannel, _RecordingChannel, ChannelContext, _SyncTaskRunner]:
agent = _FakeAgent()
a = _RecordingChannel(name="responses", path="/r")
b = _RecordingChannel(name="telegram", path="/t")
sync_runner = runner if isinstance(runner, _SyncTaskRunner) else _SyncTaskRunner()
host = AgentFrameworkHost(
target=agent,
channels=[a, b],
durable_task_runner=runner or sync_runner,
)
_ = host.app
assert a.context is not None
return host, a, b, a.context, sync_runner
def _record_identity_on(host: AgentFrameworkHost, isolation_key: str, channel: str, native_id: str) -> None:
"""Pre-seed the host's identity registry by running a request."""
host._identities.setdefault(isolation_key, {})[channel] = ChannelIdentity(channel=channel, native_id=native_id)
host._active[isolation_key] = channel
class TestDeliverResponse:
"""Delivery routing — the originating channel learns whether to render
on its own wire from the ``bool`` return; everything else
(scheduled tasks, schedule-time failures, skip reasons) lives in
the runner's own log. Tests assert the bool plus observable
state on the sync runner fake (``scheduled``, ``handler_errors``)
and on the destination channels (``pushes``)."""
async def test_originating_returns_true(self) -> None:
_, _, _, ctx, runner = _make_host_with_two_channels()
req = ChannelRequest(channel="responses", operation="op", input="x")
include_originating = await ctx.deliver_response(req, _make_reply("reply"))
assert include_originating is True
assert runner.scheduled == []
async def test_none_suppresses_everything(self) -> None:
_, _, _, ctx, runner = _make_host_with_two_channels()
req = ChannelRequest(
channel="responses",
operation="op",
input="x",
response_target=ResponseTarget.none, # type: ignore[attr-defined]
)
include_originating = await ctx.deliver_response(req, _make_reply("reply"))
assert include_originating is False
assert runner.scheduled == []
async def test_active_pushes_to_other_channel(self) -> None:
host, _a, b, ctx, runner = _make_host_with_two_channels()
# Alice was last seen on telegram.
_record_identity_on(host, "alice", "telegram", "42")
# Now she sends a message via responses; ResponseTarget.active should
# push to telegram, not back to responses.
req = ChannelRequest(
channel="responses",
operation="op",
input="x",
session=ChannelSession(isolation_key="alice"),
response_target=ResponseTarget.active, # type: ignore[attr-defined]
)
include_originating = await ctx.deliver_response(req, _make_reply("reply"))
assert include_originating is False
assert len(runner.scheduled) == 1
assert b.pushes and b.pushes[0][0].native_id == "42"
async def test_active_falls_back_to_originating_when_self(self) -> None:
host, _a, _b, ctx, runner = _make_host_with_two_channels()
_record_identity_on(host, "alice", "responses", "user:1")
req = ChannelRequest(
channel="responses",
operation="op",
input="x",
session=ChannelSession(isolation_key="alice"),
response_target=ResponseTarget.active, # type: ignore[attr-defined]
)
include_originating = await ctx.deliver_response(req, _make_reply("reply"))
assert include_originating is True
assert runner.scheduled == []
async def test_channels_with_unknown_identity_falls_back_to_originating(self) -> None:
_, _, _, ctx, runner = _make_host_with_two_channels()
# No prior identity seeded for telegram on alice.
req = ChannelRequest(
channel="responses",
operation="op",
input="x",
session=ChannelSession(isolation_key="alice"),
response_target=ResponseTarget.channel("telegram"),
)
include_originating = await ctx.deliver_response(req, _make_reply("reply"))
# Skipped at resolution → fallback to originating so the user
# still gets a reply.
assert include_originating is True
assert runner.scheduled == []
async def test_channels_with_explicit_native_id_token(self) -> None:
_, _, b, ctx, runner = _make_host_with_two_channels()
req = ChannelRequest(
channel="responses",
operation="op",
input="x",
response_target=ResponseTarget.channel("telegram:99"),
)
include_originating = await ctx.deliver_response(req, _make_reply("reply"))
assert include_originating is False
assert len(runner.scheduled) == 1
assert b.pushes[0][0].native_id == "99"
async def test_channels_originating_pseudo_includes_origin(self) -> None:
host, _a, _b, ctx, runner = _make_host_with_two_channels()
_record_identity_on(host, "alice", "telegram", "42")
req = ChannelRequest(
channel="responses",
operation="op",
input="x",
session=ChannelSession(isolation_key="alice"),
response_target=ResponseTarget.channels(["originating", "telegram"]),
)
include_originating = await ctx.deliver_response(req, _make_reply("reply"))
assert include_originating is True
assert len(runner.scheduled) == 1
async def test_channels_unknown_channel_name_falls_back(self) -> None:
_, _, _, ctx, runner = _make_host_with_two_channels()
req = ChannelRequest(
channel="responses",
operation="op",
input="x",
response_target=ResponseTarget.channel("nope"),
)
include_originating = await ctx.deliver_response(req, _make_reply("reply"))
assert include_originating is True # fallback
assert runner.scheduled == []
async def test_no_push_capability_falls_back(self) -> None:
agent = _FakeAgent()
a = _RecordingChannel(name="responses", path="/r")
b = _NoPushChannel(name="nopush", path="/n")
host = AgentFrameworkHost(target=agent, channels=[a, b])
_ = host.app
assert a.context is not None
# Pre-seed identity on the no-push channel so we get past the
# identity check and hit the ChannelPush check.
host._identities.setdefault("alice", {})["nopush"] = ChannelIdentity(channel="nopush", native_id="42")
req = ChannelRequest(
channel="responses",
operation="op",
input="x",
session=ChannelSession(isolation_key="alice"),
response_target=ResponseTarget.channel("nopush"),
)
include_originating = await a.context.deliver_response(req, _make_reply("reply"))
assert include_originating is True # fallback
async def test_all_linked_pushes_to_every_other_channel(self) -> None:
host, _a, b, ctx, runner = _make_host_with_two_channels()
# Alice on responses (originating) and telegram.
host._identities.setdefault("alice", {})
host._identities["alice"]["responses"] = ChannelIdentity(channel="responses", native_id="user:1")
host._identities["alice"]["telegram"] = ChannelIdentity(channel="telegram", native_id="42")
req = ChannelRequest(
channel="responses",
operation="op",
input="x",
session=ChannelSession(isolation_key="alice"),
response_target=ResponseTarget.all_linked, # type: ignore[attr-defined]
)
include_originating = await ctx.deliver_response(req, _make_reply("reply"))
assert include_originating is True
assert len(runner.scheduled) == 1
assert b.pushes and b.pushes[0][1].result.text == "reply"
async def test_all_linked_no_other_channels_falls_back(self) -> None:
_host, _a, _b, ctx, runner = _make_host_with_two_channels()
req = ChannelRequest(
channel="responses",
operation="op",
input="x",
session=ChannelSession(isolation_key="alice"),
response_target=ResponseTarget.all_linked, # type: ignore[attr-defined]
)
include_originating = await ctx.deliver_response(req, _make_reply("reply"))
assert include_originating is True
assert runner.scheduled == []
async def test_identities_variant_preserves_attributes(self) -> None:
"""``ResponseTarget.identities([...])`` plumbs full
:class:`ChannelIdentity` objects through resolution, preserving
``attributes`` for destination channels that need conversation/
thread metadata (Teams, Slack, Bot Framework)."""
_, _, b, ctx, runner = _make_host_with_two_channels()
ident = ChannelIdentity(
channel="telegram",
native_id="42",
attributes={"thread_id": "t1", "service_url": "https://x"},
)
req = ChannelRequest(
channel="responses",
operation="op",
input="x",
response_target=ResponseTarget.identity(ident),
)
include_originating = await ctx.deliver_response(req, _make_reply("reply"))
assert include_originating is False
assert len(runner.scheduled) == 1
# The destination identity arrived at push with attributes intact.
pushed_identity = b.pushes[0][0]
assert pushed_identity.native_id == "42"
assert dict(pushed_identity.attributes) == {"thread_id": "t1", "service_url": "https://x"}
async def test_identities_pointing_to_originating_includes_origin(self) -> None:
"""An identity whose channel matches the originating channel
folds into ``include_originating`` rather than double-delivering
via push."""
_, _, _, ctx, runner = _make_host_with_two_channels()
ident = ChannelIdentity(channel="responses", native_id="user:1")
req = ChannelRequest(
channel="responses",
operation="op",
input="x",
response_target=ResponseTarget.identities([ident]),
)
include_originating = await ctx.deliver_response(req, _make_reply("reply"))
assert include_originating is True
assert runner.scheduled == []
async def test_handler_exception_does_not_change_return_value(self) -> None:
"""When ``ChannelPush.push`` raises *inside the runner handler*
the originating channel still sees the same return value —
``DurableTaskRunner.schedule`` accepted the work, and downstream
delivery outcome is owned by the runner (it logs and retries
per the configured ``RetryPolicy``)."""
host, _a, b, ctx, runner = _make_host_with_two_channels()
b._push_raises = RuntimeError("boom") # type: ignore[attr-defined]
host._identities.setdefault("alice", {})["telegram"] = ChannelIdentity(channel="telegram", native_id="42")
req = ChannelRequest(
channel="responses",
operation="op",
input="x",
session=ChannelSession(isolation_key="alice"),
response_target=ResponseTarget.channel("telegram"),
)
include_originating = await ctx.deliver_response(req, _make_reply("reply"))
# Schedule succeeded → the return value is unaffected by a
# downstream handler failure.
assert include_originating is False
assert len(runner.scheduled) == 1
# Handler raised — runner captured the error (the real runner
# would retry it; the sync fake records it).
assert runner.handler_errors and isinstance(runner.handler_errors[0], RuntimeError)
assert str(runner.handler_errors[0]) == "boom"
async def test_schedule_exception_falls_back_to_originating(self) -> None:
"""When :meth:`DurableTaskRunner.schedule` itself raises (the
runner backend is unreachable) the destination is treated as
skipped — same outcome as any other resolution-time drop. The
host's fall-back-to-originating rule then ensures the user
still gets a reply rather than being left without one."""
host, _a, _b, ctx, runner = _make_host_with_two_channels()
runner.schedule_raises = RuntimeError("runner backend unreachable")
host._identities.setdefault("alice", {})["telegram"] = ChannelIdentity(channel="telegram", native_id="42")
req = ChannelRequest(
channel="responses",
operation="op",
input="x",
session=ChannelSession(isolation_key="alice"),
response_target=ResponseTarget.channel("telegram"),
)
include_originating = await ctx.deliver_response(req, _make_reply("reply"))
# Schedule raised → no scheduled tasks, fall back to originating.
assert runner.scheduled == []
assert include_originating is True
async def test_echo_input_pushes_user_message_then_response(self) -> None:
"""``echo_input=True`` triggers two pushes per destination,
bundled into the same scheduled task: the originating user
message first, then the agent reply. Channels downstream of a
workflow that emits to multiple channels need this to keep
their UI state coherent with the user's actual prompt."""
host, _a, b, ctx, runner = _make_host_with_two_channels()
host._identities.setdefault("alice", {})["telegram"] = ChannelIdentity(channel="telegram", native_id="42")
req = ChannelRequest(
channel="responses",
operation="op",
input="hello there",
session=ChannelSession(isolation_key="alice"),
response_target=ResponseTarget.channel("telegram", echo_input=True),
)
include_originating = await ctx.deliver_response(req, _make_reply("reply"))
assert include_originating is False
# One scheduled task per destination; the handler does echo then response inline.
assert len(runner.scheduled) == 1
_, payload = runner.scheduled[0]
assert payload["echo_result"] is not None
# Two pushes landed on the channel: echo first, then response.
assert len(b.pushes) == 2
echo_identity, echo_payload = b.pushes[0]
assert echo_identity.native_id == "42"
assert echo_payload.result.text == "hello there"
assert str(echo_payload.result.messages[0].role) == "user"
resp_identity, resp_payload = b.pushes[1]
assert resp_identity.native_id == "42"
assert resp_payload.result.text == "reply"
assert str(resp_payload.result.messages[0].role) == "assistant"
async def test_echo_input_failure_does_not_block_response(self) -> None:
"""An echo push that raises inside the handler is logged and
swallowed; the response push must still be attempted on the
same destination so the user-visible failure mode is
"response delivered without echo" rather than "no response at
all"."""
agent = _FakeAgent()
a = _RecordingChannel(name="responses", path="/r")
b = _RecordingChannel(name="telegram", path="/t")
runner = _SyncTaskRunner()
host = AgentFrameworkHost(target=agent, channels=[a, b], durable_task_runner=runner)
_ = host.app
assert a.context is not None
host._identities.setdefault("alice", {})["telegram"] = ChannelIdentity(channel="telegram", native_id="42")
# Make the FIRST push (echo) raise, but the SECOND (response) succeed.
calls = {"n": 0}
real_push = b.push
async def flaky_push(identity: ChannelIdentity, payload: HostedRunResult[Any]) -> None:
calls["n"] += 1
if calls["n"] == 1:
raise RuntimeError("echo down")
await real_push(identity, payload)
b.push = flaky_push # type: ignore[method-assign]
req = ChannelRequest(
channel="responses",
operation="op",
input="hi",
session=ChannelSession(isolation_key="alice"),
response_target=ResponseTarget.channel("telegram", echo_input=True),
)
include_originating = await a.context.deliver_response(req, _make_reply("reply"))
# Schedule succeeded; handler swallowed the echo failure and
# the response push landed on the channel.
assert include_originating is False
assert b.pushes and b.pushes[0][1].result.text == "reply"
# Handler did not raise (echo failure was swallowed inside
# the handler), so the runner saw no error.
assert runner.handler_errors == []
async def test_echo_idempotent_on_retry(self) -> None:
"""When the response push fails on a retried task, the handler
must NOT re-deliver the echo if a prior attempt already
succeeded. The ``echo_done`` cursor on the payload mapping is
the host's idempotency primitive; this test invokes the
handler directly twice with the same payload to exercise the
retry semantics."""
host, _a, b, ctx, runner = _make_host_with_two_channels()
host._identities.setdefault("alice", {})["telegram"] = ChannelIdentity(channel="telegram", native_id="42")
req = ChannelRequest(
channel="responses",
operation="op",
input="hi",
session=ChannelSession(isolation_key="alice"),
response_target=ResponseTarget.channel("telegram", echo_input=True),
)
# First scheduled invocation — echo + response both succeed.
await ctx.deliver_response(req, _make_reply("reply"))
assert len(b.pushes) == 2 # echo + response
# Simulate a retry: invoke the handler again with the same
# payload mapping (the in-process runner reuses the mapping
# across retries). After the first run ``echo_done`` was
# mutated to ``True``; the second run must skip the echo.
_, payload = runner.scheduled[0]
assert payload["echo_done"] is True
await host._handle_push_task(payload)
# Only one more push (the response) — the echo was skipped.
assert len(b.pushes) == 3
assert str(b.pushes[2][1].result.messages[0].role) == "assistant"
# --------------------------------------------------------------------------- #
# Response hook + multi-modal payload + clone-on-fan-out #
# --------------------------------------------------------------------------- #
class TestResponseHookFanOut:
async def test_response_hook_applied_per_destination(self) -> None:
"""Channels with a ``response_hook`` attribute see their hook
applied before push, with a ``ChannelResponseContext`` carrying
the destination identity, the originating request, and an
``is_echo`` flag."""
agent = _FakeAgent()
a = _RecordingChannel(name="responses", path="/r")
b = _RecordingChannel(name="telegram", path="/t")
seen: list[tuple[str, str, bool]] = []
async def telegram_hook(
result: HostedRunResult[AgentResponse],
*,
context: Any,
**_: Any,
) -> HostedRunResult[AgentResponse]:
seen.append((context.channel_name, context.destination_identity.native_id, context.is_echo))
return result.replace(
result=AgentResponse(
messages=[Message(role="assistant", contents=[Content.from_text("[hooked] " + result.result.text)])]
),
)
b.response_hook = telegram_hook # type: ignore[attr-defined]
host = AgentFrameworkHost(target=agent, channels=[a, b], durable_task_runner=_SyncTaskRunner())
_ = host.app
assert a.context is not None
host._identities.setdefault("alice", {})["telegram"] = ChannelIdentity(channel="telegram", native_id="42")
req = ChannelRequest(
channel="responses",
operation="op",
input="hi",
session=ChannelSession(isolation_key="alice"),
response_target=ResponseTarget.channel("telegram"),
)
report = await a.context.deliver_response(req, _make_reply("reply"))
assert report is False
# The pushed payload reflects the hook's transform.
assert b.pushes[0][1].result.text == "[hooked] reply"
assert seen == [("telegram", "42", False)]
async def test_response_hook_mutation_isolated_per_destination(self) -> None:
"""A hook that rebinds ``result`` on its payload must NOT affect
the payload another destination sees. The host clones the
envelope before each hook invocation so a per-destination
:meth:`HostedRunResult.replace` cannot leak across destinations."""
agent = _FakeAgent()
a = _RecordingChannel(name="responses", path="/r")
b = _RecordingChannel(name="telegram", path="/t")
c = _RecordingChannel(name="extra", path="/x")
async def hook_that_rebinds(result: HostedRunResult[AgentResponse], **_: Any) -> HostedRunResult[AgentResponse]:
# Naughty hook: rebind ``result`` to a fresh AgentResponse.
# Host's per-destination clone via ``replace()`` makes this safe
# for sibling destinations.
return result.replace(result=AgentResponse(messages=[]))
b.response_hook = hook_that_rebinds # type: ignore[attr-defined]
host = AgentFrameworkHost(target=agent, channels=[a, b, c], durable_task_runner=_SyncTaskRunner())
_ = host.app
assert a.context is not None
host._identities.setdefault("alice", {})["telegram"] = ChannelIdentity(channel="telegram", native_id="42")
host._identities["alice"]["extra"] = ChannelIdentity(channel="extra", native_id="9")
original = _make_reply("reply")
original_result_snapshot = original.result
req = ChannelRequest(
channel="responses",
operation="op",
input="hi",
session=ChannelSession(isolation_key="alice"),
response_target=ResponseTarget.channels(["telegram", "extra"]),
)
report = await a.context.deliver_response(req, original)
assert report is False
# The rebind on the telegram clone must not have touched the
# original envelope, nor the extra channel's view.
assert original.result is original_result_snapshot
# ``extra`` channel saw the original-shaped payload.
extra_push = next(p for p in c.pushes)
assert extra_push[1].result.text == "reply"
async def test_response_hook_fires_on_echo_with_is_echo_true(self) -> None:
"""When ``echo_input`` is set, the channel's response_hook fires
TWICE per destination — once for the echo (is_echo=True), once
for the response (is_echo=False)."""
agent = _FakeAgent()
a = _RecordingChannel(name="responses", path="/r")
b = _RecordingChannel(name="telegram", path="/t")
phases: list[bool] = []
async def telegram_hook(
result: HostedRunResult[AgentResponse], *, context: Any, **_: Any
) -> HostedRunResult[AgentResponse]:
phases.append(context.is_echo)
return result
b.response_hook = telegram_hook # type: ignore[attr-defined]
host = AgentFrameworkHost(target=agent, channels=[a, b], durable_task_runner=_SyncTaskRunner())
_ = host.app
assert a.context is not None
host._identities.setdefault("alice", {})["telegram"] = ChannelIdentity(channel="telegram", native_id="42")
req = ChannelRequest(
channel="responses",
operation="op",
input="hi",
session=ChannelSession(isolation_key="alice"),
response_target=ResponseTarget.channel("telegram", echo_input=True),
)
await a.context.deliver_response(req, _make_reply("reply"))
assert phases == [True, False]
# --------------------------------------------------------------------------- #
# HostedRunResult — generic typed envelope #
# --------------------------------------------------------------------------- #
@@ -1522,7 +970,7 @@ class TestBindRequestContext:
stream=True,
attributes={"response_id": "resp_stream"},
)
stream = ch.context.run_stream(req)
stream = await ch.context.run_stream(req)
# As soon as run_stream returns, the binding must already be open
# so any provider work that happens during iteration sees it.
@@ -1574,7 +1022,7 @@ class TestBoundResponseStream:
stream=True,
attributes={"response_id": "resp_get_final"},
)
stream = ch.context.run_stream(req)
stream = await ch.context.run_stream(req)
# Skip iteration and go straight to ``get_final_response``;
# the adapter must drain the inner stream itself and close
# the binding in ``finally``.
@@ -1599,7 +1047,7 @@ class TestBoundResponseStream:
stream=True,
attributes={"response_id": "resp_idem"},
)
stream = ch.context.run_stream(req)
stream = await ch.context.run_stream(req)
async for _u in stream:
pass
# Iteration's finally already closed; an explicit ``aclose``
@@ -1627,7 +1075,7 @@ class TestBoundResponseStream:
stream=True,
attributes={"response_id": "resp_abandon"},
)
stream = ch.context.run_stream(req)
stream = await ch.context.run_stream(req)
await stream.aclose() # type: ignore[attr-defined]
# Binding released without iterating.
@@ -1655,7 +1103,7 @@ class TestBoundResponseStream:
stream=True,
attributes={"response_id": "resp_getattr"},
)
stream = ch.context.run_stream(req)
stream = await ch.context.run_stream(req)
# ``with_result_hook`` is a real method on ``ResponseStream``;
# if forwarding broke this would AttributeError.
try:
@@ -1682,7 +1130,7 @@ class TestBoundResponseStream:
stream=True,
attributes={"response_id": "resp_await"},
)
stream = ch.context.run_stream(req)
stream = await ch.context.run_stream(req)
final = await stream # exercises __await__
assert final.text == "chunk-1chunk-2"
names = [n for n, _ in prov.events]
python/packages/hosting/tests/test_host_disk.py
+25 -222
View File
@@ -1,32 +1,19 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for ``state_dir`` wired through :class:`AgentFrameworkHost`."""
"""Tests for narrowed ``state_dir`` support in :class:`AgentFrameworkHost`."""
from __future__ import annotations
import asyncio
from pathlib import Path
from typing import Any
import pytest
from agent_framework_hosting import (
AgentFrameworkHost,
ChannelContext,
ChannelContribution,
ChannelIdentity,
LinkChallenge,
)
from agent_framework_hosting import AgentFrameworkHost, ChannelContext, ChannelContribution
# Skip the whole module when the optional disk extra isn't installed.
pytest.importorskip("diskcache")
# --------------------------------------------------------------------------- #
# Test helpers #
# --------------------------------------------------------------------------- #
class _AgentStub:
"""Bare-minimum SupportsAgentRun stub for host construction."""
@@ -42,65 +29,22 @@ class _ChannelStub:
return ChannelContribution()
class _NonConfigurableLinker:
async def resolve(self, _identity: ChannelIdentity) -> LinkChallenge:
return LinkChallenge("link")
class _ConfigurableLinker:
def __init__(self) -> None:
self.configured_path: Path | None = None
def configure_link_store_path(self, path: str | Path) -> None:
self.configured_path = Path(path)
async def resolve(self, _identity: ChannelIdentity) -> LinkChallenge:
return LinkChallenge("link")
def _close_host_disk(host: AgentFrameworkHost) -> None:
"""Mirror the lifespan shutdown ordering for tests that simulate restart.
The real shutdown order is ``runner.shutdown()`` → ``sessions_store.close()``;
both release their advisory file locks so a second host can take ownership.
"""
runner = host._durable_task_runner
try:
asyncio.get_event_loop().run_until_complete(runner.shutdown(timeout=1.0))
except RuntimeError:
# No running loop; spin up a throw-away one.
asyncio.run(runner.shutdown(timeout=1.0))
"""Release any session-alias store held by ``host``."""
if host._sessions_store is not None:
host._sessions_store.close()
# --------------------------------------------------------------------------- #
# state_dir=None preserves the in-memory contract #
# --------------------------------------------------------------------------- #
def test_state_dir_none_keeps_plain_dicts(tmp_path: Path) -> None:
"""No store, no sessions persistence, no files written."""
def test_state_dir_none_keeps_plain_alias_dict(tmp_path: Path) -> None:
"""No store, no alias persistence, no files written."""
host = AgentFrameworkHost(target=_AgentStub(), channels=[_ChannelStub()])
try:
assert host._sessions_store is None
assert isinstance(host._session_aliases, dict)
assert isinstance(host._active, dict)
assert isinstance(host._identities, dict)
# No accidental disk writes anywhere under tmp_path.
assert list(tmp_path.iterdir()) == []
finally:
# Nothing to close.
pass
assert host._sessions_store is None
assert isinstance(host._session_aliases, dict)
assert list(tmp_path.iterdir()) == []
# --------------------------------------------------------------------------- #
# Single string state_dir creates default subfolders #
# --------------------------------------------------------------------------- #
def test_string_state_dir_creates_subfolders(tmp_path: Path) -> None:
"""Passing a single path expands to ``runner/`` and ``sessions/``."""
def test_string_state_dir_creates_sessions_subfolder_only(tmp_path: Path) -> None:
"""Passing a single path expands to ``sessions/`` plus lazy checkpoint path."""
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub()],
@@ -108,100 +52,42 @@ def test_string_state_dir_creates_subfolders(tmp_path: Path) -> None:
)
try:
assert host._sessions_store is not None
assert (tmp_path / "runner").is_dir()
assert (tmp_path / "sessions").is_dir()
assert not (tmp_path / "runner").exists()
assert not (tmp_path / "links").exists()
# Checkpoint path is derived but not created for agent targets.
assert not (tmp_path / "checkpoints").exists()
finally:
_close_host_disk(host)
# --------------------------------------------------------------------------- #
# Per-component override via HostStatePaths-shaped dict #
# --------------------------------------------------------------------------- #
def test_per_component_paths(tmp_path: Path) -> None:
"""Dict form lets the caller route components to different roots."""
runner_dir = tmp_path / "tasks"
def test_per_component_session_path(tmp_path: Path) -> None:
"""Dict form lets callers route session aliases to a specific root."""
sessions_dir = tmp_path / "state"
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub()],
state_dir={"runner": runner_dir, "sessions": sessions_dir},
state_dir={"sessions": sessions_dir},
)
try:
assert runner_dir.is_dir()
assert sessions_dir.is_dir()
# Default subfolders should NOT exist when the caller provides
# explicit overrides.
assert not (tmp_path / "runner").is_dir() or runner_dir == (tmp_path / "runner")
assert not (tmp_path / "sessions").is_dir() or sessions_dir == (tmp_path / "sessions")
assert host._sessions_store is not None
assert host._checkpoint_location is None
finally:
_close_host_disk(host)
def test_unknown_component_key_raises(tmp_path: Path) -> None:
"""Misspelled keys should fail loudly so the user catches typos."""
@pytest.mark.parametrize("key", ["runner", "links", "active", "identities"])
def test_removed_state_dir_component_keys_raise(tmp_path: Path, key: str) -> None:
"""Obsolete follow-up components should fail loudly instead of becoming no-ops."""
with pytest.raises(ValueError, match="unknown"):
AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub()],
state_dir={"runnerr": tmp_path / "x"}, # type: ignore[dict-item]
state_dir={key: tmp_path / key}, # type: ignore[dict-item]
)
def test_links_state_path_configures_compatible_identity_linker(tmp_path: Path) -> None:
"""``state_dir['links']`` is offered to linkers that accept host-owned persistence."""
linker = _ConfigurableLinker()
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub()],
identity_linker=linker,
state_dir=tmp_path,
)
try:
assert linker.configured_path == tmp_path / "links"
finally:
_close_host_disk(host)
def test_explicit_links_state_path_without_linker_warns(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
"""Explicit ``links`` path with no linker is almost certainly dead config."""
with caplog.at_level("WARNING", logger="agent_framework.hosting"):
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub()],
state_dir={"links": tmp_path / "links"},
)
try:
assert any(
"state_dir['links']" in rec.message and "no identity_linker" in rec.message for rec in caplog.records
)
finally:
_close_host_disk(host)
def test_links_state_path_with_nonconfigurable_linker_warns(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
"""A linker that owns its persistence directly gets a clear warning."""
with caplog.at_level("WARNING", logger="agent_framework.hosting"):
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub()],
identity_linker=_NonConfigurableLinker(),
state_dir={"links": tmp_path / "links"},
)
try:
assert any(
"state_dir['links']" in rec.message and "SupportsLinkStorePath" in rec.message for rec in caplog.records
)
finally:
_close_host_disk(host)
# --------------------------------------------------------------------------- #
# Session bookkeeping survives a host restart #
# --------------------------------------------------------------------------- #
def test_session_aliases_survive_restart(tmp_path: Path) -> None:
"""Aliases written on host #1 must be visible to host #2."""
state_dir = tmp_path / "state"
@@ -219,84 +105,6 @@ def test_session_aliases_survive_restart(tmp_path: Path) -> None:
_close_host_disk(host2)
def test_active_channel_survives_restart(tmp_path: Path) -> None:
"""``_active`` must round-trip through the store."""
state_dir = tmp_path / "state"
host1 = AgentFrameworkHost(target=_AgentStub(), channels=[_ChannelStub()], state_dir=state_dir)
host1._active["user-1"] = "telegram"
host1._active["user-2"] = "responses"
_close_host_disk(host1)
host2 = AgentFrameworkHost(target=_AgentStub(), channels=[_ChannelStub()], state_dir=state_dir)
try:
assert host2._active["user-1"] == "telegram"
assert host2._active["user-2"] == "responses"
finally:
_close_host_disk(host2)
def test_identities_nested_mutation_survives_restart(tmp_path: Path) -> None:
"""Setting ``self._identities[ik][channel] = identity`` must persist.
This exercises the proxy-inner-dict ``__setitem__`` write-through path,
not just the outer-key replacement path.
"""
state_dir = tmp_path / "state"
host1 = AgentFrameworkHost(target=_AgentStub(), channels=[_ChannelStub()], state_dir=state_dir)
ident_tg = ChannelIdentity("telegram", "tg-123", {"username": "alice"})
ident_rsp = ChannelIdentity("responses", "rsp-456")
# Mirrors the host-internal path in ``_register_identity``.
host1._identities.setdefault("user-1", {})["telegram"] = ident_tg
host1._identities.setdefault("user-1", {})["responses"] = ident_rsp
host1._identities.setdefault("user-2", {})["telegram"] = ChannelIdentity("telegram", "tg-789")
_close_host_disk(host1)
host2 = AgentFrameworkHost(target=_AgentStub(), channels=[_ChannelStub()], state_dir=state_dir)
try:
u1 = host2._identities["user-1"]
assert set(u1.keys()) == {"telegram", "responses"}
assert u1["telegram"].native_id == "tg-123"
assert u1["telegram"].attributes["username"] == "alice"
assert u1["responses"].native_id == "rsp-456"
assert host2._identities["user-2"]["telegram"].native_id == "tg-789"
finally:
_close_host_disk(host2)
# --------------------------------------------------------------------------- #
# Explicit durable_task_runner + state_dir['runner'] warns #
# --------------------------------------------------------------------------- #
def test_explicit_runner_with_runner_state_warns(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
"""Caller-owned runner + state_dir['runner'] → ignore + warn."""
from agent_framework_hosting import InProcessTaskRunner
user_runner = InProcessTaskRunner()
try:
with caplog.at_level("WARNING"):
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub()],
durable_task_runner=user_runner,
allow_in_process_runner=True,
state_dir={"runner": tmp_path / "runner"},
)
assert any("state_dir['runner']" in rec.message for rec in caplog.records)
# Sessions store wasn't requested, so still None.
assert host._sessions_store is None
finally:
# user_runner has no disk state, so nothing else to clean up.
pass
# --------------------------------------------------------------------------- #
# Workflow checkpoint integration #
# --------------------------------------------------------------------------- #
def _build_simple_workflow() -> Any:
"""Build a no-op workflow for checkpoint-wiring tests."""
from tests._workflow_fixtures import build_upper_workflow
@@ -313,7 +121,6 @@ def test_single_path_state_dir_wires_workflow_checkpoints(tmp_path: Path) -> Non
state_dir=tmp_path,
)
try:
# Checkpoint location is derived from the single state_dir.
assert host._checkpoint_location == tmp_path / "checkpoints"
finally:
_close_host_disk(host)
@@ -330,7 +137,6 @@ def test_mapping_state_dir_checkpoints_key_wires_workflow_checkpoints(tmp_path:
)
try:
assert host._checkpoint_location == ckpt_dir
# No diskcache components were requested.
assert host._sessions_store is None
finally:
_close_host_disk(host)
@@ -342,9 +148,7 @@ def test_mapping_state_dir_omits_checkpoints_for_workflow(tmp_path: Path) -> Non
host = AgentFrameworkHost(
target=workflow,
channels=[_ChannelStub()],
# No 'checkpoints' key → no checkpoint persistence even though
# other components are persisted.
state_dir={"runner": tmp_path / "r", "sessions": tmp_path / "s"},
state_dir={"sessions": tmp_path / "s"},
)
try:
assert host._checkpoint_location is None
@@ -381,7 +185,6 @@ def test_state_dir_checkpoints_for_agent_target_silent_for_single_path(tmp_path:
)
try:
assert host._checkpoint_location is None
# ``checkpoints/`` subfolder is not eagerly created (no consumer).
assert not (tmp_path / "checkpoints").exists()
finally:
_close_host_disk(host)
@@ -390,7 +193,7 @@ def test_state_dir_checkpoints_for_agent_target_silent_for_single_path(tmp_path:
def test_state_dir_checkpoints_for_agent_target_warns_when_explicit(
tmp_path: Path, caplog: pytest.LogCaptureFixture
) -> None:
"""Mapping form with ``checkpoints`` + agent target → warn (dead config)."""
"""Mapping form with ``checkpoints`` + agent target → warn."""
with caplog.at_level("WARNING", logger="agent_framework.hosting"):
host = AgentFrameworkHost(
target=_AgentStub(),
python/packages/hosting/tests/test_isolation.py
+27 -6
View File
@@ -16,6 +16,7 @@ from __future__ import annotations
import asyncio
import pytest
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.routing import BaseRoute, Route
@@ -168,7 +169,22 @@ def _make_host_with_probe() -> tuple[object, _IsolationProbeChannel]:
class TestIsolationMiddlewareEndToEnd:
def test_both_headers_lifted_into_contextvar(self) -> None:
def test_headers_ignored_outside_foundry_environment(self) -> None:
host, probe = _make_host_with_probe()
with TestClient(host.app) as client: # type: ignore[attr-defined]
r = client.get(
"/probe",
headers={
ISOLATION_HEADER_USER: "alice-uid",
ISOLATION_HEADER_CHAT: "general-cid",
},
)
assert r.status_code == 200
assert r.json() == {"user": None, "chat": None, "_present": False}
assert probe.captured == [None]
def test_both_headers_lifted_into_contextvar(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("FOUNDRY_HOSTING_ENVIRONMENT", "1")
host, probe = _make_host_with_probe()
with TestClient(host.app) as client: # type: ignore[attr-defined]
r = client.get(
@@ -186,15 +202,17 @@ class TestIsolationMiddlewareEndToEnd:
assert captured.user_key == "alice-uid"
assert captured.chat_key == "general-cid"
def test_only_user_header_lifted(self) -> None:
def test_only_user_header_lifted(self, monkeypatch: pytest.MonkeyPatch) -> None:
"""One-header-only branch: the middleware still binds (chat=None)."""
monkeypatch.setenv("FOUNDRY_HOSTING_ENVIRONMENT", "1")
host, probe = _make_host_with_probe()
with TestClient(host.app) as client: # type: ignore[attr-defined]
r = client.get("/probe", headers={ISOLATION_HEADER_USER: "alice-uid"})
assert r.status_code == 200
assert r.json() == {"user": "alice-uid", "chat": None}
def test_only_chat_header_lifted(self) -> None:
def test_only_chat_header_lifted(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("FOUNDRY_HOSTING_ENVIRONMENT", "1")
host, probe = _make_host_with_probe()
with TestClient(host.app) as client: # type: ignore[attr-defined]
r = client.get("/probe", headers={ISOLATION_HEADER_CHAT: "general-cid"})
@@ -213,9 +231,10 @@ class TestIsolationMiddlewareEndToEnd:
assert r.json() == {"user": None, "chat": None, "_present": False}
assert probe.captured == [None]
def test_empty_header_value_treated_as_absent(self) -> None:
def test_empty_header_value_treated_as_absent(self, monkeypatch: pytest.MonkeyPatch) -> None:
"""A header that's present but empty must not bind an empty key —
``IsolationContext`` rejects empty strings on the read side."""
monkeypatch.setenv("FOUNDRY_HOSTING_ENVIRONMENT", "1")
host, probe = _make_host_with_probe()
with TestClient(host.app) as client: # type: ignore[attr-defined]
r = client.get(
@@ -229,10 +248,11 @@ class TestIsolationMiddlewareEndToEnd:
# Empty user header decodes to None; chat key stays bound.
assert r.json() == {"user": None, "chat": "general-cid"}
def test_contextvar_resets_after_request(self) -> None:
def test_contextvar_resets_after_request(self, monkeypatch: pytest.MonkeyPatch) -> None:
"""The middleware must call ``reset_current_isolation_keys`` in
a ``finally`` so per-request state never leaks across requests
or back into the calling thread's context."""
monkeypatch.setenv("FOUNDRY_HOSTING_ENVIRONMENT", "1")
host, probe = _make_host_with_probe()
with TestClient(host.app) as client: # type: ignore[attr-defined]
r1 = client.get("/probe", headers={ISOLATION_HEADER_USER: "alice-uid"})
@@ -245,9 +265,10 @@ class TestIsolationMiddlewareEndToEnd:
r2 = client.get("/probe")
assert r2.json() == {"user": None, "chat": None, "_present": False}
def test_concurrent_requests_get_isolated_contextvars(self) -> None:
def test_concurrent_requests_get_isolated_contextvars(self, monkeypatch: pytest.MonkeyPatch) -> None:
"""Different requests run in different async contexts; binding
from request A must NOT leak into a concurrent request B."""
monkeypatch.setenv("FOUNDRY_HOSTING_ENVIRONMENT", "1")
host, probe = _make_host_with_probe()
async def _drive() -> None:
python/packages/hosting/tests/test_runner.py
-333
View File
@@ -1,333 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for :class:`InProcessTaskRunner` and runtime-mode auto-detection."""
from __future__ import annotations
import asyncio
from collections.abc import Mapping
from typing import Any
import pytest
from agent_framework_hosting import (
AgentFrameworkHost,
ChannelContext,
ChannelContribution,
DurableTaskPayloadMode,
InProcessTaskRunner,
RetryPolicy,
TaskHandle,
)
from agent_framework_hosting._host import _detect_runtime_mode
# --------------------------------------------------------------------------- #
# Test helpers #
# --------------------------------------------------------------------------- #
class _AgentStub:
"""Bare-minimum SupportsAgentRun stub for host construction."""
async def run(self, *_args: Any, **_kwargs: Any) -> None: # pragma: no cover - unused
return None
class _ChannelStub:
name = "stub"
path = "/stub"
def contribute(self, _context: ChannelContext) -> ChannelContribution:
return ChannelContribution()
# --------------------------------------------------------------------------- #
# Runtime-mode auto-detection #
# --------------------------------------------------------------------------- #
class TestRuntimeModeDetection:
"""``_detect_runtime_mode`` is pure: tests pass a synthetic env so
they never depend on the test runner's environment. Auto-detected
mode + matched marker drive the per-host startup banner so operators
can confirm the host is running in the expected shape."""
def test_no_markers_defaults_to_long_running(self) -> None:
mode, marker = _detect_runtime_mode(env={})
assert mode == "long_running"
assert marker is None
def test_foundry_marker_selects_ephemeral(self) -> None:
mode, marker = _detect_runtime_mode(env={"FOUNDRY_HOSTING_ENVIRONMENT": "production"})
assert mode == "ephemeral"
assert marker == "FOUNDRY_HOSTING_ENVIRONMENT"
def test_azure_functions_marker_selects_ephemeral(self) -> None:
mode, marker = _detect_runtime_mode(env={"AZURE_FUNCTIONS_ENVIRONMENT": "Development"})
assert mode == "ephemeral"
assert marker == "AZURE_FUNCTIONS_ENVIRONMENT"
def test_lambda_marker_selects_ephemeral(self) -> None:
mode, marker = _detect_runtime_mode(env={"AWS_LAMBDA_FUNCTION_NAME": "my-fn"})
assert mode == "ephemeral"
assert marker == "AWS_LAMBDA_FUNCTION_NAME"
def test_empty_marker_value_ignored(self) -> None:
# Empty-string env var should not count as "set" — Foundry's
# template uses unset-or-empty as "not deployed".
mode, marker = _detect_runtime_mode(env={"FOUNDRY_HOSTING_ENVIRONMENT": ""})
assert mode == "long_running"
assert marker is None
class TestHostRuntimeMode:
"""``runtime_mode`` ctor argument overrides auto-detect; ``None``
triggers auto-detect. The detected mode is exposed via the
``runtime_mode`` property for operator inspection (and is logged at
startup via ``_log_startup``)."""
def test_explicit_long_running(self) -> None:
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub()],
runtime_mode="long_running",
)
assert host.runtime_mode == "long_running"
def test_explicit_ephemeral_with_default_runner_raises(self) -> None:
# Default runner is in-process and not durable. Ephemeral
# deployments would silently lose pushes on scale-to-zero, so
# the host refuses the combination at construction unless the
# operator opts in explicitly via ``allow_in_process_runner``.
with pytest.raises(RuntimeError, match="ephemeral"):
AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub()],
runtime_mode="ephemeral",
)
def test_explicit_ephemeral_with_in_process_opt_in_warns(self, caplog: pytest.LogCaptureFixture) -> None:
# The opt-in escape hatch keeps the old warn-and-proceed
# behaviour for local-dev / smoke-test scenarios that genuinely
# want ephemeral runtime semantics without a real durable
# backend.
with caplog.at_level("WARNING", logger="agent_framework.hosting"):
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub()],
runtime_mode="ephemeral",
allow_in_process_runner=True,
)
assert host.runtime_mode == "ephemeral"
assert any("ephemeral" in r.getMessage() and "InProcessTaskRunner" in r.getMessage() for r in caplog.records)
def test_explicit_ephemeral_with_supplied_runner_does_not_warn(self, caplog: pytest.LogCaptureFixture) -> None:
runner = InProcessTaskRunner()
with caplog.at_level("WARNING", logger="agent_framework.hosting"):
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub()],
runtime_mode="ephemeral",
durable_task_runner=runner,
)
# No warning — operator opted into a specific runner.
assert host.runtime_mode == "ephemeral"
assert host.durable_task_runner is runner
assert not any("ephemeral" in r.getMessage() for r in caplog.records)
def test_auto_detect_ephemeral_raises_without_opt_in(self, monkeypatch: pytest.MonkeyPatch) -> None:
# Auto-detected ephemeral flows through the same strict gate.
monkeypatch.setenv("FOUNDRY_HOSTING_ENVIRONMENT", "production")
with pytest.raises(RuntimeError, match="ephemeral"):
AgentFrameworkHost(target=_AgentStub(), channels=[_ChannelStub()])
def test_auto_detect_ephemeral_with_opt_in_proceeds(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("FOUNDRY_HOSTING_ENVIRONMENT", "production")
host = AgentFrameworkHost(
target=_AgentStub(),
channels=[_ChannelStub()],
allow_in_process_runner=True,
)
assert host.runtime_mode == "ephemeral"
def test_default_runner_is_in_process_task_runner(self) -> None:
host = AgentFrameworkHost(target=_AgentStub(), channels=[_ChannelStub()])
assert isinstance(host.durable_task_runner, InProcessTaskRunner)
# --------------------------------------------------------------------------- #
# InProcessTaskRunner #
# --------------------------------------------------------------------------- #
class TestInProcessTaskRunner:
async def test_schedule_runs_handler_and_records_succeeded(self) -> None:
runner = InProcessTaskRunner()
seen: list[Mapping[str, Any]] = []
async def handler(payload: Mapping[str, Any]) -> None:
seen.append(payload)
runner.register("ping", handler)
handle = await runner.schedule("ping", {"x": 1})
# ``schedule`` returns immediately; the task runs on the loop.
# Drain explicitly via ``shutdown`` to flush in-flight work,
# then assert.
await _drain(runner, handle)
assert seen == [{"x": 1}]
assert await runner.get(handle) == "succeeded"
async def test_unknown_handler_raises_keyerror(self) -> None:
runner = InProcessTaskRunner()
with pytest.raises(KeyError):
await runner.schedule("missing", {})
async def test_register_after_start_raises(self) -> None:
runner = InProcessTaskRunner()
async def noop(_p: Mapping[str, Any]) -> None:
return None
runner.register("x", noop)
handle = await runner.schedule("x", {})
await _drain(runner, handle)
# Re-registering after the runner has started scheduling is
# rejected so in-flight tasks can't have their handler swapped
# out from under them.
with pytest.raises(RuntimeError, match="register"):
runner.register("y", noop)
async def test_handler_retried_then_succeeds(self) -> None:
runner = InProcessTaskRunner()
attempts = {"n": 0}
async def flaky(_p: Mapping[str, Any]) -> None:
attempts["n"] += 1
if attempts["n"] < 3:
raise RuntimeError(f"attempt {attempts['n']}")
runner.register("flaky", flaky)
# Tight retry policy so the test doesn't sleep visibly.
policy = RetryPolicy(max_attempts=5, initial_backoff_seconds=0.001, max_backoff_seconds=0.005)
handle = await runner.schedule("flaky", {}, retry_policy=policy)
await _drain(runner, handle)
assert attempts["n"] == 3
assert await runner.get(handle) == "succeeded"
async def test_handler_failure_records_failed_after_max_attempts(self) -> None:
runner = InProcessTaskRunner()
async def always_fails(_p: Mapping[str, Any]) -> None:
raise RuntimeError("nope")
runner.register("doomed", always_fails)
policy = RetryPolicy(max_attempts=2, initial_backoff_seconds=0.001)
handle = await runner.schedule("doomed", {}, retry_policy=policy)
await _drain(runner, handle)
assert await runner.get(handle) == "failed"
async def test_shutdown_cancels_pending_tasks(self) -> None:
runner = InProcessTaskRunner()
started = asyncio.Event()
cancelled = asyncio.Event()
async def long_running(_p: Mapping[str, Any]) -> None:
started.set()
try:
# Sleep longer than the test wait so shutdown can cancel.
await asyncio.sleep(5)
except asyncio.CancelledError:
cancelled.set()
raise
runner.register("long", long_running)
handle = await runner.schedule("long", {})
await asyncio.wait_for(started.wait(), timeout=1.0)
await runner.shutdown(timeout=1.0)
assert cancelled.is_set()
assert await runner.get(handle) == "cancelled"
async def test_shutdown_grace_drain_does_not_cancel_finishing_tasks(self) -> None:
"""A short-lived task that completes within the grace window
must NOT receive a cancellation. The grace-period drain is the
graceful-shutdown contract — channels with goodbye-message
flushes rely on it."""
runner = InProcessTaskRunner()
cancelled = asyncio.Event()
completed = asyncio.Event()
async def quick(_p: Mapping[str, Any]) -> None:
try:
await asyncio.sleep(0.05)
except asyncio.CancelledError:
cancelled.set()
raise
completed.set()
runner.register("quick", quick)
handle = await runner.schedule("quick", {})
# Shutdown with a generous grace window relative to the task duration.
await runner.shutdown(timeout=1.0)
assert completed.is_set()
assert not cancelled.is_set()
assert await runner.get(handle) == "succeeded"
async def test_get_returns_none_for_unknown_handle(self) -> None:
runner = InProcessTaskRunner()
handle = TaskHandle(task_id="never-scheduled", name="x")
assert await runner.get(handle) is None
async def test_terminal_cache_evicts_oldest(self) -> None:
# Cache size of 2: drain three tasks in sequence, the first
# should age out by the time the third's terminal lands.
runner = InProcessTaskRunner(terminal_cache_size=2)
async def noop(_p: Mapping[str, Any]) -> None:
return None
runner.register("noop", noop)
h1 = await runner.schedule("noop", {})
await _drain(runner, h1)
h2 = await runner.schedule("noop", {})
await _drain(runner, h2)
h3 = await runner.schedule("noop", {})
await _drain(runner, h3)
# Oldest handle's terminal status should be evicted by now.
assert await runner.get(h1) is None
assert await runner.get(h2) == "succeeded"
assert await runner.get(h3) == "succeeded"
async def test_shutdown_is_safe_when_no_tasks_pending(self) -> None:
runner = InProcessTaskRunner()
# No-op shouldn't raise.
await runner.shutdown()
def test_payload_mode_defaults_to_object(self) -> None:
# The in-process runner passes live Python references through
# the payload — the host wires this attribute into its codec
# validator at startup. Durable adapters that persist payloads
# must override this to ``JSON`` so the host refuses to ship
# un-serialisable references.
runner = InProcessTaskRunner()
assert runner.payload_mode == DurableTaskPayloadMode.OBJECT
# --------------------------------------------------------------------------- #
# Helpers #
# --------------------------------------------------------------------------- #
async def _drain(runner: InProcessTaskRunner, handle: TaskHandle, *, timeout: float = 1.0) -> None:
"""Wait for ``handle`` to reach a terminal state.
Polls ``get`` rather than reaching into runner internals so we exercise the
public surface from the test side too.
"""
deadline = asyncio.get_event_loop().time() + timeout
while True:
status = await runner.get(handle)
if status in ("succeeded", "failed", "cancelled"):
return
if asyncio.get_event_loop().time() > deadline:
raise AssertionError(f"task {handle.task_id} did not reach terminal in {timeout}s; status={status}")
await asyncio.sleep(0.01)
python/packages/hosting/tests/test_runner_disk.py
-278
View File
@@ -1,278 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for :class:`InProcessTaskRunner` disk persistence (``state_dir``)."""
from __future__ import annotations
import asyncio
from collections.abc import Mapping
from pathlib import Path
from typing import Any
import pytest
from agent_framework_hosting import (
InProcessTaskRunner,
PushPayloadNotPicklable,
RetryPolicy,
)
# Skip the whole module if the optional diskcache dependency isn't installed.
pytest.importorskip("diskcache")
# --------------------------------------------------------------------------- #
# state_dir=None preserves today's purely in-memory contract #
# --------------------------------------------------------------------------- #
async def test_state_dir_none_is_pure_memory(tmp_path: Path) -> None:
"""No directory creation / no lock file when state_dir is omitted."""
runner = InProcessTaskRunner()
calls: list[Mapping[str, Any]] = []
async def handler(payload: Mapping[str, Any]) -> None:
calls.append(payload)
runner.register("echo", handler)
handle = await runner.schedule("echo", {"k": "v"})
# Wait for completion.
for _ in range(50):
if (await runner.get(handle)) == "succeeded":
break
await asyncio.sleep(0.01)
assert calls == [{"k": "v"}]
assert await runner.get(handle) == "succeeded"
# Confirm we didn't accidentally write to disk.
assert not (tmp_path / ".lock").exists()
await runner.shutdown()
# --------------------------------------------------------------------------- #
# Lock contention — two runners on the same dir refuse to coexist #
# --------------------------------------------------------------------------- #
async def test_two_runners_one_state_dir_raise(tmp_path: Path) -> None:
"""Second runner construction must fail loudly, not silently corrupt."""
state_dir = tmp_path / "runner"
first = InProcessTaskRunner(state_dir=state_dir)
try:
with pytest.raises(RuntimeError, match="state lock"):
InProcessTaskRunner(state_dir=state_dir)
finally:
await first.shutdown()
# --------------------------------------------------------------------------- #
# Pickle failure raises eagerly, never silently downgrades #
# --------------------------------------------------------------------------- #
async def test_unpickleable_payload_raises(tmp_path: Path) -> None:
"""Schedule must refuse payloads that can't survive a restart."""
runner = InProcessTaskRunner(state_dir=tmp_path / "runner")
async def handler(_: Mapping[str, Any]) -> None: ...
runner.register("echo", handler)
# Local lambdas / closures are the canonical unpicklable values.
with pytest.raises(PushPayloadNotPicklable):
await runner.schedule("echo", {"callback": lambda: None})
await runner.shutdown()
# --------------------------------------------------------------------------- #
# Resume — pending records replay on next process #
# --------------------------------------------------------------------------- #
async def test_pending_record_replays_on_resume(tmp_path: Path) -> None:
"""Simulate a crash: first runner schedules but never starts running."""
state_dir = tmp_path / "runner"
# Process 1 — schedule a task, then "die" before the asyncio loop runs it.
runner1 = InProcessTaskRunner(state_dir=state_dir)
blocked = asyncio.Event()
async def slow(_: Mapping[str, Any]) -> None:
# Sleep so the task is observably still in flight when we shutdown.
await blocked.wait()
runner1.register("slow", slow)
handle = await runner1.schedule("slow", {"work": 1})
# Force a hard shutdown — leaves the in-flight task in 'pending' on disk.
await runner1.shutdown(timeout=0.1)
# Process 2 — fresh runner against same state_dir, register the handler,
# call resume. We expect the persisted record to be re-scheduled.
runner2 = InProcessTaskRunner(state_dir=state_dir)
seen: list[Mapping[str, Any]] = []
async def slow_resumed(payload: Mapping[str, Any]) -> None:
seen.append(dict(payload))
runner2.register("slow", slow_resumed)
replayed = await runner2.resume()
assert replayed == 1
# Give the resumed task time to run.
for _ in range(50):
if seen:
break
await asyncio.sleep(0.01)
assert seen == [{"work": 1}]
# Status is observable via the original handle.
assert await runner2.get(handle) == "succeeded"
await runner2.shutdown()
# --------------------------------------------------------------------------- #
# echo_done cursor survives restart #
# --------------------------------------------------------------------------- #
async def test_payload_mutation_survives_restart(tmp_path: Path) -> None:
"""Handler-side payload mutations (echo_done) round-trip through disk."""
state_dir = tmp_path / "runner"
runner1 = InProcessTaskRunner(state_dir=state_dir)
# Handler sets echo_done and then blocks forever (simulating mid-flight crash).
handler_progress = asyncio.Event()
async def half_done(payload: Mapping[str, Any]) -> None:
# Mutate the payload to mark first phase complete.
payload["echo_done"] = True # type: ignore[index]
handler_progress.set()
# Sleep indefinitely so the asyncio task is still running at shutdown.
await asyncio.Event().wait()
runner1.register("two_phase", half_done)
handle = await runner1.schedule("two_phase", {"echo_done": False, "k": "v"})
await handler_progress.wait()
await runner1.shutdown(timeout=0.1)
# Process 2 — replay; the handler now sees echo_done=True from disk.
runner2 = InProcessTaskRunner(state_dir=state_dir)
observed: list[bool] = []
async def two_phase_resumed(payload: Mapping[str, Any]) -> None:
observed.append(bool(payload.get("echo_done")))
runner2.register("two_phase", two_phase_resumed)
await runner2.resume()
for _ in range(50):
if observed:
break
await asyncio.sleep(0.01)
assert observed == [True]
# And the resumed task ran to completion.
assert await runner2.get(handle) == "succeeded"
await runner2.shutdown()
# --------------------------------------------------------------------------- #
# Resume gracefully handles missing handler / corrupt entries #
# --------------------------------------------------------------------------- #
async def test_resume_with_missing_handler_marks_failed(tmp_path: Path) -> None:
"""A persisted record whose handler is no longer registered is marked failed."""
state_dir = tmp_path / "runner"
runner1 = InProcessTaskRunner(state_dir=state_dir)
async def will_be_removed(_: Mapping[str, Any]) -> None:
await asyncio.Event().wait()
runner1.register("ghost", will_be_removed)
handle = await runner1.schedule("ghost", {})
await runner1.shutdown(timeout=0.1)
# Process 2 — never registers "ghost".
runner2 = InProcessTaskRunner(state_dir=state_dir)
replayed = await runner2.resume()
assert replayed == 0
# The record is moved to terminal 'failed'.
assert await runner2.get(handle) == "failed"
await runner2.shutdown()
async def test_resume_quarantines_corrupt_entries(tmp_path: Path) -> None:
"""A non-dict on-disk entry must be quarantined, not crash resume."""
import diskcache # noqa: PLC0415 - lazy import to keep module-import cheap
state_dir = tmp_path / "runner"
state_dir.mkdir(parents=True, exist_ok=True)
# Pre-populate the cache with a junk entry.
cache = diskcache.Cache(str(state_dir))
cache.set("bad-task-id", "this is not a dict")
cache.close()
runner = InProcessTaskRunner(state_dir=state_dir)
# resume() must not raise even with a corrupt entry on disk.
replayed = await runner.resume()
assert replayed == 0
await runner.shutdown()
# The corrupt entry should have been removed.
cache2 = diskcache.Cache(str(state_dir))
assert "bad-task-id" not in cache2
cache2.close()
# --------------------------------------------------------------------------- #
# Retry attempt counter persists across resume #
# --------------------------------------------------------------------------- #
async def test_attempt_counter_persists_across_resume(tmp_path: Path) -> None:
"""A handler that crashes mid-attempt resumes with the consumed budget."""
state_dir = tmp_path / "runner"
policy = RetryPolicy(max_attempts=3, initial_backoff_seconds=0.01, backoff_multiplier=1.0)
# Process 1 — schedule, fail once, shutdown before retry settles.
runner1 = InProcessTaskRunner(state_dir=state_dir, default_retry_policy=policy)
attempts_seen_p1 = 0
async def flaky(_: Mapping[str, Any]) -> None:
nonlocal attempts_seen_p1
attempts_seen_p1 += 1
raise RuntimeError("boom-1")
runner1.register("flaky", flaky)
handle = await runner1.schedule("flaky", {})
# Let it attempt twice (waste 2 of 3 budgeted retries), then crash-shutdown.
for _ in range(50):
if attempts_seen_p1 >= 2:
break
await asyncio.sleep(0.01)
await runner1.shutdown(timeout=0.05)
# Process 2 — resume; only 1 attempt left in the budget. Confirm we don't
# re-grant the full retry budget.
runner2 = InProcessTaskRunner(state_dir=state_dir, default_retry_policy=policy)
attempts_seen_p2 = 0
async def flaky_resumed(_: Mapping[str, Any]) -> None:
nonlocal attempts_seen_p2
attempts_seen_p2 += 1
raise RuntimeError("boom-2")
runner2.register("flaky", flaky_resumed)
await runner2.resume()
# Wait for the resumed task to consume its remaining attempts and fail terminally.
for _ in range(100):
if (await runner2.get(handle)) == "failed":
break
await asyncio.sleep(0.01)
assert await runner2.get(handle) == "failed"
# Original consumed 2 attempts; we should have allowed at most max_attempts-2=1
# more in process 2.
assert attempts_seen_p2 <= 1
await runner2.shutdown()
python/packages/hosting/tests/test_types.py
-288
View File
@@ -4,63 +4,13 @@
from __future__ import annotations
from typing import Any
from agent_framework_hosting import (
ChannelContribution,
ChannelIdentity,
ChannelRequest,
ChannelResponseContext,
ChannelSession,
DurableTaskPayloadMode,
HostedRunResult,
ResponseTarget,
ResponseTargetKind,
apply_channel_response_hook,
apply_run_hook,
)
class TestResponseTarget:
def test_originating_default_singleton(self) -> None:
target = ResponseTarget.originating # type: ignore[attr-defined]
assert target.kind is ResponseTargetKind.ORIGINATING
assert target.targets == ()
def test_active_singleton(self) -> None:
target = ResponseTarget.active # type: ignore[attr-defined]
assert target.kind is ResponseTargetKind.ACTIVE
assert target.targets == ()
def test_all_linked_singleton(self) -> None:
target = ResponseTarget.all_linked # type: ignore[attr-defined]
assert target.kind is ResponseTargetKind.ALL_LINKED
def test_none_singleton(self) -> None:
target = ResponseTarget.none # type: ignore[attr-defined]
assert target.kind is ResponseTargetKind.NONE
def test_channel_builder_single(self) -> None:
target = ResponseTarget.channel("teams")
assert target.kind is ResponseTargetKind.CHANNELS
assert target.targets == ("teams",)
def test_channels_builder_list(self) -> None:
target = ResponseTarget.channels(["teams", "telegram", "originating"])
assert target.kind is ResponseTargetKind.CHANNELS
assert target.targets == ("teams", "telegram", "originating")
def test_channels_builder_accepts_tuple(self) -> None:
target = ResponseTarget.channels(("a", "b"))
assert target.targets == ("a", "b")
def test_target_is_hashable(self) -> None:
# Plain class — hashing falls back to identity, which is fine here:
# the two keys below are different instances (singleton vs builder).
d = {ResponseTarget.originating: 1, ResponseTarget.channel("t"): 2} # type: ignore[attr-defined]
assert len(d) == 2
class TestChannelRequest:
def test_required_fields_only(self) -> None:
req = ChannelRequest(channel="responses", operation="message.create", input="hi")
@@ -74,17 +24,6 @@ class TestChannelRequest:
assert req.attributes == {}
assert req.stream is False
assert req.identity is None
# Default response target is the originating singleton.
assert req.response_target.kind is ResponseTargetKind.ORIGINATING
def test_default_response_target_is_originating_singleton(self) -> None:
# Every new request shares the module-level ``originating`` singleton
# by default — instances are intended to be treated as immutable, so
# sharing is safe and avoids per-request allocation.
a = ChannelRequest(channel="a", operation="op", input="x")
b = ChannelRequest(channel="b", operation="op", input="y")
assert a.response_target is ResponseTarget.originating # type: ignore[attr-defined]
assert a.response_target is b.response_target
def test_with_session_and_identity(self) -> None:
req = ChannelRequest(
@@ -93,14 +32,12 @@ class TestChannelRequest:
input="hi",
session=ChannelSession(isolation_key="user:42"),
identity=ChannelIdentity(channel="telegram", native_id="42"),
response_target=ResponseTarget.active, # type: ignore[attr-defined]
)
assert req.session is not None
assert req.session.isolation_key == "user:42"
assert req.identity is not None
assert req.identity.channel == "telegram"
assert req.identity.native_id == "42"
assert req.response_target.kind is ResponseTargetKind.ACTIVE
class TestChannelIdentity:
@@ -111,228 +48,3 @@ class TestChannelIdentity:
def test_attributes_passthrough(self) -> None:
ident = ChannelIdentity(channel="teams", native_id="abc", attributes={"role": "user"})
assert dict(ident.attributes) == {"role": "user"}
class _DummyTarget:
"""Stand-in for the ``SupportsAgentRun | Workflow`` arg `apply_run_hook` forwards.
`apply_run_hook` doesn't introspect the target — it just forwards
it as a kwarg to the user's hook — so a bare class is enough.
"""
class _DummyChannel:
name = "dummy"
path = "/dummy"
def contribute(self, _context: Any) -> ChannelContribution:
return ChannelContribution()
class TestApplyChannelResponseHook:
async def test_originating_hook_receives_standard_context(self) -> None:
request = ChannelRequest(channel="discord", operation="message.create", input="hi")
payload = HostedRunResult("original")
captured: list[ChannelResponseContext] = []
async def hook(
result: HostedRunResult[Any],
*,
context: ChannelResponseContext,
) -> HostedRunResult[Any]:
captured.append(context)
return result.replace(result="hooked")
channel = _DummyChannel()
channel.response_hook = hook # type: ignore[attr-defined]
shaped = await apply_channel_response_hook(channel, payload, request=request, originating=True)
assert shaped.result == "hooked"
assert captured[0].request is request
assert captured[0].channel_name == "dummy"
assert captured[0].destination_identity is None
assert captured[0].originating is True
assert captured[0].is_echo is False
async def test_non_originating_hook_can_clone_before_shaping(self) -> None:
request = ChannelRequest(channel="responses", operation="message.create", input="hi")
identity = ChannelIdentity(channel="dummy", native_id="user-1")
payload = HostedRunResult("original")
seen_payloads: list[HostedRunResult[Any]] = []
seen_contexts: list[ChannelResponseContext] = []
def hook(
result: HostedRunResult[Any],
*,
context: ChannelResponseContext,
) -> HostedRunResult[Any]:
seen_payloads.append(result)
seen_contexts.append(context)
return result.replace(result="hooked")
channel = _DummyChannel()
channel.response_hook = hook # type: ignore[attr-defined]
shaped = await apply_channel_response_hook(
channel,
payload,
request=request,
destination_identity=identity,
originating=False,
is_echo=True,
clone=True,
)
assert seen_payloads[0] is not payload
assert shaped.result == "hooked"
assert seen_contexts[0].destination_identity is identity
assert seen_contexts[0].originating is False
assert seen_contexts[0].is_echo is True
async def test_missing_hook_returns_payload_or_clone(self) -> None:
request = ChannelRequest(channel="responses", operation="message.create", input="hi")
payload = HostedRunResult("original")
channel = _DummyChannel()
same = await apply_channel_response_hook(channel, payload, request=request, originating=True)
cloned = await apply_channel_response_hook(channel, payload, request=request, originating=True, clone=True)
assert same is payload
assert cloned is not payload
assert cloned.result == payload.result
class TestApplyRunHook:
"""`apply_run_hook` is the channel-side helper that invokes a
`ChannelRunHook` with the standard kwargs (`request` positional,
`target` / `protocol_request` keyword). Channels call this rather
than calling the hook directly so the convention is enforced in
one place. Cover both branching paths (sync vs async hook return)
and assert kwargs forwarding so a regression that drops `target`
or `protocol_request` is caught."""
async def test_sync_hook_returning_modified_request(self) -> None:
captured: dict[str, Any] = {}
def hook(request: ChannelRequest, **kwargs: Any) -> ChannelRequest:
# Snapshot the kwargs for the assertion below, then return a
# NEW request so we also verify the helper passes the
# replacement straight through (no merging / mutation).
captured["target"] = kwargs.get("target")
captured["protocol_request"] = kwargs.get("protocol_request")
return ChannelRequest(channel=request.channel, operation="HOOK_TOUCHED", input=request.input)
original = ChannelRequest(channel="responses", operation="op", input="hi")
target = _DummyTarget()
proto = {"raw": "payload"}
result = await apply_run_hook(hook, original, target=target, protocol_request=proto)
assert result is not original
assert result.operation == "HOOK_TOUCHED"
assert captured["target"] is target
assert captured["protocol_request"] is proto
async def test_async_hook_returning_modified_request(self) -> None:
captured: dict[str, Any] = {}
async def hook(request: ChannelRequest, **kwargs: Any) -> ChannelRequest:
captured["target"] = kwargs.get("target")
captured["protocol_request"] = kwargs.get("protocol_request")
# Return an awaitable result to exercise the async branch
# (`isinstance(result, Awaitable) → await it`).
return ChannelRequest(channel=request.channel, operation="ASYNC_HOOK", input=request.input)
original = ChannelRequest(channel="telegram", operation="op", input="hi")
target = _DummyTarget()
proto = {"update_id": 42}
result = await apply_run_hook(hook, original, target=target, protocol_request=proto)
assert result.operation == "ASYNC_HOOK"
assert captured["target"] is target
assert captured["protocol_request"] is proto
async def test_protocol_request_can_be_none(self) -> None:
"""Channels that don't have a raw protocol payload (e.g. CLI / test
harness invocations) pass ``protocol_request=None``; the helper
forwards it as-is so hooks can ``if protocol_request is None`` to
gate channel-specific logic."""
captured: dict[str, Any] = {}
async def hook(request: ChannelRequest, **kwargs: Any) -> ChannelRequest:
captured["protocol_request"] = kwargs.get("protocol_request")
captured["protocol_request_in_kwargs"] = "protocol_request" in kwargs
return request
await apply_run_hook(
hook,
ChannelRequest(channel="x", operation="op", input="hi"),
target=_DummyTarget(),
protocol_request=None,
)
assert captured["protocol_request"] is None
assert captured["protocol_request_in_kwargs"] is True
class TestDurableTaskPayloadMode:
"""``DurableTaskPayloadMode`` distinguishes object-mode (in-process,
live references) from JSON-mode (durable persistence, channel codec
required) runners. The host's startup validator uses the value to
refuse misconfigured deployments."""
def test_enum_values(self) -> None:
assert DurableTaskPayloadMode.OBJECT.value == "object"
assert DurableTaskPayloadMode.JSON.value == "json"
# Both members; no surprise additions until we ship a third
# adapter style.
assert set(DurableTaskPayloadMode) == {DurableTaskPayloadMode.OBJECT, DurableTaskPayloadMode.JSON}
class TestResponseTargetIdentities:
"""``ResponseTarget.identity``/``.identities`` carry full
:class:`ChannelIdentity` objects (incl. attributes) so destination
channels that need conversation/thread metadata (Teams, Slack, Bot
Framework) don't have to encode it through string tokens."""
def test_identity_single(self) -> None:
ident = ChannelIdentity(channel="teams", native_id="user@contoso", attributes={"tenant_id": "abc"})
target = ResponseTarget.identity(ident)
assert target.kind is ResponseTargetKind.IDENTITIES
assert len(target.target_identities) == 1
assert target.target_identities[0].channel == "teams"
assert target.target_identities[0].native_id == "user@contoso"
assert dict(target.target_identities[0].attributes) == {"tenant_id": "abc"}
def test_identities_list_preserves_attributes(self) -> None:
ident_a = ChannelIdentity(channel="teams", native_id="u1", attributes={"thread": "t1"})
ident_b = ChannelIdentity(channel="slack", native_id="u2", attributes={"channel_id": "c2"})
target = ResponseTarget.identities([ident_a, ident_b])
assert target.kind is ResponseTargetKind.IDENTITIES
assert len(target.target_identities) == 2
assert dict(target.target_identities[0].attributes) == {"thread": "t1"}
assert dict(target.target_identities[1].attributes) == {"channel_id": "c2"}
def test_identity_value_equality_matches_on_attributes(self) -> None:
# Two ``ResponseTarget.identity`` values built independently
# compare equal when the underlying ``ChannelIdentity`` content
# matches — important because tests and channel parsers use
# ``==`` on targets.
ident_a = ChannelIdentity(channel="teams", native_id="u1", attributes={"thread": "t1"})
ident_b = ChannelIdentity(channel="teams", native_id="u1", attributes={"thread": "t1"})
assert ResponseTarget.identity(ident_a) == ResponseTarget.identity(ident_b)
# Different attributes → not equal.
ident_c = ChannelIdentity(channel="teams", native_id="u1", attributes={"thread": "t2"})
assert ResponseTarget.identity(ident_a) != ResponseTarget.identity(ident_c)
def test_identity_repr_includes_targets(self) -> None:
ident = ChannelIdentity(channel="teams", native_id="u1")
rep = repr(ResponseTarget.identity(ident))
assert "ResponseTarget.identities" in rep
def test_identity_echo_input_flag(self) -> None:
ident = ChannelIdentity(channel="teams", native_id="u1")
target = ResponseTarget.identity(ident, echo_input=True)
assert target.echo_input is True