feat: add token vault support for auth0-ai-ms-agent SDK

[adam-wang-okta-public]

Description

• Adds the auth0-ai-ms-agent package — an Auth0 AI SDK adapter for the https://github.com/microsoft/agent-framework.

  This PR implements the Token Vault feature, which allows MS Agent tools to obtain access tokens for third-party APIs (Google, GitHub, Slack, etc.) on behalf of the user via Auth0's Token Vault.

Design doc

• https://oktawiki.atlassian.net/wiki/spaces/AIDEV/pages/4491870439/Auth0-ai-ms-agent+SDK+development

Public API example

  from auth0_ai_ms_agent.auth0_ai import Auth0AI

  auth0_ai = Auth0AI()

  with_google_calendar_access = auth0_ai.with_token_vault(
      connection="google-oauth2",
      scopes=["openid", "https://www.googleapis.com/auth/calendar.freebusy"],
      refresh_token=lambda *_args, **_kwargs: session["user"]["refresh_token"],
  )

  check_calendar_tool = with_google_calendar_access(
      FunctionTool(
          name="check_user_calendar",
          description="Check if the user is available on a certain date and time",
          func=tool_function,
      )
  )

Implementation details

• Auth0AI.with_token_vault() — returns a decorator that wraps an MS Agent FunctionTool with token vault authorization, delegating to TokenVaultAuthorizerBase.protect() from the core auth0-ai SDK.
• tool_wrapper — the MS Agent-specific wrapping layer that:
  • Extracts the session (injected by the framework via additional_function_arguments) and uses session.session_id as the thread_id for credential namespace resolution
  • Strips framework-injected runtime kwargs (chat_options, tools, tool_choice, options, response_format, conversation_id) before invoking the original function
  • Catches TokenVaultInterrupt and stores it in session.state["pending_interrupt"] before re-raising, so callers can inspect and redirect the user to complete authorization
  • Preserves all original FunctionTool configuration (approval_mode, max_invocations, max_invocation_exceptions, additional_properties, input schema) on the wrapped tool via FunctionTool.init introspection
  • Supports both sync and async tool functions
• Interrupt handling — when the token vault cannot provide a valid token, a TokenVaultInterrupt is raised. The caller checks session.state.get("pending_interrupt") and uses interrupt.connection, interrupt.scopes, interrupt.required_scopes, and interrupt.authorization_params to redirect the user.

Note

• original_func is called directly rather than through FunctionTool.call to avoid double-counting invocation metrics already tracked at the wrapped tool boundary.
• tool_call_id is generated per invocation (rather than forwarded from the framework) because FunctionTool.invoke() consumes it before it reaches the wrapped function.
• This SDK pins the agent-framework to ^1.0.0rc2

Testing

Manual QA

• End-to-end validation was performed using a working example (calling Google Calendar via the token vault flow). The example will be submitted in a follow-up PR once the code is cleaned up and the README is finalized.

Unit test

Unit tests are provided in tests/test_auth0_ai.py (15 tests, 100% coverage).

• Tests mock at the HTTP boundary (TokenVaultAuthorizerBase.get_access_token_impl) so that the core SDK's protect() logic runs for real, catching any interface mismatches between the adapter and the core SDK.

Test coverage includes:

• Decorator contract (returns callable, preserves tool name, description, schema, and all framework configuration)

• Authorization pass path for both sync and async tool functions

• Authorization fail path (TokenVaultInterrupt raised, session state set, original function not executed)

• Framework kwargs stripping before calling the original function

• Non-TokenVaultInterrupt exceptions propagate unchanged without setting session state

• RuntimeError raised when no session is provided

• To run tests:

  cd packages/auth0-ai-ms-agent
  poetry install
  poetry run pytest tests/ --cov=auth0_ai_ms_agent --cov-report=term-missing -v

Checklist

• I have added documentation for new/changed functionality in this PR or in auth0.com/docs
• All active GitHub checks for tests, formatting, and security are passing
• The correct base branch is being used, if not the default branch

References

• MS agent framework public doc
• MS agent framework repo
• Jira ticket
  • https://auth0team.atlassian.net/browse/AIDX-487
  • https://auth0team.atlassian.net/browse/AIDX-488

[agupta-ghub]

Quickly reviewed it with Claude

PR Review

The core logic is sound — framework kwargs stripping, context building, protect_fn delegation, and sync/async handling all follow the
patterns established by the LangChain and LlamaIndex adapters. A few issues need addressing before merge, one of which is a publishing
blocker.

---

🔴 Critical (blocking)

auth0-ai path dependency will break PyPI publishing

pyproject.toml puts the local path dep in main dependencies:

[tool.poetry.dependencies]
auth0-ai = { path = "../auth0-ai", develop = true }  # won't resolve from PyPI

Anyone who runs pip install auth0-ai-ms-agent will get a broken package. The LangChain adapter handles this correctly — version pin in
main deps, path override in dev deps:

[tool.poetry.dependencies]
auth0-ai = "^1.0.x"

[tool.poetry.group.dev.dependencies]
auth0-ai = { path = "../auth0-ai", develop = true }

---
🟠 Major

Dead openfga-sdk dependency

openfga-sdk = "^0.9.5" is listed but there is no FGA code in this package. The other adapters include it because they ship FGA support —
this one is token-vault only. Remove it to avoid unnecessary transitive dependencies for all users.

---
TokenVaultAuthorizer incorrectly inherits from ABC

token_vault_authorizer.py:11 — TokenVaultAuthorizer is a concrete class, it's directly instantiated in auth0_ai.py. Declaring it abstract
is contradictory. TokenVaultAuthorizerBase is Generic, not ABC, so this isn't even inherited. The boilerplate __init__ that only calls
super() can also be removed entirely.

# current
class TokenVaultAuthorizer(TokenVaultAuthorizerBase, ABC):
    def __init__(self, params, auth0=None):
        super().__init__(params, auth0)

# suggested
class TokenVaultAuthorizer(TokenVaultAuthorizerBase):
    pass  # or just remove __init__ entirely

---
inspect._empty is a private API

tool_wrapper.py:45:

# current — uses private implementation detail
if value is not None or param.default is inspect._empty:

# fix — use the public API
if value is not None or param.default is inspect.Parameter.empty:

---
Non-serializable exception stored in session state

tool_wrapper.py:92:

session.state["pending_interrupt"] = e  # stores a live Python exception object

If AgentSession.state is ever persisted between turns (Redis, DB, over the wire) — which is the point of a session in multi-turn
conversations — serializing a Python exception will fail. The data callers actually need (connection, scopes, required_scopes,
authorization_params) is all on the interrupt object; the exception wrapper is not necessary. The README's isinstance(interrupt,
TokenVaultInterrupt) check also couples callers to never persisting state. At minimum this constraint should be clearly documented;
ideally the stored value is a plain serializable dict.

---
pending_interrupt is never cleared on subsequent runs

After a TokenVaultInterrupt, session.state["pending_interrupt"] is written but never cleaned up. If the user completes authorization and
the agent runs successfully on the next turn, the stale interrupt is still present. Any caller checking
session.state.get("pending_interrupt") after a successful run will see a false positive.

Simple fix — clear it at the start of each invocation before calling protect_fn:

session.state.pop("pending_interrupt", None)

---
🟡 Minor

**params: TokenVaultAuthorizerParams type annotation is misleading

auth0_ai.py:27 — **params: T means "each kwarg value has type T", not "these kwargs together conform to T's shape".
TokenVaultAuthorizerParams is a regular class, not a TypedDict. The annotation doesn't give callers useful IDE autocomplete and is
semantically incorrect. Either annotate as **kwargs: Any or explicitly declare each parameter mirroring the
TokenVaultAuthorizerParams.__init__ signature.

---
_schema_supplied private attribute is fragile

tool_wrapper.py:96:

schema_or_model = tool.parameters() if getattr(tool, "_schema_supplied", False) else input_model

_schema_supplied is a private attribute on an RC-stage FunctionTool. The getattr default of False means if this attribute is renamed or
removed in a future framework version, this silently falls back to input_model with no error — just a quiet schema regression. Deserves a
comment explaining why it's needed and what the silent fallback means, so it's easy to audit on framework upgrades.

---
_FRAMEWORK_KWARGS risk going stale silently

The set is a deny-list of framework-injected kwargs derived from a specific line in the RC framework. When the framework graduates from RC
 and adds/removes injected kwargs, this set won't update automatically. If a new framework kwarg isn't listed here it leaks into
original_func as an unexpected argument. The comment linking to the source line is good — worth also noting explicitly that this is a
maintenance point that needs revisiting on each framework version bump.

---
Missing pytest asyncio mode configuration

With pytest-asyncio = "^0.25.0", omitting asyncio_mode in pyproject.toml produces deprecation warnings. Add:

[tool.pytest.ini_options]
asyncio_mode = "auto"

---
Missing newline at end of token_vault_authorizer.py

---
🔵 Test coverage gaps

- No test for stale pending_interrupt being cleared on a successful retry — would directly catch the issue raised above.
- tool_call_id uniqueness is untested — the PR explicitly notes this is generated per-invocation as a design decision; worth asserting
that two successive invocations produce distinct tool_call_id values in their contexts.
- No test for the non-_schema_supplied path — test_returned_function_tool_preserves_schema_supplied_input_model covers the schema dict
path but not the pydantic model path.

[adam-wang-okta-public]

Comments were addressed with code update in a new commit. Complete replies to comments as below.

#1 — PyPI path dependency (blocking)
Fixed. auth0-ai is now pinned as "^1.0.2" in [tool.poetry.dependencies] and the local path override moved to [tool.poetry.group.dev.dependencies] only.

---

#2 — openfga-sdk dead dependency
FGA (Fine-Grained Authorization) is a planned feature for this SDK. The README already lists it under "coming soon". Keeping the dependency now.

---

#3 — Incorrect ABC + redundant init
Fixed. Removed ABC from the inheritance — TokenVaultAuthorizer is a concrete class that implements authorizer() and is directly instantiated, so ABC was misleading. Also removed the init that only called super().init unchanged, which is redundant in Python (MRO handles it automatically).

---

#4 — inspect._empty private API
Fixed. Changed to the public inspect.Parameter.empty.

---

#5 — AgentSession.state serialization
Not a concern in this context. agent.session is injected by the MS Agent Framework as a runtime parameter to the tool call — it is never persisted. It's a live in-memory object for the duration of the invocation. The external store support in auth0-ai core is for credentials, which is a separate concern entirely. Storing a live TokenVaultInterrupt in session.state is safe.

---

#6 — Stale interrupt never cleared on retry
Fixed. The success path now calls session.state.pop("pending_interrupt", None) before returning, clearing any interrupt left over from a previous failed attempt. Added test_stale_interrupt_is_cleared_on_successful_retry to cover this scenario — it pre-seeds session.state["pending_interrupt"] and asserts it is gone after a successful run.

---

#7 — **params: TokenVaultAuthorizerParams misleading
Unlike typescript, Python annotations is mainly for document purpose, and aren't enforced at runtime or compile time. We're intentionally keeping TokenVaultAuthorizerParams here as inline documentation — it signals to callers exactly what structure the kwargs should conform to.

---

#8 — Fragile _schema_supplied private attribute
Acknowledged. The getattr(tool, "_schema_supplied", False) pattern means if the framework ever removes this attribute the code silently falls back to input_model, which is the safe default. The reference link in _FRAMEWORK_KWARGS establishes the pattern of tracking internal framework details — we'll apply the same vigilance here on framework upgrades.

---

#9 — _FRAMEWORK_KWARGS silently going stale
Agreed this is a maintenance risk. The comment on line 9 already links directly to the framework source where these kwargs are injected, making it straightforward to diff on upgrades. We'll keep this as a checklist item when bumping agent-framework.

---

#10 — Missing pytest asyncio config
All async tests use explicit @pytest.mark.asyncio markers, which is correct for pytest-asyncio strict mode (the default in 0.25+). No asyncio_mode = "auto" config is needed.

---

#11 — Missing trailing newline
Fixed.

---

Test gaps — tool_call_id uniqueness + non-_schema_supplied path
tool_call_id uniqueness is guaranteed by uuid.uuid4() — testing it would be testing the standard library, not our logic. The non-_schema_supplied path is exercised implicitly by every test that wraps a FunctionTool without a schema dict — test_returned_function_tool_preserves_configuration being the clearest example.

Thanks for the review. @agupta-ghub

[agupta-ghub]

Thanks for the thorough responses, most items are resolved. Just my 2 cents :

---

openfga-sdk dependency

"Coming soon" is not a sufficient reason to ship a transitive dependency to all users today. Every pip install auth0-ai-ms-agent will pull in the FGA dependency tree even though there is zero FGA code in this package. The right time to add the dep is when the FGA code ships, in a follow-up PR alongside the feature. Do we know when it be ready? is it before or after this SDK ships?

params: TokenVaultAuthorizerParams annotation
If the goal is documentation, the docstring already does that job. For the annotation, use kwargs: Any. If you want real type safety and autocomplete, use kwargs: Unpack[...] with a TypedDict (available from typing since Python 3.11, which is already your minimum). But params: TokenVaultAuthorizerParams doesn't achieve either goal correctly.

[adam-wang-okta-public]

1. FGA is removed in this PR, but for clarity, this PR merge will not ship a new release. This repo still needs manual release. Both token vault and FGA are in scope of this ms-agent SDK release. That means the planned auth0-ai-ms-agent SDK will not release until the FGA, (plus async auth) is built in.

2. Unpack[TypedDict] is the right solution, but TokenVaultAuthorizerParams is a Generic class with complex Union types — implementing Unpack requires a parallel TypedDict that mirrors its __init__ signature, creating duplication that needs to stay in sync. The clean fix is to refactor TokenVaultAuthorizerParams to a TypedDict in the core auth0-ai package, which would benefit all adapters (langchain, llamaindex) and is worth its own PR. Changing to **kwargs: Any here loses the documentation value, so keeping **params: TokenVaultAuthorizerParams until the proper Unpack refactor lands.

@agupta-ghub

[priley86]

I mainly reviewed this update for consistency with the existing packages, w/o the example posted here. Here's what I'm seeing w/ Claude / Opus 4.6:

## PR Review: `auth0-ai-ms-agent` Package Consistency

I compared `packages/auth0-ai-ms-agent` against the three existing packages (`auth0-ai`, `auth0-ai-langchain`, `auth0-ai-llamaindex`) and found the following inconsistencies:

### 1. Missing `LICENSE` file

All three existing packages include an Apache 2.0 `LICENSE` file at the package root. The new `auth0-ai-ms-agent` package does **not** have one, even though its `pyproject.toml` declares `license = "apache-2.0"`.

### 2. Missing `tests/__init__.py`

All three existing packages have an empty `tests/__init__.py`. The new package is missing it. While this may work in practice, it's inconsistent and could cause issues with some test discovery configurations.

### 3. Missing README badges

The `auth0-ai-langchain` and `auth0-ai-llamaindex` packages include PyPI release, download, and license badges at the top of their READMEs. The `auth0-ai-ms-agent` README has none.

### 4. Python version constraint is different

| Package | Constraint |
|---|---|
| `auth0-ai`, `auth0-ai-langchain`, `auth0-ai-llamaindex` | `python = "^3.11"` (>=3.11, <4.0) |
| `auth0-ai-ms-agent` | `python = ">=3.11,<3.14"` |

The ms-agent constraint is tighter (capped at <3.14). This may be intentional due to `agent-framework` requirements, but worth confirming.

### 5. `TokenVaultAuthorizer` does not inherit `ABC` or override `_handle_authorization_interrupts`

In both langchain and llamaindex:

`class TokenVaultAuthorizer(TokenVaultAuthorizerBase, ABC):`

In ms-agent:

`class TokenVaultAuthorizer(TokenVaultAuthorizerBase):`

The `ABC` mixin is dropped. Additionally, the langchain authorizer overrides `_handle_authorization_interrupts` to convert to a LangGraph interrupt. The ms-agent version does **not** override it — it falls through to the base class default (`raise err`). This may be fine for the MS Agent Framework's interrupt model (which uses `session.state["pending_interrupt"]` instead), but the lack of `ABC` is a minor inconsistency.

### 6. Test dependency differences

- `auth0-ai-ms-agent` includes `pytest-cov` and `pytest-randomly` in its `[tool.poetry.group.test.dependencies]`, which the langchain and llamaindex packages do not.
- The base `auth0-ai` package puts pytest under `[tool.poetry.group.dev.dependencies]` rather than a separate `test` group.

Minor, but inconsistent across packages.

### 7. Top-level `__init__.py` exports are different in style

The ms-agent `__init__.py` exports `Auth0AI`, `TokenVaultAuthorizer`, `get_credentials_from_token_vault`, and `get_access_token_from_token_vault` at the top level. The langchain and llamaindex packages only export `FGARetriever` from their top-level `__init__.py`, with `Auth0AI` and token vault symbols accessed via submodules. This gives ms-agent a flatter public API surface.

### 8. No `async_authorization` or `fga` submodules

The langchain and llamaindex packages both have `async_authorization/` and `fga/` submodules. The ms-agent package has neither (the README notes these as "coming soon"). Not a bug — just a scope gap to be aware of for feature parity tracking.

---

### Summary

| Issue | Severity |
|---|---|
| Missing `LICENSE` file | **High** — should be added before merge |
| Missing `tests/__init__.py` | **Medium** — consistency + potential test discovery issue |
| Missing README badges | **Low** — cosmetic |
| Python version constraint difference | **Low** — confirm if intentional |
| Dropped `ABC` mixin on `TokenVaultAuthorizer` | **Low** — worth aligning |
| Test dependency inconsistencies | **Low** — cosmetic |
| Flatter `__init__.py` exports | **Low** — style preference |
| Missing `async_authorization`/`fga` modules | **Info** — known scope gap |

From my view here, i'm really only concerned with:

1.) I think a LICENSE file would be nice here so when we go to publish to PyPI it's included.
4.) Python version constraint - we should check for certainty about this restriction / consistency
5.) The ABC mixin is applied in the other packages now. Do you feel it's overkill, unnecessary to keep? Or should we be consistent there?

The other comments are largely unrelated to this PR, and would seemingly be addressed with future updates.

[adam-wang-okta-public]

Thanks for the check @priley86 .

1. License file added.
2. Python version in SDK package is updated to ^3.11. The original >=3.11,<3.14 was driven by a transitive dependency: agent-framework (umbrella) → agent-framework-redis → redisvl, which declares python <3.14. This PR switches from agent-framework to agent-framework-core (the SDK only needs FunctionTool), which has no Python upper bound — so that restriction disappears, plus a smaller install size. We did test Python 3.14 and it currently fails to install, but the cause is entirely outside our dependency tree: agent-framework-core, openai, and mcp all require pydantic >=2, which pulls in pydantic-core — a Rust extension built with PyO3 0.24, which only supports up to Python 3.13. This is a known ecosystem-wide gap; every major AI framework library in Python is blocked on the same thing.

Best practice is ^3.11 rather than hardcoding <3.14, because:

• The PyO3/pydantic gap is temporary and will be resolved upstream — when it is, ^3.11 just works with no release needed from us
• Hardcoding <3.14 would require us to cut a release purely to bump a constraint we never actually owned
• It's consistent with auth0-ai-langchain and auth0-ai-llamaindex

3. ABC removal in intentional. Removed intentionally. ABC is only meaningful when a class declares @abstractmethod methods that subclasses must implement. Neither TokenVaultAuthorizerBase nor TokenVaultAuthorizer has any abstract methods — _handle_authorization_interrupts is a concrete method with a default implementation (raise err) on the base class.
4. Readme badge added.
5. __init__.py is actually not needed for pytest.
6. Do not quite understand the benefits for the top-level __init__.py in Langchain and LLamaindex SDK only exporting FGARetriever..lol.
7. FGA and async auth module will be added in follow PRs.