mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: add agent-framework-hosting core package (#5638)
* feat(hosting): add agent-framework-hosting core package
New ``agent-framework-hosting`` package implementing ADR 0026 / SPEC-002:
the channel-neutral host that lets a single ``Agent`` (or ``Workflow``)
fan out across multiple wire protocols ("channels") behind one Starlette
ASGI app.
Surface (re-exported from ``agent_framework_hosting``):
- ``AgentFrameworkHost`` — wraps a hostable target, mounts channels onto
an ASGI app, owns per-isolation-key ``AgentSession`` reuse, threads
request context (``response_id`` / ``previous_response_id``) into
context providers via an ``ExitStack`` of ``bind_request_context``
calls, and exposes an opt-in Hypercorn ``serve()`` helper (extra
``[serve]``).
- ``Channel`` protocol + ``ChannelContribution`` — the surface a channel
package implements (routes, lifespans, identity hooks, …).
- ``ChannelRequest`` / ``ChannelSession`` / ``ChannelIdentity`` /
``ChannelPush`` / ``ChannelCommand[Context]`` / ``ChannelRunHook`` /
``ChannelStreamTransformHook`` / ``DeliveryReport`` /
``HostedRunResult`` / ``ResponseTarget`` / ``ResponseTargetKind`` /
``apply_run_hook`` — channel-side dataclasses + helpers.
- ``IsolationKeys`` + ``ISOLATION_HEADER_USER`` / ``..._CHAT`` +
``get/set/reset_current_isolation_keys`` — the host's ASGI middleware
reads the ``x-agent-{user,chat}-isolation-key`` headers off each
inbound request and exposes them to the agent stack via a
``ContextVar`` so storage-side providers (e.g.
``FoundryHostedAgentHistoryProvider``) can apply per-tenant
partitioning without channels having to forward anything.
Includes 45 unit tests covering the host, channel contributions,
isolation contextvar, and shared types. Registers the package in
``python/pyproject.toml`` ``[tool.uv.sources]`` and adds the matching
pyright ``executionEnvironments`` entry for tests.
Hypercorn is an optional dependency (``[serve]`` extra); the soft import
in ``serve()`` is annotated for pyright since it isn't on the default
install.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* fix(hosting): address PR-2 review comments
Source-code changes
- _suppress_already_consumed: narrow contract — RuntimeError now logs
at WARNING with exc_info; non-RuntimeError still logs at exception().
Docstring clarifies that any non-clean teardown is observable.
- _BoundResponseStream: add aclose() and route __await__ through
get_final_response() so the binding is always released — fixes
contextvar leak when channels abandon the stream or use the
await-the-stream convenience.
- Lifespan: aggregate startup/shutdown callback errors; every callback
runs, all failures are logged with their qualname, and the first
error is re-raised so Starlette still aborts boot.
- _build_run_kwargs: switch session-cache write to dict.setdefault so
concurrent racers cannot orphan a session if create_session ever
yields.
- _deliver_response: introduce DeliveryReport.failed for push outages
vs explicit "no link" drops; an outage no longer triggers an
originating fallback so the channel can decide degraded behaviour.
Test additions
- tests/test_isolation.py (new): full coverage of IsolationKeys, the
contextvar helpers, header constants, and end-to-end ASGI
middleware lift / reset / passthrough.
- tests/test_host.py: TestBindRequestContext, TestBoundResponseStream
(aclose / __await__ / __getattr__ forwarding / double-close
idempotency), TestWrapInputListMessages (list[Message] LAST
precedence), TestLifespanAggregation (startup + shutdown).
- tests/test_types.py: TestApplyRunHook (sync/async/None), and
TestDeliveryReport (new failed field).
- Updated test_push_exception_marks_skipped ->
test_push_exception_lands_in_failed_no_fallback to match the new
delivery contract.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* fix(hosting): address PR-2 round-2 review comments
- Refactor workflow checkpoint restoration into shared helpers
(_restore_workflow_checkpoint for blocking; the streaming sibling
drains the rehydration stream) so the blocking and streaming paths
rehydrate identically — clarifies the previously inline _maybe_restore
by hoisting the pattern next to the blocking call site.
- Document that blocking workflow output is text-only by design;
richer modalities ride the streaming AgentResponseUpdate channel,
which preserves all content parts.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* review: address PR-4 _host.py round 2 feedback
These review comments were filed on PR-4 (#5640) but target lines that
live in the hosting-core package (PR-2 / #5638), so the fixes land here
and PR-4's stack will pick them up on rebase.
- _suppress_already_consumed: narrow the RuntimeError catch to the two
documented benign messages (`Inner stream not available`, `Event loop
is closed`); any other RuntimeError now logs at ERROR with a full
traceback so executor bugs / runner-context state errors / checkpoint
RuntimeErrors during the post-run flush no longer masquerade as
benign cleanup noise. Still no propagation (we're in an
async-generator finally during teardown) — see the docstring.
- _restore_workflow_checkpoint{,_streaming}: log a WARNING when a
non-None latest checkpoint drains to zero events, so a stale or
partially-written checkpoint_id surfaces as an operator signal
instead of a silent state-loss.
(The `deliver_response` "no destinations resolvable" vs "every
destination errored" concern raised in 3198268038 is already addressed
by the existing `failed` vs `skipped` distinction surfaced through
`DeliveryReport.failed` — see lines 1080-1102 and the
`DeliveryReport` docstring.)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* fix(hosting): reject path-traversal patterns in checkpoint isolation_key
The host's `_resolve_checkpoint_storage` joined `request.session.isolation_key`
directly into the configured `checkpoint_location`. The key is caller-
controlled — sourced from inbound headers (`x-agent-{user,chat}-isolation-key`
injected by the Foundry runtime), from channel-supplied derivations such as
`telegram:<chat_id>` / `entra:<oid>`, or from values set by a channel
`run_hook`. A value like `../../../etc/foo` or an absolute path would let
the resulting checkpoint directory escape the configured root (CWE-22).
This matches the path-traversal class fixed upstream in #5851 for the
foundry_hosting checkpoint storage.
New `_checkpoint_path_for_isolation_key(root, isolation_key)` helper:
- Uses a denylist (not allowlist) so legitimate namespaced keys
(`telegram:42`, `entra:abc-def`) continue to pass through unmodified.
- Rejects path separators (`/`, `\`), NUL, all-dot reductions (`.`, `..`,
`...`, ...), absolute paths (`os.path.isabs`), and drive-letter prefixes
(`os.path.splitdrive` plus an explicit `^[A-Za-z]:` check so payloads
crafted on a POSIX host still fail closed if the resulting directory
ever round-trips to Windows storage).
- After joining, resolves both sides and verifies
`target.is_relative_to(root)` as defence-in-depth.
`_resolve_checkpoint_storage` now logs a WARNING and returns `None` for
invalid keys rather than crashing the request — checkpointing is best-
effort and we prefer dropping it to letting one malformed key abort an
otherwise valid agent run.
Tests:
- `TestCheckpointPathForIsolationKey` exercises the helper directly with
legitimate keys (alphanumeric, `:`-namespaced, dotted, 200-char), all
rejected traversal patterns from #5851's MSRC repro list, and
non-string input.
- `TestHostWorkflowCheckpointingPathTraversal` verifies the end-to-end
request path: a traversal key (`../escape`) and an in-key separator
(`evil/sub`) both produce a successful agent response with no files
written under `checkpoint_location`, and the traversal case logs a
WARNING citing `isolation_key`.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* fix(hosting): address PR-2 round-3 review feedback + add response hooks
Round-3 review comment fixes:
- _types.py: drop the _EMPTY_MAPPING sentinel; ChannelIdentity.attributes
uses plain dict() as the default — simpler, no extra symbol to track.
- _host.py: drop the local `import asyncio` + `from typing import cast as
_cast` inside `serve()`; rely on the module-level imports.
- _host.py: switch `_log_incoming` to structured `extra={...}` payloads
for both INFO and DEBUG so log aggregators get queryable fields.
- _host.py: delete `_flat_context_providers` and stop descending into a
`.providers` attribute. Aggregator providers (AggregateContextProvider /
ContextProviderBase) are responsible for forwarding `response_context`
to their children themselves; the host treats whatever
`agent.context_providers` exposes as the final, flat list.
- _host.py: stop collapsing agent / workflow output to text. `_invoke`
forwards `AgentResponse.messages` (and `raw_response`) on the
`HostedRunResult`. `_invoke_workflow` builds a per-event message list
via a new `_workflow_output_to_messages` helper that preserves
AgentResponse / AgentResponseUpdate / Message / Content branches and
falls back to text only for arbitrary objects.
- _host.py: `_workflow_event_to_update` carries Content payloads through
unchanged so multi-modal workflow outputs (images, function-call
metadata, ...) survive into channels.
New features (per design discussion in the PR thread):
- HostedRunResult: rebuilt around `messages: list[Message]` with
`.text` / `.contents` as projections, a `raw_response` slot for the
underlying AgentResponse, and a `replace(messages=..., raw_response=...)`
clone helper used by the delivery layer for per-destination isolation.
The `HostedRunResult(text="...")` ctor is preserved as a back-compat
shim that synthesises a single assistant text message.
- ResponseTarget: gain `echo_input: bool = False` (also exposed on
`.channel(name, *, echo_input=...)` / `.channels([...], *, echo_input=...)`).
When set, the host pushes the originating user message to each
non-originating destination before the agent reply. Channels can
filter or transform echoes via their response_hook.
- DeliveryReport: add `echoed` / `echo_failed` tuples to surface
per-destination outcomes of the new echo phase. Echo failures do not
abort the corresponding response push on the same destination.
- ChannelResponseHook + ChannelResponseContext + apply_response_hook:
duck-typed `response_hook` attribute on channels for per-destination
post-processing. Receives a clone of the HostedRunResult and a
context carrying the request, channel name, destination identity,
originating flag, and `is_echo` phase flag. Channels stay
modality-aware (text-only wires flatten via the hook; card-capable
channels render structured contents directly).
- _deliver_response: clone-before-hook fan-out so a hook mutating one
channel's payload cannot leak into another destination's view.
Tests:
- Update _FakeAgentResponse to expose `.messages` (single assistant text
message synthesised from `text`) so existing tests pass unchanged on
the new multi-modal _invoke path.
- Replace the obsolete `test_bind_descends_one_level_into_providers_attribute`
with a regression guard asserting the host does NOT descend into
`.providers` (matches new contract).
- New tests for HostedRunResult multi-modal preservation, echo_input
fan-out with success + failure, response_hook applied per destination,
per-destination mutation isolation, and is_echo phase observability.
Docs:
- spec 002: rewrite Canonical flow with the new input → run_hook → host
→ target → wrap → per-destination clone → response_hook → push
pipeline; document multi-modality contract and per-destination
cloning; add `echo_input` row to ResponseTarget table; rewrite
HostedRunResult/HostedStreamResult row; add ChannelResponseHook /
ChannelResponseContext / apply_response_hook table; log decisions
Q28 (no host-side text collapse), Q29 (duck-typed response_hook),
Q30 (opt-in `echo_input` on ResponseTarget).
- ADR 0026: add ChannelResponseHook + multi-modality bullets;
surface `echo_input` on the ResponseTarget bullet.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* fix(hosting): drop HostedRunResult(text=...) back-compat shim; use from_text()
Pre-release cleanup — no released callers to break, so consolidate on one
canonical entry point plus a classmethod for the ergonomic
single-text-message case:
- HostedRunResult.__init__ takes ``messages`` positionally (required); no
more ``text=`` kwarg overload, no more "synthesise an empty message
when no args" path.
- New HostedRunResult.from_text(text, *, role="assistant", raw_response=None)
classmethod for the common "wrap a single text content as one message"
case (tests, channels emitting plain strings, the echo-input phase
wrapping a user's text turn).
- ``_build_echo_payload`` uses ``HostedRunResult.from_text(raw, role="user")``
for the ``str`` and fallback branches; the other branches use the plain
ctor with explicit ``Message`` lists.
- Tests rewritten to use ``from_text("reply")`` everywhere
``HostedRunResult(text="reply")`` appeared. Added an explicit
``test_from_text_role_kwarg_overrides_default`` regression guard.
- spec 002: HostedRunResult row updated to describe the
``from_text(text, *, role="assistant")`` classmethod instead of the
removed back-compat shim.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* refactor(hosting-core): reshape HostedRunResult into generic typed envelope
Replace the flattened multi-modal HostedRunResult (carrying
messages/raw_response/.text projections) with a typed generic
envelope around the target's full-fidelity output:
class HostedRunResult(Generic[TResult]):
result: TResult
session: AgentSession | None
- Agent targets produce HostedRunResult[AgentResponse]; channels
read result.messages, result.text, result.value, result.response_id,
result.usage_details directly off the underlying response.
- Workflow targets produce HostedRunResult[WorkflowRunResult];
channels iterate result.get_outputs() and inspect
result.get_final_state() themselves (the host no longer collapses
workflow outputs onto a synthesised message list).
- The echo-input phase synthesises a HostedRunResult[AgentResponse]
wrapping the user's turn so the same per-destination delivery
machinery applies.
- replace() is now {result, session} only; the host's clone is
shallow — channels that need to mutate result itself are
responsible for their own deep copy.
Rationale: the earlier shape pre-shaped target output (collapsing
workflows onto a Message list, losing per-executor outputs, final
state, and structured value affordances). Carrying the target output
unchanged keeps the host modality-agnostic, gives channel authors
static typing where they want it, and removes 30+ lines of
host-side projection helpers.
Also updates ADR 0026 + spec 002 (Q3, Q28, Q29 amended; new Q31
captures the generic-envelope decision and rationale).
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* docs(hosting-core): document echo vs response distinction for push channels
The host already encodes the echo-vs-response phase via the
underlying Message.role on the pushed HostedRunResult:
- echo phase: payload.result.messages[*].role == "user"
- response phase: payload.result.messages[*].role == "assistant"
Both pushes go through the same ChannelPush.push(identity, payload)
entry point. Channels distinguish either by inspecting role (which
works for any push-capable channel) or — when a response_hook is
wired — by branching on ChannelResponseContext.is_echo directly.
Expand the ChannelPush Protocol docstring to make this discoverable
for channel implementers (esp. chat bots that cannot impersonate
the user on their wire and need to render echoes as quoted /
prefixed blocks rather than as bot replies).
Mirror the explanation into the spec's echo_input section.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* docs(hosting-core): fix quickstart to use current Agent API
ChatAgent was renamed to Agent and the preferred construction pattern
is client.as_agent(...). Also drop the sibling channel import so the
snippet imports only modules declared as dependencies of this package;
point readers at the sibling packages instead.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* test(hosting-core): drop redundant @pytest.mark.asyncio decorators
asyncio_mode = "auto" is configured in pyproject.toml, so individual
@pytest.mark.asyncio decorators are unnecessary.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* docs(hosting): add authorization profiles + IdentityAllowlist seam to ADR/spec
Composes `require_link` + `allowlist` into three named profiles (open,
forced-link, allowlist) with the allowlist itself keyed on either the
channel-native id (pre-link) or a verified IdP claim (post-link), plus
`AnyOf`/`AllOf` combinators for mixed setups. Lifts the design into
an explicit host seam (`host.authorize(...)` → `AuthorizationOutcome`
of `Allowed` / `LinkRequired` / `Denied`) instead of leaving each
channel to roll its own.
Key contract bits:
- Tri-state `AllowlistDecision` (ALLOW / DENY / ABSTAIN) so claim-based
lists can ABSTAIN until claims are available without composition
silently flipping that into DENY.
- `AuthorizationContext` carries explicit `phase` + `claim_source`
so allowlists can tell pre-link from post-link without overloading
`verified_claims is None`.
- Channel-side `allowlist: ... | Literal["inherit"] | None` with an
explicit inheritance sentinel, so the host-level `default_allowlist`
is opt-out, not opt-in.
- Construction-time validator rejects silent-deny configurations
(`LinkedClaimAllowlist` without a claim source) with a typed
`ChannelConfigurationError`.
- Group-chat denial mirrors the existing `LinkChallenge` DM-redirect
pattern; only the redacted `user_message` reaches the wire,
structured `log_details` stay in telemetry.
Ships in two waves: the Protocol + `NativeIdAllowlist` + config
validator land with the next core PR ahead of the linker; the full
pipeline + `LinkedClaimAllowlist` enforcement land with the
`IdentityLinker` core PR.
Updates: ADR 0026 (summary bullet + conceptual-API table row + resolved
Q16), spec 002 (new req #22, renumbered v1 fast-follow #23..#29 and
stretch #30..#31, new "Authorization profiles and the IdentityAllowlist
seam" subsection, inbound-ownership row, resolved Q32, follow-up entry).
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* feat(hosting): add DurableTaskRunner seam + runtime_mode auto-detect
Introduces the explicit long-running vs ephemeral runtime distinction
and a generic DurableTaskRunner Protocol that owns non-originating
push dispatch — collapsing the previous deliveries[] per-destination
state machine, SupportsDeliveryTracking provider capability, and
Foundry update_item service ask down to a single immutable
intended_targets[] write on the message.
Spec / ADR:
- New §"Runtime modes" with auto-detect markers + defaults matrix.
- Rewrites §"Delivery tracking" → §"Intended targets + durable
delivery": intent-only on the message, operational state lives in
the runner.
- New §"Durable task runner" defining DurableTaskRunner / RetryPolicy
/ TaskHandle / TaskStatus.
- Drops §SupportsDeliveryTracking and §Foundry update_item gap.
- Resolved Qs: 12, 18, 21, 26 revised; new 17/18/19 (ADR) and
33/34/35 (spec).
Code:
- New _runner.py with InProcessTaskRunner (asyncio + bounded retry,
bounded terminal-status cache, register-after-start guard,
shutdown drain).
- _host.py: runtime_mode + durable_task_runner ctor params;
auto-detect via FOUNDRY_HOSTING_ENVIRONMENT /
AZURE_FUNCTIONS_ENVIRONMENT / AWS_LAMBDA_FUNCTION_NAME;
HOSTING_PUSH_TASK_NAME handler registered eagerly so
_deliver_response can be called outside the lifespan;
_handle_push_task does echo-then-response inline per destination;
_deliver_response now schedules one task per destination via the
runner (DeliveryReport.pushed = scheduled; .failed = schedule-time
outage only).
- _types.py: new DurableTaskRunner Protocol + RetryPolicy /
TaskHandle / TaskStatus; DeliveryReport drops echoed /
echo_failed (echo outcome owned by the runner).
- __init__.py exports the new public surface.
Tests: 132 passing, 90% coverage. New test_runner.py covers
InProcessTaskRunner success/retry/terminal-failure/cancellation/
register-after-start, runtime-mode auto-detect with synthetic env,
and the warning-on-ephemeral-without-runner path. test_host.py
delivery tests use a sync runner fake for deterministic assertions
and validate the new "schedule succeeded vs runner backend
unreachable" semantics.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* feat(hosting): rubber-duck round-5 — strict ephemeral, codec seam, allowlist Wave-1, drop DeliveryReport
Adopts the rubber-duck-approved package of changes from the round-5
review of PR #5638 (modulo DeliveryReport.failed — the value type is
removed entirely now that durable delivery covers the failure
surface, per user direction).
Code:
- Drop DeliveryReport value type; host-internal _deliver_response
returns bool. Failure observability is now logs (in-process) /
runner backend (durable adapters).
- Strict ephemeral default: ephemeral runtime_mode with the default
in-process runner raises RuntimeError; opt-in via
allow_in_process_runner=True (warns).
- ChannelPushCodec Protocol + DurableTaskPayloadMode enum +
_validate_runner_codec_pairing so JSON-mode runners can be safely
paired with channels via codecs; _handle_push_task accepts both
object- and JSON-envelope shapes.
- ResponseTarget.identity(...) / .identities([...]) builders +
IDENTITIES kind for explicit caller-supplied recipients; field
rename identities → _target_identities (private) with a
target_identities property to resolve the classmethod collision.
- Intent-only audit: _annotate_intended_targets writes
hosting.intended_targets / skipped_targets / includes_originating /
originating_channel onto assistant messages — single immutable
write per the runner-owned operational-state model.
- InProcessTaskRunner: 2-phase drain on shutdown
(shutdown_grace_seconds, default 5.0) so a clean shutdown does not
abandon work mid-retry; payload_mode = OBJECT class-level.
- Echo idempotency: _handle_push_task tracks an echo_done cursor on
runner-owned task state so a retry that fires after the echo
phase succeeded does not double-echo.
Wave-1 authorization seam (full landing):
- New _authorization.py with AllowlistDecision tri-state,
AuthorizationContext, IdentityAllowlist Protocol, AllowAll /
NativeIdAllowlist (with async loader cache + channel-scope ABSTAIN) /
LinkedClaimAllowlist (raise-until-Wave-2) / AnyOfAllowlists /
AllOfAllowlists / CallableAllowlist built-ins, Allowed /
LinkRequired / Denied outcomes, ChannelConfigurationError.
- Host(default_allowlist=..., identity_linker=...) + per-channel
allowlist parameter with 'inherit' / None semantics.
- _validate_channel_authorization enforces all three rules at
construction: claim-source requirement, linker presence for
require_link=True (elevated from no-op — must not ship
unenforced), and NativeIdAllowlist(channel=...) typo detection.
Combinator-walking via _flatten_allowlists catches nested
misconfigs.
- host.authorize(...) for the native-id pipeline: open path returns
Allowed with auto-issued <channel>:<native_id> isolation key (or
the existing key when the identity has been seen); ABSTAIN on a
claim-required allowlist maps to
Denied(reason_code='allowlist_requires_link') until Wave 2 wires
the linker to convert it to LinkRequired.
Spec / ADR:
- docs/specs/002-python-hosting-channels.md: Wave-1 status updated
to reflect the linker-presence rule elevation and the
host.authorize landing; new sub-sections (codec contract, drain,
echo cursor); Qs 18 / 21 DeliveryReport references purged; new
resolved Qs 36–40 covering the strict-ephemeral default, codec
contract, DeliveryReport removal, echo cursor, and drain.
- docs/decisions/0026-hosting-channels.md: Q12 DeliveryReport
reference purged; Q16 updated to reflect Wave-1 landing; new
resolved Qs 20 (codec contract) + 21 (strict ephemeral / drain /
echo cursor).
Tests:
- New tests/test_authorization.py (35 cases) covering every Wave-1
built-in, the three validator rules, combinator decision
semantics, and host.authorize across open / allow / deny /
abstain-with-claim-dep / abstain-without-claim-dep paths plus
existing-key reuse and verified-claims propagation.
- tests/test_host.py: TestDeliverResponse rewritten for the bool
return + runner.scheduled-count assertions; new tests for
IDENTITIES variant + echo idempotency.
- tests/test_runner.py: strict-ephemeral now expects RuntimeError;
allow_in_process_runner opt-in tests; shutdown drain test;
payload_mode default test.
- tests/test_types.py: TestDeliveryReport removed; new
TestDurableTaskPayloadMode + TestResponseTargetIdentities.
Validation: 178 tests pass, 91% coverage, fmt + lint + pyright +
mypy clean.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* docs(hosting): add mermaid flow diagrams to ADR, spec, README
Insert the 10 hosting flow diagrams reviewed in
python/.user/hosting-diagrams.md into the public docs:
- README: runtime topology (1a) + cross-link to the spec for the
richer set.
- ADR: runtime topology, channel contribution shape, and authorization
decision (1a, 1b, 3) at the end of 'Conceptual API shape'.
- Spec: all 10 diagrams — 1a/1b at the top of API Surface, 2 in
Canonical flow, 3 in Authorization profiles, 4-7 in Scenarios 6-8,
8 in Codec contract, 9 in Echo idempotency, 10 in Scenario 9.
Doc-only; no API or behaviour change.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* feat(hosting): add opt-in disk persistence via state_dir
Long-running hosts (always-on container, single-VM bot, local dev) lose
state on every restart today. Add an opt-in disk persistence layer under
a new `state_dir` constructor parameter on `AgentFrameworkHost` that
survives process restarts without taking on a heavyweight database
dependency.
Backed by `diskcache` (installed via the new `[disk]` optional extra).
An OS-level advisory file lock guarantees single-owner semantics so two
hosts pointed at the same directory cannot double-execute scheduled
pushes.
What persists when `state_dir` is set:
- Pending durable-task records — scheduled-but-not-yet-completed pushes
replay on the next host startup via `InProcessTaskRunner.resume()`.
Records that crashed mid-attempt resume with the already-consumed
retry budget (no full-budget re-grant).
- `_session_aliases` — per-isolation-key session-id rewrites.
- `_active` — most-recently-active channel per isolation key.
- `_identities` — `ChannelIdentity` rows for fan-out targeting,
including nested mutations of the form
`self._identities[ik][channel] = identity`.
The `state_dir` parameter accepts any of:
- `None` — today's purely in-memory behaviour.
- `str` / `PathLike` — single root; host auto-creates `runner/` and
`sessions/` subfolders.
- `HostStatePaths` TypedDict / plain mapping — per-component overrides
routed to different roots. Unknown keys raise `ValueError` to surface
typos early.
Unpicklable push payloads raise `PushPayloadNotPicklable` eagerly from
`schedule()` so issues surface at the call site rather than on the
next restart. Corrupt on-disk records are quarantined-and-logged; the
runner never crashes on resume.
Live `AgentSession` objects stay in memory and are rehydrated lazily
by the history provider on the next turn.
- New modules: `_persistence.py` (lock + normalisation),
`_state_store.py` (session-bookkeeping store).
- Runner rewrite: 4-state model (`pending` / `succeeded` / `failed`
/ `cancelled`); the transient `running` state was a bug that caused
resume to skip records that crashed mid-handler.
- New tests: `test_runner_disk.py` (8 tests), `test_host_disk.py` (8
tests). 194 passed total. pyright + mypy + ruff clean.
- README: new "Optional disk persistence" section with code samples.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* feat(hosting): add checkpoints to state_dir + fix host docstring
Three related polish changes on top of the disk-persistence landing:
1. Extend `state_dir` to cover workflow checkpoints. Adds
`checkpoints` as a third `HostStatePaths` key. Single-path form
(`state_dir="/foo"`) now also auto-derives `/foo/checkpoints/`
for workflow targets (equivalent to passing
`checkpoint_location="/foo/checkpoints"`). The mapping form lets
workflow callers opt out by omitting the key, or route checkpoints
to a different volume.
Conflict / precedence rules:
* Explicit `checkpoint_location` always wins over the state_dir
derived path; a warning surfaces the double-config.
* Single-path `state_dir` + non-Workflow target → checkpoints path
silently ignored (no eager directory creation either).
* Mapping form with `checkpoints` + non-Workflow target → warn
(almost certainly dead config).
* Derived path with a workflow that already has its own
`checkpoint_storage` → same `RuntimeError` as the explicit
parameter triggers, so ownership stays unambiguous.
Checkpoint persistence uses `FileCheckpointStorage` from the
framework core — no extra dependency. Only `runner` and
`sessions` require the `[disk]` extra.
2. Move `AgentFrameworkHost.__init__` parameter docs from `Args:` to
`Keyword Args:` for every parameter after the `*`. Only `target`
remains under `Args:`. Brings the docstring in line with the
actual signature (the params have always been keyword-only).
3. `HostStatePaths` already existed as a TypedDict but did not cover
`checkpoints`; updated to document the new key with the same
per-attribute docstring style as `runner` / `sessions` so editors
can surface help on the keys.
Validation: 201 tests pass (was 194; +7 checkpoint integration tests
in test_host_disk.py). pyright + mypy + ruff + bandit clean.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* feat(hosting): add core IdentityLinker authorization seam
Fold the core IdentityLinker pieces into the hosting-core PR so the
authorization surface no longer has a deferred Wave-2 placeholder.
Provider-specific linkers (for example Entra OAuth helpers) can now plug
into core without core depending on an IdP SDK.
Core additions:
- Add LinkChallenge, LinkedIdentity, LinkResolution, and IdentityLinker.
IdentityLinker.resolve(identity) is a single-call decision that returns
either a linked identity with verified claims or a challenge the channel
can render.
- Enable LinkedClaimAllowlist end-to-end. It now abstains pre-link and
allows/denies post-link against verified claims, including multi-valued
claims such as groups.
- Add AuthPolicy factories for common allowlist shapes.
- Extend Allowed with verified_claims and claim_source for audit/telemetry
without requiring callers to re-derive how the decision was made.
Host behavior:
- identity_linker is now typed as IdentityLinker | None.
- authorize() supports open, native-id, forced-link, and linked-claim
profiles end-to-end.
- require_link=True resolves via the linker and returns LinkRequired when
the identity is not linked.
- claim-based allowlists use channel-emitted verified_claims when present,
or linker-resolved claims otherwise.
- authorize() remains decision-only and does not mutate _identities/_active;
identity registry writes remain on the actual request execution path.
Docs/tests:
- Remove Wave-1/Wave-2 language from core/spec/ADR surfaces touched here.
- Update the spec/ADR to describe the core linker seam and provider-specific
linker packages.
- Add authorization tests for linker challenges, linked identities, linked
claim allowlists, channel-emitted claims, AuthPolicy factories, and the
no-mutation contract.
Validation: 214 tests pass, pyright/mypy/ruff clean.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* feat(hosting): add link-store path to state_dir
Identity linking introduces host-adjacent state that needs the same state_dir treatment as runner, session, and checkpoint state. Add a links component to the host state paths so applications and linker packages have a typed, discoverable persistence location.
Changes:
- Extend HostStatePaths with links and include it in state_dir normalization (state_dir/links/ for the single-path form).
- Add SupportsLinkStorePath, an optional protocol for identity linkers that accept a host-provided link-store path.
- AgentFrameworkHost now offers state_dir links to compatible linkers, warns when an explicit links path is supplied without a linker, and warns when the configured linker manages persistence directly instead of implementing SupportsLinkStorePath.
- Update README and spec text to document the link-store component and clarify that concrete linkers still own the storage format.
- Add disk-state tests for compatible, missing, and non-configurable linkers.
Validation: 217 tests pass, pyright/mypy/ruff clean.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---------
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
committed by
GitHub
Unverified
GitHub
parent
e666cdc7c8
commit
0cb9b52a4b
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) Microsoft Corporation.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE
|
||||
@@ -0,0 +1,192 @@
|
||||
# agent-framework-hosting
|
||||
|
||||
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, …).
|
||||
|
||||
The base package contains only the 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.
|
||||
|
||||
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).
|
||||
|
||||
## Install
|
||||
|
||||
```bash
|
||||
pip install agent-framework-hosting agent-framework-hosting-responses
|
||||
# or with uvicorn 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)
|
||||
pip install "agent-framework-hosting[disk]"
|
||||
```
|
||||
|
||||
## Quickstart
|
||||
|
||||
```python
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
from agent_framework_hosting import AgentFrameworkHost, Channel
|
||||
|
||||
agent = OpenAIChatClient().as_agent(name="Assistant")
|
||||
|
||||
# Add channels from sibling packages, e.g. `agent-framework-hosting-responses`
|
||||
# exposes a `ResponsesChannel` that serves the OpenAI Responses API.
|
||||
channels: list[Channel] = []
|
||||
|
||||
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.).
|
||||
|
||||
## Optional disk persistence (`state_dir`)
|
||||
|
||||
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.
|
||||
|
||||
```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.
|
||||
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",
|
||||
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.
|
||||
|
||||
@@ -0,0 +1,143 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Multi-channel hosting for Microsoft Agent Framework agents.
|
||||
|
||||
Serve a single agent target through one or more **channels** — pluggable
|
||||
adapters that expose the target over different transports such as the
|
||||
OpenAI Responses API, Microsoft Teams, Telegram, and others. The base
|
||||
package contains only the channel-neutral plumbing; concrete channels
|
||||
ship in their own packages (``agent-framework-hosting-responses``,
|
||||
``agent-framework-hosting-telegram``, …) so users install only what
|
||||
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 ._isolation import (
|
||||
ISOLATION_HEADER_CHAT,
|
||||
ISOLATION_HEADER_USER,
|
||||
IsolationKeys,
|
||||
get_current_isolation_keys,
|
||||
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,
|
||||
HostedRunResult,
|
||||
HostStatePaths,
|
||||
PushPayloadNotPicklable,
|
||||
PushPayloadNotSerializable,
|
||||
ResponseTarget,
|
||||
ResponseTargetKind,
|
||||
RetryPolicy,
|
||||
TaskHandle,
|
||||
TaskStatus,
|
||||
apply_response_hook,
|
||||
apply_run_hook,
|
||||
)
|
||||
|
||||
try:
|
||||
__version__ = importlib.metadata.version(__name__)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
__version__ = "0.0.0"
|
||||
|
||||
__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",
|
||||
"HostStatePaths",
|
||||
"HostedRunResult",
|
||||
"IdentityAllowlist",
|
||||
"IdentityLinker",
|
||||
"InProcessTaskRunner",
|
||||
"IsolationKeys",
|
||||
"LinkChallenge",
|
||||
"LinkRequired",
|
||||
"LinkResolution",
|
||||
"LinkedClaimAllowlist",
|
||||
"LinkedIdentity",
|
||||
"NativeIdAllowlist",
|
||||
"PushPayloadNotPicklable",
|
||||
"PushPayloadNotSerializable",
|
||||
"ResponseTarget",
|
||||
"ResponseTargetKind",
|
||||
"RetryPolicy",
|
||||
"RuntimeMode",
|
||||
"SupportsLinkStorePath",
|
||||
"TaskHandle",
|
||||
"TaskStatus",
|
||||
"__version__",
|
||||
"apply_response_hook",
|
||||
"apply_run_hook",
|
||||
"get_current_isolation_keys",
|
||||
"logger",
|
||||
"reset_current_isolation_keys",
|
||||
"set_current_isolation_keys",
|
||||
]
|
||||
@@ -0,0 +1,485 @@
|
||||
# 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",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,76 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Per-request isolation keys read from inbound HTTP headers.
|
||||
|
||||
The Foundry Hosted Agents runtime injects two well-known headers on every
|
||||
request it forwards to the user's container:
|
||||
|
||||
* ``x-agent-user-isolation-key`` — opaque per-user partition key
|
||||
* ``x-agent-chat-isolation-key`` — opaque per-conversation partition key
|
||||
|
||||
When the headers are present we are running inside (or being driven by) the
|
||||
Foundry runtime; when they are absent we are running in plain local dev. The
|
||||
host installs an ASGI middleware in :meth:`AgentFrameworkHost._build_app`
|
||||
that reads both headers off every inbound HTTP request and pushes them into
|
||||
the :data:`current_isolation_keys` contextvar for the duration of the
|
||||
request, then resets it. Providers that need partition-aware storage (most
|
||||
notably ``FoundryHostedAgentHistoryProvider``) read the contextvar via
|
||||
:func:`get_current_isolation_keys` and apply the keys to their backend
|
||||
calls — so app authors don't have to wire any middleware themselves and
|
||||
channels stay free of Foundry-specific header knowledge.
|
||||
|
||||
The contextvar holds a plain :class:`IsolationKeys` mapping; conversion to
|
||||
provider-specific types (e.g. Foundry's ``IsolationContext``) happens at
|
||||
the consuming provider so this module has no provider dependencies.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextvars import ContextVar, Token
|
||||
|
||||
__all__ = [
|
||||
"ISOLATION_HEADER_CHAT",
|
||||
"ISOLATION_HEADER_USER",
|
||||
"IsolationKeys",
|
||||
"current_isolation_keys",
|
||||
"get_current_isolation_keys",
|
||||
"reset_current_isolation_keys",
|
||||
"set_current_isolation_keys",
|
||||
]
|
||||
|
||||
|
||||
ISOLATION_HEADER_USER = "x-agent-user-isolation-key"
|
||||
ISOLATION_HEADER_CHAT = "x-agent-chat-isolation-key"
|
||||
|
||||
|
||||
class IsolationKeys:
|
||||
"""Per-request Foundry isolation keys lifted off the inbound headers."""
|
||||
|
||||
def __init__(self, user_key: str | None = None, chat_key: str | None = None) -> None:
|
||||
self.user_key = user_key
|
||||
self.chat_key = chat_key
|
||||
|
||||
@property
|
||||
def is_empty(self) -> bool:
|
||||
return self.user_key is None and self.chat_key is None
|
||||
|
||||
|
||||
current_isolation_keys: ContextVar[IsolationKeys | None] = ContextVar(
|
||||
"agent_framework_hosting_isolation_keys",
|
||||
default=None,
|
||||
)
|
||||
|
||||
|
||||
def get_current_isolation_keys() -> IsolationKeys | None:
|
||||
"""Return the isolation keys bound to the current request, if any."""
|
||||
return current_isolation_keys.get()
|
||||
|
||||
|
||||
def set_current_isolation_keys(keys: IsolationKeys | None) -> Token[IsolationKeys | None]:
|
||||
"""Bind ``keys`` to the current async context and return a reset token."""
|
||||
return current_isolation_keys.set(keys)
|
||||
|
||||
|
||||
def reset_current_isolation_keys(token: Token[IsolationKeys | None]) -> None:
|
||||
"""Restore the isolation contextvar to its prior value."""
|
||||
current_isolation_keys.reset(token)
|
||||
@@ -0,0 +1,195 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""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.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import os
|
||||
import sys
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
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")
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
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 "
|
||||
"is not installed. Install the disk extra: "
|
||||
"`pip install 'agent-framework-hosting[disk]'`."
|
||||
) from exc
|
||||
return diskcache
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
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":
|
||||
import msvcrt
|
||||
|
||||
try:
|
||||
msvcrt.locking(fh.fileno(), msvcrt.LK_NBLCK, 1)
|
||||
except OSError as exc:
|
||||
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."
|
||||
) from exc
|
||||
else:
|
||||
import fcntl
|
||||
|
||||
try:
|
||||
fcntl.flock(fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
except OSError as exc:
|
||||
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."
|
||||
) from exc
|
||||
except RuntimeError:
|
||||
raise
|
||||
except Exception:
|
||||
fh.close()
|
||||
raise
|
||||
return fh
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
if handle is None:
|
||||
return
|
||||
with contextlib.suppress(Exception): # close errors are not actionable
|
||||
handle.close()
|
||||
|
||||
|
||||
def normalize_state_dir(
|
||||
state_dir: str | os.PathLike[str] | HostStatePaths | Mapping[str, str | os.PathLike[str]] | None,
|
||||
) -> 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`.
|
||||
"""
|
||||
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."
|
||||
)
|
||||
for name in _KNOWN_COMPONENTS:
|
||||
raw_value: Any = state_dir.get(name)
|
||||
if raw_value is None:
|
||||
result[name] = None
|
||||
continue
|
||||
if isinstance(raw_value, (str, os.PathLike)):
|
||||
result[name] = Path(os.fspath(raw_value))
|
||||
else:
|
||||
raise TypeError(f"state_dir[{name!r}] must be a str or PathLike — got {type(raw_value).__name__}")
|
||||
return result
|
||||
|
||||
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",
|
||||
]
|
||||
@@ -0,0 +1,751 @@
|
||||
# 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"]
|
||||
@@ -0,0 +1,402 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Disk-backed wrappers for the host's in-memory state dicts.
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Any, TypeVar, cast
|
||||
|
||||
from ._persistence import (
|
||||
acquire_state_dir_lock,
|
||||
load_diskcache,
|
||||
release_state_dir_lock,
|
||||
)
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
def __init__(self, sessions_dir: str | os.PathLike[str]) -> None:
|
||||
self._sessions_dir: Path = Path(os.fspath(sessions_dir))
|
||||
diskcache = load_diskcache()
|
||||
self._lock_handle: Any = acquire_state_dir_lock(self._sessions_dir)
|
||||
try:
|
||||
self._cache: Any = diskcache.Cache(str(self._sessions_dir))
|
||||
except Exception:
|
||||
release_state_dir_lock(self._lock_handle)
|
||||
self._lock_handle = None
|
||||
raise
|
||||
|
||||
@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 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.
|
||||
"""
|
||||
if self._cache is not None:
|
||||
try:
|
||||
self._cache.close()
|
||||
except Exception: # pragma: no cover - close errors aren't actionable
|
||||
logger.exception("SessionsStateStore: 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
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
store: SessionsStateStore,
|
||||
key_prefix: str,
|
||||
initial: Mapping[str, _V] | None = None,
|
||||
) -> None:
|
||||
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)
|
||||
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
|
||||
|
||||
def __setitem__(self, key: str, value: _V) -> None:
|
||||
super().__setitem__(key, value)
|
||||
try:
|
||||
self._store.cache.set(self._prefix + key, value)
|
||||
except Exception: # pragma: no cover - cache write failures aren't actionable
|
||||
logger.exception("SessionsStateStore: failed to persist %s%s", self._prefix, key)
|
||||
|
||||
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 - cache write failures aren't actionable
|
||||
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.
|
||||
value: _V = super().pop(key, *args)
|
||||
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)
|
||||
return value
|
||||
|
||||
def clear(self) -> None:
|
||||
keys = list(self.keys())
|
||||
super().clear()
|
||||
cache = self._store.cache
|
||||
for k in keys:
|
||||
try:
|
||||
del cache[self._prefix + k]
|
||||
except KeyError:
|
||||
pass
|
||||
except Exception: # pragma: no cover
|
||||
logger.exception("SessionsStateStore: failed to evict %s%s during clear", self._prefix, k)
|
||||
|
||||
def update( # type: ignore[override]
|
||||
self,
|
||||
other: Mapping[str, _V] | None = None,
|
||||
/,
|
||||
**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.
|
||||
if other is not None:
|
||||
for k in other:
|
||||
self[k] = other[k]
|
||||
for k, v in kwargs.items():
|
||||
self[k] = v
|
||||
|
||||
|
||||
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",
|
||||
]
|
||||
@@ -0,0 +1,915 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
# ``ChannelRequest`` is the only intentional dataclass here (callers use
|
||||
# ``dataclasses.replace`` on it in run hooks). The other types are plain
|
||||
# Python classes by preference, so the "could be a dataclass" lint is muted
|
||||
# at the file level.
|
||||
# ruff: noqa: B903
|
||||
|
||||
"""Channel-neutral request envelope and channel protocol types.
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
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, runtime_checkable
|
||||
|
||||
from agent_framework import (
|
||||
AgentResponse,
|
||||
AgentResponseUpdate,
|
||||
AgentRunInputs,
|
||||
ResponseStream,
|
||||
SupportsAgentRun,
|
||||
Workflow,
|
||||
)
|
||||
from starlette.routing import BaseRoute
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._host import ChannelContext
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Channel-neutral request envelope
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class ChannelSession:
|
||||
"""Channel-supplied session hint.
|
||||
|
||||
The host turns this into an ``AgentSession`` keyed by ``isolation_key`` so
|
||||
every distinct end user gets their own context-provider state (e.g. one
|
||||
``FileHistoryProvider`` JSONL file per user).
|
||||
"""
|
||||
|
||||
def __init__(self, isolation_key: str | None = None) -> None:
|
||||
self.isolation_key = isolation_key
|
||||
|
||||
|
||||
class ChannelIdentity:
|
||||
"""Channel-native identity the host sees on each 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.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channel: str,
|
||||
native_id: str,
|
||||
attributes: Mapping[str, Any] | None = None,
|
||||
) -> None:
|
||||
self.channel = channel
|
||||
self.native_id = native_id
|
||||
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.
|
||||
|
||||
Kept as a dataclass so app authors can use ``dataclasses.replace(...)`` in
|
||||
run hooks to produce a modified envelope without re-listing every field.
|
||||
"""
|
||||
|
||||
channel: str
|
||||
operation: str # e.g. "message.create", "command.invoke"
|
||||
input: AgentRunInputs
|
||||
session: ChannelSession | None = None
|
||||
options: Mapping[str, Any] | None = None
|
||||
session_mode: str = "auto" # "auto" | "required" | "disabled"
|
||||
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:
|
||||
"""A discoverable command a channel exposes to its users (e.g. ``/reset``)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
description: str,
|
||||
handle: Callable[[ChannelCommandContext], Awaitable[None]],
|
||||
) -> None:
|
||||
self.name = name
|
||||
self.description = description
|
||||
self.handle = handle
|
||||
|
||||
|
||||
class ChannelCommandContext:
|
||||
"""Context passed to a :class:`ChannelCommand` handler."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
request: ChannelRequest,
|
||||
reply: Callable[[str], Awaitable[None]],
|
||||
) -> None:
|
||||
self.request = request
|
||||
self.reply = reply
|
||||
|
||||
|
||||
_EMPTY_ROUTES: tuple[BaseRoute, ...] = ()
|
||||
_EMPTY_COMMANDS: tuple[ChannelCommand, ...] = ()
|
||||
_EMPTY_LIFECYCLE: tuple[Callable[[], Awaitable[None]], ...] = ()
|
||||
|
||||
|
||||
class ChannelContribution:
|
||||
"""Routes, commands, and lifecycle hooks a channel contributes to the host."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
routes: Sequence[BaseRoute] = _EMPTY_ROUTES,
|
||||
commands: Sequence[ChannelCommand] = _EMPTY_COMMANDS,
|
||||
on_startup: Sequence[Callable[[], Awaitable[None]]] = _EMPTY_LIFECYCLE,
|
||||
on_shutdown: Sequence[Callable[[], Awaitable[None]]] = _EMPTY_LIFECYCLE,
|
||||
) -> None:
|
||||
self.routes = routes
|
||||
self.commands = commands
|
||||
self.on_startup = on_startup
|
||||
self.on_shutdown = on_shutdown
|
||||
|
||||
|
||||
class _Unset:
|
||||
"""Sentinel for ``HostedRunResult.replace`` overrides.
|
||||
|
||||
Distinguishes "caller did not pass this kwarg" from "caller passed
|
||||
``None`` explicitly" — needed because ``session`` is ``None`` in
|
||||
many envelopes and we want the no-arg call to preserve it.
|
||||
"""
|
||||
|
||||
|
||||
_UNSET = _Unset()
|
||||
|
||||
|
||||
TResult = TypeVar("TResult")
|
||||
|
||||
|
||||
class HostedRunResult(Generic[TResult]):
|
||||
r"""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.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
result: TResult,
|
||||
*,
|
||||
session: Any | None = None,
|
||||
) -> None:
|
||||
self.result = result
|
||||
self.session = session
|
||||
|
||||
def replace(
|
||||
self,
|
||||
*,
|
||||
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.
|
||||
"""
|
||||
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.
|
||||
"""
|
||||
|
||||
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."""
|
||||
|
||||
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."""
|
||||
|
||||
|
||||
# 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[
|
||||
[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.
|
||||
"""
|
||||
|
||||
name: str
|
||||
path: str # default mount path (e.g. "/responses"); use "" to mount routes at the app root
|
||||
|
||||
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: ...
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 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_response_hook",
|
||||
"apply_run_hook",
|
||||
]
|
||||
@@ -0,0 +1,110 @@
|
||||
[project]
|
||||
name = "agent-framework-hosting"
|
||||
description = "Multi-channel hosting for Microsoft Agent Framework agents."
|
||||
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
version = "1.0.0a260424"
|
||||
license-files = ["LICENSE"]
|
||||
urls.homepage = "https://aka.ms/agent-framework"
|
||||
urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
|
||||
urls.release_notes = "https://github.com/microsoft/agent-framework/releases?q=tag%3Apython-1&expanded=true"
|
||||
urls.issues = "https://github.com/microsoft/agent-framework/issues"
|
||||
classifiers = [
|
||||
"License :: OSI Approved :: MIT License",
|
||||
"Development Status :: 3 - Alpha",
|
||||
"Intended Audience :: Developers",
|
||||
"Programming Language :: Python :: 3",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
"Programming Language :: Python :: 3.13",
|
||||
"Programming Language :: Python :: 3.14",
|
||||
"Typing :: Typed",
|
||||
]
|
||||
dependencies = [
|
||||
"agent-framework-core>=1.2.0,<2",
|
||||
"starlette>=0.37",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
serve = [
|
||||
"hypercorn>=0.17",
|
||||
]
|
||||
disk = [
|
||||
"diskcache>=5.6",
|
||||
]
|
||||
|
||||
[tool.uv]
|
||||
prerelease = "if-necessary-or-explicit"
|
||||
environments = [
|
||||
"sys_platform == 'darwin'",
|
||||
"sys_platform == 'linux'",
|
||||
"sys_platform == 'win32'"
|
||||
]
|
||||
|
||||
[tool.uv-dynamic-versioning]
|
||||
fallback-version = "0.0.0"
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
testpaths = 'tests'
|
||||
addopts = "-ra -q -r fEX"
|
||||
asyncio_mode = "auto"
|
||||
asyncio_default_fixture_loop_scope = "function"
|
||||
filterwarnings = []
|
||||
timeout = 120
|
||||
markers = [
|
||||
"integration: marks tests as integration tests that require external services",
|
||||
]
|
||||
|
||||
[tool.ruff]
|
||||
extend = "../../pyproject.toml"
|
||||
|
||||
[tool.coverage.run]
|
||||
omit = [
|
||||
"**/__init__.py"
|
||||
]
|
||||
|
||||
[tool.pyright]
|
||||
extends = "../../pyproject.toml"
|
||||
include = ["agent_framework_hosting"]
|
||||
exclude = ['tests']
|
||||
|
||||
[tool.mypy]
|
||||
plugins = ['pydantic.mypy']
|
||||
strict = true
|
||||
python_version = "3.10"
|
||||
ignore_missing_imports = true
|
||||
disallow_untyped_defs = true
|
||||
no_implicit_optional = true
|
||||
check_untyped_defs = true
|
||||
warn_return_any = true
|
||||
show_error_codes = true
|
||||
warn_unused_ignores = false
|
||||
disallow_incomplete_defs = true
|
||||
disallow_untyped_decorators = true
|
||||
|
||||
[tool.bandit]
|
||||
targets = ["agent_framework_hosting"]
|
||||
exclude_dirs = ["tests"]
|
||||
|
||||
[tool.poe]
|
||||
executor.type = "uv"
|
||||
include = "../../shared_tasks.toml"
|
||||
|
||||
[tool.poe.tasks.mypy]
|
||||
help = "Run MyPy for this package."
|
||||
cmd = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_hosting"
|
||||
|
||||
[tool.poe.tasks.test]
|
||||
help = "Run the default unit test suite for this package."
|
||||
cmd = 'pytest -m "not integration" --cov=agent_framework_hosting --cov-report=term-missing:skip-covered tests'
|
||||
|
||||
[build-system]
|
||||
requires = ["flit-core >= 3.11,<4.0"]
|
||||
build-backend = "flit_core.buildapi"
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"httpx>=0.28.1",
|
||||
]
|
||||
@@ -0,0 +1,43 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Workflow fixtures for hosting tests.
|
||||
|
||||
Defined in a module that does not use ``from __future__ import annotations``
|
||||
because the workflow handler validation reflects on real annotation objects
|
||||
rather than stringified forms.
|
||||
"""
|
||||
|
||||
from agent_framework import Executor, Workflow, WorkflowBuilder, WorkflowContext, handler
|
||||
|
||||
|
||||
class _UpperExecutor(Executor):
|
||||
@handler
|
||||
async def handle(self, text: str, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.yield_output(text.upper())
|
||||
|
||||
|
||||
class _EchoExecutor(Executor):
|
||||
@handler
|
||||
async def handle(self, text: str, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.yield_output(text)
|
||||
|
||||
|
||||
def build_upper_workflow() -> Workflow:
|
||||
return WorkflowBuilder(start_executor=_UpperExecutor(id="upper")).build()
|
||||
|
||||
|
||||
def build_echo_workflow() -> Workflow:
|
||||
return WorkflowBuilder(start_executor=_EchoExecutor(id="echo")).build()
|
||||
|
||||
|
||||
class _MultiChunkExecutor(Executor):
|
||||
"""Yields three separate ``output`` events so streaming has something to chew on."""
|
||||
|
||||
@handler
|
||||
async def handle(self, text: str, ctx: WorkflowContext[str]) -> None:
|
||||
for chunk in (f"{text}-1", f"{text}-2", f"{text}-3"):
|
||||
await ctx.yield_output(chunk)
|
||||
|
||||
|
||||
def build_multi_chunk_workflow() -> Workflow:
|
||||
return WorkflowBuilder(start_executor=_MultiChunkExecutor(id="multi")).build()
|
||||
@@ -0,0 +1,580 @@
|
||||
# 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"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,424 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for ``state_dir`` wired through :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,
|
||||
)
|
||||
|
||||
# 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."""
|
||||
|
||||
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()
|
||||
|
||||
|
||||
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))
|
||||
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."""
|
||||
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
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 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/``."""
|
||||
host = AgentFrameworkHost(
|
||||
target=_AgentStub(),
|
||||
channels=[_ChannelStub()],
|
||||
state_dir=tmp_path,
|
||||
)
|
||||
try:
|
||||
assert host._sessions_store is not None
|
||||
assert (tmp_path / "runner").is_dir()
|
||||
assert (tmp_path / "sessions").is_dir()
|
||||
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"
|
||||
sessions_dir = tmp_path / "state"
|
||||
host = AgentFrameworkHost(
|
||||
target=_AgentStub(),
|
||||
channels=[_ChannelStub()],
|
||||
state_dir={"runner": runner_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")
|
||||
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."""
|
||||
with pytest.raises(ValueError, match="unknown"):
|
||||
AgentFrameworkHost(
|
||||
target=_AgentStub(),
|
||||
channels=[_ChannelStub()],
|
||||
state_dir={"runnerr": tmp_path / "x"}, # 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"
|
||||
|
||||
host1 = AgentFrameworkHost(target=_AgentStub(), channels=[_ChannelStub()], state_dir=state_dir)
|
||||
host1._session_aliases["user-1"] = "sess-abc"
|
||||
host1._session_aliases["user-2"] = "sess-def"
|
||||
_close_host_disk(host1)
|
||||
|
||||
host2 = AgentFrameworkHost(target=_AgentStub(), channels=[_ChannelStub()], state_dir=state_dir)
|
||||
try:
|
||||
assert host2._session_aliases["user-1"] == "sess-abc"
|
||||
assert host2._session_aliases["user-2"] == "sess-def"
|
||||
finally:
|
||||
_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
|
||||
|
||||
return build_upper_workflow()
|
||||
|
||||
|
||||
def test_single_path_state_dir_wires_workflow_checkpoints(tmp_path: Path) -> None:
|
||||
"""``state_dir="/foo"`` + workflow target → ``/foo/checkpoints/`` is used."""
|
||||
workflow = _build_simple_workflow()
|
||||
host = AgentFrameworkHost(
|
||||
target=workflow,
|
||||
channels=[_ChannelStub()],
|
||||
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)
|
||||
|
||||
|
||||
def test_mapping_state_dir_checkpoints_key_wires_workflow_checkpoints(tmp_path: Path) -> None:
|
||||
"""``state_dir={"checkpoints": ...}`` + workflow target → that path is used."""
|
||||
workflow = _build_simple_workflow()
|
||||
ckpt_dir = tmp_path / "ck"
|
||||
host = AgentFrameworkHost(
|
||||
target=workflow,
|
||||
channels=[_ChannelStub()],
|
||||
state_dir={"checkpoints": ckpt_dir},
|
||||
)
|
||||
try:
|
||||
assert host._checkpoint_location == ckpt_dir
|
||||
# No diskcache components were requested.
|
||||
assert host._sessions_store is None
|
||||
finally:
|
||||
_close_host_disk(host)
|
||||
|
||||
|
||||
def test_mapping_state_dir_omits_checkpoints_for_workflow(tmp_path: Path) -> None:
|
||||
"""Mapping form lets workflow callers opt out of checkpoint persistence."""
|
||||
workflow = _build_simple_workflow()
|
||||
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"},
|
||||
)
|
||||
try:
|
||||
assert host._checkpoint_location is None
|
||||
finally:
|
||||
_close_host_disk(host)
|
||||
|
||||
|
||||
def test_explicit_checkpoint_location_wins_over_state_dir(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""``checkpoint_location`` + ``state_dir`` → explicit param wins + warn."""
|
||||
workflow = _build_simple_workflow()
|
||||
explicit = tmp_path / "explicit-ck"
|
||||
with caplog.at_level("WARNING", logger="agent_framework.hosting"):
|
||||
host = AgentFrameworkHost(
|
||||
target=workflow,
|
||||
channels=[_ChannelStub()],
|
||||
checkpoint_location=explicit,
|
||||
state_dir=tmp_path,
|
||||
)
|
||||
try:
|
||||
assert host._checkpoint_location == explicit
|
||||
assert any(
|
||||
"state_dir['checkpoints']" in rec.message and "checkpoint_location" in rec.message for rec in caplog.records
|
||||
)
|
||||
finally:
|
||||
_close_host_disk(host)
|
||||
|
||||
|
||||
def test_state_dir_checkpoints_for_agent_target_silent_for_single_path(tmp_path: Path) -> None:
|
||||
"""Single-path state_dir + agent target → no checkpoint, no warning."""
|
||||
host = AgentFrameworkHost(
|
||||
target=_AgentStub(),
|
||||
channels=[_ChannelStub()],
|
||||
state_dir=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)
|
||||
|
||||
|
||||
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)."""
|
||||
with caplog.at_level("WARNING", logger="agent_framework.hosting"):
|
||||
host = AgentFrameworkHost(
|
||||
target=_AgentStub(),
|
||||
channels=[_ChannelStub()],
|
||||
state_dir={"checkpoints": tmp_path / "ck"},
|
||||
)
|
||||
try:
|
||||
assert host._checkpoint_location is None
|
||||
assert any(
|
||||
"state_dir['checkpoints']" in rec.message and "not a Workflow" in rec.message for rec in caplog.records
|
||||
)
|
||||
finally:
|
||||
_close_host_disk(host)
|
||||
|
||||
|
||||
def test_state_dir_checkpoints_conflicts_with_workflow_own_storage(tmp_path: Path) -> None:
|
||||
"""Derived checkpoint path triggers the same conflict guard as explicit."""
|
||||
from agent_framework import InMemoryCheckpointStorage, WorkflowBuilder
|
||||
|
||||
from tests._workflow_fixtures import _UpperExecutor
|
||||
|
||||
workflow = WorkflowBuilder(
|
||||
start_executor=_UpperExecutor(id="upper"),
|
||||
checkpoint_storage=InMemoryCheckpointStorage(),
|
||||
).build()
|
||||
with pytest.raises(RuntimeError, match="already has checkpoint storage"):
|
||||
AgentFrameworkHost(
|
||||
target=workflow,
|
||||
channels=[_ChannelStub()],
|
||||
state_dir=tmp_path,
|
||||
)
|
||||
@@ -0,0 +1,282 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for the per-request isolation contextvar surface in
|
||||
:mod:`agent_framework_hosting._isolation`.
|
||||
|
||||
The isolation keys are the ONLY seam Foundry-aware providers use to
|
||||
find partition keys, and the host's ASGI middleware lifts them off the
|
||||
two well-known headers on every inbound HTTP request. A regression
|
||||
that drops the lookup, mistypes a header name, or fails to reset the
|
||||
contextvar would silently misroute writes / leak per-request state
|
||||
across requests, with zero unit-test signal — so cover the surface
|
||||
fully here.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import JSONResponse
|
||||
from starlette.routing import BaseRoute, Route
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from agent_framework_hosting import (
|
||||
Channel,
|
||||
ChannelContext,
|
||||
ChannelContribution,
|
||||
IsolationKeys,
|
||||
get_current_isolation_keys,
|
||||
reset_current_isolation_keys,
|
||||
set_current_isolation_keys,
|
||||
)
|
||||
from agent_framework_hosting._isolation import ( # pyright: ignore[reportPrivateUsage]
|
||||
ISOLATION_HEADER_CHAT,
|
||||
ISOLATION_HEADER_USER,
|
||||
current_isolation_keys,
|
||||
)
|
||||
|
||||
|
||||
class TestIsolationKeys:
|
||||
def test_defaults_to_none_pair(self) -> None:
|
||||
keys = IsolationKeys()
|
||||
assert keys.user_key is None
|
||||
assert keys.chat_key is None
|
||||
assert keys.is_empty is True
|
||||
|
||||
def test_partial_with_only_user_is_not_empty(self) -> None:
|
||||
keys = IsolationKeys(user_key="alice")
|
||||
assert keys.user_key == "alice"
|
||||
assert keys.chat_key is None
|
||||
assert keys.is_empty is False
|
||||
|
||||
def test_partial_with_only_chat_is_not_empty(self) -> None:
|
||||
keys = IsolationKeys(chat_key="general")
|
||||
assert keys.is_empty is False
|
||||
|
||||
def test_full_pair_is_not_empty(self) -> None:
|
||||
keys = IsolationKeys(user_key="alice", chat_key="general")
|
||||
assert keys.is_empty is False
|
||||
|
||||
|
||||
class TestContextVarHelpers:
|
||||
def test_default_is_none(self) -> None:
|
||||
# Each test gets a fresh contextvar value because pytest runs
|
||||
# tests in fresh contexts. ``get`` returns the default.
|
||||
assert get_current_isolation_keys() is None
|
||||
|
||||
def test_set_and_get_round_trip(self) -> None:
|
||||
token = set_current_isolation_keys(IsolationKeys(user_key="alice", chat_key="general"))
|
||||
try:
|
||||
current = get_current_isolation_keys()
|
||||
assert current is not None
|
||||
assert current.user_key == "alice"
|
||||
assert current.chat_key == "general"
|
||||
finally:
|
||||
reset_current_isolation_keys(token)
|
||||
# Reset restores prior value (None in the default context).
|
||||
assert get_current_isolation_keys() is None
|
||||
|
||||
def test_set_with_none_clears(self) -> None:
|
||||
outer = set_current_isolation_keys(IsolationKeys(user_key="alice"))
|
||||
try:
|
||||
inner = set_current_isolation_keys(None)
|
||||
try:
|
||||
assert get_current_isolation_keys() is None
|
||||
finally:
|
||||
reset_current_isolation_keys(inner)
|
||||
# Reset surfaces the outer value again.
|
||||
current = get_current_isolation_keys()
|
||||
assert current is not None
|
||||
assert current.user_key == "alice"
|
||||
finally:
|
||||
reset_current_isolation_keys(outer)
|
||||
|
||||
def test_module_level_contextvar_is_the_same_instance(self) -> None:
|
||||
"""Direct contextvar access (used by the ASGI middleware) and the
|
||||
public `get_current_isolation_keys()` helper read from the SAME
|
||||
underlying contextvar. A regression that introduced a second
|
||||
contextvar would silently break the middleware → provider hop."""
|
||||
token = current_isolation_keys.set(IsolationKeys(user_key="bob"))
|
||||
try:
|
||||
via_helper = get_current_isolation_keys()
|
||||
assert via_helper is not None
|
||||
assert via_helper.user_key == "bob"
|
||||
finally:
|
||||
current_isolation_keys.reset(token)
|
||||
|
||||
|
||||
class TestHeaderConstants:
|
||||
"""The two header names are part of the public contract — they
|
||||
match the ones the Foundry Hosted Agents runtime stamps on every
|
||||
inbound request. A typo here would silently misroute partition
|
||||
writes."""
|
||||
|
||||
def test_user_header_value(self) -> None:
|
||||
assert ISOLATION_HEADER_USER == "x-agent-user-isolation-key"
|
||||
|
||||
def test_chat_header_value(self) -> None:
|
||||
assert ISOLATION_HEADER_CHAT == "x-agent-chat-isolation-key"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# End-to-end: ASGI middleware lifts the headers into the contextvar.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class _IsolationProbeChannel:
|
||||
"""A minimal Channel that exposes a single GET route which captures
|
||||
the contextvar value INSIDE the request and returns it as JSON.
|
||||
|
||||
Tests use this to exercise the full middleware → contextvar →
|
||||
handler hop end-to-end.
|
||||
"""
|
||||
|
||||
name = "probe"
|
||||
path = ""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.captured: list[IsolationKeys | None] = []
|
||||
|
||||
async def _handler(_request: Request) -> JSONResponse:
|
||||
keys = get_current_isolation_keys()
|
||||
self.captured.append(keys)
|
||||
payload = (
|
||||
{"user": keys.user_key, "chat": keys.chat_key}
|
||||
if keys is not None
|
||||
else {"user": None, "chat": None, "_present": False}
|
||||
)
|
||||
return JSONResponse(payload)
|
||||
|
||||
self._routes: list[BaseRoute] = [Route("/probe", _handler)]
|
||||
|
||||
def contribute(self, _context: ChannelContext) -> ChannelContribution:
|
||||
return ChannelContribution(routes=self._routes)
|
||||
|
||||
|
||||
def _make_host_with_probe() -> tuple[object, _IsolationProbeChannel]:
|
||||
from agent_framework_hosting import AgentFrameworkHost
|
||||
|
||||
class _NoopAgent:
|
||||
async def run(self, *_args: object, **_kwargs: object) -> object: # pragma: no cover - never called
|
||||
raise RuntimeError("not invoked")
|
||||
|
||||
probe = _IsolationProbeChannel()
|
||||
assert isinstance(probe, Channel)
|
||||
host = AgentFrameworkHost(target=_NoopAgent(), channels=[probe]) # type: ignore[arg-type]
|
||||
return host, probe
|
||||
|
||||
|
||||
class TestIsolationMiddlewareEndToEnd:
|
||||
def test_both_headers_lifted_into_contextvar(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": "alice-uid", "chat": "general-cid"}
|
||||
assert len(probe.captured) == 1
|
||||
captured = probe.captured[0]
|
||||
assert captured is not None
|
||||
assert captured.user_key == "alice-uid"
|
||||
assert captured.chat_key == "general-cid"
|
||||
|
||||
def test_only_user_header_lifted(self) -> None:
|
||||
"""One-header-only branch: the middleware still binds (chat=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"})
|
||||
assert r.status_code == 200
|
||||
assert r.json() == {"user": "alice-uid", "chat": None}
|
||||
|
||||
def test_only_chat_header_lifted(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_CHAT: "general-cid"})
|
||||
assert r.status_code == 200
|
||||
assert r.json() == {"user": None, "chat": "general-cid"}
|
||||
|
||||
def test_no_headers_keeps_contextvar_none(self) -> None:
|
||||
"""Local-dev path: with neither header present the middleware is
|
||||
a no-op and the contextvar stays at its default ``None`` —
|
||||
providers see "no isolation" and route to the in-memory
|
||||
fallback rather than picking up stale per-request state."""
|
||||
host, probe = _make_host_with_probe()
|
||||
with TestClient(host.app) as client: # type: ignore[attr-defined]
|
||||
r = client.get("/probe")
|
||||
assert r.status_code == 200
|
||||
assert r.json() == {"user": None, "chat": None, "_present": False}
|
||||
assert probe.captured == [None]
|
||||
|
||||
def test_empty_header_value_treated_as_absent(self) -> None:
|
||||
"""A header that's present but empty must not bind an empty key —
|
||||
``IsolationContext`` rejects empty strings on the read side."""
|
||||
host, probe = _make_host_with_probe()
|
||||
with TestClient(host.app) as client: # type: ignore[attr-defined]
|
||||
r = client.get(
|
||||
"/probe",
|
||||
headers={
|
||||
ISOLATION_HEADER_USER: "",
|
||||
ISOLATION_HEADER_CHAT: "general-cid",
|
||||
},
|
||||
)
|
||||
assert r.status_code == 200
|
||||
# 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:
|
||||
"""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."""
|
||||
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"})
|
||||
assert r1.status_code == 200
|
||||
# Reading the contextvar OUTSIDE the request scope must see
|
||||
# the default — not the value the prior request bound.
|
||||
assert get_current_isolation_keys() is None
|
||||
# And a follow-up request without headers gets a clean
|
||||
# ``None`` rather than inheriting alice-uid.
|
||||
r2 = client.get("/probe")
|
||||
assert r2.json() == {"user": None, "chat": None, "_present": False}
|
||||
|
||||
def test_concurrent_requests_get_isolated_contextvars(self) -> None:
|
||||
"""Different requests run in different async contexts; binding
|
||||
from request A must NOT leak into a concurrent request B."""
|
||||
host, probe = _make_host_with_probe()
|
||||
|
||||
async def _drive() -> None:
|
||||
# Run two requests in parallel asyncio tasks against the
|
||||
# same TestClient and assert their captures don't bleed
|
||||
# into each other.
|
||||
async def _hit(user_key: str) -> dict[str, str | None]:
|
||||
with TestClient(host.app) as client: # type: ignore[attr-defined]
|
||||
r = client.get("/probe", headers={ISOLATION_HEADER_USER: user_key})
|
||||
return r.json() # type: ignore[no-any-return]
|
||||
|
||||
r_alice, r_bob = await asyncio.gather(_hit("alice-uid"), _hit("bob-uid"))
|
||||
assert r_alice == {"user": "alice-uid", "chat": None}
|
||||
assert r_bob == {"user": "bob-uid", "chat": None}
|
||||
|
||||
asyncio.run(_drive())
|
||||
|
||||
|
||||
class TestNonHttpScopesPassThrough:
|
||||
"""The middleware intentionally only inspects ``http`` scopes;
|
||||
lifespan / websocket scopes are forwarded untouched. A regression
|
||||
that touched lifespan scopes here would crash boot."""
|
||||
|
||||
async def test_lifespan_scope_does_not_consult_headers(self) -> None:
|
||||
# The TestClient context manager exercises the lifespan scope
|
||||
# implicitly; if the middleware tried to decode headers on a
|
||||
# non-http scope this would raise. Exercise it without binding
|
||||
# any contextvar work.
|
||||
host, _probe = _make_host_with_probe()
|
||||
with TestClient(host.app): # type: ignore[attr-defined]
|
||||
# Just enter / exit; no requests.
|
||||
pass
|
||||
@@ -0,0 +1,333 @@
|
||||
# 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)
|
||||
@@ -0,0 +1,278 @@
|
||||
# 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()
|
||||
@@ -0,0 +1,252 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for the channel-neutral envelope types in :mod:`agent_framework_hosting._types`."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from agent_framework_hosting import (
|
||||
ChannelIdentity,
|
||||
ChannelRequest,
|
||||
ChannelSession,
|
||||
DurableTaskPayloadMode,
|
||||
ResponseTarget,
|
||||
ResponseTargetKind,
|
||||
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")
|
||||
assert req.channel == "responses"
|
||||
assert req.operation == "message.create"
|
||||
assert req.input == "hi"
|
||||
assert req.session is None
|
||||
assert req.options is None
|
||||
assert req.session_mode == "auto"
|
||||
assert req.metadata == {}
|
||||
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(
|
||||
channel="telegram",
|
||||
operation="message.create",
|
||||
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:
|
||||
def test_attributes_default_empty_mapping(self) -> None:
|
||||
ident = ChannelIdentity(channel="teams", native_id="abc")
|
||||
assert dict(ident.attributes) == {}
|
||||
|
||||
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 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
|
||||
Reference in New Issue
Block a user