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.

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.