diff --git a/agentclaw/app/agent_claw_main/main.py b/agentclaw/app/agent_claw_main/main.py index 2320094..72f4481 100644 --- a/agentclaw/app/agent_claw_main/main.py +++ b/agentclaw/app/agent_claw_main/main.py @@ -94,8 +94,8 @@ def write_workspace_file(path: str, content: str) -> str: @tool -def connect_google_account() -> str: - """Generate a Google OAuth authorization URL for the current user to connect their Google account. +def connect_google_account(label: str = 'primary') -> str: + """Connect a Google account with a custom label (e.g. 'work', 'personal'). Defaults to 'primary'. Use this when the user wants to connect Google Workspace (Gmail, Calendar, Drive, etc.) or when Google tools fail due to missing credentials.""" if not OAUTH_START_URL: @@ -103,8 +103,18 @@ def connect_google_account() -> str: actor_id = _current_actor_id if not actor_id: return 'Cannot determine actor_id for OAuth flow.' - url = f'{OAUTH_START_URL}?actor_id={actor_id}' - return f'Please open this URL to connect your Google account:\n{url}\n\nAfter authorizing, Google Workspace tools (Gmail, Calendar, Drive) will be available.' + url = f'{OAUTH_START_URL}?actor_id={actor_id}&label={label}' + return f'Please open this URL to connect your Google account as "{label}":\n{url}\n\nAfter authorizing, Google Workspace tools (Gmail, Calendar, Drive) will be available.' + + +@tool +def list_google_accounts() -> str: + """List all connected Google accounts and their labels.""" + accounts = _gws._current_google_accounts + if not accounts: + return 'No Google accounts connected. Use connect_google_account to add one.' + parts = [f'{label} ({email})' for label, email in accounts.items()] + return 'Connected Google accounts: ' + ', '.join(parts) @tool @@ -145,11 +155,9 @@ def manage_service(action: str, service: str, config: dict | None = None) -> str return 'service name is required.' if not config: return 'config dict is required for enroll.' - # Validate known services if service == 'home_assistant': if 'url' not in config or 'token' not in config: return 'home_assistant config requires "url" and "token" keys.' - # Update in-memory config immediately for this session set_ha_config(config['url'], config['token']) table.update_item( Key={'actor_id': actor_id}, @@ -245,17 +253,21 @@ async def main(payload: dict, context): ha_cfg = services.get('home_assistant', {}) set_ha_config(ha_cfg.get('url', ''), ha_cfg.get('token', '')) + # Sync google_accounts to google_workspace module + google_accounts = user_profile.get('google_accounts', {}) + _gws._current_google_accounts = google_accounts + # Build system prompt — base cached, user context injected per-invocation user_context = '' if user_profile: name = user_profile.get('display_name', '') username = user_profile.get('telegram_username', '') - google_email = user_profile.get('google_email', '') user_context = f'Name: {name}' if username: user_context += f'\nTelegram username: @{username}' - if google_email: - user_context += f'\nGoogle account: {google_email}' + if google_accounts: + acct_list = ', '.join(f'{label} ({email})' for label, email in google_accounts.items()) + user_context += f'\nGoogle accounts: {acct_list}' else: user_context += '\nGoogle account: not connected (use connect_google_account tool to connect)' enrolled = list(services.keys()) @@ -282,7 +294,7 @@ async def main(payload: dict, context): ) base_tools = [web_search, web_fetch, read_workspace_file, write_workspace_file, - home_assistant, connect_google_account, + home_assistant, connect_google_account, list_google_accounts, manage_service, schedule_reminder, list_reminders, cancel_reminder, list_calendars, get_calendar_events, list_gmail_messages, get_gmail_message] diff --git a/agentclaw/app/agent_claw_main/tools/google_workspace.py b/agentclaw/app/agent_claw_main/tools/google_workspace.py index 6daf721..00e9a1b 100644 --- a/agentclaw/app/agent_claw_main/tools/google_workspace.py +++ b/agentclaw/app/agent_claw_main/tools/google_workspace.py @@ -4,21 +4,26 @@ Mirrors workspace-mcp (gcalendar/gmail) logic using google-api-python-client dir since workspace-mcp tool functions require FastMCP request context and cannot be called outside an MCP server. -Credential secret: agent-claw/google-credentials/{actor_id.replace(':', '-')} -Contains: token, refresh_token, token_uri, client_id, client_secret, scopes +Credential secrets: agent-claw/google-credentials/{safe_actor_id}/{label} +Backward compat: agent-claw/google-credentials/{safe_actor_id} (treated as "primary") """ import json +import time import traceback import boto3 -import httplib2 from strands import tool from google.oauth2.credentials import Credentials from google.auth.transport.requests import Request from googleapiclient.discovery import build from datetime import datetime, timezone, timedelta -_HTTP_TIMEOUT = 15 _sm = None +# Cache: actor_id -> (timestamp, {label: Credentials}) +_creds_cache: dict[str, tuple[float, dict[str, Credentials]]] = {} + +# Set per-invocation by main.py +_current_actor_id: str = '' +_current_google_accounts: dict = {} # {label: email} from DynamoDB def _secrets(): @@ -28,24 +33,24 @@ def _secrets(): return _sm -def _get_creds(actor_id: str) -> Credentials: - from datetime import datetime - secret_name = 'agent-claw/google-credentials/' + actor_id.replace(':', '-') - print(f'[google] fetching creds actor={actor_id}') - data = json.loads(_secrets().get_secret_value(SecretId=secret_name)['SecretString']) +def _actor_id(): + return _current_actor_id + + +def _load_creds_from_secret(secret_name: str) -> Credentials: + """Load, optionally refresh, and return Credentials from a named secret.""" + sm = _secrets() + data = json.loads(sm.get_secret_value(SecretId=secret_name)['SecretString']) expiry_str = data.get('expiry') + expiry = None if expiry_str: exp_aware = datetime.fromisoformat(expiry_str.replace('Z', '+00:00')) - expiry = exp_aware.replace(tzinfo=None) # google-auth uses naive UTC datetimes - else: - expiry = None + expiry = exp_aware.replace(tzinfo=None) stored_scopes = data.get('scopes', []) api_scopes = [s for s in stored_scopes if s.startswith('https://')] if stored_scopes else None - # Fix stored scopes if they contain OIDC scopes if stored_scopes and any(s in stored_scopes for s in ['openid', 'email', 'profile']): data['scopes'] = api_scopes - _secrets().put_secret_value(SecretId=secret_name, SecretString=json.dumps(data)) - print('[google] fixed stored scopes: removed OIDC scopes') + sm.put_secret_value(SecretId=secret_name, SecretString=json.dumps(data)) creds = Credentials( token=data.get('token'), refresh_token=data.get('refresh_token'), @@ -55,51 +60,96 @@ def _get_creds(actor_id: str) -> Credentials: scopes=api_scopes, expiry=expiry, ) - print(f'[google] creds loaded, expired={creds.expired}') if (creds.expired or not creds.valid) and creds.refresh_token: - print('[google] refreshing token') creds.refresh(Request()) data['token'] = creds.token if creds.expiry: data['expiry'] = creds.expiry.isoformat() - _secrets().put_secret_value(SecretId=secret_name, SecretString=json.dumps(data)) - print('[google] token refreshed and saved') + sm.put_secret_value(SecretId=secret_name, SecretString=json.dumps(data)) return creds -def _actor_id(): - # Read from module-level var set by main.py per invocation - # DO NOT use 'import main as _main' — it re-runs main.py including app.run() which hangs - return _current_actor_id +def _load_all_creds(actor_id: str) -> dict[str, Credentials]: + """Load all labeled credentials for actor_id, with 5-min TTL cache.""" + now = time.time() + if actor_id in _creds_cache: + ts, cached = _creds_cache[actor_id] + if now - ts < 300: + return cached + safe = actor_id.replace(':', '-').replace('/', '-') + prefix = f'agent-claw/google-credentials/{safe}/' + sm = _secrets() + result: dict[str, Credentials] = {} -# Set per-invocation by main.py before any tool call -_current_actor_id: str = '' + try: + paginator = sm.get_paginator('list_secrets') + for page in paginator.paginate(Filters=[{'Key': 'name', 'Values': [prefix]}]): + for secret in page.get('SecretList', []): + name = secret['Name'] + label = name[len(prefix):] + if not label or '/' in label: + continue + try: + result[label] = _load_creds_from_secret(name) + print(f'[google] loaded creds actor={actor_id} label={label}') + except Exception as e: + print(f'[google] failed to load label={label}: {e}') + except Exception as e: + print(f'[google] list_secrets failed: {e}') + + # Backward compat: flat secret path + if not result: + flat = f'agent-claw/google-credentials/{safe}' + try: + result['primary'] = _load_creds_from_secret(flat) + print(f'[google] loaded creds from flat path actor={actor_id}') + except Exception: + pass + + _creds_cache[actor_id] = (now, result) + return result def _svc(api: str, version: str, creds: Credentials): - # Standard google-auth pattern: pass credentials= directly to build(). - # google.oauth2.credentials.Credentials is natively supported by googleapiclient. - # (creds.authorize() is oauth2client only — not available here) return build(api, version, credentials=creds, cache_discovery=False) +def _get_creds_for_label(all_creds: dict[str, Credentials], label: str | None): + """Return {label: creds} filtered by label, or all if label is None.""" + if label: + if label not in all_creds: + return {} + return {label: all_creds[label]} + return all_creds + + @tool -def list_calendars() -> str: - """List all Google Calendars for the current user.""" +def list_calendars(account_label: str = None) -> str: + """List all Google Calendars for the current user. + + Args: + account_label: Optional account label (e.g. 'work', 'personal'). Lists all accounts if omitted. + """ try: - creds = _get_creds(_actor_id()) - result = _svc('calendar', 'v3', creds).calendarList().list().execute() - items = result.get('items', []) - if not items: - return 'No calendars found.' - return '\n'.join( - f'- "{c.get("summary", "")}"{" (Primary)" if c.get("primary") else ""} (ID: {c["id"]})' - for c in items - ) + all_creds = _load_all_creds(_actor_id()) + if not all_creds: + return 'No Google accounts connected. Use connect_google_account to add one.' + creds_map = _get_creds_for_label(all_creds, account_label) + if not creds_map: + return f'No account with label "{account_label}" found.' + multi = len(creds_map) > 1 + parts = [] + for label, creds in creds_map.items(): + items = _svc('calendar', 'v3', creds).calendarList().list().execute().get('items', []) + lines = [ + f'{"[" + label + "] " if multi else ""}- "{c.get("summary", "")}"{" (Primary)" if c.get("primary") else ""} (ID: {c["id"]})' + for c in items + ] + parts.append('\n'.join(lines) if lines else f'{"[" + label + "] " if multi else ""}No calendars found.') + return '\n'.join(parts) except Exception as e: - tb = traceback.format_exc() - print(f'[google] list_calendars error: {e}\n{tb}') + print(f'[google] list_calendars error: {e}\n{traceback.format_exc()}') return f'Error listing calendars: {e}' @@ -111,6 +161,7 @@ def get_calendar_events( time_max: str = '', max_results: int = 25, query: str = '', + account_label: str = None, ) -> str: """Get upcoming Google Calendar events. @@ -121,10 +172,16 @@ def get_calendar_events( time_max: End of time range in RFC3339 format (optional) max_results: Maximum events to return (default: 25) query: Keyword search within event fields (optional) + account_label: Optional account label (e.g. 'work', 'personal'). Queries all accounts if omitted. """ try: - creds = _get_creds(_actor_id()) - svc = _svc('calendar', 'v3', creds) + all_creds = _load_all_creds(_actor_id()) + if not all_creds: + return 'No Google accounts connected.' + creds_map = _get_creds_for_label(all_creds, account_label) + if not creds_map: + return f'No account with label "{account_label}" found.' + multi = len(creds_map) > 1 now = datetime.now(timezone.utc) params = { 'calendarId': calendar_id, @@ -136,92 +193,118 @@ def get_calendar_events( } if query: params['q'] = query - events = svc.events().list(**params).execute().get('items', []) - if not events: - return f'No events found in calendar "{calendar_id}".' - lines = [] - for e in events: - start = e['start'].get('dateTime', e['start'].get('date', '')) - end = e['end'].get('dateTime', e['end'].get('date', '')) - eid = e.get('id', '') - lines.append(f'- "{e.get("summary", "No Title")}" (Starts: {start}, Ends: {end}) ID: {eid}') - return f'Retrieved {len(events)} events from "{calendar_id}":\n' + '\n'.join(lines) + parts = [] + for label, creds in creds_map.items(): + events = _svc('calendar', 'v3', creds).events().list(**params).execute().get('items', []) + if not events: + parts.append(f'{"[" + label + "] " if multi else ""}No events found in calendar "{calendar_id}".') + continue + lines = [] + for e in events: + start = e['start'].get('dateTime', e['start'].get('date', '')) + end = e['end'].get('dateTime', e['end'].get('date', '')) + prefix = f'[{label}] ' if multi else '' + lines.append(f'{prefix}- "{e.get("summary", "No Title")}" (Starts: {start}, Ends: {end}) ID: {e.get("id", "")}') + parts.append(f'Retrieved {len(events)} events{" [" + label + "]" if multi else ""} from "{calendar_id}":\n' + '\n'.join(lines)) + return '\n\n'.join(parts) except Exception as e: - tb = traceback.format_exc() - print(f'[google] get_calendar_events error: {e}\n{tb}') + print(f'[google] get_calendar_events error: {e}\n{traceback.format_exc()}') return f'Error fetching calendar events: {e}' @tool -def list_gmail_messages(max_results: int = 10, query: str = 'in:inbox') -> str: +def list_gmail_messages(max_results: int = 10, query: str = 'in:inbox', account_label: str = None) -> str: """List Gmail messages. Args: max_results: Maximum number of messages to return (default: 10) query: Gmail search query (default: 'in:inbox') + account_label: Optional account label (e.g. 'work', 'personal'). Lists all accounts if omitted. """ try: - creds = _get_creds(_actor_id()) - svc = _svc('gmail', 'v1', creds) - result = svc.users().messages().list(userId='me', q=query, maxResults=max_results).execute() - messages = result.get('messages', []) - if not messages: - return 'No messages found.' - lines = [] - for m in messages: - msg = svc.users().messages().get( - userId='me', id=m['id'], format='metadata', - metadataHeaders=['Subject', 'From', 'Date'] - ).execute() - h = {hdr['name']: hdr['value'] for hdr in msg.get('payload', {}).get('headers', [])} - lines.append(f"id={m['id']} | {h.get('Date', '')} | From: {h.get('From', '')} | {h.get('Subject', '(no subject)')}") - next_token = result.get('nextPageToken') - out = '\n'.join(lines) - if next_token: - out += f'\n(more results available)' - return out + all_creds = _load_all_creds(_actor_id()) + if not all_creds: + return 'No Google accounts connected.' + creds_map = _get_creds_for_label(all_creds, account_label) + if not creds_map: + return f'No account with label "{account_label}" found.' + multi = len(creds_map) > 1 + parts = [] + for label, creds in creds_map.items(): + svc = _svc('gmail', 'v1', creds) + result = svc.users().messages().list(userId='me', q=query, maxResults=max_results).execute() + messages = result.get('messages', []) + if not messages: + parts.append(f'{"[" + label + "] " if multi else ""}No messages found.') + continue + lines = [] + for m in messages: + msg = svc.users().messages().get( + userId='me', id=m['id'], format='metadata', + metadataHeaders=['Subject', 'From', 'Date'] + ).execute() + h = {hdr['name']: hdr['value'] for hdr in msg.get('payload', {}).get('headers', [])} + prefix = f'[{label}] ' if multi else '' + lines.append(f"{prefix}id={m['id']} | {h.get('Date', '')} | From: {h.get('From', '')} | {h.get('Subject', '(no subject)')}") + if result.get('nextPageToken'): + lines.append(f'{"[" + label + "] " if multi else ""}(more results available)') + parts.append('\n'.join(lines)) + return '\n'.join(parts) except Exception as e: - tb = traceback.format_exc() - print(f'[google] list_gmail_messages error: {e}\n{tb}') + print(f'[google] list_gmail_messages error: {e}\n{traceback.format_exc()}') return f'Error listing Gmail messages: {e}' @tool -def get_gmail_message(message_id: str, body_format: str = 'text') -> str: +def get_gmail_message(message_id: str, body_format: str = 'text', account_label: str = None) -> str: """Get the full content of a Gmail message by ID. Args: message_id: The Gmail message ID body_format: 'text' (default), 'html', or 'raw' + account_label: Optional account label. Tries all accounts if omitted. """ try: - creds = _get_creds(_actor_id()) - svc = _svc('gmail', 'v1', creds) - meta = svc.users().messages().get( - userId='me', id=message_id, format='metadata', - metadataHeaders=['Subject', 'From', 'To', 'Cc', 'Date'] - ).execute() - h = {hdr['name']: hdr['value'] for hdr in meta.get('payload', {}).get('headers', [])} - - if body_format == 'raw': - import base64 - raw = svc.users().messages().get(userId='me', id=message_id, format='raw').execute() - body = base64.urlsafe_b64decode(raw.get('raw', '') + '==').decode('utf-8', errors='replace') - else: - full = svc.users().messages().get(userId='me', id=message_id, format='full').execute() - body = _extract_body(full.get('payload', {}), prefer_html=(body_format == 'html')) - - lines = [ - f"From: {h.get('From', '')}", - f"To: {h.get('To', '')}", - f"Date: {h.get('Date', '')}", - f"Subject: {h.get('Subject', '')}", - '', - body, - ] - if h.get('Cc'): - lines.insert(3, f"Cc: {h['Cc']}") - return '\n'.join(lines) + all_creds = _load_all_creds(_actor_id()) + if not all_creds: + return 'No Google accounts connected.' + creds_map = _get_creds_for_label(all_creds, account_label) + if not creds_map: + return f'No account with label "{account_label}" found.' + multi = len(creds_map) > 1 + for label, creds in creds_map.items(): + try: + svc = _svc('gmail', 'v1', creds) + meta = svc.users().messages().get( + userId='me', id=message_id, format='metadata', + metadataHeaders=['Subject', 'From', 'To', 'Cc', 'Date'] + ).execute() + h = {hdr['name']: hdr['value'] for hdr in meta.get('payload', {}).get('headers', [])} + if body_format == 'raw': + import base64 + raw = svc.users().messages().get(userId='me', id=message_id, format='raw').execute() + body = base64.urlsafe_b64decode(raw.get('raw', '') + '==').decode('utf-8', errors='replace') + else: + full = svc.users().messages().get(userId='me', id=message_id, format='full').execute() + body = _extract_body(full.get('payload', {}), prefer_html=(body_format == 'html')) + prefix = f'[{label}]\n' if multi else '' + lines = [ + f"{prefix}From: {h.get('From', '')}", + f"To: {h.get('To', '')}", + f"Date: {h.get('Date', '')}", + f"Subject: {h.get('Subject', '')}", + '', + body, + ] + if h.get('Cc'): + lines.insert(3, f"Cc: {h['Cc']}") + return '\n'.join(lines) + except Exception as e: + if multi: + print(f'[google] get_gmail_message label={label} not found: {e}') + continue + return f'Error fetching Gmail message: {e}' + return f'Message {message_id} not found in any connected account.' except Exception as e: return f'Error fetching Gmail message: {e}' @@ -237,17 +320,14 @@ def _extract_body(payload: dict, prefer_html: bool = False) -> str: return base64.urlsafe_b64decode(data + '==').decode('utf-8', errors='replace') if data else '' parts = payload.get('parts', []) - # First pass: preferred type for part in parts: if part.get('mimeType') == target: data = part.get('body', {}).get('data', '') return base64.urlsafe_b64decode(data + '==').decode('utf-8', errors='replace') if data else '' - # Second pass: fallback type for part in parts: if part.get('mimeType') == fallback: data = part.get('body', {}).get('data', '') return base64.urlsafe_b64decode(data + '==').decode('utf-8', errors='replace') if data else '' - # Recurse into multipart for part in parts: text = _extract_body(part, prefer_html) if text: diff --git a/src/lambdas/agent-runner/handler.py b/src/lambdas/agent-runner/handler.py index 4a82d1f..7ba13fc 100644 --- a/src/lambdas/agent-runner/handler.py +++ b/src/lambdas/agent-runner/handler.py @@ -185,7 +185,7 @@ def handler(event, context): 'user_profile': { 'display_name': user_profile.get('display_name', actor_id), 'telegram_username': user_profile.get('telegram_username', ''), - 'google_email': user_profile.get('google_email', ''), + 'google_accounts': user_profile.get('google_accounts', {'primary': user_profile['google_email']} if user_profile.get('google_email') else {}), 'allowed': user_profile.get('allowed', True), 'services': user_profile.get('enrolled_services', user_profile.get('services', {})), }, diff --git a/src/lambdas/oauth-handler/handler.py b/src/lambdas/oauth-handler/handler.py index df24ce4..3012b68 100644 --- a/src/lambdas/oauth-handler/handler.py +++ b/src/lambdas/oauth-handler/handler.py @@ -2,12 +2,10 @@ Google OAuth handler Lambda. Routes: - GET /oauth/start?actor_id=telegram:123 → redirect to Google OAuth consent + GET /oauth/start?actor_id=telegram:123&label=work → redirect to Google OAuth consent GET /oauth/callback?code=...&state=... → exchange code, store tokens, update DynamoDB """ import base64 -import hashlib -import hmac import json import os import time @@ -52,9 +50,9 @@ def get_oauth_client() -> tuple[str, str]: return secret['client_id'], secret['client_secret'] -def actor_id_to_secret_name(actor_id: str) -> str: +def actor_id_to_secret_name(actor_id: str, label: str = 'primary') -> str: safe = actor_id.replace(':', '-').replace('/', '-') - return f'agent-claw/google-credentials/{safe}' + return f'agent-claw/google-credentials/{safe}/{label}' def _redirect(url: str) -> dict: @@ -82,11 +80,14 @@ def handle_start(params: dict) -> dict: if not actor_id: return _html('
Connected {user_email} to your agent account.
' + f'Connected {user_email} as "{label}".
' f'You can close this window and return to Telegram.
' )