mirror of
https://github.com/Sea-Haven-Industries/open-swe.git
synced 2026-09-30 06:53:14 +00:00
feat: add user-scoped Notion MCP OAuth (#1593)
* feat: add user-scoped Notion MCP OAuth Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com> * fix: refresh Notion token per tool call Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com> * fix: normalize Notion wrapper response format Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com> --------- Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com>
This commit is contained in:
Johannes du Plessis
GitHub
parent
3a0e2b4672
commit
860aee48ee
10 changed files with 1212 additions and 9 deletions
324
agent/dashboard/notion_oauth.py
Normal file
324
agent/dashboard/notion_oauth.py
Normal file
|
|
@ -0,0 +1,324 @@
|
|||
"""Notion MCP OAuth helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import os
|
||||
import secrets
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
from urllib.parse import urlencode, urlparse
|
||||
|
||||
import httpx
|
||||
from langgraph_sdk import get_client
|
||||
|
||||
from ..encryption import decrypt_token, encrypt_token
|
||||
|
||||
NOTION_MCP_URL = "https://mcp.notion.com/mcp"
|
||||
NOTION_STATE_COOKIE_NAME = "osw_notion_oauth_state"
|
||||
NOTION_OAUTH_FLOW_NAMESPACE: list[str] = ["notion_oauth_flows"]
|
||||
|
||||
_NOTION_HOST = "mcp.notion.com"
|
||||
_PROTECTED_RESOURCE_METADATA_URL = "https://mcp.notion.com/.well-known/oauth-protected-resource"
|
||||
_AUTHORIZATION_SERVER_METADATA_PATH = "/.well-known/oauth-authorization-server"
|
||||
_HTTP_TIMEOUT = httpx.Timeout(15.0, connect=5.0)
|
||||
|
||||
|
||||
class NotionOAuthError(Exception):
|
||||
"""Notion OAuth endpoint error."""
|
||||
|
||||
def __init__(self, status_code: int, detail: str, *, error_code: str | None = None) -> None:
|
||||
super().__init__(detail)
|
||||
self.status_code = status_code
|
||||
self.detail = detail
|
||||
self.error_code = error_code
|
||||
|
||||
|
||||
def _client():
|
||||
return get_client()
|
||||
|
||||
|
||||
def _is_notion_https_url(url: str) -> bool:
|
||||
parsed = urlparse(url)
|
||||
return parsed.scheme == "https" and parsed.hostname == _NOTION_HOST
|
||||
|
||||
|
||||
def _require_notion_https_url(url: str, label: str) -> str:
|
||||
if not _is_notion_https_url(url):
|
||||
raise NotionOAuthError(502, f"invalid Notion OAuth {label}")
|
||||
return url
|
||||
|
||||
|
||||
def _metadata_url(auth_server_url: str) -> str:
|
||||
parsed = urlparse(_require_notion_https_url(auth_server_url, "authorization server"))
|
||||
return f"{parsed.scheme}://{parsed.netloc}{_AUTHORIZATION_SERVER_METADATA_PATH}"
|
||||
|
||||
|
||||
def _base64url(data: bytes) -> str:
|
||||
return base64.urlsafe_b64encode(data).decode().rstrip("=")
|
||||
|
||||
|
||||
def generate_code_verifier() -> str:
|
||||
return _base64url(secrets.token_bytes(32))
|
||||
|
||||
|
||||
def code_challenge_for_verifier(verifier: str) -> str:
|
||||
return _base64url(hashlib.sha256(verifier.encode()).digest())
|
||||
|
||||
|
||||
def build_notion_authorize_url(
|
||||
*,
|
||||
authorization_endpoint: str,
|
||||
client_id: str,
|
||||
redirect_uri: str,
|
||||
code_challenge: str,
|
||||
state: str,
|
||||
) -> str:
|
||||
_require_notion_https_url(authorization_endpoint, "authorization endpoint")
|
||||
params = {
|
||||
"response_type": "code",
|
||||
"client_id": client_id,
|
||||
"redirect_uri": redirect_uri,
|
||||
"state": state,
|
||||
"code_challenge": code_challenge,
|
||||
"code_challenge_method": "S256",
|
||||
"prompt": "consent",
|
||||
}
|
||||
return f"{authorization_endpoint}?{urlencode(params)}"
|
||||
|
||||
|
||||
async def discover_notion_oauth_metadata() -> dict[str, Any]:
|
||||
"""Discover Notion MCP OAuth endpoints."""
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=_HTTP_TIMEOUT) as client:
|
||||
protected_resource = await client.get(_PROTECTED_RESOURCE_METADATA_URL)
|
||||
if not protected_resource.is_success:
|
||||
raise _oauth_error_from_response(
|
||||
protected_resource,
|
||||
"Notion OAuth protected resource discovery failed",
|
||||
)
|
||||
resource_metadata = protected_resource.json()
|
||||
auth_servers = resource_metadata.get("authorization_servers")
|
||||
if not isinstance(auth_servers, list) or not auth_servers:
|
||||
raise NotionOAuthError(
|
||||
502, "Notion OAuth discovery returned no authorization server"
|
||||
)
|
||||
auth_server = auth_servers[0]
|
||||
if not isinstance(auth_server, str):
|
||||
raise NotionOAuthError(
|
||||
502,
|
||||
"Notion OAuth discovery returned invalid authorization server",
|
||||
)
|
||||
|
||||
metadata_response = await client.get(_metadata_url(auth_server))
|
||||
if not metadata_response.is_success:
|
||||
raise _oauth_error_from_response(
|
||||
metadata_response,
|
||||
"Notion OAuth authorization server discovery failed",
|
||||
)
|
||||
metadata = metadata_response.json()
|
||||
except httpx.RequestError as exc:
|
||||
raise NotionOAuthError(503, "Notion OAuth discovery failed") from exc
|
||||
|
||||
if not isinstance(metadata, dict):
|
||||
raise NotionOAuthError(502, "Notion OAuth discovery returned invalid metadata")
|
||||
for key in ("authorization_endpoint", "token_endpoint", "registration_endpoint"):
|
||||
value = metadata.get(key)
|
||||
if not isinstance(value, str) or not value:
|
||||
raise NotionOAuthError(502, f"Notion OAuth discovery missing {key}")
|
||||
_require_notion_https_url(value, key)
|
||||
return metadata
|
||||
|
||||
|
||||
async def register_notion_oauth_client(
|
||||
metadata: dict[str, Any],
|
||||
*,
|
||||
redirect_uri: str,
|
||||
) -> dict[str, Any]:
|
||||
"""Register this deployment as a Notion MCP OAuth client."""
|
||||
registration_endpoint = metadata.get("registration_endpoint")
|
||||
if not isinstance(registration_endpoint, str):
|
||||
raise NotionOAuthError(502, "Notion OAuth metadata missing registration endpoint")
|
||||
_require_notion_https_url(registration_endpoint, "registration endpoint")
|
||||
body: dict[str, Any] = {
|
||||
"client_name": os.environ.get("NOTION_MCP_CLIENT_NAME", "Open SWE"),
|
||||
"redirect_uris": [redirect_uri],
|
||||
"grant_types": ["authorization_code", "refresh_token"],
|
||||
"response_types": ["code"],
|
||||
"token_endpoint_auth_method": "none",
|
||||
}
|
||||
client_uri = os.environ.get("DASHBOARD_BASE_URL", "").strip()
|
||||
if client_uri:
|
||||
body["client_uri"] = client_uri
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=_HTTP_TIMEOUT) as client:
|
||||
response = await client.post(
|
||||
registration_endpoint,
|
||||
headers={"Accept": "application/json", "Content-Type": "application/json"},
|
||||
json=body,
|
||||
)
|
||||
except httpx.RequestError as exc:
|
||||
raise NotionOAuthError(503, "Notion OAuth client registration failed") from exc
|
||||
if not response.is_success:
|
||||
raise _oauth_error_from_response(response, "Notion OAuth client registration failed")
|
||||
data = response.json()
|
||||
if not isinstance(data, dict) or not isinstance(data.get("client_id"), str):
|
||||
raise NotionOAuthError(502, "Notion OAuth client registration missing client_id")
|
||||
return data
|
||||
|
||||
|
||||
async def store_notion_oauth_flow(
|
||||
login: str,
|
||||
nonce_hash: str,
|
||||
*,
|
||||
redirect_uri: str,
|
||||
state: str,
|
||||
) -> str:
|
||||
"""Create and store a short-lived Notion OAuth flow."""
|
||||
metadata = await discover_notion_oauth_metadata()
|
||||
client_info = await register_notion_oauth_client(metadata, redirect_uri=redirect_uri)
|
||||
verifier = generate_code_verifier()
|
||||
authorization_endpoint = str(metadata["authorization_endpoint"])
|
||||
client_id = str(client_info["client_id"])
|
||||
client_secret = (
|
||||
client_info.get("client_secret")
|
||||
if isinstance(client_info.get("client_secret"), str)
|
||||
else None
|
||||
)
|
||||
value = {
|
||||
"login": login,
|
||||
"encrypted_code_verifier": encrypt_token(verifier),
|
||||
"client_id": client_id,
|
||||
"encrypted_client_secret": encrypt_token(client_secret) if client_secret else None,
|
||||
"token_endpoint": str(metadata["token_endpoint"]),
|
||||
"redirect_uri": redirect_uri,
|
||||
"created_at": datetime.now(UTC).isoformat(),
|
||||
}
|
||||
await _client().store.put_item([*NOTION_OAUTH_FLOW_NAMESPACE, login], nonce_hash, value)
|
||||
return build_notion_authorize_url(
|
||||
authorization_endpoint=authorization_endpoint,
|
||||
client_id=client_id,
|
||||
redirect_uri=redirect_uri,
|
||||
code_challenge=code_challenge_for_verifier(verifier),
|
||||
state=state,
|
||||
)
|
||||
|
||||
|
||||
async def pop_notion_oauth_flow(login: str, nonce_hash: str) -> dict[str, Any] | None:
|
||||
"""Read and delete a pending Notion OAuth flow."""
|
||||
namespace = [*NOTION_OAUTH_FLOW_NAMESPACE, login]
|
||||
try:
|
||||
item = await _client().store.get_item(namespace, nonce_hash)
|
||||
except httpx.HTTPStatusError as exc:
|
||||
if exc.response.status_code == 404:
|
||||
return None
|
||||
raise
|
||||
try:
|
||||
await _client().store.delete_item(namespace, nonce_hash)
|
||||
except httpx.HTTPStatusError as exc:
|
||||
if exc.response.status_code != 404:
|
||||
raise
|
||||
if item is None:
|
||||
return None
|
||||
value = item.get("value") if isinstance(item, dict) else getattr(item, "value", None)
|
||||
if not isinstance(value, dict):
|
||||
return None
|
||||
encrypted_code_verifier = value.pop("encrypted_code_verifier", "")
|
||||
if encrypted_code_verifier:
|
||||
value["code_verifier"] = decrypt_token(encrypted_code_verifier)
|
||||
encrypted_client_secret = value.pop("encrypted_client_secret", "")
|
||||
if encrypted_client_secret:
|
||||
value["client_secret"] = decrypt_token(encrypted_client_secret) or None
|
||||
return value
|
||||
|
||||
|
||||
def _oauth_error_from_response(response: httpx.Response, fallback: str) -> NotionOAuthError:
|
||||
error_code = None
|
||||
detail = fallback
|
||||
try:
|
||||
data = response.json()
|
||||
except ValueError:
|
||||
data = None
|
||||
if isinstance(data, dict):
|
||||
raw_error = data.get("error")
|
||||
if isinstance(raw_error, str):
|
||||
error_code = raw_error
|
||||
raw_description = data.get("error_description")
|
||||
description = raw_description if isinstance(raw_description, str) else raw_error
|
||||
detail = f"{fallback}: {description}"
|
||||
elif response.text:
|
||||
detail = f"{fallback}: {response.text[:200]}"
|
||||
return NotionOAuthError(response.status_code, detail, error_code=error_code)
|
||||
|
||||
|
||||
async def exchange_notion_code(code: str, flow: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Exchange a Notion OAuth code for tokens."""
|
||||
token_endpoint = flow.get("token_endpoint")
|
||||
client_id = flow.get("client_id")
|
||||
redirect_uri = flow.get("redirect_uri")
|
||||
code_verifier = flow.get("code_verifier")
|
||||
if not all(
|
||||
isinstance(v, str) and v for v in (token_endpoint, client_id, redirect_uri, code_verifier)
|
||||
):
|
||||
raise NotionOAuthError(400, "stored Notion OAuth flow is incomplete")
|
||||
_require_notion_https_url(token_endpoint, "token endpoint")
|
||||
body = {
|
||||
"grant_type": "authorization_code",
|
||||
"code": code,
|
||||
"client_id": client_id,
|
||||
"redirect_uri": redirect_uri,
|
||||
"code_verifier": code_verifier,
|
||||
}
|
||||
client_secret = flow.get("client_secret")
|
||||
if isinstance(client_secret, str) and client_secret:
|
||||
body["client_secret"] = client_secret
|
||||
return await _request_token(token_endpoint, body, fallback="Notion OAuth token exchange failed")
|
||||
|
||||
|
||||
async def refresh_notion_access_token(
|
||||
*,
|
||||
refresh_token: str,
|
||||
token_endpoint: str,
|
||||
client_id: str,
|
||||
client_secret: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Refresh a Notion OAuth access token."""
|
||||
_require_notion_https_url(token_endpoint, "token endpoint")
|
||||
body = {
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": refresh_token,
|
||||
"client_id": client_id,
|
||||
}
|
||||
if client_secret:
|
||||
body["client_secret"] = client_secret
|
||||
return await _request_token(token_endpoint, body, fallback="Notion OAuth token refresh failed")
|
||||
|
||||
|
||||
async def _request_token(
|
||||
token_endpoint: str, body: dict[str, str], *, fallback: str
|
||||
) -> dict[str, Any]:
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=_HTTP_TIMEOUT) as client:
|
||||
response = await client.post(
|
||||
token_endpoint,
|
||||
headers={
|
||||
"Accept": "application/json",
|
||||
"Content-Type": "application/x-www-form-urlencoded",
|
||||
"User-Agent": "OpenSWE-Notion-MCP/1.0",
|
||||
},
|
||||
data=body,
|
||||
)
|
||||
except httpx.RequestError as exc:
|
||||
raise NotionOAuthError(503, f"{fallback}: network error") from exc
|
||||
if not response.is_success:
|
||||
raise _oauth_error_from_response(response, fallback)
|
||||
data = response.json()
|
||||
if not isinstance(data, dict) or not isinstance(data.get("access_token"), str):
|
||||
raise NotionOAuthError(502, f"{fallback}: missing access_token")
|
||||
return data
|
||||
|
||||
|
||||
def is_reauth_required_error(exc: BaseException) -> bool:
|
||||
return isinstance(exc, NotionOAuthError) and exc.error_code == "invalid_grant"
|
||||
|
|
@ -35,6 +35,13 @@ from .enabled_repos import (
|
|||
from .eval_jobs import (
|
||||
get_reviewer_eval_status,
|
||||
)
|
||||
from .notion_oauth import (
|
||||
NOTION_STATE_COOKIE_NAME,
|
||||
NotionOAuthError,
|
||||
exchange_notion_code,
|
||||
pop_notion_oauth_flow,
|
||||
store_notion_oauth_flow,
|
||||
)
|
||||
from .oauth import (
|
||||
COOKIE_NAME,
|
||||
SESSION_TTL_SECONDS,
|
||||
|
|
@ -146,8 +153,11 @@ from .thread_api import (
|
|||
from .user_credentials import (
|
||||
CurrentsCredentialsUpdate,
|
||||
connect_currents,
|
||||
connect_notion,
|
||||
disconnect_currents,
|
||||
disconnect_notion,
|
||||
get_currents_status,
|
||||
get_notion_status,
|
||||
)
|
||||
from .user_mappings import (
|
||||
delete_mapping,
|
||||
|
|
@ -290,6 +300,26 @@ def _clear_slack_state_cookie(response: Response) -> None:
|
|||
)
|
||||
|
||||
|
||||
def _set_notion_state_cookie(response: Response, nonce: str) -> None:
|
||||
secure, _ = _cookie_security()
|
||||
response.set_cookie(
|
||||
key=NOTION_STATE_COOKIE_NAME,
|
||||
value=nonce,
|
||||
max_age=STATE_TTL_SECONDS,
|
||||
httponly=True,
|
||||
secure=secure,
|
||||
samesite="lax",
|
||||
path="/dashboard/api/notion",
|
||||
)
|
||||
|
||||
|
||||
def _clear_notion_state_cookie(response: Response) -> None:
|
||||
secure, _ = _cookie_security()
|
||||
response.delete_cookie(
|
||||
NOTION_STATE_COOKIE_NAME, path="/dashboard/api/notion", samesite="lax", secure=secure
|
||||
)
|
||||
|
||||
|
||||
@router.get("/auth/login")
|
||||
async def auth_login(
|
||||
request: Request,
|
||||
|
|
@ -436,6 +466,89 @@ async def disconnect_my_currents(
|
|||
return status.get("currents", {"connected": False})
|
||||
|
||||
|
||||
@router.get("/my-credentials/notion")
|
||||
async def get_my_notion_status(
|
||||
session: dict[str, Any] = _SESSION_DEP,
|
||||
) -> dict[str, Any]:
|
||||
status = await get_notion_status(session["sub"])
|
||||
return status.get("notion", {"connected": False})
|
||||
|
||||
|
||||
@router.delete("/my-credentials/notion")
|
||||
async def disconnect_my_notion(
|
||||
session: dict[str, Any] = _SESSION_DEP,
|
||||
) -> dict[str, Any]:
|
||||
status = await disconnect_notion(session["sub"])
|
||||
return status.get("notion", {"connected": False})
|
||||
|
||||
|
||||
@router.get("/notion/login")
|
||||
async def notion_login(
|
||||
session: dict[str, Any] = _SESSION_DEP,
|
||||
) -> RedirectResponse:
|
||||
redirect_uri = f"{_api_base_url()}/dashboard/api/notion/callback"
|
||||
nonce = new_state_nonce()
|
||||
nonce_hash = hash_state_nonce(nonce)
|
||||
state = issue_state(
|
||||
redirect_to=f"{_frontend_base_url()}/my-settings",
|
||||
nonce_hash=nonce_hash,
|
||||
)
|
||||
try:
|
||||
url = await store_notion_oauth_flow(
|
||||
session["sub"],
|
||||
nonce_hash,
|
||||
redirect_uri=redirect_uri,
|
||||
state=state,
|
||||
)
|
||||
except NotionOAuthError as exc:
|
||||
raise HTTPException(exc.status_code, exc.detail) from exc
|
||||
response = RedirectResponse(url, status_code=302)
|
||||
_set_notion_state_cookie(response, nonce)
|
||||
return response
|
||||
|
||||
|
||||
@router.get("/notion/callback")
|
||||
async def notion_callback(
|
||||
request: Request,
|
||||
state: str,
|
||||
code: str | None = None,
|
||||
error: str | None = None,
|
||||
error_description: str | None = None,
|
||||
session: dict[str, Any] = _SESSION_DEP,
|
||||
) -> RedirectResponse:
|
||||
state_payload = decode_state(state)
|
||||
nonce_hash = state_payload.get("nonce_hash")
|
||||
cookie_nonce = request.cookies.get(NOTION_STATE_COOKIE_NAME)
|
||||
if (
|
||||
not isinstance(nonce_hash, str)
|
||||
or not cookie_nonce
|
||||
or not hmac.compare_digest(hash_state_nonce(cookie_nonce), nonce_hash)
|
||||
):
|
||||
raise HTTPException(400, "oauth state mismatch — please retry")
|
||||
|
||||
flow = await pop_notion_oauth_flow(session["sub"], nonce_hash)
|
||||
if flow is None:
|
||||
raise HTTPException(400, "oauth flow expired — please retry")
|
||||
if error:
|
||||
detail = error_description or error
|
||||
raise HTTPException(400, f"Notion OAuth failed: {detail}")
|
||||
if not code:
|
||||
raise HTTPException(400, "Notion OAuth callback missing code")
|
||||
|
||||
try:
|
||||
token_data = await exchange_notion_code(code, flow)
|
||||
await connect_notion(session["sub"], token_data, flow)
|
||||
except NotionOAuthError as exc:
|
||||
raise HTTPException(exc.status_code, exc.detail) from exc
|
||||
except ValueError as exc:
|
||||
raise HTTPException(400, str(exc)) from exc
|
||||
|
||||
redirect_to = sanitize_redirect_to(state_payload.get("redirect_to")) or _frontend_base_url()
|
||||
response = RedirectResponse(redirect_to, status_code=302)
|
||||
_clear_notion_state_cookie(response)
|
||||
return response
|
||||
|
||||
|
||||
@router.get("/slack/login")
|
||||
async def slack_login(
|
||||
_session: dict[str, Any] = _SESSION_DEP,
|
||||
|
|
|
|||
|
|
@ -1,27 +1,27 @@
|
|||
"""Per-user third-party service credentials (Currents.dev).
|
||||
|
||||
Credentials are encrypted at rest with :mod:`agent.encryption` and stored in a
|
||||
dedicated LangGraph Store namespace, keyed by the user's GitHub login. The
|
||||
sandbox never holds these keys — they feed server-side read-only tools.
|
||||
"""
|
||||
"""Per-user third-party service credentials."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import UTC, datetime
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Any
|
||||
|
||||
from langgraph_sdk import get_client
|
||||
from pydantic import BaseModel, field_validator
|
||||
|
||||
from ..encryption import decrypt_token, encrypt_token
|
||||
from .notion_oauth import is_reauth_required_error, refresh_notion_access_token
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
USER_CREDENTIALS_NAMESPACE: list[str] = ["user_credentials"]
|
||||
CURRENTS_KEY = "currents"
|
||||
NOTION_KEY = "notion"
|
||||
|
||||
CURRENTS_API_BASE = "https://api.currents.dev/v1"
|
||||
_NOTION_TOKEN_EXPIRY_SKEW_SECONDS = 300
|
||||
|
||||
|
||||
def _client():
|
||||
|
|
@ -68,6 +68,248 @@ class CurrentsCredentialsUpdate(BaseModel):
|
|||
return v.strip()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class NotionCredentials:
|
||||
access_token: str
|
||||
refresh_token: str | None
|
||||
token_endpoint: str
|
||||
client_id: str
|
||||
client_secret: str | None = None
|
||||
|
||||
|
||||
def _expires_at_from_response(data: dict[str, Any], *, field: str = "expires_in") -> str | None:
|
||||
raw = data.get(field)
|
||||
if not isinstance(raw, int | float) or raw <= 0:
|
||||
return None
|
||||
return (datetime.now(UTC) + timedelta(seconds=int(raw))).isoformat()
|
||||
|
||||
|
||||
def _token_expired(
|
||||
expires_at: str | None, *, skew_seconds: int = _NOTION_TOKEN_EXPIRY_SKEW_SECONDS
|
||||
) -> bool:
|
||||
if not isinstance(expires_at, str) or not expires_at:
|
||||
return False
|
||||
try:
|
||||
exp = datetime.fromisoformat(expires_at.replace("Z", "+00:00"))
|
||||
if exp.tzinfo is None:
|
||||
exp = exp.replace(tzinfo=UTC)
|
||||
except ValueError:
|
||||
return False
|
||||
return datetime.now(UTC) + timedelta(seconds=skew_seconds) >= exp
|
||||
|
||||
|
||||
async def get_notion_status(login: str) -> dict[str, Any]:
|
||||
"""Return a redacted view of the user's Notion MCP connection."""
|
||||
notion = await _get_provider(login, NOTION_KEY)
|
||||
return {
|
||||
"notion": {
|
||||
"connected": True,
|
||||
"token_expires_at": notion.get("token_expires_at"),
|
||||
"updated_at": notion.get("updated_at"),
|
||||
}
|
||||
if notion
|
||||
else {"connected": False},
|
||||
}
|
||||
|
||||
|
||||
def _notion_record_from_response(
|
||||
data: dict[str, Any],
|
||||
*,
|
||||
existing: dict[str, Any] | None = None,
|
||||
client_id: str | None = None,
|
||||
client_secret: str | None = None,
|
||||
token_endpoint: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
existing = existing or {}
|
||||
access_token = data.get("access_token")
|
||||
if not isinstance(access_token, str) or not access_token:
|
||||
raise ValueError("Notion OAuth response missing access_token")
|
||||
refresh_token = data.get("refresh_token")
|
||||
record: dict[str, Any] = {
|
||||
"encrypted_access_token": encrypt_token(access_token),
|
||||
"client_id": client_id or existing.get("client_id", ""),
|
||||
"token_endpoint": token_endpoint or existing.get("token_endpoint", ""),
|
||||
"updated_at": datetime.now(UTC).isoformat(),
|
||||
}
|
||||
token_type = data.get("token_type")
|
||||
if isinstance(token_type, str):
|
||||
record["token_type"] = token_type
|
||||
scope = data.get("scope")
|
||||
if isinstance(scope, str):
|
||||
record["scope"] = scope
|
||||
token_expires_at = _expires_at_from_response(data)
|
||||
if token_expires_at:
|
||||
record["token_expires_at"] = token_expires_at
|
||||
elif existing.get("token_expires_at"):
|
||||
record["token_expires_at"] = existing["token_expires_at"]
|
||||
refresh_token_expires_at = _expires_at_from_response(data, field="refresh_token_expires_in")
|
||||
if refresh_token_expires_at:
|
||||
record["refresh_token_expires_at"] = refresh_token_expires_at
|
||||
elif existing.get("refresh_token_expires_at"):
|
||||
record["refresh_token_expires_at"] = existing["refresh_token_expires_at"]
|
||||
if isinstance(refresh_token, str) and refresh_token:
|
||||
record["encrypted_refresh_token"] = encrypt_token(refresh_token)
|
||||
elif existing.get("encrypted_refresh_token"):
|
||||
record["encrypted_refresh_token"] = existing["encrypted_refresh_token"]
|
||||
if client_secret:
|
||||
record["encrypted_client_secret"] = encrypt_token(client_secret)
|
||||
elif existing.get("encrypted_client_secret"):
|
||||
record["encrypted_client_secret"] = existing["encrypted_client_secret"]
|
||||
return record
|
||||
|
||||
|
||||
async def connect_notion(login: str, data: dict[str, Any], flow: dict[str, Any]) -> dict[str, Any]:
|
||||
client_id = flow.get("client_id")
|
||||
token_endpoint = flow.get("token_endpoint")
|
||||
if not isinstance(client_id, str) or not isinstance(token_endpoint, str):
|
||||
raise ValueError("stored Notion OAuth flow is incomplete")
|
||||
client_secret = (
|
||||
flow.get("client_secret") if isinstance(flow.get("client_secret"), str) else None
|
||||
)
|
||||
await _put_provider(
|
||||
login,
|
||||
NOTION_KEY,
|
||||
_notion_record_from_response(
|
||||
data,
|
||||
client_id=client_id,
|
||||
client_secret=client_secret,
|
||||
token_endpoint=token_endpoint,
|
||||
),
|
||||
)
|
||||
return await get_notion_status(login)
|
||||
|
||||
|
||||
async def disconnect_notion(login: str) -> dict[str, Any]:
|
||||
await _delete_provider(login, NOTION_KEY)
|
||||
return await get_notion_status(login)
|
||||
|
||||
|
||||
def _decrypt_notion_access_token(record: dict[str, Any]) -> str | None:
|
||||
token = decrypt_token(record.get("encrypted_access_token", ""))
|
||||
return token or None
|
||||
|
||||
|
||||
def _decrypt_notion_refresh_token(record: dict[str, Any]) -> str | None:
|
||||
token = decrypt_token(record.get("encrypted_refresh_token", ""))
|
||||
return token or None
|
||||
|
||||
|
||||
def _decrypt_notion_client_secret(record: dict[str, Any]) -> str | None:
|
||||
token = decrypt_token(record.get("encrypted_client_secret", ""))
|
||||
return token or None
|
||||
|
||||
|
||||
_notion_refresh_locks: dict[str, asyncio.Lock] = {}
|
||||
|
||||
|
||||
def _notion_refresh_lock(login: str) -> asyncio.Lock:
|
||||
lock = _notion_refresh_locks.get(login)
|
||||
if lock is None:
|
||||
lock = asyncio.Lock()
|
||||
_notion_refresh_locks[login] = lock
|
||||
return lock
|
||||
|
||||
|
||||
async def _refresh_stored_notion_token(
|
||||
login: str,
|
||||
record: dict[str, Any],
|
||||
) -> tuple[str | None, bool]:
|
||||
refresh_token = _decrypt_notion_refresh_token(record)
|
||||
token_endpoint = record.get("token_endpoint")
|
||||
client_id = record.get("client_id")
|
||||
if not refresh_token or not isinstance(token_endpoint, str) or not isinstance(client_id, str):
|
||||
return None, False
|
||||
try:
|
||||
data = await refresh_notion_access_token(
|
||||
refresh_token=refresh_token,
|
||||
token_endpoint=token_endpoint,
|
||||
client_id=client_id,
|
||||
client_secret=_decrypt_notion_client_secret(record),
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("Notion token refresh failed for %s", login, exc_info=True)
|
||||
return None, is_reauth_required_error(exc)
|
||||
await _put_provider(login, NOTION_KEY, _notion_record_from_response(data, existing=record))
|
||||
access_token = data.get("access_token")
|
||||
return (access_token if isinstance(access_token, str) else None), False
|
||||
|
||||
|
||||
async def get_notion_credentials(
|
||||
login: str, *, force_refresh: bool = False
|
||||
) -> NotionCredentials | None:
|
||||
"""Return a valid Notion MCP credential set for a user."""
|
||||
record = await _get_provider(login, NOTION_KEY)
|
||||
if not record:
|
||||
return None
|
||||
access_token = _decrypt_notion_access_token(record)
|
||||
if not access_token:
|
||||
return None
|
||||
if not force_refresh and not _token_expired(record.get("token_expires_at")):
|
||||
return NotionCredentials(
|
||||
access_token=access_token,
|
||||
refresh_token=_decrypt_notion_refresh_token(record),
|
||||
token_endpoint=record.get("token_endpoint", ""),
|
||||
client_id=record.get("client_id", ""),
|
||||
client_secret=_decrypt_notion_client_secret(record),
|
||||
)
|
||||
if not _decrypt_notion_refresh_token(record):
|
||||
return None
|
||||
async with _notion_refresh_lock(login):
|
||||
record = await _get_provider(login, NOTION_KEY)
|
||||
if not record:
|
||||
return None
|
||||
access_token = _decrypt_notion_access_token(record)
|
||||
if not access_token:
|
||||
return None
|
||||
if not force_refresh and not _token_expired(record.get("token_expires_at")):
|
||||
return NotionCredentials(
|
||||
access_token=access_token,
|
||||
refresh_token=_decrypt_notion_refresh_token(record),
|
||||
token_endpoint=record.get("token_endpoint", ""),
|
||||
client_id=record.get("client_id", ""),
|
||||
client_secret=_decrypt_notion_client_secret(record),
|
||||
)
|
||||
refreshed, refresh_token_dead = await _refresh_stored_notion_token(login, record)
|
||||
if refreshed:
|
||||
refreshed_record = await _get_provider(login, NOTION_KEY) or record
|
||||
return NotionCredentials(
|
||||
access_token=refreshed,
|
||||
refresh_token=_decrypt_notion_refresh_token(refreshed_record),
|
||||
token_endpoint=refreshed_record.get("token_endpoint", ""),
|
||||
client_id=refreshed_record.get("client_id", ""),
|
||||
client_secret=_decrypt_notion_client_secret(refreshed_record),
|
||||
)
|
||||
if refresh_token_dead:
|
||||
latest = await _get_provider(login, NOTION_KEY)
|
||||
if latest and latest.get("encrypted_refresh_token") != record.get(
|
||||
"encrypted_refresh_token"
|
||||
):
|
||||
latest_access_token = _decrypt_notion_access_token(latest)
|
||||
if latest_access_token:
|
||||
return NotionCredentials(
|
||||
access_token=latest_access_token,
|
||||
refresh_token=_decrypt_notion_refresh_token(latest),
|
||||
token_endpoint=latest.get("token_endpoint", ""),
|
||||
client_id=latest.get("client_id", ""),
|
||||
client_secret=_decrypt_notion_client_secret(latest),
|
||||
)
|
||||
logger.info("Dropping dead Notion authorization for %s; reconnect required", login)
|
||||
await disconnect_notion(login)
|
||||
return None
|
||||
return NotionCredentials(
|
||||
access_token=access_token,
|
||||
refresh_token=_decrypt_notion_refresh_token(record),
|
||||
token_endpoint=record.get("token_endpoint", ""),
|
||||
client_id=record.get("client_id", ""),
|
||||
client_secret=_decrypt_notion_client_secret(record),
|
||||
)
|
||||
|
||||
|
||||
async def get_notion_access_token(login: str) -> str | None:
|
||||
credentials = await get_notion_credentials(login)
|
||||
return credentials.access_token if credentials else None
|
||||
|
||||
|
||||
async def get_currents_status(login: str) -> dict[str, Any]:
|
||||
"""Return a redacted, dashboard-safe view of the user's Currents key."""
|
||||
currents = await _get_provider(login, CURRENTS_KEY)
|
||||
|
|
|
|||
101
agent/integrations/notion_mcp.py
Normal file
101
agent/integrations/notion_mcp.py
Normal file
|
|
@ -0,0 +1,101 @@
|
|||
"""Server-side Notion tools backed by Notion's hosted MCP server."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.tools import BaseTool
|
||||
|
||||
from ..dashboard.notion_oauth import NOTION_MCP_URL
|
||||
from ..dashboard.user_credentials import get_notion_access_token
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_MCP_TIMEOUT_SECONDS = 30.0
|
||||
|
||||
|
||||
async def _build_mcp_tools(access_token: str) -> list[BaseTool]:
|
||||
from langchain_mcp_adapters.client import MultiServerMCPClient
|
||||
|
||||
client = MultiServerMCPClient(
|
||||
{
|
||||
"notion": {
|
||||
"transport": "streamable_http",
|
||||
"url": NOTION_MCP_URL,
|
||||
"headers": {
|
||||
"Authorization": f"Bearer {access_token}",
|
||||
},
|
||||
"timeout": timedelta(seconds=_MCP_TIMEOUT_SECONDS),
|
||||
}
|
||||
}
|
||||
)
|
||||
return await client.get_tools()
|
||||
|
||||
|
||||
async def _fresh_mcp_tool(login: str, tool_name: str) -> BaseTool:
|
||||
access_token = await get_notion_access_token(login)
|
||||
if not access_token:
|
||||
raise RuntimeError(
|
||||
"Notion MCP authorization unavailable; reconnect Notion in Profile Settings"
|
||||
)
|
||||
tools = await _build_mcp_tools(access_token)
|
||||
for tool in tools:
|
||||
if tool.name == tool_name:
|
||||
return tool
|
||||
raise RuntimeError(f"Notion MCP tool {tool_name!r} is no longer available")
|
||||
|
||||
|
||||
def _tool_input(args: tuple[Any, ...], kwargs: dict[str, Any]) -> str | dict[str, Any]:
|
||||
if args and kwargs:
|
||||
raise TypeError("Notion MCP tool received both positional and keyword input")
|
||||
if not args:
|
||||
return kwargs
|
||||
if len(args) == 1 and isinstance(args[0], str):
|
||||
return args[0]
|
||||
if len(args) == 1 and isinstance(args[0], dict):
|
||||
return args[0]
|
||||
raise TypeError("Notion MCP tool received invalid positional input")
|
||||
|
||||
|
||||
class _RefreshingNotionMCPTool(BaseTool):
|
||||
login: str
|
||||
mcp_tool_name: str
|
||||
|
||||
def _run(self, *args: Any, **kwargs: Any) -> Any:
|
||||
try:
|
||||
asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
return asyncio.run(self._arun(*args, **kwargs))
|
||||
raise RuntimeError("Notion MCP tools must be called asynchronously")
|
||||
|
||||
async def _arun(self, *args: Any, **kwargs: Any) -> Any:
|
||||
tool = await _fresh_mcp_tool(self.login, self.mcp_tool_name)
|
||||
return await tool.ainvoke(_tool_input(args, kwargs))
|
||||
|
||||
|
||||
def _refreshing_tool(login: str, tool: BaseTool) -> BaseTool:
|
||||
return _RefreshingNotionMCPTool(
|
||||
name=tool.name,
|
||||
description=tool.description,
|
||||
args_schema=tool.args_schema,
|
||||
response_format="content",
|
||||
login=login,
|
||||
mcp_tool_name=tool.name,
|
||||
)
|
||||
|
||||
|
||||
async def load_notion_tools(login: str) -> list[BaseTool]:
|
||||
"""Return Notion MCP tools for a connected user."""
|
||||
access_token = await get_notion_access_token(login)
|
||||
if not access_token:
|
||||
return []
|
||||
try:
|
||||
tools = await _build_mcp_tools(access_token)
|
||||
except Exception: # noqa: BLE001
|
||||
logger.warning("Failed to load Notion MCP tools", exc_info=True)
|
||||
return []
|
||||
logger.info("Loaded %d Notion MCP tool(s) for %s", len(tools), login)
|
||||
return [_refreshing_tool(login, tool) for tool in tools]
|
||||
|
|
@ -51,6 +51,7 @@ from .integrations.currents_tools import load_currents_tools
|
|||
from .integrations.datadog_mcp import load_datadog_tools
|
||||
from .integrations.langsmith import _configure_github_proxy
|
||||
from .integrations.langsmith_tools import load_langsmith_tools
|
||||
from .integrations.notion_mcp import load_notion_tools
|
||||
from .middleware import (
|
||||
ModelFallbackMiddleware,
|
||||
PlanModeMiddleware,
|
||||
|
|
@ -739,12 +740,18 @@ async def get_agent(config: RunnableConfig) -> Pregel:
|
|||
corridor_tools = await _load_corridor_mcp_tools()
|
||||
|
||||
currents_tools: list[Any] = []
|
||||
notion_tools: list[Any] = []
|
||||
if profile_login:
|
||||
try:
|
||||
currents_tools = await load_currents_tools(profile_login)
|
||||
except Exception:
|
||||
logger.warning("Failed to load Currents tools", exc_info=True)
|
||||
currents_tools = []
|
||||
try:
|
||||
notion_tools = await load_notion_tools(profile_login)
|
||||
except Exception:
|
||||
logger.warning("Failed to load Notion tools", exc_info=True)
|
||||
notion_tools = []
|
||||
|
||||
logger.info("Returning agent with sandbox for thread %s", thread_id)
|
||||
main_model = make_model(model_id, **model_kwargs)
|
||||
|
|
@ -785,6 +792,7 @@ async def get_agent(config: RunnableConfig) -> Pregel:
|
|||
*corridor_tools,
|
||||
*observability_tools,
|
||||
*currents_tools,
|
||||
*notion_tools,
|
||||
],
|
||||
subagents=[_general_purpose_subagent(subagent_model)],
|
||||
backend=backend_factory,
|
||||
|
|
|
|||
111
tests/test_notion_oauth.py
Normal file
111
tests/test_notion_oauth.py
Normal file
|
|
@ -0,0 +1,111 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import pytest
|
||||
from cryptography.fernet import Fernet
|
||||
|
||||
from agent.dashboard import notion_oauth as no
|
||||
|
||||
|
||||
class _FakeStore:
|
||||
def __init__(self) -> None:
|
||||
self.items: dict[tuple[tuple[str, ...], str], dict[str, Any]] = {}
|
||||
|
||||
async def get_item(self, namespace: list[str], key: str):
|
||||
value = self.items.get((tuple(namespace), key))
|
||||
return {"value": value} if value is not None else None
|
||||
|
||||
async def put_item(self, namespace: list[str], key: str, value: dict[str, Any]) -> None:
|
||||
self.items[(tuple(namespace), key)] = value
|
||||
|
||||
async def delete_item(self, namespace: list[str], key: str) -> None:
|
||||
self.items.pop((tuple(namespace), key), None)
|
||||
|
||||
|
||||
class _FakeClient:
|
||||
def __init__(self, store: _FakeStore) -> None:
|
||||
self.store = store
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def fake_store(monkeypatch: pytest.MonkeyPatch) -> _FakeStore:
|
||||
store = _FakeStore()
|
||||
monkeypatch.setattr(no, "_client", lambda: _FakeClient(store))
|
||||
monkeypatch.setenv("TOKEN_ENCRYPTION_KEY", Fernet.generate_key().decode())
|
||||
return store
|
||||
|
||||
|
||||
def test_code_challenge_matches_rfc7636_vector() -> None:
|
||||
verifier = "dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk"
|
||||
assert no.code_challenge_for_verifier(verifier) == "E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM"
|
||||
|
||||
|
||||
def test_build_notion_authorize_url() -> None:
|
||||
url = no.build_notion_authorize_url(
|
||||
authorization_endpoint="https://mcp.notion.com/authorize",
|
||||
client_id="cid",
|
||||
redirect_uri="https://example.com/dashboard/api/notion/callback",
|
||||
code_challenge="challenge",
|
||||
state="state-token",
|
||||
)
|
||||
parsed = urlparse(url)
|
||||
query = parse_qs(parsed.query)
|
||||
assert parsed.netloc == "mcp.notion.com"
|
||||
assert parsed.path == "/authorize"
|
||||
assert query["response_type"] == ["code"]
|
||||
assert query["client_id"] == ["cid"]
|
||||
assert query["code_challenge"] == ["challenge"]
|
||||
assert query["code_challenge_method"] == ["S256"]
|
||||
assert query["prompt"] == ["consent"]
|
||||
|
||||
|
||||
def test_build_notion_authorize_url_rejects_other_hosts() -> None:
|
||||
with pytest.raises(no.NotionOAuthError):
|
||||
no.build_notion_authorize_url(
|
||||
authorization_endpoint="https://example.com/authorize",
|
||||
client_id="cid",
|
||||
redirect_uri="https://example.com/callback",
|
||||
code_challenge="challenge",
|
||||
state="state-token",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_and_pop_notion_oauth_flow(
|
||||
fake_store: _FakeStore,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
no,
|
||||
"discover_notion_oauth_metadata",
|
||||
AsyncMock(
|
||||
return_value={
|
||||
"authorization_endpoint": "https://mcp.notion.com/authorize",
|
||||
"token_endpoint": "https://mcp.notion.com/token",
|
||||
"registration_endpoint": "https://mcp.notion.com/register",
|
||||
}
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
no,
|
||||
"register_notion_oauth_client",
|
||||
AsyncMock(return_value={"client_id": "cid", "client_secret": "secret"}),
|
||||
)
|
||||
monkeypatch.setattr(no, "generate_code_verifier", lambda: "verifier")
|
||||
|
||||
url = await no.store_notion_oauth_flow(
|
||||
"alice",
|
||||
"nonce-hash",
|
||||
redirect_uri="https://example.com/dashboard/api/notion/callback",
|
||||
state="state-token",
|
||||
)
|
||||
assert parse_qs(urlparse(url).query)["client_id"] == ["cid"]
|
||||
|
||||
flow = await no.pop_notion_oauth_flow("alice", "nonce-hash")
|
||||
assert flow is not None
|
||||
assert flow["code_verifier"] == "verifier"
|
||||
assert flow["client_secret"] == "secret"
|
||||
assert await no.pop_notion_oauth_flow("alice", "nonce-hash") is None
|
||||
|
|
@ -1,12 +1,14 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import Literal
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from langchain_core.tools import StructuredTool
|
||||
|
||||
from agent import server
|
||||
from agent.dashboard.team_credentials import DatadogCredentials, LangSmithCredentials
|
||||
from agent.integrations import datadog_mcp, langsmith_tools
|
||||
from agent.integrations import datadog_mcp, langsmith_tools, notion_mcp
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -36,6 +38,103 @@ async def test_load_datadog_tools_returns_tools() -> None:
|
|||
assert await datadog_mcp.load_datadog_tools() == sentinel
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_notion_tools_empty_when_not_connected() -> None:
|
||||
with patch.object(notion_mcp, "get_notion_access_token", AsyncMock(return_value=None)):
|
||||
assert await notion_mcp.load_notion_tools("alice") == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_notion_tools_degrades_on_error() -> None:
|
||||
with (
|
||||
patch.object(notion_mcp, "get_notion_access_token", AsyncMock(return_value="tok")),
|
||||
patch.object(notion_mcp, "_build_mcp_tools", AsyncMock(side_effect=RuntimeError("boom"))),
|
||||
):
|
||||
assert await notion_mcp.load_notion_tools("alice") == []
|
||||
|
||||
|
||||
def _notion_tool_for_token(
|
||||
token: str,
|
||||
response_format: Literal["content", "content_and_artifact"] = "content",
|
||||
) -> StructuredTool:
|
||||
async def notion_search(query: str):
|
||||
"""Search Notion."""
|
||||
content = {"query": query, "token": token}
|
||||
if response_format == "content_and_artifact":
|
||||
return content, {"artifact_token": token}
|
||||
return content
|
||||
|
||||
return StructuredTool.from_function(
|
||||
coroutine=notion_search,
|
||||
name="notion_search",
|
||||
description="Search Notion",
|
||||
response_format=response_format,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_notion_tools_returns_wrappers() -> None:
|
||||
discovered = _notion_tool_for_token("initial-token")
|
||||
with (
|
||||
patch.object(notion_mcp, "get_notion_access_token", AsyncMock(return_value="tok")),
|
||||
patch.object(notion_mcp, "_build_mcp_tools", AsyncMock(return_value=[discovered])),
|
||||
):
|
||||
tools = await notion_mcp.load_notion_tools("alice")
|
||||
assert len(tools) == 1
|
||||
assert tools[0].name == "notion_search"
|
||||
assert tools[0].description == discovered.description
|
||||
assert tools[0].args_schema == discovered.args_schema
|
||||
assert tools[0].response_format == "content"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_notion_wrapper_normalizes_content_and_artifact_tools() -> None:
|
||||
get_token = AsyncMock(side_effect=["initial-token", "fresh-token"])
|
||||
build_tools = AsyncMock(
|
||||
side_effect=lambda token: [_notion_tool_for_token(token, "content_and_artifact")]
|
||||
)
|
||||
with (
|
||||
patch.object(notion_mcp, "get_notion_access_token", get_token),
|
||||
patch.object(notion_mcp, "_build_mcp_tools", build_tools),
|
||||
):
|
||||
tools = await notion_mcp.load_notion_tools("alice")
|
||||
assert tools[0].response_format == "content"
|
||||
result = await tools[0].ainvoke({"query": "roadmap"})
|
||||
assert result == {"query": "roadmap", "token": "fresh-token"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_notion_wrapper_refreshes_token_at_call_time() -> None:
|
||||
get_token = AsyncMock(side_effect=["initial-token", "fresh-token"])
|
||||
build_tools = AsyncMock(side_effect=lambda token: [_notion_tool_for_token(token)])
|
||||
with (
|
||||
patch.object(notion_mcp, "get_notion_access_token", get_token),
|
||||
patch.object(notion_mcp, "_build_mcp_tools", build_tools),
|
||||
):
|
||||
tools = await notion_mcp.load_notion_tools("alice")
|
||||
result = await tools[0].ainvoke({"query": "roadmap"})
|
||||
assert result == {"query": "roadmap", "token": "fresh-token"}
|
||||
assert get_token.await_count == 2
|
||||
assert [call.args[0] for call in build_tools.await_args_list] == [
|
||||
"initial-token",
|
||||
"fresh-token",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_notion_wrapper_fails_when_token_missing_at_call_time() -> None:
|
||||
get_token = AsyncMock(side_effect=["initial-token", None])
|
||||
build_tools = AsyncMock(return_value=[_notion_tool_for_token("initial-token")])
|
||||
with (
|
||||
patch.object(notion_mcp, "get_notion_access_token", get_token),
|
||||
patch.object(notion_mcp, "_build_mcp_tools", build_tools),
|
||||
):
|
||||
tools = await notion_mcp.load_notion_tools("alice")
|
||||
with pytest.raises(RuntimeError, match="Notion MCP authorization unavailable"):
|
||||
await tools[0].ainvoke({"query": "roadmap"})
|
||||
assert build_tools.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_langsmith_tools_empty_when_not_connected() -> None:
|
||||
with patch.object(langsmith_tools, "get_langsmith_credentials", AsyncMock(return_value=None)):
|
||||
|
|
|
|||
|
|
@ -1,12 +1,15 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from cryptography.fernet import Fernet
|
||||
from pydantic import ValidationError
|
||||
|
||||
from agent.dashboard import user_credentials as uc
|
||||
from agent.dashboard.notion_oauth import NotionOAuthError
|
||||
from agent.dashboard.user_credentials import CurrentsCredentialsUpdate
|
||||
|
||||
|
||||
|
|
@ -93,3 +96,107 @@ async def test_currents_status_when_not_connected(fake_store: _FakeStore) -> Non
|
|||
@pytest.mark.asyncio
|
||||
async def test_get_currents_api_key_none_when_not_connected(fake_store: _FakeStore) -> None:
|
||||
assert await uc.get_currents_api_key("nobody") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_notion_roundtrip_and_redaction(fake_store: _FakeStore) -> None:
|
||||
status = await uc.connect_notion(
|
||||
"alice",
|
||||
{
|
||||
"access_token": "notion-access-1234",
|
||||
"refresh_token": "notion-refresh",
|
||||
"expires_in": 3600,
|
||||
"token_type": "Bearer",
|
||||
},
|
||||
{
|
||||
"client_id": "client-id",
|
||||
"client_secret": "client-secret",
|
||||
"token_endpoint": "https://mcp.notion.com/token",
|
||||
},
|
||||
)
|
||||
assert status["notion"]["connected"] is True
|
||||
|
||||
record = fake_store.items[(("user_credentials", "alice"), "notion")]
|
||||
assert record["encrypted_access_token"] != "notion-access-1234"
|
||||
assert record["encrypted_refresh_token"] != "notion-refresh"
|
||||
assert record["encrypted_client_secret"] != "client-secret"
|
||||
|
||||
creds = await uc.get_notion_credentials("alice")
|
||||
assert creds is not None
|
||||
assert creds.access_token == "notion-access-1234"
|
||||
assert creds.refresh_token == "notion-refresh"
|
||||
assert creds.client_id == "client-id"
|
||||
assert creds.client_secret == "client-secret"
|
||||
|
||||
after = await uc.disconnect_notion("alice")
|
||||
assert after["notion"]["connected"] is False
|
||||
assert await uc.get_notion_credentials("alice") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_notion_refresh_rotates_tokens(fake_store: _FakeStore) -> None:
|
||||
await uc.connect_notion(
|
||||
"alice",
|
||||
{
|
||||
"access_token": "old-access",
|
||||
"refresh_token": "old-refresh",
|
||||
"expires_in": 3600,
|
||||
},
|
||||
{
|
||||
"client_id": "client-id",
|
||||
"client_secret": "client-secret",
|
||||
"token_endpoint": "https://mcp.notion.com/token",
|
||||
},
|
||||
)
|
||||
record = fake_store.items[(("user_credentials", "alice"), "notion")]
|
||||
record["token_expires_at"] = (datetime.now(UTC) - timedelta(minutes=1)).isoformat()
|
||||
|
||||
with patch.object(
|
||||
uc,
|
||||
"refresh_notion_access_token",
|
||||
new_callable=AsyncMock,
|
||||
return_value={
|
||||
"access_token": "new-access",
|
||||
"refresh_token": "new-refresh",
|
||||
"expires_in": 3600,
|
||||
},
|
||||
) as refresh:
|
||||
creds = await uc.get_notion_credentials("alice")
|
||||
|
||||
assert creds is not None
|
||||
assert creds.access_token == "new-access"
|
||||
assert creds.refresh_token == "new-refresh"
|
||||
refresh.assert_awaited_once_with(
|
||||
refresh_token="old-refresh",
|
||||
token_endpoint="https://mcp.notion.com/token",
|
||||
client_id="client-id",
|
||||
client_secret="client-secret",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_notion_invalid_grant_disconnects(fake_store: _FakeStore) -> None:
|
||||
await uc.connect_notion(
|
||||
"alice",
|
||||
{
|
||||
"access_token": "old-access",
|
||||
"refresh_token": "old-refresh",
|
||||
"expires_in": 3600,
|
||||
},
|
||||
{
|
||||
"client_id": "client-id",
|
||||
"token_endpoint": "https://mcp.notion.com/token",
|
||||
},
|
||||
)
|
||||
record = fake_store.items[(("user_credentials", "alice"), "notion")]
|
||||
record["token_expires_at"] = (datetime.now(UTC) - timedelta(minutes=1)).isoformat()
|
||||
|
||||
with patch.object(
|
||||
uc,
|
||||
"refresh_notion_access_token",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=NotionOAuthError(400, "dead", error_code="invalid_grant"),
|
||||
):
|
||||
assert await uc.get_notion_credentials("alice") is None
|
||||
|
||||
assert await uc.get_notion_status("alice") == {"notion": {"connected": False}}
|
||||
|
|
|
|||
|
|
@ -202,6 +202,12 @@ export interface CurrentsConnectBody {
|
|||
api_key: string
|
||||
}
|
||||
|
||||
export interface NotionCredentialStatus {
|
||||
connected: boolean
|
||||
token_expires_at?: string | null
|
||||
updated_at?: string | null
|
||||
}
|
||||
|
||||
export interface UserMapping {
|
||||
github_login: string
|
||||
work_email: string
|
||||
|
|
@ -612,6 +618,12 @@ export const api = {
|
|||
request<CurrentsCredentialStatus>("/my-credentials/currents", {
|
||||
method: "DELETE",
|
||||
}),
|
||||
getMyNotionStatus: () =>
|
||||
request<NotionCredentialStatus>("/my-credentials/notion"),
|
||||
disconnectNotion: () =>
|
||||
request<NotionCredentialStatus>("/my-credentials/notion", {
|
||||
method: "DELETE",
|
||||
}),
|
||||
listEnabledReviewRepos: () =>
|
||||
request<{ repos: Array<string> }>("/enabled-review-repos"),
|
||||
setEnabledReviewRepo: (full_name: string, enabled: boolean) =>
|
||||
|
|
@ -685,3 +697,7 @@ export function loginUrl(redirectTo?: string): string {
|
|||
export function slackConnectUrl(): string {
|
||||
return `${API_BASE}/dashboard/api/slack/login`
|
||||
}
|
||||
|
||||
export function notionConnectUrl(): string {
|
||||
return `${API_BASE}/dashboard/api/notion/login`
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import { Navigate, createFileRoute, useNavigate } from "@tanstack/react-router"
|
|||
import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"
|
||||
import { useState } from "react"
|
||||
import { IoLogoSlack } from "react-icons/io5"
|
||||
import { SiNotion } from "react-icons/si"
|
||||
|
||||
import type { CurrentsConnectBody, SessionUser } from "@/lib/api"
|
||||
import { AppShell, SettingsRow, SettingsSection } from "@/components/AppShell"
|
||||
|
|
@ -16,7 +17,7 @@ import {
|
|||
} from "@/components/ui/select"
|
||||
import { Skeleton } from "@/components/ui/skeleton"
|
||||
import { Switch } from "@/components/ui/switch"
|
||||
import { api, slackConnectUrl } from "@/lib/api"
|
||||
import { api, notionConnectUrl, slackConnectUrl } from "@/lib/api"
|
||||
import {
|
||||
buildProfileUpdate,
|
||||
useOptions,
|
||||
|
|
@ -265,6 +266,85 @@ function CurrentsCredentialsSection() {
|
|||
)
|
||||
}
|
||||
|
||||
function NotionCredentialsSection() {
|
||||
const qc = useQueryClient()
|
||||
const creds = useQuery({
|
||||
queryKey: ["myNotion"],
|
||||
queryFn: api.getMyNotionStatus,
|
||||
})
|
||||
const [connecting, setConnecting] = useState(false)
|
||||
const [error, setError] = useState<string | null>(null)
|
||||
|
||||
const onSuccess = () => {
|
||||
qc.invalidateQueries({ queryKey: ["myNotion"] })
|
||||
setError(null)
|
||||
}
|
||||
const onError = (e: Error) => setError(e.message)
|
||||
|
||||
const disconnect = useMutation({
|
||||
mutationFn: () => api.disconnectNotion(),
|
||||
onSuccess,
|
||||
onError,
|
||||
})
|
||||
|
||||
const connected = creds.data?.connected
|
||||
const connect = () => {
|
||||
setConnecting(true)
|
||||
void qc.invalidateQueries({ queryKey: ["myNotion"] })
|
||||
window.location.assign(notionConnectUrl())
|
||||
}
|
||||
|
||||
return (
|
||||
<SettingsSection
|
||||
title="Notion"
|
||||
description="Connect your Notion account through Notion MCP so agent runs can use Notion tools with your workspace permissions. OAuth tokens are encrypted at rest and scoped to your account only."
|
||||
>
|
||||
<SettingsRow
|
||||
label="Notion MCP"
|
||||
description={
|
||||
connected
|
||||
? "Connected. Reconnect if your Notion authorization expires or workspace access changes."
|
||||
: "Not connected. Sign in with Notion to authorize the hosted Notion MCP server."
|
||||
}
|
||||
control={
|
||||
<div className="flex items-center gap-2">
|
||||
<span
|
||||
className={cn(
|
||||
"rounded-full px-2 py-0.5 text-[10px] font-medium",
|
||||
connected
|
||||
? "bg-primary/10 text-primary"
|
||||
: "bg-muted text-muted-foreground"
|
||||
)}
|
||||
>
|
||||
{connected ? "Connected" : "Not connected"}
|
||||
</span>
|
||||
{connected ? (
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
onClick={() => disconnect.mutate()}
|
||||
disabled={disconnect.isPending}
|
||||
>
|
||||
Disconnect
|
||||
</Button>
|
||||
) : (
|
||||
<Button
|
||||
size="sm"
|
||||
onClick={connect}
|
||||
disabled={connecting || creds.isLoading}
|
||||
>
|
||||
<SiNotion className="size-4" />
|
||||
{connecting ? "Redirecting…" : "Connect Notion"}
|
||||
</Button>
|
||||
)}
|
||||
</div>
|
||||
}
|
||||
/>
|
||||
{error && <p className="px-4 pb-3 text-xs text-destructive">{error}</p>}
|
||||
</SettingsSection>
|
||||
)
|
||||
}
|
||||
|
||||
function MySettingsPage() {
|
||||
const session = useSession()
|
||||
const qc = useQueryClient()
|
||||
|
|
@ -368,6 +448,8 @@ function MySettingsPage() {
|
|||
|
||||
<CurrentsCredentialsSection />
|
||||
|
||||
<NotionCredentialsSection />
|
||||
|
||||
<SettingsSection title="Account">
|
||||
<SettingsRow
|
||||
label="Sign out"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue