diff --git a/backend/apps/tools/oauth/OAuthService.py b/backend/apps/tools/oauth/OAuthService.py new file mode 100644 index 00000000..876149d0 --- /dev/null +++ b/backend/apps/tools/oauth/OAuthService.py @@ -0,0 +1,258 @@ +"""OAuth service — token exchange, refresh, disconnect, flow initiation. + +Pure business logic with no HTTP/FastAPI dependencies. +""" + +import base64 +import hashlib +import logging +import os +import secrets +import time +from typing import Any, Optional, Tuple +from urllib.parse import urlencode + +import httpx +from pydantic import BaseModel, Field +from typing import Dict +from backend.apps.tools.oauth.OAUTH_PROVIDERS.OAuthProvider import OAuthProvider +from backend.apps.tools.oauth.OAUTH_PROVIDERS.OAUTH_PROVIDERS import OAUTH_PROVIDERS +from backend.apps.tools.shared_utils.ToolDefinition import ToolDefinition +from backend.core.db.PydanticStore import PydanticStore +from typeguard import typechecked +from backend.ports import BACKEND_DEV_PORT + +logger = logging.getLogger(__name__) + + +class OAuthService(BaseModel): + store: PydanticStore[ToolDefinition] + pending_oauth: Dict[str, str] = Field(default_factory=dict) + pending_pkce: Dict[str, str] = Field(default_factory=dict) + + @typechecked + def p_redirect_uri(self) -> str: + port = os.environ.get("OPENSWARM_PORT", str(BACKEND_DEV_PORT)) + return f"http://localhost:{port}/api/tools/oauth/callback" + + @typechecked + async def start_flow(self, tool_id: str) -> str: + """Build the authorization URL and stash pending state. + + Returns the full auth URL the client should redirect to. + Raises ValueError if the provider's client ID env var is unset. + """ + tool: ToolDefinition = self.store.load(tool_id) + provider: OAuthProvider = OAUTH_PROVIDERS[tool.oauth_provider] + + client_id: str = os.environ.get(provider.client_id_env, "") + if not client_id: + raise ValueError(f"{provider.client_id_env} not set in backend .env") + + assert tool.oauth_provider is not None + + provider_key: str = tool.oauth_provider + state: str = f"{provider_key}:{tool_id}" + self.pending_oauth[state] = tool_id + + params: Dict[str, str] = { + "client_id": client_id, + "redirect_uri": self._redirect_uri(), + "response_type": "code", + "state": state, + **provider.extra_auth_params, + } + if provider.scopes: + params["scope"] = " ".join(provider.scopes) + + if provider.pkce_required: + code_verifier: str = secrets.token_urlsafe(64) + code_challenge: str = ( + base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest()) + .rstrip(b"=") + .decode() + ) + params["code_challenge"] = code_challenge + params["code_challenge_method"] = "S256" + self.pending_pkce[state] = code_verifier + + return f"{provider.auth_url}?{urlencode(params)}" + + @typechecked + async def handle_callback(self, code: str, state: str) -> Tuple[str, ToolDefinition]: + """Exchange the authorization code for tokens and persist them. + + Returns (tool_id, updated_tool) on success. + Raises LookupError if the state token is unknown. + """ + tool_id: Optional[str] = self.pending_oauth.pop(state, None) + if not tool_id: + alt_key: str = state.split(":")[-1] if ":" in state else state + tool_id = self.pending_oauth.pop(alt_key, None) + if not tool_id: + raise LookupError("Invalid OAuth state") + + tool: ToolDefinition = self.store.load(tool_id) + provider: OAuthProvider = OAUTH_PROVIDERS[tool.oauth_provider] + + client_id: str = os.environ.get(provider.client_id_env, "") + client_secret: str = os.environ.get(provider.client_secret_env, "") + + token_data: Dict[str, str] = { + "code": code, + "redirect_uri": self._redirect_uri(), + "grant_type": "authorization_code", + } + headers: Dict[str, str] = {} + + if provider.token_auth_method == "basic": + creds: str = base64.b64encode(f"{client_id}:{client_secret}".encode()).decode() + headers["Authorization"] = f"Basic {creds}" + elif provider.token_auth_method == "basic_json": + creds: str = base64.b64encode(f"{client_id}:{client_secret}".encode()).decode() + headers["Authorization"] = f"Basic {creds}" + headers["Content-Type"] = "application/json" + else: + token_data["client_id"] = client_id + token_data["client_secret"] = client_secret + + if tool.oauth_provider == "github": + headers["Accept"] = "application/json" + + code_verifier: Optional[str] = self.pending_pkce.pop(state, None) + if code_verifier: + token_data["code_verifier"] = code_verifier + + async with httpx.AsyncClient(timeout=15.0) as client: + if provider.token_auth_method == "basic_json": + resp = await client.post(provider.token_url, json=token_data, headers=headers) + else: + resp = await client.post(provider.token_url, data=token_data, headers=headers) + + if resp.status_code != 200: + logger.warning("OAuth token exchange failed: %s", resp.text) + raise RuntimeError(resp.text) + + tokens: Dict[str, Any] = resp.json() + + access_token: str = tokens.get("access_token", "") + if provider.token_response_path and not access_token: + obj: Dict[str, Any] = tokens + for part in provider.token_response_path.split("."): + obj = obj.get(part, {}) if isinstance(obj, dict) else "" + if isinstance(obj, str) and obj: + access_token = obj + + tool.oauth_tokens: Dict[str, Any] = { + "access_token": access_token, + "refresh_token": tokens.get("refresh_token", ""), + "token_expiry": time.time() + tokens.get("expires_in", 3600), + } + + for response_path, env_var in provider.extra_token_fields.items(): + obj_val: Dict[str, Any] = tokens + for part in response_path.split("."): + obj_val = obj_val.get(part, "") if isinstance(obj_val, dict) else "" + if obj_val: + tool.oauth_tokens[env_var] = str(obj_val) + + tool.auth_status = "connected" + + if access_token and provider.userinfo_url: + tool.connected_account_email = await self.p_fetch_userinfo( + provider.userinfo_url, provider.userinfo_field, access_token, + label=tool.oauth_provider or "google", + ) + + if (tool.oauth_provider or "google") == "notion" and not tool.connected_account_email: + workspace_name = tokens.get("workspace_name") + if workspace_name: + tool.connected_account_email = workspace_name + + self.store.save(tool) + return tool_id, tool + + @typechecked + async def disconnect(self, tool_id: str) -> ToolDefinition: + """Revoke the access token (best-effort) and clear stored credentials.""" + tool: ToolDefinition = self.store.load(tool_id) + access_token = tool.oauth_tokens.get("access_token") + + if access_token: + provider: OAuthProvider = OAUTH_PROVIDERS[tool.oauth_provider] + revoke_url: Optional[str] = provider.revoke_url + assert revoke_url is not None, "Revoke URL is required" + try: + async with httpx.AsyncClient(timeout=10.0) as client: + await client.post( + revoke_url, + params={"token": access_token}, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + ) + except Exception as e: + logger.warning("Failed to revoke token for tool %s: %s", tool.id, e) + + tool.oauth_tokens = {} + tool.auth_status = "configured" + tool.connected_account_email = None + self.store.save(tool) + return tool + + @typechecked + async def refresh_token(self, tool: ToolDefinition) -> Optional[str]: + """Refresh an expired OAuth token. Returns the fresh access_token or None. + + Mutates the tool in-place and saves to the store on success. + """ + if tool.auth_type != "oauth2": + return None + refresh_tok = tool.oauth_tokens.get("refresh_token") + if not refresh_tok: + return None + expiry = tool.oauth_tokens.get("token_expiry", 0) + if time.time() < expiry - 60: + return tool.oauth_tokens.get("access_token") + + provider: OAuthProvider = OAUTH_PROVIDERS[tool.oauth_provider] + client_id: str = os.environ.get(provider.client_id_env, "") + client_secret: str = os.environ.get(provider.client_secret_env, "") + if not client_id or not client_secret: + return None + + try: + async with httpx.AsyncClient(timeout=15.0) as client: + resp = await client.post(provider.token_url, data={ + "client_id": client_id, + "client_secret": client_secret, + "refresh_token": refresh_tok, + "grant_type": "refresh_token", + }) + if resp.status_code == 200: + data = resp.json() + new_token = data["access_token"] + tool.oauth_tokens["access_token"] = new_token + tool.oauth_tokens["token_expiry"] = time.time() + data.get("expires_in", 3600) + + if not tool.connected_account_email and provider.userinfo_url: + tool.connected_account_email = await self.p_fetch_userinfo( + provider.userinfo_url, provider.userinfo_field, new_token, + ) + + self.store.save(tool) + return new_token + except Exception as e: + logger.warning("OAuth token refresh failed for tool %s: %s", tool.id, e) + return None + + @typechecked + async def p_fetch_userinfo( + self, url: str, field: str, access_token: str, *, label: str = "", + ) -> Optional[str]: + try: + async with httpx.AsyncClient(timeout=10.0) as client: + resp = await client.get(url, headers={"Authorization": f"Bearer {access_token}"}) + if resp.status_code == 200: + return resp.json().get(field) + except Exception as e: + logger.warning("Failed to fetch userinfo%s: %s", f" for {label}" if label else "", e) + return None \ No newline at end of file diff --git a/backend/apps/tools/oauth/oauth.py b/backend/apps/tools/oauth/oauth.py deleted file mode 100644 index b6f1f551..00000000 --- a/backend/apps/tools/oauth/oauth.py +++ /dev/null @@ -1,266 +0,0 @@ -"""OAuth flow logic — callback, start, disconnect, refresh. - -The tool store is injected via set_store() from the tools sub-app -to avoid circular imports. -""" - -import base64 -import hashlib -import logging -import os -import secrets -import time -from typing import Any, Optional -from urllib.parse import urlencode - -import httpx -from fastapi import HTTPException, Query -from fastapi.responses import HTMLResponse - -from backend.apps.tools.oauth.OAUTH_PROVIDERS.OAUTH_PROVIDERS import OAUTH_PROVIDERS -from backend.core.db.PydanticStore import PydanticStore -from backend.apps.tools.shared_utils.ToolDefinition import ToolDefinition -from backend.ports import BACKEND_DEV_PORT - - -logger = logging.getLogger(__name__) - -_pending_oauth: dict[str, str] = {} -_pending_pkce: dict[str, str] = {} - -_store: Optional[PydanticStore[ToolDefinition]] = None - - -def set_store(store: PydanticStore[ToolDefinition]) -> None: - global _store - _store = store - - -def _get_store() -> PydanticStore[ToolDefinition]: - assert _store is not None, "OAuth store not initialized — call set_store() first" - return _store - - -async def oauth_callback(code: str = Query(...), state: str = Query("")) -> HTMLResponse: - tool_id = _pending_oauth.pop(state, None) - if not tool_id: - tool_id = _pending_oauth.pop(state.split(":")[-1] if ":" in state else state, None) - if not tool_id: - return HTMLResponse("
{resp.text}",
- status_code=400,
- )
-
- tokens = resp.json()
-
- access_token = tokens.get("access_token", "")
- if provider.token_response_path and not access_token:
- obj: Any = tokens
- for part in provider.token_response_path.split("."):
- obj = obj.get(part, {}) if isinstance(obj, dict) else ""
- if isinstance(obj, str) and obj:
- access_token = obj
-
- tool.oauth_tokens = {
- "access_token": access_token,
- "refresh_token": tokens.get("refresh_token", ""),
- "token_expiry": time.time() + tokens.get("expires_in", 3600),
- }
-
- for response_path, env_var in provider.extra_token_fields.items():
- obj_val: Any = tokens
- for part in response_path.split("."):
- obj_val = obj_val.get(part, "") if isinstance(obj_val, dict) else ""
- if obj_val:
- tool.oauth_tokens[env_var] = str(obj_val)
-
- tool.auth_status = "connected"
-
- if access_token and provider.userinfo_url:
- try:
- async with httpx.AsyncClient(timeout=10.0) as info_client:
- info_resp = await info_client.get(
- provider.userinfo_url,
- headers={"Authorization": f"Bearer {access_token}"},
- )
- if info_resp.status_code == 200:
- tool.connected_account_email = info_resp.json().get(provider.userinfo_field)
- except Exception as e:
- logger.warning(f"Failed to fetch userinfo for {tool.oauth_provider or 'google'}: {e}")
-
- if (tool.oauth_provider or "google") == "notion" and not tool.connected_account_email:
- workspace_name = tokens.get("workspace_name")
- if workspace_name:
- tool.connected_account_email = workspace_name
-
- store.save(tool)
-
- return HTMLResponse(
- ""
- 'You can close this window.
' - "" - "" - ) - - -async def oauth_start(tool_id: str) -> dict: - store = _get_store() - tool = store.load(tool_id) - provider = OAUTH_PROVIDERS[tool.oauth_provider] - - client_id = os.environ.get(provider.client_id_env, "") - if not client_id: - raise HTTPException(status_code=400, detail=f"{provider.client_id_env} not set in backend .env") - - port = os.environ.get("OPENSWARM_PORT", str(BACKEND_DEV_PORT)) - redirect_uri = f"http://localhost:{port}/api/tools/oauth/callback" - provider_key = tool.oauth_provider or "google" - state = f"{provider_key}:{tool_id}" - - _pending_oauth[state] = tool_id - - params: dict[str, str] = { - "client_id": client_id, - "redirect_uri": redirect_uri, - "response_type": "code", - "state": state, - **provider.extra_auth_params, - } - if provider.scopes: - params["scope"] = " ".join(provider.scopes) - - if provider.pkce_required: - code_verifier = secrets.token_urlsafe(64) - code_challenge = base64.urlsafe_b64encode( - hashlib.sha256(code_verifier.encode()).digest() - ).rstrip(b"=").decode() - params["code_challenge"] = code_challenge - params["code_challenge_method"] = "S256" - _pending_pkce[state] = code_verifier - - auth_url = f"{provider.auth_url}?{urlencode(params)}" - return {"auth_url": auth_url} - - -async def oauth_disconnect(tool_id: str) -> dict: - store = _get_store() - tool = store.load(tool_id) - access_token = tool.oauth_tokens.get("access_token") - - if access_token: - provider = OAUTH_PROVIDERS[tool.oauth_provider] - revoke_url = provider.revoke_url or "https://oauth2.googleapis.com/revoke" - try: - async with httpx.AsyncClient(timeout=10.0) as client: - await client.post( - revoke_url, - params={"token": access_token}, - headers={"Content-Type": "application/x-www-form-urlencoded"}, - ) - except Exception as e: - logger.warning(f"Failed to revoke token for tool {tool.id}: {e}") - - tool.oauth_tokens = {} - tool.auth_status = "configured" - tool.connected_account_email = None - store.save(tool) - return {"ok": True, "tool": tool.model_dump()} - - -async def refresh_oauth_token(tool: ToolDefinition) -> Optional[str]: - """Refresh an expired OAuth token. Returns the fresh access_token or None. - - Mutates the tool in-place and saves to the store if refresh succeeds. - """ - if tool.auth_type != "oauth2": - return None - refresh_token = tool.oauth_tokens.get("refresh_token") - if not refresh_token: - return None - expiry = tool.oauth_tokens.get("token_expiry", 0) - if time.time() < expiry - 60: - return tool.oauth_tokens.get("access_token") - - provider = OAUTH_PROVIDERS[tool.oauth_provider] - client_id = os.environ.get(provider.client_id_env, "") - client_secret = os.environ.get(provider.client_secret_env, "") - if not client_id or not client_secret: - return None - - try: - async with httpx.AsyncClient(timeout=15.0) as client: - resp = await client.post(provider.token_url, data={ - "client_id": client_id, - "client_secret": client_secret, - "refresh_token": refresh_token, - "grant_type": "refresh_token", - }) - if resp.status_code == 200: - data = resp.json() - new_token = data["access_token"] - tool.oauth_tokens["access_token"] = new_token - tool.oauth_tokens["token_expiry"] = time.time() + data.get("expires_in", 3600) - - if not tool.connected_account_email and provider.userinfo_url: - try: - async with httpx.AsyncClient(timeout=10.0) as info_client: - info_resp = await info_client.get( - provider.userinfo_url, - headers={"Authorization": f"Bearer {new_token}"}, - ) - if info_resp.status_code == 200: - tool.connected_account_email = info_resp.json().get(provider.userinfo_field) - except Exception: - pass - - _get_store().save(tool) - return new_token - except Exception as e: - logger.warning(f"OAuth token refresh failed for tool {tool.id}: {e}") - return None diff --git a/backend/apps/tools/shared_utils/ToolDefinition.py b/backend/apps/tools/shared_utils/ToolDefinition.py index 83268bc3..aa1ecd89 100644 --- a/backend/apps/tools/shared_utils/ToolDefinition.py +++ b/backend/apps/tools/shared_utils/ToolDefinition.py @@ -1,9 +1,10 @@ +from typing_extensions import Dict from pydantic import BaseModel, Field from typing import Optional, Any from uuid import uuid4 from backend.core.tools.shared_structs.TOOL_PERMISSIONS import TOOL_PERMISSIONS - +# TODO: better type specing of this whole class, also we may not even need this class???? class ToolDefinition(BaseModel): model_config = {"extra": "ignore"} @@ -11,13 +12,13 @@ class ToolDefinition(BaseModel): name: str description: str = "" command: str = "" - mcp_config: dict[str, Any] = Field(default_factory=dict) - credentials: dict[str, str] = Field(default_factory=dict) + mcp_config: Dict[str, Any] = Field(default_factory=dict) + credentials: Dict[str, str] = Field(default_factory=dict) auth_type: str = "none" auth_status: str = "none" oauth_provider: Optional[str] = None - oauth_tokens: dict[str, Any] = Field(default_factory=dict) - tool_permissions: dict[str, TOOL_PERMISSIONS] = Field(default_factory=dict) - tool_descriptions: dict[str, str] = Field(default_factory=dict) # tool_descriptions[tool_name] = tool_description + oauth_tokens: Dict[str, Any] = Field(default_factory=dict) + tool_permissions: Dict[str, TOOL_PERMISSIONS] = Field(default_factory=dict) + tool_descriptions: Dict[str, str] = Field(default_factory=dict) # tool_descriptions[tool_name] = tool_description connected_account_email: Optional[str] = None enabled: bool = True diff --git a/backend/apps/tools/tools.py b/backend/apps/tools/tools.py index 79e45264..740be85e 100644 --- a/backend/apps/tools/tools.py +++ b/backend/apps/tools/tools.py @@ -1,5 +1,6 @@ """Tools sub-app — CRUD for user-installed MCP tools, builtin permissions, and discovery.""" +from typing_extensions import List from claude_agent_sdk.types import McpServerConfig from pydantic import BaseModel, Field import json @@ -9,7 +10,8 @@ import time from contextlib import asynccontextmanager from typing import Any, Optional -from fastapi import HTTPException +from fastapi import HTTPException, Query +from fastapi.responses import HTMLResponse from backend.config.Apps import SubApp from backend.config.paths import DB_ROOT @@ -18,13 +20,14 @@ from backend.apps.tools.shared_utils.ToolDefinition import ToolDefinition from backend.apps.tools.discover_tools.discover_tools import discover_tools from backend.apps.tools.discover_tools.DiscoveryError import DiscoveryError, DiscoveryConfigError from backend.apps.tools.tool_definition_to_mcp_tool.tool_definition_to_mcp_tool import tool_definition_to_mcp_tool -from backend.apps.tools.oauth.oauth import refresh_oauth_token, oauth_callback, oauth_start, oauth_disconnect, set_store +from backend.apps.tools.oauth.OAuthService import OAuthService from backend.apps.tools.oauth.OAUTH_PROVIDERS.OAUTH_PROVIDERS import OAUTH_PROVIDERS from backend.apps.tools.builtin_tools import BUILTIN_TOOLS from backend.core.tools.shared_structs.TOOL_PERMISSIONS import TOOL_PERMISSIONS from backend.core.tools.shared_structs.Toolkit import Toolkit from backend.core.tools.shared_structs.Tool import Tool from backend.core.tools.shared_structs.MCP_Tool import MCP_Tool +from typeguard import typechecked logger = logging.getLogger(__name__) @@ -39,11 +42,13 @@ TOOL_STORE: PydanticStore[ToolDefinition] = PydanticStore[ToolDefinition]( not_found_detail="Tool not found", ) +OAUTH_SERVICE: Optional[OAuthService] = None @asynccontextmanager async def tools_lifespan(): + global OAUTH_SERVICE os.makedirs(TOOLS_DIR, exist_ok=True) - set_store(TOOL_STORE) + OAUTH_SERVICE = OAuthService(store=TOOL_STORE) yield @@ -167,7 +172,8 @@ async def discover(tool_id: str) -> dict: tool = TOOL_STORE.load(tool_id) if tool.auth_type == "oauth2" and tool.auth_status == "connected": - refreshed = await refresh_oauth_token(tool) + assert OAUTH_SERVICE is not None, "OAuthService not initialized" + refreshed = await OAUTH_SERVICE.refresh_token(tool) if not refreshed and tool.oauth_tokens.get("access_token"): expiry = tool.oauth_tokens.get("token_expiry", 0) if isinstance(expiry, (int, float)) and time.time() >= expiry - 60: @@ -219,7 +225,7 @@ async def load_user_toolkit() -> Optional[Toolkit]: Returns None if no valid tools could be converted. """ - mcp_tools: list[Tool] = [] + mcp_tools: List[Tool] = [] for td in TOOL_STORE.load_all(): if not td.mcp_config or not td.enabled: continue @@ -246,9 +252,51 @@ async def load_user_toolkit() -> Optional[Toolkit]: ) # --------------------------------------------------------------------------- -# OAuth routes +# OAuth # --------------------------------------------------------------------------- -tools.router.add_api_route("/oauth/callback", oauth_callback, methods=["GET"]) -tools.router.add_api_route("/{tool_id}/oauth/start", oauth_start, methods=["POST"]) -tools.router.add_api_route("/{tool_id}/oauth/disconnect", oauth_disconnect, methods=["POST"]) +@tools.router.get("/oauth/callback") +async def oauth_callback(code: str = Query(...), state: str = Query("")) -> HTMLResponse: + try: + assert OAUTH_SERVICE is not None, "OAuthService not initialized" + tool_id, _tool = await OAUTH_SERVICE.handle_callback(code, state) + except LookupError: + return HTMLResponse( + "{e}",
+ status_code=400,
+ )
+
+ return HTMLResponse(
+ ""
+ 'You can close this window.
' + "" + "" + ) + + +@tools.router.post("/{tool_id}/oauth/start") +async def oauth_start(tool_id: str) -> dict: + try: + assert OAUTH_SERVICE is not None, "OAuthService not initialized" + auth_url = await OAUTH_SERVICE.start_flow(tool_id) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) + return {"auth_url": auth_url} + + +@tools.router.post("/{tool_id}/oauth/disconnect") +async def oauth_disconnect(tool_id: str) -> dict: + assert OAUTH_SERVICE is not None, "OAuthService not initialized" + tool = await OAUTH_SERVICE.disconnect(tool_id) + return {"ok": True, "tool": tool.model_dump()}