daniel/agent-claw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Multi-account Google support with user labels

Browse Source
This commit is contained in:
daniel
2026-05-09 11:21:37 -05:00
parent 01b258579b
commit 38d828ef74
4 changed files with 239 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/lambdas/agent-runner/handler.py
View File

@@ -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', {})),
},

50
src/lambdas/oauth-handler/handler.py
View File

@@ -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('<h1>Missing actor_id</h1>', 400)
label = params.get('label', 'primary')
client_id, _ = get_oauth_client()
redirect_uri = os.environ['OAUTH_REDIRECT_URI']
# Encode actor_id in state (base64 to keep URL-safe)
state = base64.urlsafe_b64encode(actor_id.encode()).decode().rstrip('=')
# Encode actor_id + label in state (JSON → base64)
state_data = json.dumps({'a': actor_id, 'l': label})
state = base64.urlsafe_b64encode(state_data.encode()).decode().rstrip('=')
auth_url = (
'https://accounts.google.com/o/oauth2/v2/auth?'
@@ -113,12 +114,20 @@ def handle_callback(params: dict) -> dict:
if not code or not state:
return _html('<h1>Missing code or state</h1>', 400)
# Decode actor_id from state
# Decode actor_id + label from state
try:
padding = 4 - len(state) % 4
actor_id = base64.urlsafe_b64decode(state + '=' * padding).decode()
state_data = json.loads(base64.urlsafe_b64decode(state + '=' * padding).decode())
actor_id = state_data['a']
label = state_data.get('l', 'primary')
except Exception:
return _html('<h1>Invalid state</h1>', 400)
# Backward compat: old state was just base64(actor_id)
try:
padding = 4 - len(state) % 4
actor_id = base64.urlsafe_b64decode(state + '=' * padding).decode()
label = 'primary'
except Exception:
return _html('<h1>Invalid state</h1>', 400)
client_id, client_secret = get_oauth_client()
redirect_uri = os.environ['OAUTH_REDIRECT_URI']
@@ -155,7 +164,6 @@ def handle_callback(params: dict) -> dict:
pass
if not user_email:
# Fallback: call userinfo endpoint
try:
access_token = tokens.get('access_token', '')
req2 = urllib.request.Request(
@@ -184,23 +192,24 @@ def handle_callback(params: dict) -> dict:
time.gmtime(time.time() + int(tokens['expires_in']))
)
# Store in Secrets Manager
secret_name = actor_id_to_secret_name(actor_id)
# Store in Secrets Manager at labeled path
secret_name = actor_id_to_secret_name(actor_id, label)
sm = get_sm()
try:
sm.create_secret(Name=secret_name, SecretString=json.dumps(creds))
except sm.exceptions.ResourceExistsException:
sm.put_secret_value(SecretId=secret_name, SecretString=json.dumps(creds))
print(f'[oauth] Stored credentials for actor={actor_id} email={user_email}')
print(f'[oauth] Stored credentials for actor={actor_id} label={label} email={user_email}')
# Update DynamoDB users table with google_email
# Update DynamoDB: merge into google_accounts map
table_name = os.environ.get('USERS_TABLE_NAME', '')
if table_name and actor_id:
try:
get_ddb().Table(table_name).update_item(
Key={'actor_id': actor_id},
UpdateExpression='SET google_email = :e',
ExpressionAttributeValues={':e': user_email},
UpdateExpression='SET google_accounts = if_not_exists(google_accounts, :empty), google_accounts.#label = :email',
ExpressionAttributeNames={'#label': label},
ExpressionAttributeValues={':email': user_email, ':empty': {}},
)
except Exception as e:
print(f'[oauth] DynamoDB update failed: {e}')
@@ -211,10 +220,7 @@ def handle_callback(params: dict) -> dict:
if bot_token_arn and actor_id.startswith('telegram:'):
chat_id = actor_id.split(':', 1)[1]
bot_token = get_sm().get_secret_value(SecretId=bot_token_arn)['SecretString']
tg_text = (
f'✅ Google account connected!\n\n'
f'{user_email} is now linked. You can now ask me about your Gmail, Calendar, and Drive.'
)
tg_text = f'✅ Connected {user_email} as "{label}"'
tg_payload = json.dumps({'chat_id': chat_id, 'text': tg_text}).encode()
tg_req = urllib.request.Request(
f'https://api.telegram.org/bot{bot_token}/sendMessage',
@@ -227,6 +233,6 @@ def handle_callback(params: dict) -> dict:
return _html(
f'<h1>✅ Google account connected!</h1>'
f'<p>Connected <b>{user_email}</b> to your agent account.</p>'
f'<p>Connected <b>{user_email}</b> as "<b>{label}</b>".</p>'
f'<p>You can close this window and return to Telegram.</p>'
)