diff --git a/agents/providers/__init__.py b/agents/providers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/agents/providers/core.py b/agents/providers/core.py new file mode 100644 index 0000000..14a7b08 --- /dev/null +++ b/agents/providers/core.py @@ -0,0 +1,111 @@ +""" +Core Context Providers +====================== + +Always-on providers: web, workspace, CRM, knowledge. +These don't require external credentials — they work out of the box. +""" + +from os import getenv +from pathlib import Path + +from agno.context.database import DatabaseContextProvider +from agno.context.mode import ContextMode +from agno.context.web.parallel import ParallelBackend +from agno.context.web.parallel_mcp import ParallelMCPBackend +from agno.context.web.provider import WebContextProvider +from agno.context.wiki import FileSystemBackend, GitBackend, WikiContextProvider +from agno.context.workspace import WorkspaceContextProvider +from agno.tools.workspace import DEFAULT_EXCLUDE_PATTERNS +from agno.utils.log import log_info, log_warning + +from agents.instructions import CRM_READ, CRM_WRITE, KNOWLEDGE_READ, KNOWLEDGE_WRITE +from app.settings import default_model +from db import SCHEMA, get_readonly_engine, get_sql_engine + +REPO_ROOT = Path(__file__).resolve().parents[2] +KNOWLEDGE_PATH = REPO_ROOT / "knowledge" + + +def create_web_provider() -> WebContextProvider: + """Web search and page reading via Parallel or MCP backend.""" + model = default_model() + if getenv("PARALLEL_API_KEY"): + return WebContextProvider(backend=ParallelBackend(), model=model) + return WebContextProvider(backend=ParallelMCPBackend(), model=model) + + +def create_workspace_provider() -> WorkspaceContextProvider: + """Filesystem context for the context repo itself. + + mode=tools exposes the read tools (list_files / search_content / read_file) + straight to the main context agent instead of behind a nested sub-agent, so a + codebase question is answered in the agent's own turn (one pass, bounded by its + tool_call_limit) rather than paying a full sub-agent round-trip per file read. + + agno's defaults already exclude .env*, .git, caches, etc. Also keep Google + credential files out: in local dev compose mounts the repo at /app, so + without this the owner's own agent could read the minted OAuth token (or a + stray key file) back through read_file. + """ + return WorkspaceContextProvider( + root=REPO_ROOT, + model=default_model(), + mode=ContextMode.tools, + exclude_patterns=[*DEFAULT_EXCLUDE_PATTERNS, "*_token.json", "google-service-account.json"], + ) + + +def create_crm_provider() -> DatabaseContextProvider: + """The CRM — the structured database, read + write over the `crm` schema. + + The tuned instructions know the managed table shape + (projects/meetings/reminders/notes/contacts), rendered from the schema spec. + """ + return DatabaseContextProvider( + id="crm", + name="CRM", + sql_engine=get_sql_engine(), + readonly_engine=get_readonly_engine(), + schema=SCHEMA, + read_instructions=CRM_READ, + write_instructions=CRM_WRITE, + model=default_model(), + ) + + +def create_knowledge_provider() -> WikiContextProvider: + """The knowledge base — read + write knowledge, organized folder-per-spec. + + Filesystem-backed by default. Set `KNOWLEDGE_REPO_URL` AND `KNOWLEDGE_GITHUB_TOKEN` + to switch to `GitBackend` for durable storage with an audit trail. + """ + repo_url = getenv("KNOWLEDGE_REPO_URL", "").strip() + github_token = getenv("KNOWLEDGE_GITHUB_TOKEN", "").strip() + + backend: FileSystemBackend | GitBackend + if repo_url and github_token: + backend = GitBackend( + repo_url=repo_url, + github_token=github_token, + branch=getenv("KNOWLEDGE_BRANCH", "main"), + local_path=getenv("KNOWLEDGE_LOCAL_PATH") or None, + ) + log_info(f"Knowledge base: GitBackend ({repo_url})") + else: + if repo_url or github_token: + log_warning( + "Knowledge base: KNOWLEDGE_REPO_URL and KNOWLEDGE_GITHUB_TOKEN must both be set " + "to enable GitBackend; falling back to FileSystemBackend." + ) + KNOWLEDGE_PATH.mkdir(parents=True, exist_ok=True) + backend = FileSystemBackend(path=KNOWLEDGE_PATH) + + return WikiContextProvider( + id="knowledge", + name="Knowledge Base", + backend=backend, + read_instructions=KNOWLEDGE_READ, + write_instructions=KNOWLEDGE_WRITE, + model=default_model(), + ) diff --git a/agents/providers/google.py b/agents/providers/google.py new file mode 100644 index 0000000..0d7ec58 --- /dev/null +++ b/agents/providers/google.py @@ -0,0 +1,296 @@ +""" +Google Context Providers +======================== + +Gmail and Calendar provider factories with shared auth config. + +When ``GOOGLE_TOKEN_ENCRYPTION_KEY`` is set, OAuth tokens are encrypted +and stored in PostgreSQL instead of file-based token paths. The shared +``AuthConfig`` consolidates scopes across Gmail + Calendar, so a single +OAuth consent covers both services. + +Agno's Google toolkits handle all credential resolution automatically: +1. Check shared auth cache (already authenticated by another toolkit) +2. Load from DB (encrypted if key set) +3. Fallback to file (local dev) +4. Interactive OAuth (first-time setup) + +The precheck functions below validate tokens BEFORE spinning up a sub-agent, +avoiding wasted work when auth is dead. +""" + +import asyncio +import json +from os import getenv +from pathlib import Path +from typing import TYPE_CHECKING + +from agno.utils.log import log_debug, log_warning + +from agents.instructions import CALENDAR_READ, GMAIL_READ +from app.settings import default_model + +if TYPE_CHECKING: + from agno.context.provider import ContextProvider + from agno.tools.google.auth import AuthConfig + +# Repo root for default token paths +REPO_ROOT = Path(__file__).resolve().parents[2] + +# Shared auth config — lazily initialized +_google_auth_config: "AuthConfig | None" = None + + +def google_configured() -> bool: + """True when the Gmail/Calendar OAuth client is configured. + + Set ``GOOGLE_CLIENT_ID`` + ``GOOGLE_CLIENT_SECRET`` and mint the consent + tokens once with ``scripts/google_mint_tokens.py`` — see ``docs/GOOGLE.md``. + """ + return bool(getenv("GOOGLE_CLIENT_ID") and getenv("GOOGLE_CLIENT_SECRET")) + + +def get_google_auth() -> "AuthConfig | None": + """Get shared Google AuthConfig with DB storage and encryption. + + Lazily initialized to avoid import overhead when Google isn't configured. + Shared across all Google toolkits so OAuth scopes consolidate into one + consent screen and credentials are cached across providers. + + Token storage priority: + 1. DB with encryption (production) — set GOOGLE_TOKEN_ENCRYPTION_KEY + 2. DB without encryption (not recommended) — set encrypt_tokens=False + 3. File fallback (local dev) — uses token_path from each provider + """ + global _google_auth_config + if _google_auth_config is not None: + return _google_auth_config + + if not google_configured(): + return None + + try: + from agno.tools.google.auth import AuthConfig + + from db import get_postgres_db + + db = get_postgres_db() + encryption_key = getenv("GOOGLE_TOKEN_ENCRYPTION_KEY") + + # 5 min timeout for API calls — context's sub-agents can take time + http_timeout = float(getenv("GOOGLE_API_TIMEOUT", "300")) + + _google_auth_config = AuthConfig( + db=db, + token_encryption_key=encryption_key, + http_timeout=http_timeout, + ) + + # Pre-register all scopes BEFORE any toolkit is instantiated. + # Toolkits are lazy-loaded (created on first query, not at provider init), + # so if Gmail runs first, OAuth would only have Gmail scopes unless we + # pre-register everything here. + _google_auth_config.register_scopes([ + # Gmail + "https://www.googleapis.com/auth/gmail.readonly", + "https://www.googleapis.com/auth/gmail.modify", + "https://www.googleapis.com/auth/gmail.compose", + # Calendar + "https://www.googleapis.com/auth/calendar.readonly", + "https://www.googleapis.com/auth/calendar", + ]) + + if encryption_key: + log_debug("Google auth: DB storage + encryption enabled") + else: + log_debug("Google auth: DB storage enabled (no encryption key)") + + return _google_auth_config + except ImportError: + log_warning("Google auth: AuthConfig not available (google-api-python-client not installed)") + return None + + +def gmail_token_path() -> str: + """Where the Gmail OAuth token cache lives (``GMAIL_TOKEN_FILE`` or repo root). + + The single source of truth for this path: the provider reads it, the mint + script (``scripts/google_mint_tokens.py``) writes it, and the entrypoint's + base64 materialization restores it on deploys that don't keep files. + """ + return getenv("GMAIL_TOKEN_FILE") or str(REPO_ROOT / "gmail_token.json") + + +def calendar_token_path() -> str: + """Where the Calendar OAuth token cache lives (``CALENDAR_TOKEN_FILE`` or repo root).""" + return getenv("CALENDAR_TOKEN_FILE") or str(REPO_ROOT / "calendar_token.json") + + +def create_gmail_provider() -> "ContextProvider | None": + """Gmail — read + draft. ``update_gmail`` only ever creates a draft; it + never sends, so it is *not* an act tool and needs no approval gate (a draft + is private and reversible — you review and send from Gmail). + + Imported lazily: the google client libraries are optional, and the + registry's try/except treats a missing import as "provider not available" + instead of taking the app down. + """ + if not google_configured(): + return None + from agno.context.gmail import GmailContextProvider + from agno.tools.google.gmail import GmailTools + + class _DraftOnlyGmail(GmailContextProvider): + """Lock the Gmail write surface to drafts — it can never send. + + Agno's Gmail write sub-agent already drafts by default; we override the + toolkit hook to drop every outward-send tool, making drafts-only a hard + guarantee rather than a prompt convention. + + To let @context send for you instead: use ``GmailContextProvider`` + directly (drop this subclass) and add ``update_gmail`` to ``ACT_TOOLS`` + so every send pauses for your approval, like the calendar. The + implications + steps are in ``docs/GOOGLE.md``. + """ + + def _build_write_toolkit(self) -> GmailTools: + toolkit = super()._build_write_toolkit() + for name in ("send_email", "send_email_reply", "send_draft"): + toolkit.functions.pop(name, None) + toolkit.async_functions.pop(name, None) + return toolkit + + auth = get_google_auth() + return _DraftOnlyGmail( + auth=auth, + model=default_model(), + write=True, + token_path=gmail_token_path(), + read_instructions=GMAIL_READ, + ) + + +def create_calendar_provider() -> "ContextProvider | None": + """Google Calendar — read + write. ``update_calendar`` is approval-gated. + + Uses shared ``AuthConfig`` with Gmail for consolidated OAuth scopes + and encrypted DB token storage. + """ + if not google_configured(): + return None + from agno.context.calendar import GoogleCalendarContextProvider + + auth = get_google_auth() + return GoogleCalendarContextProvider( + auth=auth, + model=default_model(), + write=True, + token_path=calendar_token_path(), + read_instructions=CALENDAR_READ, + ) + + +# --------------------------------------------------------------------------- +# Token validation helpers (for precheck in tool hardening) +# --------------------------------------------------------------------------- + + +def google_token_usable_from_file(token_path: str) -> bool: + """True iff a file-based Google OAuth token is valid or can be refreshed.""" + p = Path(token_path) + if not p.exists(): + return False + try: + from google.auth.transport.requests import Request + from google.oauth2.credentials import Credentials + + creds = Credentials.from_authorized_user_file(str(p)) + except Exception: + return False + if creds.valid: + return True + if creds.expired and creds.refresh_token: + try: + creds.refresh(Request()) + except Exception: + return False + try: + p.write_text(creds.to_json()) + except Exception: + pass + return bool(creds.valid) + return False + + +def google_token_usable_from_db() -> bool: + """True iff a DB-stored Google OAuth token is valid or can be refreshed.""" + auth = get_google_auth() + if auth is None: + return False + + # Check in-memory cache first — populated by actual queries via _resolve_creds() + if auth.creds and auth.creds.valid: + return True + + if auth.db is None: + return False + try: + from agno.utils.encryption import decrypt_dict, is_encrypted + from google.auth.transport.requests import Request + from google.oauth2.credentials import Credentials + + row = auth.db.get_auth_token("google", None, "google") + if not row: + return False + token_data = row.get("token_data") + if not token_data: + return False + if is_encrypted(token_data): + token_data = decrypt_dict(token_data, key=auth.token_encryption_key) + creds = Credentials.from_authorized_user_info(token_data, row.get("granted_scopes") or []) + except Exception: + return False + if creds.valid: + # Cache valid creds so subsequent prechecks return immediately + auth.creds = creds + return True + if creds.expired and creds.refresh_token: + try: + creds.refresh(Request()) + if creds.valid: + # Cache refreshed creds — matches what _resolve_creds() does + auth.creds = creds + return True + except Exception: + return False + return False + + +def google_token_precheck(provider_id: str): + """Build a precheck for time-boxed query tools that skips Google reads on a dead token. + + Returns an async callable that yields ``None`` when the token is usable, or a one-line + "skipped" chunk to short-circuit before the sub-agent spins up. The token check runs + off the loop and is itself bounded, so a hung refresh can't stall the run either. + + Checks DB-stored tokens first (when AuthConfig is configured), falling back to file. + """ + token_path = gmail_token_path() if provider_id == "gmail" else calendar_token_path() + + async def _precheck(): + try: + # 1. Check DB-stored token first (preferred when AuthConfig is configured) + auth = get_google_auth() + if auth is not None and auth.db is not None: + usable = await asyncio.wait_for(asyncio.to_thread(google_token_usable_from_db), timeout=8) + if usable: + return None + # 2. Fall back to file-based token + usable = await asyncio.wait_for(asyncio.to_thread(google_token_usable_from_file, token_path), timeout=8) + except Exception: + usable = False + if usable: + return None + return json.dumps({"error": f"{provider_id} is unavailable right now (auth needs refresh) — skipped"}) + + return _precheck diff --git a/agents/providers/hardening.py b/agents/providers/hardening.py new file mode 100644 index 0000000..1578bc5 --- /dev/null +++ b/agents/providers/hardening.py @@ -0,0 +1,109 @@ +""" +Tool Hardening +============== + +Time-boxing and prechecks for provider tools. + +A rundown fans `use_context` out to several provider sub-agents back to back, and +agno puts no timeout around each one. We time-box every read here so a slow source +degrades to a one-line "skipped" and the rest of the brief still lands. + +Google reads get an extra guard: on a dead OAuth token we skip before spinning the +sub-agent, which also avoids agno's interactive browser-auth fallback (wrong on a +headless server). +""" + +import asyncio +import contextlib +import inspect +import json + +from agno.run import RunContext +from agno.tools import tool + + +def timeout_error(label: str, timeout: float) -> str: + """JSON error chunk for a timed-out tool.""" + return json.dumps({"error": f"{label} timed out after {int(timeout)}s — skipped"}) + + +async def _drain_into(queue: asyncio.Queue, sentinel: object, make_call) -> None: + """Producer task: run a provider tool and push each chunk onto ``queue``. + + A provider ``query_*`` entrypoint returns a coroutine; awaiting it yields either an + async generator of streamed events or a finished value. Running this in its own + task means a timeout cancels only the task — never the consumer or the calling + agent's tool flow — so a slow source can't corrupt the outer stream. + """ + try: + res = make_call() + if inspect.iscoroutine(res): + res = await res + if inspect.isasyncgen(res): + async for chunk in res: + await queue.put(chunk) + else: + await queue.put(res) + except asyncio.CancelledError: + raise + except Exception as exc: + await queue.put(json.dumps({"error": f"{type(exc).__name__}: {exc}"})) + finally: + await queue.put(sentinel) + + +async def bounded_tool_call(make_call, timeout: float, label: str): + """Yield a provider tool's chunks under a total wall-clock ``timeout``. + + The tool runs as an isolated producer task feeding a queue. On timeout we emit one + error chunk (the providers' own ``{"error": ...}`` shape) and cancel the producer. + The remaining budget also caps inter-chunk stalls, not just the total. + """ + queue: asyncio.Queue = asyncio.Queue() + sentinel = object() + task = asyncio.create_task(_drain_into(queue, sentinel, make_call)) + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + try: + while True: + remaining = deadline - loop.time() + if remaining <= 0: + yield timeout_error(label, timeout) + return + try: + item = await asyncio.wait_for(queue.get(), timeout=remaining) + except (asyncio.TimeoutError, TimeoutError): + yield timeout_error(label, timeout) + return + if item is sentinel: + return + yield item + finally: + if not task.done(): + task.cancel() + with contextlib.suppress(BaseException): + await task + + +def time_boxed_query_tool(original, timeout: float, precheck=None): + """Wrap a provider ``query_*`` tool so its sub-agent run is time-boxed. + + Same name + description; the explicit ``question`` / ``run_context`` signature keeps + agno's schema inference and run_context injection unchanged. The optional ``precheck`` + (an async callable) runs first: if it returns a chunk, we yield that and skip the + sub-agent — the Google guard uses it to short-circuit on a dead token. + """ + raw = original.entrypoint + label = original.name + + @tool(name=original.name, description=original.description) + async def _query(question: str, run_context: RunContext | None = None): + if precheck is not None: + skip = await precheck() + if skip is not None: + yield skip + return + async for chunk in bounded_tool_call(lambda: raw(question=question, run_context=run_context), timeout, label): + yield chunk + + return _query diff --git a/agents/providers/slack.py b/agents/providers/slack.py new file mode 100644 index 0000000..064751b --- /dev/null +++ b/agents/providers/slack.py @@ -0,0 +1,32 @@ +""" +Slack Context Provider +====================== + +Slack read + write. `query_slack` reads channels/DMs; `update_slack` posts. +""" + +from os import getenv + +from agno.context.slack import SlackContextProvider + +from agents.instructions import SLACK_READ +from app.settings import default_model + + +def create_slack_provider() -> SlackContextProvider | None: + """Slack — read + write. + + Note: `search.messages` needs a *user* token (`xoxp-`, scope `search:read`); a bot + token returns `not_allowed_token_type`. Agno hard-codes `enable_search_messages=True` + with no user-token slot, so search errors out and the read falls back to + channel/thread history. Pass a user token here to restore it. + """ + if not getenv("SLACK_BOT_TOKEN"): + return None + return SlackContextProvider( + model=default_model(), + read=True, + write=True, + read_instructions=SLACK_READ, + stream_sub_agent_events=False, + ) diff --git a/agents/sources.py b/agents/sources.py index 42363a8..f280672 100644 --- a/agents/sources.py +++ b/agents/sources.py @@ -1,56 +1,47 @@ """ -Context's Provider Registry -=========================== +Context Provider Registry +========================= -Wiring for the context providers available to Context. The structured database (`crm`), the knowledge base (`knowledge`), the workspace, and web are always on; Slack, Gmail, and Calendar are added when their credentials are set. +Wiring for the context providers available to @context. The structured database +(`crm`), the knowledge base (`knowledge`), the workspace, and web are always on; +Slack, Gmail, and Calendar are added when their credentials are set. -Each provider exposes at most two tools to the main agent — `query_` and `update_` — so the tool surface stays linear at 2N as sources grow. +Each provider exposes at most two tools to the main agent — `query_` and +`update_` — so the tool surface stays linear at 2N as sources grow. -`ACT_TOOLS` is the canonical list of tools that act on the outside world *as the owner* and so are approval-gated (the run pauses for the owner's explicit OK before they execute). Only `update_calendar` qualifies. Two writes are deliberately excluded: `update_gmail` only ever drafts (never sends — private and reversible), and `update_slack` is ordinary messaging. Filing into your own store is frictionless; only the sensitive outward action is gated. See `docs/SECURITY.md`. +Provider factories live in `agents/providers/`. This module handles: +- Registry lifecycle (create, get, setup, close) +- Tool hardening (time-boxing, prechecks) +- Introspection (status, logging) + +`ACT_TOOLS` gates tools that act on the outside world as the owner. Only +`update_calendar` qualifies. Two writes are deliberately excluded: `update_gmail` +only ever drafts (never sends), and `update_slack` is ordinary messaging. +See `docs/SECURITY.md`. """ import asyncio -import contextlib -import inspect import json from concurrent.futures import ThreadPoolExecutor from os import getenv -from pathlib import Path -from agno.context.database import DatabaseContextProvider -from agno.context.mode import ContextMode from agno.context.provider import ContextProvider -from agno.context.slack import SlackContextProvider -from agno.context.web.parallel import ParallelBackend -from agno.context.web.parallel_mcp import ParallelMCPBackend -from agno.context.web.provider import WebContextProvider -from agno.context.wiki import FileSystemBackend, GitBackend, WikiContextProvider -from agno.context.workspace import WorkspaceContextProvider from agno.run import RunContext from agno.tools import tool -from agno.tools.workspace import DEFAULT_EXCLUDE_PATTERNS from agno.utils.log import log_info, log_warning -from agents.instructions import ( - CALENDAR_READ, - CRM_READ, - CRM_WRITE, - GMAIL_READ, - KNOWLEDGE_READ, - KNOWLEDGE_WRITE, - SLACK_READ, +from agents.providers.core import ( + create_crm_provider, + create_knowledge_provider, + create_web_provider, + create_workspace_provider, ) -from app.settings import backbone_query_timeout, default_model, provider_query_timeout -from db import SCHEMA, get_readonly_engine, get_sql_engine - -# Workspace root for the always-on filesystem context. Hardcoded to the context repo so @context can answer questions about its own codebase out of the box. -REPO_ROOT = Path(__file__).resolve().parents[1] - -# Knowledge-base root - where @context stores files. Filesystem-backed by default; set KNOWLEDGE_REPO_URL + KNOWLEDGE_GITHUB_TOKEN to switch to GitBackend at startup for durable storage with an audit trail. -KNOWLEDGE_PATH = REPO_ROOT / "knowledge" +from agents.providers.google import create_calendar_provider, create_gmail_provider, google_token_precheck +from agents.providers.hardening import time_boxed_query_tool +from agents.providers.slack import create_slack_provider +from app.settings import backbone_query_timeout, provider_query_timeout # Tools that act on the outside world as the owner → approval-gated by gate_act_tools. -# Only the calendar; gmail is draft-only, slack is ordinary messaging (see module docstring). ACT_TOOLS: frozenset[str] = frozenset({"update_calendar"}) @@ -80,15 +71,16 @@ def gate_act_tools(tools: list) -> list: def create_context_providers() -> list[ContextProvider]: """Build the registered context providers from env and cache them. - Optional builders are wrapped in try/except so one bad config doesn't take the whole registry down. + Optional builders are wrapped in try/except so one bad config doesn't take + the whole registry down. """ configured: list[ContextProvider] = [ - _create_web_provider(), - _create_workspace_provider(), - _create_crm_provider(), - _create_knowledge_provider(), + create_web_provider(), + create_workspace_provider(), + create_crm_provider(), + create_knowledge_provider(), ] - for factory in (_create_slack_provider, _create_gmail_provider, _create_calendar_provider): + for factory in (create_slack_provider, create_gmail_provider, create_calendar_provider): try: provider = factory() except Exception as exc: @@ -153,341 +145,11 @@ async def close_context_providers() -> None: # --------------------------------------------------------------------------- -# Context Providers -# --------------------------------------------------------------------------- - - -def _create_web_provider() -> WebContextProvider: - model = default_model() - if getenv("PARALLEL_API_KEY"): - return WebContextProvider(backend=ParallelBackend(), model=model) - return WebContextProvider(backend=ParallelMCPBackend(), model=model) - - -def _create_workspace_provider() -> WorkspaceContextProvider: - # mode=tools exposes the read tools (list_files / search_content / read_file) - # straight to the main context agent instead of behind a nested sub-agent, so a - # codebase question is answered in the agent's own turn (one pass, bounded by its - # tool_call_limit) rather than paying a full sub-agent round-trip per file read. - # The usage guidance lives in OWNER_GUIDE (the agent never sees provider - # instructions); no per-source time-box is needed (see BACKBONE_SOURCES). - # - # agno's defaults already exclude .env*, .git, caches, etc. Also keep Google - # credential files out: in local dev compose mounts the repo at /app, so - # without this the owner's own agent could read the minted OAuth token (or a - # stray key file) back through read_file. (The image is clean — .dockerignore - # excludes them — this covers the mounted-dev case.) - return WorkspaceContextProvider( - root=REPO_ROOT, - model=default_model(), - mode=ContextMode.tools, - exclude_patterns=[*DEFAULT_EXCLUDE_PATTERNS, "*_token.json", "google-service-account.json"], - ) - - -def _create_crm_provider() -> DatabaseContextProvider: - """The CRM — the structured database, read + write over the `crm` schema. - - The tuned instructions know the managed table shape (projects/meetings/reminders/notes/contacts), rendered from the schema spec. - """ - return DatabaseContextProvider( - id="crm", - name="CRM", - sql_engine=get_sql_engine(), - readonly_engine=get_readonly_engine(), - schema=SCHEMA, - read_instructions=CRM_READ, - write_instructions=CRM_WRITE, - model=default_model(), - ) - - -def _create_knowledge_provider() -> WikiContextProvider: - """The knowledge base — read + write knowledge, organized folder-per-spec. - - Filesystem-backed by default. Set `KNOWLEDGE_REPO_URL` AND `KNOWLEDGE_GITHUB_TOKEN` — ideally pointing at your specs repo — to switch to `GitBackend` for durable storage with an audit trail. Optional knobs: `KNOWLEDGE_BRANCH` (default `main`), `KNOWLEDGE_LOCAL_PATH`. - """ - repo_url = getenv("KNOWLEDGE_REPO_URL", "").strip() - github_token = getenv("KNOWLEDGE_GITHUB_TOKEN", "").strip() - - backend: FileSystemBackend | GitBackend - if repo_url and github_token: - backend = GitBackend( - repo_url=repo_url, - github_token=github_token, - branch=getenv("KNOWLEDGE_BRANCH", "main"), - local_path=getenv("KNOWLEDGE_LOCAL_PATH") or None, - ) - log_info(f"Knowledge base: GitBackend ({repo_url})") - else: - if repo_url or github_token: - log_warning( - "Knowledge base: KNOWLEDGE_REPO_URL and KNOWLEDGE_GITHUB_TOKEN must both be set " - "to enable GitBackend; falling back to FileSystemBackend." - ) - KNOWLEDGE_PATH.mkdir(parents=True, exist_ok=True) - backend = FileSystemBackend(path=KNOWLEDGE_PATH) - - return WikiContextProvider( - id="knowledge", - name="Knowledge Base", - backend=backend, - read_instructions=KNOWLEDGE_READ, - write_instructions=KNOWLEDGE_WRITE, - model=default_model(), - ) - - -def _create_slack_provider() -> SlackContextProvider | None: - """Slack — read + write. `query_slack` reads channels/DMs; `update_slack` posts. - - Note: `search.messages` needs a *user* token (`xoxp-`, scope `search:read`); a bot - token returns `not_allowed_token_type`. Agno hard-codes `enable_search_messages=True` - with no user-token slot, so search errors out and the read falls back to - channel/thread history. Pass a user token here to restore it. - """ - if not getenv("SLACK_BOT_TOKEN"): - return None - return SlackContextProvider(model=default_model(), read=True, write=True, read_instructions=SLACK_READ) - - -def _google_configured() -> bool: - """True when the Gmail/Calendar OAuth client is configured. - - Set ``GOOGLE_CLIENT_ID`` + ``GOOGLE_CLIENT_SECRET`` and mint the consent - tokens once with ``scripts/google_mint_tokens.py`` — see ``docs/GOOGLE.md``. - """ - return bool(getenv("GOOGLE_CLIENT_ID") and getenv("GOOGLE_CLIENT_SECRET")) - - -def gmail_token_path() -> str: - """Where the Gmail OAuth token cache lives (``GMAIL_TOKEN_FILE`` or repo root). - - The single source of truth for this path: the provider reads it, the mint - script (``scripts/google_mint_tokens.py``) writes it, and the entrypoint's - base64 materialization restores it on deploys that don't keep files. - """ - return getenv("GMAIL_TOKEN_FILE") or str(REPO_ROOT / "gmail_token.json") - - -def calendar_token_path() -> str: - """Where the Calendar OAuth token cache lives (``CALENDAR_TOKEN_FILE`` or repo root).""" - return getenv("CALENDAR_TOKEN_FILE") or str(REPO_ROOT / "calendar_token.json") - - -def _create_gmail_provider() -> ContextProvider | None: - """Gmail — read + draft. ``update_gmail`` only ever creates a draft; it - never sends, so it is *not* an act tool and needs no approval gate (a draft - is private and reversible — you review and send from Gmail). - - Imported lazily: the google client libraries are optional, and the - registry's try/except treats a missing import as "provider not available" - instead of taking the app down. - """ - if not _google_configured(): - return None - from agno.context.gmail import GmailContextProvider - from agno.tools.google.gmail import GmailTools - - class _DraftOnlyGmail(GmailContextProvider): - """Lock the Gmail write surface to drafts — it can never send. - - Agno's Gmail write sub-agent already drafts by default; we override the - toolkit hook to drop every outward-send tool, making drafts-only a hard - guarantee rather than a prompt convention. - - To let @context send for you instead: use ``GmailContextProvider`` - directly (drop this subclass) and add ``update_gmail`` to ``ACT_TOOLS`` - so every send pauses for your approval, like the calendar. The - implications + steps are in ``docs/GOOGLE.md``. - """ - - def _build_write_toolkit(self) -> GmailTools: - toolkit = super()._build_write_toolkit() - # Strip the send tools; keep create_draft_email / update_draft. - for name in ("send_email", "send_email_reply", "send_draft"): - toolkit.functions.pop(name, None) - toolkit.async_functions.pop(name, None) - return toolkit - - return _DraftOnlyGmail( - model=default_model(), - write=True, - token_path=gmail_token_path(), - read_instructions=GMAIL_READ, - ) - - -def _create_calendar_provider() -> ContextProvider | None: - """Google Calendar — read + write. ``update_calendar`` is approval-gated.""" - if not _google_configured(): - return None - from agno.context.calendar import GoogleCalendarContextProvider - - return GoogleCalendarContextProvider( - model=default_model(), - write=True, - token_path=calendar_token_path(), - read_instructions=CALENDAR_READ, - ) - - -# --------------------------------------------------------------------------- -# Owner tool hardening +# Tool hardening # --------------------------------------------------------------------------- -# -# A rundown fans `use_context` out to several provider sub-agents back to back, and -# agno puts no timeout around each one. We time-box every read here so a slow source -# degrades to a one-line "skipped" and the rest of the brief still lands. -# -# Google reads get an extra guard: on a dead OAuth token we skip before spinning the -# sub-agent, which also avoids agno's interactive browser-auth fallback (wrong on a -# headless server). The token check uses only the public google.oauth2 API. - - -def _timeout_error(label: str, timeout: float) -> str: - return json.dumps({"error": f"{label} timed out after {int(timeout)}s — skipped"}) - - -async def _drain_into(queue: asyncio.Queue, sentinel: object, make_call) -> None: - """Producer task: run a provider tool and push each chunk onto ``queue``. - - A provider ``query_*`` entrypoint returns a coroutine; awaiting it yields either an - async generator of streamed events or a finished value. Running this in its own - task means a timeout cancels only the task — never the consumer or the calling - agent's tool flow — so a slow source can't corrupt the outer stream. - """ - try: - res = make_call() - if inspect.iscoroutine(res): - res = await res - if inspect.isasyncgen(res): - async for chunk in res: - await queue.put(chunk) - else: - await queue.put(res) - except asyncio.CancelledError: - raise - except Exception as exc: - await queue.put(json.dumps({"error": f"{type(exc).__name__}: {exc}"})) - finally: - await queue.put(sentinel) - - -async def _bounded_tool_call(make_call, timeout: float, label: str): - """Yield a provider tool's chunks under a total wall-clock ``timeout``. - - The tool runs as an isolated producer task feeding a queue. On timeout we emit one - error chunk (the providers' own ``{"error": ...}`` shape) and cancel the producer. - The remaining budget also caps inter-chunk stalls, not just the total. - """ - queue: asyncio.Queue = asyncio.Queue() - sentinel = object() - task = asyncio.create_task(_drain_into(queue, sentinel, make_call)) - loop = asyncio.get_running_loop() - deadline = loop.time() + timeout - try: - while True: - remaining = deadline - loop.time() - if remaining <= 0: - yield _timeout_error(label, timeout) - return - try: - item = await asyncio.wait_for(queue.get(), timeout=remaining) - except (asyncio.TimeoutError, TimeoutError): - yield _timeout_error(label, timeout) - return - if item is sentinel: - return - yield item - finally: - if not task.done(): - task.cancel() - with contextlib.suppress(BaseException): - await task - - -def _time_boxed_query_tool(original, timeout: float, precheck=None): - """Wrap a provider ``query_*`` tool so its sub-agent run is time-boxed. - - Same name + description; the explicit ``question`` / ``run_context`` signature keeps - agno's schema inference and run_context injection unchanged. The optional ``precheck`` - (an async callable) runs first: if it returns a chunk, we yield that and skip the - sub-agent — the Google guard uses it to short-circuit on a dead token. - """ - raw = original.entrypoint - label = original.name - - @tool(name=original.name, description=original.description) - async def _query(question: str, run_context: RunContext | None = None): - if precheck is not None: - skip = await precheck() - if skip is not None: - yield skip - return - async for chunk in _bounded_tool_call(lambda: raw(question=question, run_context=run_context), timeout, label): - yield chunk - - return _query - - -def _google_token_usable(token_path: str) -> bool: - """True iff a Google OAuth token is valid or can be refreshed without a browser. - - Refreshes and persists an expired-but-refreshable token in place, so the provider's - sub-agent then loads a valid one. Never triggers interactive auth — a dead token - just returns False. - """ - p = Path(token_path) - if not p.exists(): - return False - try: - from google.auth.transport.requests import Request - from google.oauth2.credentials import Credentials - - creds = Credentials.from_authorized_user_file(str(p)) - except Exception: - return False - if creds.valid: - return True - if creds.expired and creds.refresh_token: - try: - creds.refresh(Request()) - except Exception: - return False - try: - p.write_text(creds.to_json()) - except Exception: - pass - return bool(creds.valid) - return False - - -def _google_token_precheck(provider_id: str): - """Build a precheck for `_time_boxed_query_tool` that skips Google reads on a dead token. - - Returns an async callable that yields ``None`` when the token is usable, or a one-line - "skipped" chunk to short-circuit before the sub-agent spins up. The token check runs - off the loop and is itself bounded, so a hung refresh can't stall the run either. - """ - token_path = gmail_token_path() if provider_id == "gmail" else calendar_token_path() - - async def _precheck(): - try: - usable = await asyncio.wait_for(asyncio.to_thread(_google_token_usable, token_path), timeout=8) - except Exception: - usable = False - if usable: - return None - return json.dumps({"error": f"{provider_id} is unavailable right now (auth needs refresh) — skipped"}) - - return _precheck - # Backbone read sources — the brief's spine. They get a longer per-source budget -# than best-effort sources (see backbone_query_timeout) so they reliably land in the -# concurrent fan-out, where best-effort sources still skip fast. Just the CRM today; -# the inbound queue (`rundown`) isn't a query_* sub-agent, so it isn't time-boxed. +# than best-effort sources so they reliably land in the concurrent fan-out. BACKBONE_SOURCES: frozenset[str] = frozenset({"crm"}) @@ -509,8 +171,8 @@ def owner_provider_tools() -> list: tools.append(t) continue timeout = backbone if ctx.id in BACKBONE_SOURCES else best_effort - precheck = _google_token_precheck(ctx.id) if ctx.id in ("gmail", "calendar") else None - tools.append(_time_boxed_query_tool(t, timeout, precheck)) + precheck = google_token_precheck(ctx.id) if ctx.id in ("gmail", "calendar") else None + tools.append(time_boxed_query_tool(t, timeout, precheck)) return tools diff --git a/app/main.py b/app/main.py index cd42d6a..df2181d 100644 --- a/app/main.py +++ b/app/main.py @@ -3,6 +3,10 @@ ================== """ +from dotenv import load_dotenv + +load_dotenv() # ruff: noqa: E402 + from contextlib import asynccontextmanager from os import getenv from pathlib import Path diff --git a/app/settings.py b/app/settings.py index 457cc71..cab64a3 100644 --- a/app/settings.py +++ b/app/settings.py @@ -89,7 +89,7 @@ def use_context_timeout() -> float: message instead of hanging the client. Keep it under the client's own tool timeout (e.g. Claude Code's ``MCP_TOOL_TIMEOUT``) so our message wins, not a dead stream. """ - return _float_env("USE_CONTEXT_TIMEOUT", 55.0) + return _float_env("USE_CONTEXT_TIMEOUT", 300.0) def provider_query_timeout() -> float: @@ -98,7 +98,7 @@ def provider_query_timeout() -> float: A slow source degrades to a one-line "skipped" and the rest of the brief still lands. Smaller than ``use_context_timeout`` so several can skip within one budget. """ - return _float_env("PROVIDER_TIMEOUT", 20.0) + return _float_env("PROVIDER_TIMEOUT", 180.0) def backbone_query_timeout() -> float: diff --git a/db/url.py b/db/url.py index 9fe252b..728bf07 100644 --- a/db/url.py +++ b/db/url.py @@ -8,7 +8,20 @@ def build_db_url() -> str: - """Build database URL from environment variables.""" + """Build database URL from environment variables. + + Supports three formats (checked in order): + 1. DATABASE_PUBLIC_URL — Railway's public proxy URL (for local scripts) + 2. DATABASE_URL — Railway's internal URL (for deployed services) + 3. DB_* env vars — explicit config (docker compose default) + """ + # Railway URLs (postgresql:// -> postgresql+psycopg://) + for var in ("DATABASE_PUBLIC_URL", "DATABASE_URL"): + url = getenv(var) + if url and url.startswith("postgresql://"): + return url.replace("postgresql://", "postgresql+psycopg://", 1) + + # Explicit config driver = getenv("DB_DRIVER", "postgresql+psycopg") user = getenv("DB_USER", "context") password = quote(getenv("DB_PASS", "context"), safe="") diff --git a/docs/GOOGLE.md b/docs/GOOGLE.md index 87dcbee..2d76abf 100644 --- a/docs/GOOGLE.md +++ b/docs/GOOGLE.md @@ -30,24 +30,26 @@ Out of the box, Google puts a new app in **"Testing,"** and in that mode the log ### 3. Add the credentials to `.env` ```bash +# Google OAuth GOOGLE_CLIENT_ID=***.apps.googleusercontent.com GOOGLE_CLIENT_SECRET=*** -GOOGLE_PROJECT_ID=your-project-id + +# Token encryption (required) — generate with: +# python scripts/google_mint_tokens.py --generate-key +GOOGLE_TOKEN_ENCRYPTION_KEY=*** ``` ### 4. Connect your account -The consent screen opens a browser, so this runs on your machine, not in the container. From the repo root, with the venv active (`./scripts/venv_setup.sh` if you don't have one): +The OAuth consent screen opens a browser, so this step runs on your machine (not in the container). From the repo root, with the venv active (`./scripts/venv_setup.sh` if you don't have one): ```bash python scripts/google_mint_tokens.py ``` -It reads `.env` (and `.env.production`, if you have one) and opens a browser to connect each service you haven't linked yet — Gmail, then Calendar; approve both — writing `gmail_token.json` / `calendar_token.json` at the repo root. - -It prints the Google account each token is connected to, so you can confirm it's the right one. The dev container picks the token files up through the existing `.:/app` mount. +It opens a browser to connect Gmail + Calendar. Approve both, and the encrypted token is saved to PostgreSQL. -Pass `--force` to re-mint from scratch or to connect a different account. +Pass `--force` to re-connect from scratch or to switch accounts. ### 5. Restart @@ -57,11 +59,15 @@ docker compose up -d ## Deploying (Railway) -The tokens are stored in a file on the local filesystem. These files are gitignored and are not part of the image. +Tokens are encrypted and stored in PostgreSQL, so they persist across deploys automatically. + +To mint directly to your production database: -But for production, we store the tokens as environment variables that are decoded by the entrypoint at startup. +```bash +railway run python scripts/google_mint_tokens.py +``` -The `google_mint_tokens.py` script sets this up for you: once the tokens are written it base64s them into `GMAIL_TOKEN_JSON_B64` and `CALENDAR_TOKEN_JSON_B64` in `.env.production`, then offers to run `./scripts/railway/env-sync.sh` to push them to Railway. It only writes to `.env.production` — local dev reads the token files directly through the `.:/app` mount, so `.env` never needs them. +This injects your Railway env vars (including `DB_*`), opens the browser locally for OAuth, and saves the encrypted token to your prod DB. ## Verify @@ -88,10 +94,10 @@ and it's booked, decline and nothing happens. Off by default on purpose — drafting keeps you in control and means nothing goes out unread. If you do want @context to send directly: -1. In [`agents/sources.py`](../agents/sources.py), in `_create_gmail_provider`, +1. In [`agents/providers/google.py`](../agents/providers/google.py), in `create_gmail_provider`, stop stripping the send tools (the `_DraftOnlyGmail` subclass) — or use `GmailContextProvider` directly. -2. Add `update_gmail` to `ACT_TOOLS` in the same file, so every send **pauses +2. Add `update_gmail` to `ACT_TOOLS` in [`agents/sources.py`](../agents/sources.py), so every send **pauses for your approval** the same way calendar changes do — never ship sending ungated. @@ -104,17 +110,15 @@ approvals queue. credentials are set; a misconfig is logged and skipped (the app still starts). Check the startup logs for a `_create_gmail_provider failed` / `_create_calendar_provider failed` warning, and confirm `.env` has - `GOOGLE_CLIENT_ID` + `GOOGLE_CLIENT_SECRET`. + `GOOGLE_CLIENT_ID`, `GOOGLE_CLIENT_SECRET`, and `GOOGLE_TOKEN_ENCRYPTION_KEY`. - **Access stopped working after about a week.** The app is still in "Testing" — do step 2 (Internal for Workspace, or Publish for personal), then re-run - `python scripts/google_mint_tokens.py --force` (it re-mints, rewrites the - base64, and offers the Railway sync). + `python scripts/google_mint_tokens.py --force`. - **`access_blocked` / "app isn't verified" during connect.** On a personal account, add your address as a **test user** on the consent screen, or finish publishing it (step 2). Clicking past the unverified notice is expected for your own app. -- **A token got revoked.** Re-run the mint script with `--force` — it re-mints, - rewrites the base64, and offers the Railway sync. +- **A token got revoked.** Re-run the mint script with `--force`. ## Scope it down diff --git a/example.env b/example.env index 5d9b6ba..ebe0542 100644 --- a/example.env +++ b/example.env @@ -75,22 +75,16 @@ OPENAI_API_KEY=sk-*** # Gmail + Google Calendar # Lets @context read your inbox and calendar, draft emails (it never sends — # you send from Gmail), and propose calendar changes (they go to your approvals -# queue). A few minutes to set up — full guide in docs/GOOGLE.md. +# queue). Full guide in docs/GOOGLE.md. # -# 1. Create a Google OAuth client (with the Gmail + Calendar APIs on). -# 2. Mint the consent tokens once: python scripts/google_mint_tokens.py +# Setup: +# 1. Create a Google OAuth client (enable Gmail + Calendar APIs) +# 2. Generate encryption key: python scripts/google_mint_tokens.py --generate-key +# 3. Mint the token: python scripts/google_mint_tokens.py # --------------------------------------------------------------------------- # GOOGLE_CLIENT_ID= # GOOGLE_CLIENT_SECRET= -# GOOGLE_PROJECT_ID= -# -# Token caches (written by the mint script). Default to gmail_token.json / -# calendar_token.json at the repo root; override the path, or ship the minted -# tokens as base64 so they survive a deploy (the entrypoint restores them). -# GMAIL_TOKEN_FILE= -# CALENDAR_TOKEN_FILE= -# GMAIL_TOKEN_JSON_B64= -# CALENDAR_TOKEN_JSON_B64= +# GOOGLE_TOKEN_ENCRYPTION_KEY= # --------------------------------------------------------------------------- # Scheduled digests (Slack) diff --git a/requirements.txt b/requirements.txt index f48a72a..b23bf9b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,26 +1,26 @@ # This file was autogenerated by uv via the following command: # ./scripts/generate_requirements.sh upgrade -agno==2.6.16 +agno==2.6.20 aiohappyeyeballs==2.6.2 aiohttp==3.14.1 aiosignal==1.4.0 annotated-doc==0.0.4 annotated-types==0.7.0 -anyio==4.13.0 +anyio==4.14.1 attrs==26.1.0 authlib==1.7.2 beartype==0.22.9 burner-redis==0.1.7 cachetools==7.1.4 -certifi==2026.5.20 +certifi==2026.6.17 cffi==2.0.0 charset-normalizer==3.4.7 -click==8.4.1 +click==8.4.2 cloudpickle==3.1.2 croniter==6.2.2 cronsim==2.7 cryptography==49.0.0 -cyclopts==4.18.0 +cyclopts==4.19.0 detect-installer==0.1.0 diskcache==5.6.3 distro==1.9.0 @@ -29,25 +29,25 @@ docstring-parser==0.18.0 email-validator==2.3.0 exceptiongroup==1.3.1 fakeredis==2.34.1 -fastapi==0.137.1 -fastapi-cli==0.0.24 -fastapi-cloud-cli==0.20.0 +fastapi==0.138.1 +fastapi-cli==0.0.27 +fastapi-cloud-cli==0.21.0 fastar==0.11.0 fastmcp==2.14.7 frozenlist==1.8.0 gitdb==4.0.12 gitpython==3.1.50 google-api-core==2.31.0 -google-api-python-client==2.197.0 -google-auth==2.54.0 +google-api-python-client==2.198.0 +google-auth==2.55.1 google-auth-httplib2==0.4.0 google-auth-oauthlib==1.4.0 googleapis-common-protos==1.75.0 h11==0.16.0 h2==4.3.0 -hpack==4.1.0 +hpack==4.2.0 httpcore==1.0.9 -httplib2==0.31.2 +httplib2==0.32.0 httptools==0.8.0 httpx==0.28.1 httpx-sse==0.4.3 @@ -67,21 +67,21 @@ keyring==25.7.0 lupa==2.8 markdown-it-py==4.2.0 markupsafe==3.0.3 -mcp==1.27.2 +mcp==1.28.1 mdurl==0.1.2 more-itertools==11.1.0 multidict==6.7.1 -numpy==2.4.6 +numpy==2.5.0 oauthlib==3.3.1 -openai==2.41.1 +openai==2.44.0 openapi-pydantic==0.5.1 openinference-instrumentation==0.1.53 -openinference-instrumentation-agno==0.1.37 +openinference-instrumentation-agno==0.1.38 openinference-semantic-conventions==0.1.30 -opentelemetry-api==1.42.1 -opentelemetry-instrumentation==0.63b1 -opentelemetry-sdk==1.42.1 -opentelemetry-semantic-conventions==0.63b1 +opentelemetry-api==1.43.0 +opentelemetry-instrumentation==0.64b0 +opentelemetry-sdk==1.43.0 +opentelemetry-semantic-conventions==0.64b0 packaging==26.2 parallel-web==1.1.0 pathable==0.6.0 @@ -102,7 +102,7 @@ pycparser==3.0 pydantic==2.13.4 pydantic-core==2.46.4 pydantic-extra-types==2.11.1 -pydantic-settings==2.14.1 +pydantic-settings==2.14.2 pydocket==0.22.0 pygments==2.20.0 pyjwt==2.13.0 @@ -114,16 +114,16 @@ python-json-logger==4.1.0 python-multipart==0.0.32 pytz==2026.2 pyyaml==6.0.3 -redis==8.0.0 +redis==8.0.1 referencing==0.37.0 requests==2.34.2 requests-oauthlib==2.0.0 rich==15.0.0 -rich-rst==2.0.1 +rich-rst==2.0.2 rich-toolkit==0.20.1 rignore==0.7.6 rpds-py==2026.5.1 -sentry-sdk==2.62.0 +sentry-sdk==2.63.0 shellingham==1.5.4 six==1.17.0 slack-sdk==3.42.0 @@ -131,10 +131,10 @@ smmap==5.0.3 sniffio==1.3.1 sortedcontainers==2.4.0 sqlalchemy==2.0.51 -sse-starlette==3.4.4 +sse-starlette==3.4.5 starlette==1.3.1 -tqdm==4.68.2 -typer==0.26.7 +tqdm==4.68.3 +typer==0.26.8 typing-extensions==4.15.0 typing-inspection==0.4.2 uncalled-for==0.3.2 @@ -144,5 +144,5 @@ uvicorn==0.49.0 uvloop==0.22.1 watchfiles==1.2.0 websockets==16.0 -wrapt==2.2.1 +wrapt==2.2.2 yarl==1.24.2 diff --git a/scripts/google_mint_tokens.py b/scripts/google_mint_tokens.py old mode 100755 new mode 100644 index a87df8d..0278ad7 --- a/scripts/google_mint_tokens.py +++ b/scripts/google_mint_tokens.py @@ -1,259 +1,136 @@ #!/usr/bin/env python3 """ -@context Google token minter — connects your Gmail + Calendar. - -The OAuth consent flow opens a browser, so it can't run in the container — you -mint the tokens once on your machine and the dev container picks them up through -the .:/app mount. This script is the one-command version of that: it loads your -env, checks the OAuth client is configured, mints a token file for Gmail and one -for Calendar with exactly the scopes @context's providers use, and writes the -tokens back as base64 so a Railway deploy can restore them. - -Usage (from the repo root, with the venv active — ./scripts/venv_setup.sh): - - python scripts/google_mint_tokens.py # mint what's missing - python scripts/google_mint_tokens.py --force # re-mint even if tokens exist - -Prereqs (in .env or .env.production — see docs/GOOGLE.md): - GOOGLE_CLIENT_ID, GOOGLE_CLIENT_SECRET, GOOGLE_PROJECT_ID - -What it does: - - Loads creds from .env and .env.production (real env > .env > .env.production). - - Mints GMAIL_TOKEN_FILE / CALENDAR_TOKEN_FILE (repo root by default) — gitignored. - - Prints which Google account each token authorized, so you can confirm it. - - Upserts GMAIL_TOKEN_JSON_B64 / CALENDAR_TOKEN_JSON_B64 into .env.production - for the entrypoint to restore on deploy (dev reads the token files directly, - so .env never needs them). - - Offers to run ./scripts/railway/env-sync.sh so the deploy picks them up. +Mint Google OAuth tokens — encrypted, DB-backed. + +Opens a browser for OAuth consent, saves the encrypted token to PostgreSQL. +Requires GOOGLE_TOKEN_ENCRYPTION_KEY to be set (tokens are always encrypted). + +Usage: + python scripts/google_mint_tokens.py # mint if no token exists + python scripts/google_mint_tokens.py --force # re-mint (delete + mint) + python scripts/google_mint_tokens.py --generate-key # print a new encryption key + +Setup: + 1. Create OAuth credentials at https://console.cloud.google.com + (Enable Gmail + Calendar APIs, create Desktop app credentials) + 2. Add to .env: + GOOGLE_CLIENT_ID=... + GOOGLE_CLIENT_SECRET=... + GOOGLE_TOKEN_ENCRYPTION_KEY= + 3. Run this script to mint the token """ import argparse -import base64 -import subprocess import sys from pathlib import Path REPO_ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(REPO_ROOT)) -# The env files we read creds from, in precedence order (earlier wins; real -# environment beats both) — so dev creds in .env beat deploy creds in .env.production. -# The base64 tokens are written back only to .env.production (see main()): they -# exist solely to survive a baked deploy image, and in dev the container reads the -# token files through the .:/app mount, so .env never needs them. -ENV_FILES = (".env", ".env.production") - -# The scopes the Gmail + Calendar providers use. Minted as the union of read + -# write so the read and write sub-agents can share one token file per service. -# Kept in sync with docs/GOOGLE.md. -GMAIL_SCOPES = [ - "https://www.googleapis.com/auth/gmail.readonly", - "https://www.googleapis.com/auth/gmail.modify", - "https://www.googleapis.com/auth/gmail.compose", -] -CALENDAR_SCOPES = [ - "https://www.googleapis.com/auth/calendar.readonly", - "https://www.googleapis.com/auth/calendar", -] - - -def _load_dotenv() -> None: - """Populate os.environ from .env and .env.production (without overriding real env). - - A tiny, dependency-free KEY=VALUE parser so you don't have to `set -a; source` - first. Precedence: real environment > .env > .env.production — so dev creds in - .env win over deploy creds in .env.production, and an exported value beats both. - """ - from os import environ - - for name in ENV_FILES: - env_file = REPO_ROOT / name - if not env_file.exists(): - continue - for raw in env_file.read_text().splitlines(): - line = raw.strip() - if not line or line.startswith("#") or "=" not in line: - continue - key, _, value = line.partition("=") - key = key.strip() - value = value.strip().strip('"').strip("'") - if key and key not in environ: - environ[key] = value - - -def _managed_key(line: str, keys: dict[str, str]) -> str | None: - """The managed key a line sets — matching both `KEY=...` and a commented `# KEY=`. - - This lets the minter fill the placeholders that ship commented-out in - example.env (e.g. `# GMAIL_TOKEN_JSON_B64=`) in place, rather than appending a - duplicate beside them. Prose comments never match: the key has to sit - immediately before the `=`, and only the keys we're writing are considered. - """ - body = line.lstrip() - if body.startswith("#"): - body = body.lstrip("#").lstrip() - if "=" not in body: - return None - candidate = body.split("=", 1)[0].strip() - return candidate if candidate in keys else None - - -def _upsert_env(path: Path, updates: dict[str, str]) -> None: - """Set KEY=VALUE lines in an env file, idempotently. - - Fills the first existing line for each key in place — whether it's already - active (`KEY=old`) or a commented placeholder (`# KEY=`) — so re-runs and the - commented placeholders shipped in example.env never leave duplicates. Any - later line for an already-filled key is dropped; keys with no line at all are - appended under a labeled comment. - """ - lines = path.read_text().splitlines() if path.exists() else [] - pending = dict(updates) - out: list[str] = [] - for line in lines: - key = _managed_key(line, updates) - if key is None: - out.append(line) - elif key in pending: - out.append(f"{key}={pending.pop(key)}") # fill the first slot we find - # else: a later duplicate / placeholder for an already-filled key — drop it - if pending: - if out and out[-1].strip(): - out.append("") - out.append("# Gmail/Calendar OAuth tokens (base64) — restored by scripts/entrypoint.sh on deploy") - out.extend(f"{key}={value}" for key, value in pending.items()) - path.write_text("\n".join(out) + "\n") - - -def _gmail_account(token_path: Path) -> str | None: - """The Google address this Gmail token authorized (None if it can't be read).""" - try: - from google.oauth2.credentials import Credentials - from googleapiclient.discovery import build - - creds = Credentials.from_authorized_user_file(str(token_path), GMAIL_SCOPES) - service = build("gmail", "v1", credentials=creds, cache_discovery=False) - return service.users().getProfile(userId="me").execute().get("emailAddress") - except Exception: - return None - - -def _calendar_account(token_path: Path) -> str | None: - """The Google address this Calendar token authorized (the primary calendar id).""" - try: - from google.oauth2.credentials import Credentials - from googleapiclient.discovery import build - - creds = Credentials.from_authorized_user_file(str(token_path), CALENDAR_SCOPES) - service = build("calendar", "v3", credentials=creds, cache_discovery=False) - return service.calendarList().get(calendarId="primary").execute().get("id") - except Exception: - return None - - -def _mint(label: str, token_path: Path, trigger, account, force: bool) -> bool: - """Mint one service's token (or reuse an existing one). Returns True on success. - - `trigger()` runs the OAuth flow (opens the browser on first use and writes the - token file); `account(path)` reads back which Google account it authorized. - """ - if token_path.exists() and not force: - print(f"→ {label}: token already at {token_path} (skipping; --force to re-mint).") - else: - if force and token_path.exists(): - token_path.unlink() - print(f"→ {label}: opening browser for consent...") - trigger() - if not token_path.exists(): - print(f" {label} token was not written to {token_path} — check the client credentials.") - return False - print(f" wrote {token_path}") - - who = account(token_path) - print( - f" connected account: {who}" if who else " (couldn't read the connected account — token may need a re-mint)" - ) - return True +def generate_key() -> str: + from agno.utils.encryption import generate_encryption_key + + return generate_encryption_key() -def main() -> int: - from os import getenv - parser = argparse.ArgumentParser(description="Mint @context's Gmail + Calendar OAuth tokens.") - parser.add_argument("--force", action="store_true", help="re-mint even if a token file already exists") +def main() -> int: + parser = argparse.ArgumentParser(description="Mint Google OAuth tokens (encrypted, DB-backed)") + parser.add_argument("--force", action="store_true", help="Delete existing token and re-mint") + parser.add_argument("--generate-key", action="store_true", help="Generate an encryption key and exit") args = parser.parse_args() - _load_dotenv() + if args.generate_key: + key = generate_key() + print(f"GOOGLE_TOKEN_ENCRYPTION_KEY={key}") + print("\nAdd this to your .env file.") + return 0 + + from dotenv import load_dotenv + + load_dotenv() - missing = [k for k in ("GOOGLE_CLIENT_ID", "GOOGLE_CLIENT_SECRET") if not getenv(k)] - if missing: - print(f"Missing required env: {', '.join(missing)}.") - print(f"Set the OAuth client in {' or '.join(ENV_FILES)} — see docs/GOOGLE.md.") + import os + + from agents.providers.google import get_google_auth, google_configured + + # 1. Check OAuth credentials + if not google_configured(): + print("ERROR: Set GOOGLE_CLIENT_ID + GOOGLE_CLIENT_SECRET in .env") + print(" See docs/GOOGLE.md for setup instructions.") return 1 - from agents.sources import calendar_token_path, gmail_token_path + # 2. Check encryption key + encryption_key = os.getenv("GOOGLE_TOKEN_ENCRYPTION_KEY") + if not encryption_key: + print("ERROR: GOOGLE_TOKEN_ENCRYPTION_KEY not set.") + print(" Generate one with: python scripts/google_mint_tokens.py --generate-key") + return 1 - gmail_token = Path(gmail_token_path()) - calendar_token = Path(calendar_token_path()) + # 3. Create auth config + auth = get_google_auth() + if not auth: + print("ERROR: Failed to create AuthConfig") + return 1 - print("Minting Google tokens. Your browser opens for each one that isn't already minted.\n") + if not auth.db: + print("ERROR: Database not configured. Check DB_* env vars.") + return 1 + + print(f"Scopes: {len(auth.scopes)} (Gmail + Calendar)") + print(f"DB: {auth.db.id}") + print(f"Encryption: enabled") + print() - # Importing here keeps the missing-env message fast and dependency-light. - from agno.tools.google.calendar import GoogleCalendarTools + # 4. Check existing token + row = auth.db.get_auth_token("google", None, "google") + if row and not args.force: + scopes = row.get("granted_scopes", []) + print(f"Token already exists with {len(scopes)} scopes:") + for s in scopes: + print(f" - {s.split('/')[-1]}") + print("\nUse --force to delete and re-mint.") + return 0 + + if row and args.force: + from db.url import db_url + from sqlalchemy import create_engine, text + + engine = create_engine(db_url) + with engine.begin() as conn: + conn.execute(text("DELETE FROM ai.agno_auth_tokens WHERE provider = 'google'")) + print("Deleted existing token.\n") + + # 5. Trigger OAuth from agno.tools.google.gmail import GmailTools - ok = _mint( - "Gmail", - gmail_token, - lambda: GmailTools(token_path=str(gmail_token), scopes=GMAIL_SCOPES).get_latest_emails(1), - _gmail_account, - args.force, - ) - if not ok: + print("Opening browser for OAuth consent...") + print("(Grant access to Gmail + Calendar)\n") + + gmail = GmailTools(auth=auth) + result = gmail.get_latest_emails(count=1) + + if "error" in result.lower(): + print(f"FAILED: {result}") return 1 - ok = _mint( - "Calendar", - calendar_token, - lambda: GoogleCalendarTools(token_path=str(calendar_token), scopes=CALENDAR_SCOPES).list_events(limit=1), - _calendar_account, - args.force, - ) - if not ok: + # 6. Verify token was saved + row = auth.db.get_auth_token("google", None, "google") + if not row: + print("FAILED: Token not saved to DB") return 1 - # Write the tokens back as base64 so a Railway deploy can restore them (the - # token files are gitignored + .dockerignore'd and don't survive a redeploy). - # This is deploy-only: in dev the container reads the token files through the - # .:/app mount, so .env never needs these — we only touch .env.production. - b64 = { - "GMAIL_TOKEN_JSON_B64": base64.b64encode(gmail_token.read_bytes()).decode(), - "CALENDAR_TOKEN_JSON_B64": base64.b64encode(calendar_token.read_bytes()).decode(), - } - prod = REPO_ROOT / ".env.production" + from agno.utils.encryption import is_encrypted - print() - if prod.exists(): - _upsert_env(prod, b64) - print("Wrote GMAIL_TOKEN_JSON_B64 + CALENDAR_TOKEN_JSON_B64 to .env.production.") - else: - print("Dev needs nothing more — the container reads the token files via the .:/app mount.") - print("base64 is only for a Railway deploy. Create .env.production and re-run, or add:") - for key, value in b64.items(): - print(f" {key}={value}") - - print("\nDone. Restart to pick up the tokens: docker compose up -d") - - # Offer the Railway sync when we wrote the deploy file and the script is there. - sync = REPO_ROOT / "scripts" / "railway" / "env-sync.sh" - if prod.exists() and sync.exists(): - try: - answer = input("\nPush to Railway now with ./scripts/railway/env-sync.sh? [y/N] ") - except EOFError: - answer = "" - if answer.strip().lower().startswith("y"): - return subprocess.run(["bash", str(sync)], cwd=str(REPO_ROOT)).returncode - print("Skipped. Run ./scripts/railway/env-sync.sh yourself when you're ready to deploy.") + token_data = row.get("token_data", {}) + if not is_encrypted(token_data): + print("WARNING: Token saved but NOT encrypted. Check encryption key.") + return 1 + + scopes = row.get("granted_scopes", []) + print(f"SUCCESS: Encrypted token saved to DB with {len(scopes)} scopes:") + for s in scopes: + print(f" - {s.split('/')[-1]}") return 0 diff --git a/scripts/railway/google_auth.py b/scripts/railway/google_auth.py new file mode 100755 index 0000000..7df40d1 --- /dev/null +++ b/scripts/railway/google_auth.py @@ -0,0 +1,193 @@ +#!/usr/bin/env python3 +""" +Mint Google OAuth tokens to Railway DB. + +Connects to Railway PostgreSQL via TCP proxy, runs OAuth flow, saves encrypted +tokens. Creates TCP proxy if needed via Railway GraphQL API. + +Usage: + python scripts/railway/google_auth.py # mint if no token exists + python scripts/railway/google_auth.py --force # re-mint (delete + mint) + +Prerequisites: + - Railway CLI installed and logged in (`railway login`) + - Project deployed (`./scripts/railway/up.sh`) + - .env with GOOGLE_CLIENT_ID, GOOGLE_CLIENT_SECRET, GOOGLE_TOKEN_ENCRYPTION_KEY +""" + +import argparse +import json +import os +import subprocess +import sys +from pathlib import Path + +import requests + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT)) + + +def run(cmd: str) -> str: + return subprocess.run(cmd, shell=True, capture_output=True, text=True).stdout.strip() + + +def gql(query: str, token: str) -> dict: + return requests.post( + "https://backboard.railway.com/graphql/v2", + json={"query": query}, + headers={"Authorization": f"Bearer {token}"}, + ).json() + + +def main() -> int: + parser = argparse.ArgumentParser(description="Mint Google OAuth tokens to Railway DB") + parser.add_argument("--force", action="store_true", help="Delete existing token and re-mint") + args = parser.parse_args() + + from dotenv import load_dotenv + + os.chdir(REPO_ROOT) + load_dotenv() + + # 1. Check prerequisites + if not os.getenv("GOOGLE_CLIENT_ID") or not os.getenv("GOOGLE_CLIENT_SECRET"): + print("ERROR: GOOGLE_CLIENT_ID and GOOGLE_CLIENT_SECRET required in .env") + return 1 + + if not os.getenv("GOOGLE_TOKEN_ENCRYPTION_KEY"): + print("ERROR: GOOGLE_TOKEN_ENCRYPTION_KEY not set.") + print(" Generate: python scripts/google_mint_tokens.py --generate-key") + return 1 + + # 2. Railway project + print("Connecting to Railway...\n") + + project_json = run("railway status --json") + if not project_json: + print("ERROR: No Railway project linked. Run: railway link") + return 1 + + project = json.loads(project_json).get("name", "") + print(f" Project: {project}") + + pgvars_json = run("railway variables --service pgvector --json") + if not pgvars_json: + print("ERROR: pgvector service not found. Run: ./scripts/railway/up.sh") + return 1 + + pgvars = json.loads(pgvars_json) + service_id = pgvars.get("RAILWAY_SERVICE_ID", "") + env_id = pgvars.get("RAILWAY_ENVIRONMENT_ID", "") + + # 3. Railway access token + config_path = Path.home() / ".railway" / "config.json" + with open(config_path) as f: + access_token = json.load(f).get("user", {}).get("accessToken", "") + + if not service_id or not access_token: + print("ERROR: Missing Railway credentials. Run: railway login") + return 1 + + # 4. TCP proxy (query or create) + query = f'{{ tcpProxies(serviceId: "{service_id}", environmentId: "{env_id}") {{ domain proxyPort applicationPort }} }}' + result = gql(query, access_token) + + proxy = None + for p in result.get("data", {}).get("tcpProxies", []): + if p.get("applicationPort") == 5432: + proxy = f"{p['domain']}:{p['proxyPort']}" + break + + if not proxy: + mutation = f'mutation {{ tcpProxyCreate(input: {{ serviceId: "{service_id}", environmentId: "{env_id}", applicationPort: 5432 }}) {{ domain proxyPort }} }}' + data = gql(mutation, access_token).get("data", {}).get("tcpProxyCreate", {}) + if not data.get("domain"): + print(f"ERROR: Failed to create TCP proxy") + return 1 + proxy = f"{data['domain']}:{data['proxyPort']}" + print(f" TCP: {proxy} (created)") + else: + print(f" TCP: {proxy}") + + # 5. Database connection + db_user = pgvars.get("POSTGRES_USER", "context") + db_pass = pgvars.get("POSTGRES_PASSWORD", "context") + db_name = pgvars.get("POSTGRES_DB", "context") + db_url = f"postgresql+psycopg://{db_user}:{db_pass}@{proxy}/{db_name}" + + from sqlalchemy import create_engine, text + + try: + engine = create_engine(db_url) + with engine.connect() as conn: + conn.execute(text("SELECT 1")) + print(" Database: connected\n") + except Exception as e: + print(f"ERROR: Database connection failed: {e}") + return 1 + + # 6. Auth config + os.environ["DATABASE_PUBLIC_URL"] = db_url + + from agents.providers.google import get_google_auth + + auth = get_google_auth() + if not auth or not auth.db: + print("ERROR: Failed to create auth config") + return 1 + + print(f"Scopes: {len(auth.scopes)} (Gmail + Calendar)") + print(f"DB: {auth.db.id}") + print("Encryption: enabled\n") + + # 7. Check existing token + row = auth.db.get_auth_token("google", None, "google") + if row and not args.force: + scopes = row.get("granted_scopes", []) + print(f"Token already exists with {len(scopes)} scopes:") + for s in scopes: + print(f" - {s.split('/')[-1]}") + print("\nUse --force to delete and re-mint.") + return 0 + + if row and args.force: + with engine.begin() as conn: + conn.execute(text("DELETE FROM ai.agno_auth_tokens WHERE provider = 'google'")) + print("Deleted existing token.\n") + + # 8. OAuth flow + from agno.tools.google.gmail import GmailTools + + print("Opening browser for OAuth consent...") + print("(Grant access to Gmail + Calendar)\n") + + gmail = GmailTools(auth=auth) + result = gmail.get_latest_emails(count=1) + + if "error" in result.lower(): + print(f"FAILED: {result}") + return 1 + + # 9. Verify + row = auth.db.get_auth_token("google", None, "google") + if not row: + print("FAILED: Token not saved to DB") + return 1 + + from agno.utils.encryption import is_encrypted + + if not is_encrypted(row.get("token_data", {})): + print("WARNING: Token saved but NOT encrypted.") + return 1 + + scopes = row.get("granted_scopes", []) + print(f"SUCCESS: Encrypted token saved to Railway DB with {len(scopes)} scopes:") + for s in scopes: + print(f" - {s.split('/')[-1]}") + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())