From 490bb3c5c5a9176ac763127c53f5f9a15a17625d Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 5 Jun 2026 14:14:32 +0200 Subject: [PATCH] refactor(agents): extract shared Google OAuth helper from gmail connector build_credentials/get_token_encryption are Google-OAuth helpers used by both the Gmail and Calendar connector tools. They lived inside gmail/tools/_helpers.py, forcing calendar -> gmail coupling. Move them to a neutral connector-level module (connectors/google_auth.py); gmail/_helpers.py re-exports them under the legacy private names so existing gmail tools are untouched, and calendar now imports the shared module directly. --- .../calendar/tools/search_events.py | 4 +- .../connectors/gmail/tools/_helpers.py | 65 +++++-------------- .../subagents/connectors/google_auth.py | 59 +++++++++++++++++ 3 files changed, 76 insertions(+), 52 deletions(-) create mode 100644 surfsense_backend/app/agents/chat/multi_agent_chat/subagents/connectors/google_auth.py diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/connectors/calendar/tools/search_events.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/connectors/calendar/tools/search_events.py index 68189a99f..cf9a015cf 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/connectors/calendar/tools/search_events.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/connectors/calendar/tools/search_events.py @@ -5,8 +5,8 @@ from langchain_core.tools import tool from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select -from app.agents.chat.multi_agent_chat.subagents.connectors.gmail.tools._helpers import ( - _build_credentials, +from app.agents.chat.multi_agent_chat.subagents.connectors.google_auth import ( + build_credentials as _build_credentials, ) from app.db import SearchSourceConnector, SearchSourceConnectorType diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/connectors/gmail/tools/_helpers.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/connectors/gmail/tools/_helpers.py index 5a467e328..12d984352 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/connectors/gmail/tools/_helpers.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/connectors/gmail/tools/_helpers.py @@ -1,61 +1,26 @@ -"""Shared helpers for Gmail connector tools. +"""Gmail-specific helpers for the Gmail connector tools. -Credential construction (``_build_credentials``) is also reused by the -Calendar connector tools, since both are Google OAuth backed. +Google OAuth credential construction lives in +``app.agents.chat.multi_agent_chat.subagents.connectors.google_auth`` (shared +with the Calendar connector). It is re-exported here under the legacy private +names so the existing Gmail tools keep importing it from this module. """ from __future__ import annotations -from datetime import datetime from typing import Any -from app.db import SearchSourceConnector +from app.agents.chat.multi_agent_chat.subagents.connectors.google_auth import ( + build_credentials as _build_credentials, + get_token_encryption as _get_token_encryption, +) -_token_encryption_cache: object | None = None - - -def _get_token_encryption(): - global _token_encryption_cache - if _token_encryption_cache is None: - from app.config import config - from app.utils.oauth_security import TokenEncryption - - if not config.SECRET_KEY: - raise RuntimeError("SECRET_KEY not configured for token decryption.") - _token_encryption_cache = TokenEncryption(config.SECRET_KEY) - return _token_encryption_cache - - -def _build_credentials(connector: SearchSourceConnector): - """Build Google OAuth Credentials from a connector's stored config. - - Handles both native OAuth connectors (with encrypted tokens) and - Composio-backed connectors. Shared by Gmail and Calendar tools. - """ - from app.utils.google_credentials import COMPOSIO_GOOGLE_CONNECTOR_TYPES - - if connector.connector_type in COMPOSIO_GOOGLE_CONNECTOR_TYPES: - raise ValueError("Composio connectors must use Composio tool execution.") - - from google.oauth2.credentials import Credentials - - cfg = dict(connector.config) - if cfg.get("_token_encrypted"): - enc = _get_token_encryption() - for key in ("token", "refresh_token", "client_secret"): - if cfg.get(key): - cfg[key] = enc.decrypt_token(cfg[key]) - - exp = (cfg.get("expiry") or "").replace("Z", "") - return Credentials( - token=cfg.get("token"), - refresh_token=cfg.get("refresh_token"), - token_uri=cfg.get("token_uri"), - client_id=cfg.get("client_id"), - client_secret=cfg.get("client_secret"), - scopes=cfg.get("scopes", []), - expiry=datetime.fromisoformat(exp) if exp else None, - ) +__all__ = [ + "_build_credentials", + "_format_gmail_summary", + "_get_token_encryption", + "_gmail_headers", +] def _gmail_headers(message: dict[str, Any]) -> dict[str, str]: diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/connectors/google_auth.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/connectors/google_auth.py new file mode 100644 index 000000000..6eb60ef2a --- /dev/null +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/connectors/google_auth.py @@ -0,0 +1,59 @@ +"""Google OAuth credential construction shared across Google connectors. + +Both the Gmail and Calendar connector tools are Google OAuth backed and build +``google.oauth2.credentials.Credentials`` from a stored ``SearchSourceConnector`` +the same way. This module is the single owner of that logic so neither connector +has to import the other. +""" + +from __future__ import annotations + +from datetime import datetime + +from app.db import SearchSourceConnector + +_token_encryption_cache: object | None = None + + +def get_token_encryption(): + global _token_encryption_cache + if _token_encryption_cache is None: + from app.config import config + from app.utils.oauth_security import TokenEncryption + + if not config.SECRET_KEY: + raise RuntimeError("SECRET_KEY not configured for token decryption.") + _token_encryption_cache = TokenEncryption(config.SECRET_KEY) + return _token_encryption_cache + + +def build_credentials(connector: SearchSourceConnector): + """Build Google OAuth Credentials from a connector's stored config. + + Handles both native OAuth connectors (with encrypted tokens) and + Composio-backed connectors. Shared by Gmail and Calendar tools. + """ + from app.utils.google_credentials import COMPOSIO_GOOGLE_CONNECTOR_TYPES + + if connector.connector_type in COMPOSIO_GOOGLE_CONNECTOR_TYPES: + raise ValueError("Composio connectors must use Composio tool execution.") + + from google.oauth2.credentials import Credentials + + cfg = dict(connector.config) + if cfg.get("_token_encrypted"): + enc = get_token_encryption() + for key in ("token", "refresh_token", "client_secret"): + if cfg.get(key): + cfg[key] = enc.decrypt_token(cfg[key]) + + exp = (cfg.get("expiry") or "").replace("Z", "") + return Credentials( + token=cfg.get("token"), + refresh_token=cfg.get("refresh_token"), + token_uri=cfg.get("token_uri"), + client_id=cfg.get("client_id"), + client_secret=cfg.get("client_secret"), + scopes=cfg.get("scopes", []), + expiry=datetime.fromisoformat(exp) if exp else None, + )