336 lines
14 KiB
Python
336 lines
14 KiB
Python
"""Google Calendar and Gmail tools — credentials injected from Secrets Manager per call.
|
|
|
|
Mirrors workspace-mcp (gcalendar/gmail) logic using google-api-python-client directly,
|
|
since workspace-mcp tool functions require FastMCP request context and cannot be called
|
|
outside an MCP server.
|
|
|
|
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
|
|
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
|
|
|
|
_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():
|
|
global _sm
|
|
if _sm is None:
|
|
_sm = boto3.client('secretsmanager', region_name='us-east-1')
|
|
return _sm
|
|
|
|
|
|
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)
|
|
stored_scopes = data.get('scopes', [])
|
|
api_scopes = [s for s in stored_scopes if s.startswith('https://')] if stored_scopes else None
|
|
if stored_scopes and any(s in stored_scopes for s in ['openid', 'email', 'profile']):
|
|
data['scopes'] = api_scopes
|
|
sm.put_secret_value(SecretId=secret_name, SecretString=json.dumps(data))
|
|
creds = Credentials(
|
|
token=data.get('token'),
|
|
refresh_token=data.get('refresh_token'),
|
|
token_uri=data.get('token_uri', 'https://oauth2.googleapis.com/token'),
|
|
client_id=data.get('client_id'),
|
|
client_secret=data.get('client_secret'),
|
|
scopes=api_scopes,
|
|
expiry=expiry,
|
|
)
|
|
if (creds.expired or not creds.valid) and creds.refresh_token:
|
|
creds.refresh(Request())
|
|
data['token'] = creds.token
|
|
if creds.expiry:
|
|
data['expiry'] = creds.expiry.isoformat()
|
|
sm.put_secret_value(SecretId=secret_name, SecretString=json.dumps(data))
|
|
return creds
|
|
|
|
|
|
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] = {}
|
|
|
|
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}')
|
|
|
|
# Always include flat secret as 'primary' if not already in result
|
|
if 'primary' not in 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):
|
|
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(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:
|
|
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:
|
|
print(f'[google] list_calendars error: {e}\n{traceback.format_exc()}')
|
|
return f'Error listing calendars: {e}'
|
|
|
|
|
|
@tool
|
|
def get_calendar_events(
|
|
calendar_id: str = 'primary',
|
|
days_ahead: int = 7,
|
|
time_min: str = '',
|
|
time_max: str = '',
|
|
max_results: int = 25,
|
|
query: str = '',
|
|
account_label: str = None,
|
|
) -> str:
|
|
"""Get upcoming Google Calendar events.
|
|
|
|
Args:
|
|
calendar_id: Calendar ID (default: 'primary')
|
|
days_ahead: Days ahead to fetch when time_min/time_max not specified (default: 7)
|
|
time_min: Start of time range in RFC3339 format (optional)
|
|
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:
|
|
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,
|
|
'timeMin': time_min or now.isoformat().replace('+00:00', 'Z'),
|
|
'timeMax': time_max or (now + timedelta(days=days_ahead)).isoformat().replace('+00:00', 'Z'),
|
|
'maxResults': max_results,
|
|
'singleEvents': True,
|
|
'orderBy': 'startTime',
|
|
}
|
|
if query:
|
|
params['q'] = query
|
|
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:
|
|
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', 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:
|
|
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:
|
|
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', 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:
|
|
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}'
|
|
|
|
|
|
def _extract_body(payload: dict, prefer_html: bool = False) -> str:
|
|
import base64
|
|
mime = payload.get('mimeType', '')
|
|
target = 'text/html' if prefer_html else 'text/plain'
|
|
fallback = 'text/plain' if prefer_html else 'text/html'
|
|
|
|
if mime == target:
|
|
data = payload.get('body', {}).get('data', '')
|
|
return base64.urlsafe_b64decode(data + '==').decode('utf-8', errors='replace') if data else ''
|
|
|
|
parts = payload.get('parts', [])
|
|
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 ''
|
|
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 ''
|
|
for part in parts:
|
|
text = _extract_body(part, prefer_html)
|
|
if text:
|
|
return text
|
|
return ''
|