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

packages/auth0-ai-ms-agent/README.md

@@ -0,0 +1,149 @@
+# Auth0 AI for Microsoft Agent Framework
+
+`auth0-ai-ms-agent` is an SDK for building secure AI-powered applications using [Auth0](https://www.auth0.ai/) and [Microsoft Agent Framework](https://github.com/microsoft/agent-framework).
+
+## Installation
+
+> ⚠️ **WARNING**: `auth0-ai-ms-agent` is currently **under heavy development**. We strictly follow [Semantic Versioning (SemVer)](https://semver.org/), meaning all **breaking changes will only occur in major versions**. However, please note that during this early phase, **major versions may be released frequently** as the API evolves. We recommend locking versions when using this in production.
+
+```bash
+pip install auth0-ai-ms-agent
+```
+
+## Features
+
+- **Token Vault**: OAuth-based authorization for calling third-party APIs (GitHub, Slack, Google Calendar, etc.)
+- **Async Authorization** _(coming soon)_: CIBA-based user approval workflows
+- **Fine-Grained Authorization** _(coming soon)_: Integration with Okta FGA for document and tool-level permissions
+
+## Calling APIs On User's Behalf
+
+The `Auth0AI.with_token_vault` function exchanges a user's refresh token (or access token) for a Token Vault access token that is valid to call a third-party API.
+
+Full Example of [Calling APIs On User's Behalf](../../../examples/calling-apis/ms-agent-examples/).
+
+### Basic Usage
+
+1. Define a tool with the proper authorizer:
+
+```python
+from auth0_ai_ms_agent.auth0_ai import Auth0AI
+from auth0_ai_ms_agent.token_vault import get_credentials_from_token_vault
+from agent_framework import FunctionTool
+from datetime import datetime
+
+# If not provided, Auth0 settings will be read from env variables: `AUTH0_DOMAIN`, `AUTH0_CLIENT_ID`, and `AUTH0_CLIENT_SECRET`
+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"],
+    # Optional:
+    # login_hint="user@example.com",
+    # authorization_params={"ui_locales": "en"}
+    # store=InMemoryStore(),
+)
+
+def tool_function(date: datetime):
+    credentials = get_credentials_from_token_vault()
+    # Call Google API using credentials["access_token"]
+
+check_calendar_tool = with_google_calendar_access(
+    FunctionTool(
+        name="check_user_calendar",
+        description="Use this function to check if the user is available on a certain date and time",
+        func=tool_function,
+    )
+)
+```
+
+2. Pass the agent session to the tool at runtime. The MS Agent Framework injects the session automatically when you configure it via `additional_function_arguments`:
+
+```python
+result = await agent.run(
+    "Am I free on Friday at 10am?",
+    options={
+        "additional_function_arguments": {
+            "session": session,
+        }
+    }
+)
+```
+
+3. Handle interruptions properly. If the tool does not have access to the user's calendar, it will raise a `TokenVaultInterrupt`. See [Handling Interrupts](#handling-interrupts).
+
+### Additional Authorization Parameters
+
+The `authorization_params` parameter is optional and can be used to pass additional authorization parameters needed to connect an account (e.g., `ui_locales`).
+If you need `login_hint` during the token exchange, pass it via the top-level `login_hint` argument.
+
+```python
+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"],
+    login_hint="user@example.com",
+    authorization_params={"ui_locales": "en"}
+)
+```
+
+## Handling Interrupts
+
+When authorization is required, the tool raises a `TokenVaultInterrupt` and stores it in `session.state["pending_interrupt"]`. Your application should check for this after each agent run and redirect the user to complete authorization.
+
+```python
+from auth0_ai_ms_agent.token_vault import TokenVaultInterrupt
+
+result = await agent.run(
+    "Am I free on Friday at 10am?",
+    options={
+        "additional_function_arguments": {
+            "session": session,
+        }
+    }
+)
+
+interrupt = session.state.get("pending_interrupt")
+if isinstance(interrupt, TokenVaultInterrupt):
+    # Redirect user to complete authorization
+    # interrupt.connection, interrupt.scopes, interrupt.required_scopes, interrupt.authorization_params are available
+    pass
+```
+
+## Configuration
+
+### Environment Variables
+
+```bash
+# Auth0 Configuration
+AUTH0_DOMAIN=your-tenant.auth0.com
+AUTH0_CLIENT_ID=your_client_id
+AUTH0_CLIENT_SECRET=your_client_secret
+```
+
+### Token Vault Setup
+
+1. Configure Auth0 connections (GitHub, Slack, Google, etc.)
+2. Enable Token Vault in Auth0 dashboard
+3. Obtain user's refresh token or access token through Auth0 login
+4. Pass the token resolver to `with_token_vault`
+
+## Examples
+
+See the [examples directory](../../../examples/) for complete working examples:
+
+- [Calling APIs](../../../examples/calling-apis/ms-agent-examples/) - Token Vault with GitHub, Slack, and Google Calendar
+
+---
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: light)" srcset="https://cdn.auth0.com/website/sdks/logos/auth0_light_mode.png"   width="150">
+    <source media="(prefers-color-scheme: dark)" srcset="https://cdn.auth0.com/website/sdks/logos/auth0_dark_mode.png" width="150">
+    <img alt="Auth0 Logo" src="https://cdn.auth0.com/website/sdks/logos/auth0_light_mode.png" width="150">
+  </picture>
+</p>
+<p align="center">Auth0 is an easy to implement, adaptable authentication and authorization platform. To learn more checkout <a href="https://auth0.com/why-auth0">Why Auth0?</a></p>
+<p align="center">
+This project is licensed under the Apache 2.0 license. See the <a href="https://github.com/auth0/auth0-ai-python/blob/main/LICENSE"> LICENSE</a> file for more info.</p>

packages/auth0-ai-ms-agent/auth0_ai_ms_agent/__init__.py

@@ -0,0 +1,6 @@
+from auth0_ai_ms_agent.auth0_ai import Auth0AI as Auth0AI
+from auth0_ai_ms_agent.token_vault import (
+    TokenVaultAuthorizer as TokenVaultAuthorizer,
+    get_credentials_from_token_vault as get_credentials_from_token_vault,
+    get_access_token_from_token_vault as get_access_token_from_token_vault,
+)

packages/auth0-ai-ms-agent/auth0_ai_ms_agent/auth0_ai.py

@@ -0,0 +1,74 @@
+"""Auth0 AI for Microsoft Agent Framework.
+
+Provides decorators to secure MS Agent tools using Auth0 authorization flows.
+"""
+
+from typing import Callable, Optional
+from agent_framework import FunctionTool
+from auth0_ai.authorizers.token_vault_authorizer import TokenVaultAuthorizerParams
+from auth0_ai.authorizers.types import Auth0ClientParams
+from auth0_ai_ms_agent.token_vault.token_vault_authorizer import TokenVaultAuthorizer
+
+
+class Auth0AI:
+    """Provides decorators to secure MS Agent tools using Auth0 authorization flows."""
+
+    def __init__(self, auth0: Optional[Auth0ClientParams] = None):
+        """Initializes the Auth0AI instance.
+
+        Args:
+            auth0 (Optional[Auth0ClientParams]): Parameters for the Auth0 client.
+                If not provided, values will be automatically read from environment
+                variables: `AUTH0_DOMAIN`, `AUTH0_CLIENT_ID`, and `AUTH0_CLIENT_SECRET`.
+        """
+        self.auth0 = auth0
+
+    def with_token_vault(self, **params: TokenVaultAuthorizerParams) -> Callable[[FunctionTool], FunctionTool]:
+        """Enables a tool to obtain an access token from a Token Vault identity provider (e.g., Google, Azure AD).
+
+        The token can then be used within the tool to call third-party APIs on behalf of the user.
+
+        Args:
+            **params: Parameters defined in `TokenVaultAuthorizerParams`.
+
+        Returns:
+            Callable[[FunctionTool], FunctionTool]: A decorator to wrap an MS Agent FunctionTool.
+
+        Example:
+            ```python
+            from auth0_ai_ms_agent.auth0_ai import Auth0AI
+            from auth0_ai_ms_agent.token_vault import get_credentials_from_token_vault
+            from agent_framework import FunctionTool
+            from datetime import datetime
+
+            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"],
+            )
+
+            def tool_function(date: datetime):
+                credentials = get_credentials_from_token_vault()
+                # Call Google API using credentials["access_token"]
+
+            check_calendar_tool = with_google_calendar_access(
+                FunctionTool(
+                    name="check_user_calendar",
+                    description="Use this function to check if the user is available on a certain date and time",
+                    func=tool_function,
+                )
+            )
+
+            # Pass session when running the agent:
+            # result = await agent.run(
+            #     "Am I free on Friday at 10am?",
+            #     options={"additional_function_arguments": {"session": session}},
+            # )
+            ```
+        """
+        authorizer = TokenVaultAuthorizer(
+            TokenVaultAuthorizerParams(**params), self.auth0
+        )
+        return authorizer.authorizer()

packages/auth0-ai-ms-agent/auth0_ai_ms_agent/token_vault/__init__.py

@@ -0,0 +1,10 @@
+from auth0_ai.interrupts.token_vault_interrupt import (
+    TokenVaultError as TokenVaultError,
+    TokenVaultInterrupt as TokenVaultInterrupt
+)
+
+from auth0_ai.authorizers.token_vault_authorizer import (
+    get_credentials_from_token_vault as get_credentials_from_token_vault,
+    get_access_token_from_token_vault as get_access_token_from_token_vault
+)
+from auth0_ai_ms_agent.token_vault.token_vault_authorizer import TokenVaultAuthorizer as TokenVaultAuthorizer

packages/auth0-ai-ms-agent/auth0_ai_ms_agent/token_vault/token_vault_authorizer.py

@@ -0,0 +1,23 @@
+from abc import ABC
+from agent_framework import FunctionTool
+
+from auth0_ai.authorizers.token_vault_authorizer import (
+    TokenVaultAuthorizerBase,
+    TokenVaultAuthorizerParams
+)
+from auth0_ai.authorizers.types import Auth0ClientParams
+from auth0_ai_ms_agent.utils.tool_wrapper import tool_wrapper
+
+class TokenVaultAuthorizer(TokenVaultAuthorizerBase, ABC):
+    def __init__(
+        self,
+        params: TokenVaultAuthorizerParams,
+        auth0: Auth0ClientParams = None,
+    ):
+        super().__init__(params, auth0)
+
+    def authorizer(self):
+        def wrap_tool(tool: FunctionTool) -> FunctionTool:
+            return tool_wrapper(tool, self.protect)
+
+        return wrap_tool

packages/auth0-ai-ms-agent/auth0_ai_ms_agent/utils/tool_wrapper.py

@@ -0,0 +1,119 @@
+import inspect
+import uuid
+from typing import Any, Callable
+from agent_framework import FunctionTool
+from auth0_ai.interrupts.token_vault_interrupt import TokenVaultInterrupt
+
+# Framework-injected runtime kwargs forwarded by the agent to tool functions.
+# These must be stripped before invoking the original user-defined function.
+# Reference: https://github.com/microsoft/agent-framework/blob/main/python/packages/core/agent_framework/_tools.py#L514-L529
+_FRAMEWORK_KWARGS = frozenset({
+    "chat_options", "tools", "tool_choice", "options", "response_format", "conversation_id",
+})
+
+
+def _build_function_tool_kwargs(tool: FunctionTool) -> dict[str, Any]:
+    """Build constructor kwargs by copying configuration from an existing FunctionTool.
+
+    Introspects FunctionTool.__init__ to forward all named parameters except func,
+    name, description, and input_model (handled separately by the caller). None values
+    are omitted for optional parameters to avoid overriding framework defaults.
+
+    Args:
+        tool: The source FunctionTool to copy configuration from.
+
+    Returns:
+        A dict of kwargs suitable for unpacking into the FunctionTool constructor.
+    """
+    init_signature = inspect.signature(FunctionTool)
+    init_kwargs: dict[str, Any] = {}
+    for param_name, param in init_signature.parameters.items():
+        if param_name in {"self", "func", "name", "description", "input_model"}:
+            continue
+        if param.kind not in (param.POSITIONAL_OR_KEYWORD, param.KEYWORD_ONLY):
+            continue
+        if hasattr(tool, param_name):
+            value = getattr(tool, param_name)
+            if value is not None or param.default is inspect._empty:
+                init_kwargs[param_name] = value
+    return init_kwargs
+
+
+def tool_wrapper(tool: FunctionTool, protect_fn: Callable) -> FunctionTool:
+    """Wrap an MS Agent FunctionTool with an Auth0 authorization protect function.
+
+    The returned FunctionTool preserves the original tool's name, description, input
+    schema, and all framework configuration (approval_mode, max_invocations, etc.).
+    At runtime the wrapped function:
+    - Extracts and validates the session injected by the agent framework
+    - Strips framework-injected kwargs before forwarding to the original function
+    - Builds the authorization context (thread_id, tool_call_id, tool_name)
+    - Delegates to protect_fn for credential acquisition and validation
+    - Stores any TokenVaultInterrupt in session.state["pending_interrupt"] before re-raising
+
+    Args:
+        tool: The FunctionTool to wrap.
+        protect_fn: The bound protect method from a TokenVaultAuthorizerBase instance.
+
+    Returns:
+        A new FunctionTool with authorization applied, preserving all original configuration.
+    """
+    original_func = tool.func
+    tool_name = tool.name
+    tool_description = tool.description
+    input_model = tool.input_model
+
+    async def wrapped_func(**kwargs: Any):
+        # The agent framework injects `session` (AgentSession) into runtime kwargs
+        # for tools that accept **kwargs. session.session_id is used as thread_id
+        # for credential namespace resolution in the token vault authorizer.
+        session = kwargs.pop("session", None)
+        for key in _FRAMEWORK_KWARGS:
+            kwargs.pop(key, None)
+
+        if session is None:
+            raise RuntimeError(
+                f"[{tool_name}] A session is required to record tool state. "
+                "Pass session via options={'additional_function_arguments': {'session': session}}."
+            )
+
+        thread_id = session.session_id
+
+        # tool_call_id is consumed by FunctionTool.invoke() before reaching this
+        # function, so a unique ID is generated per invocation instead.
+        # Reference: https://github.com/microsoft/agent-framework/blob/main/python/packages/core/agent_framework/_tools.py#L466-L470
+        tool_call_id = str(uuid.uuid4())
+
+        def get_context(*_, **__):
+            return {
+                "thread_id": thread_id,
+                "tool_call_id": tool_call_id,
+                "tool_name": tool_name,
+            }
+
+        async def execute_fn(*_, **__):
+            # original_func is called directly rather than through FunctionTool.__call__
+            # to avoid double-counting framework metrics (invocation_count,
+            # invocation_exception_count) already tracked at the wrapped tool boundary.
+            if inspect.iscoroutinefunction(original_func):
+                return await original_func(**kwargs)
+            else:
+                return original_func(**kwargs)
+
+        try:
+            return await protect_fn(get_context, execute_fn)(**kwargs)
+        except TokenVaultInterrupt as e:
+            session.state["pending_interrupt"] = e
+            raise
+
+    schema_or_model = tool.parameters() if getattr(tool, "_schema_supplied", False) else input_model
+
+    init_kwargs = _build_function_tool_kwargs(tool)
+
+    return FunctionTool(
+        func=wrapped_func,
+        name=tool_name,
+        description=tool_description,
+        input_model=schema_or_model,
+        **init_kwargs,
+    )