mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-09-11 12:17:45 +02:00
[Haik]: cleaned up oauth logic so theres a single OAuthService class that handles logic only, then all endpoints and stores are handled in the tools subapp. Now gonna rename the oauth folder and fix imports
This commit is contained in:
haikdc
@@ -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
|
||||
@@ -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("<html><body><h2>Invalid OAuth state</h2></body></html>", status_code=400)
|
||||
|
||||
store = _get_store()
|
||||
tool = store.load(tool_id)
|
||||
provider = OAUTH_PROVIDERS[tool.oauth_provider]
|
||||
|
||||
client_id = os.environ.get(provider.client_id_env, "")
|
||||
client_secret = os.environ.get(provider.client_secret_env, "")
|
||||
port = os.environ.get("OPENSWARM_PORT", str(BACKEND_DEV_PORT))
|
||||
redirect_uri = f"http://localhost:{port}/api/tools/oauth/callback"
|
||||
|
||||
token_data: dict[str, str] = {
|
||||
"code": code, "redirect_uri": redirect_uri, "grant_type": "authorization_code",
|
||||
}
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
if provider.token_auth_method == "basic":
|
||||
creds = base64.b64encode(f"{client_id}:{client_secret}".encode()).decode()
|
||||
headers["Authorization"] = f"Basic {creds}"
|
||||
elif provider.token_auth_method == "basic_json":
|
||||
creds = 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 or "google") == "github":
|
||||
headers["Accept"] = "application/json"
|
||||
|
||||
code_verifier = _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(f"OAuth token exchange failed: {resp.text}")
|
||||
return HTMLResponse(
|
||||
f"<html><body><h2>Token exchange failed</h2><pre>{resp.text}</pre></body></html>",
|
||||
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(
|
||||
"<html><body>"
|
||||
'<h2 style="font-family:sans-serif;color:#22c55e">Connected successfully!</h2>'
|
||||
'<p style="font-family:sans-serif;color:#666">You can close this window.</p>'
|
||||
"<script>"
|
||||
"if (window.opener) window.opener.postMessage({type:'oauth_complete', tool_id:'" + tool_id + "'}, '*');"
|
||||
"setTimeout(() => window.close(), 1500);"
|
||||
"</script>"
|
||||
"</body></html>"
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
"<html><body><h2>Invalid OAuth state</h2></body></html>",
|
||||
status_code=400,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
return HTMLResponse(
|
||||
f"<html><body><h2>Token exchange failed</h2><pre>{e}</pre></body></html>",
|
||||
status_code=400,
|
||||
)
|
||||
|
||||
return HTMLResponse(
|
||||
"<html><body>"
|
||||
'<h2 style="font-family:sans-serif;color:#22c55e">Connected successfully!</h2>'
|
||||
'<p style="font-family:sans-serif;color:#666">You can close this window.</p>'
|
||||
"<script>"
|
||||
"if (window.opener) window.opener.postMessage({type:'oauth_complete', tool_id:'"
|
||||
+ tool_id
|
||||
+ "'}, '*');"
|
||||
"setTimeout(() => window.close(), 1500);"
|
||||
"</script>"
|
||||
"</body></html>"
|
||||
)
|
||||
|
||||
|
||||
@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()}
|
||||
|
||||
Reference in New Issue
Block a user