multi-tenant Phase 2: per-user Google OAuth
- workspace-mcp: add proxy.py (port 8080) that reads X-Actor-Id header,
fetches per-user Google credentials from Secrets Manager, writes creds
file, sets USER_GOOGLE_EMAIL, proxies to workspace-mcp on port 8081
- workspace-mcp: update bootstrap to start workspace-mcp on 8081 + proxy on 8080
- workspace-mcp: update Dockerfile to include proxy.py
- oauth-handler Lambda: new Lambda with /oauth/start + /oauth/callback
routes; exchanges Google auth code, stores tokens in Secrets Manager
at agent-claw/google-credentials/{actor_id_safe}, updates DynamoDB
- CDK: add OAuthHandler Lambda + GET /oauth/start + /oauth/callback routes
- CDK: remove shared google-workspace-credentials secret; add per-user
secret IAM grants (agent-claw/google-credentials/*) for workspace-mcp
role, runtime1 role, and oauth-handler role
- CDK: output OAuthStartUrl + OAuthRedirectUri
- agent-runner: pass google_email in user_profile payload
- main.py: pass actor_id as X-Actor-Id header in workspace-mcp MCP calls;
skip workspace-mcp if user has no google_email; add connect_google_account
tool that generates OAuth URL for the current user
- main.py: include google_email in user_context for system prompt
- agentcore.json: add OAUTH_START_URL env var for agent runtime
This commit is contained in:
@@ -9,11 +9,13 @@ RUN pip install workspace-mcp==1.20.3 boto3 --quiet
|
||||
# Copy bootstrap and helper scripts
|
||||
COPY bootstrap /var/task/bootstrap
|
||||
COPY fetch_credentials.py /var/task/fetch_credentials.py
|
||||
COPY proxy.py /var/task/proxy.py
|
||||
RUN chmod +x /var/task/bootstrap
|
||||
|
||||
# Lambda Web Adapter config
|
||||
# Lambda Web Adapter config — proxy listens on 8080, workspace-mcp on 8081
|
||||
ENV AWS_LAMBDA_EXEC_WRAPPER=/opt/bootstrap
|
||||
ENV PORT=8080
|
||||
ENV PROXY_PORT=8080
|
||||
ENV READINESS_CHECK_PATH=/health
|
||||
|
||||
CMD ["/var/task/bootstrap"]
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
#!/bin/bash
|
||||
# Lambda Web Adapter bootstrap for workspace-mcp
|
||||
# Dependencies are in /opt/python (Lambda layer)
|
||||
# Lambda Web Adapter bootstrap for workspace-mcp (multi-tenant)
|
||||
# Proxy on port 8080 (Lambda Web Adapter entry), workspace-mcp on port 8081
|
||||
|
||||
set -e
|
||||
|
||||
@@ -14,14 +14,27 @@ export HOME=/tmp
|
||||
export WORKSPACE_MCP_LOG_DIR=/tmp
|
||||
export GOOGLE_WORKSPACE_MCP_CREDENTIALS_DIR=/tmp/workspace_mcp_credentials
|
||||
|
||||
echo "[workspace-mcp] Fetching Google credentials..." >&2
|
||||
$PYTHON /var/task/fetch_credentials.py
|
||||
echo "[workspace-mcp] Fetching default Google credentials..." >&2
|
||||
$PYTHON /var/task/fetch_credentials.py || true # non-fatal: per-user creds loaded by proxy
|
||||
|
||||
if [ -f /tmp/workspace-mcp-env ]; then
|
||||
source /tmp/workspace-mcp-env
|
||||
fi
|
||||
|
||||
echo "[workspace-mcp] Starting on port $PORT..." >&2
|
||||
exec $PYTHON /opt/python/bin/workspace-mcp \
|
||||
echo "[workspace-mcp] Starting workspace-mcp on port 8081..." >&2
|
||||
PORT=8081 $PYTHON /opt/python/bin/workspace-mcp \
|
||||
--transport streamable-http \
|
||||
--tool-tier extended
|
||||
--tool-tier extended &
|
||||
WMCP_PID=$!
|
||||
|
||||
# Wait for workspace-mcp to be ready
|
||||
for i in $(seq 1 30); do
|
||||
if curl -sf http://127.0.0.1:8081/health > /dev/null 2>&1; then
|
||||
echo "[workspace-mcp] Ready on port 8081" >&2
|
||||
break
|
||||
fi
|
||||
sleep 0.5
|
||||
done
|
||||
|
||||
echo "[workspace-mcp] Starting credential proxy on port $PORT..." >&2
|
||||
exec $PYTHON /var/task/proxy.py
|
||||
|
||||
143
src/lambdas/workspace-mcp/proxy.py
Normal file
143
src/lambdas/workspace-mcp/proxy.py
Normal file
@@ -0,0 +1,143 @@
|
||||
"""
|
||||
Multi-tenant credential proxy for workspace-mcp.
|
||||
|
||||
Sits on port 8080 (Lambda Web Adapter entry point).
|
||||
Reads X-Actor-Id header, fetches per-user Google credentials from Secrets Manager,
|
||||
writes them to /tmp/workspace_mcp_credentials/{email}.json, sets USER_GOOGLE_EMAIL,
|
||||
then proxies the request to workspace-mcp on port 8081.
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
import urllib.request
|
||||
import urllib.error
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
|
||||
import boto3
|
||||
|
||||
_sm = None
|
||||
_sm_lock = threading.Lock()
|
||||
|
||||
# Cache: actor_id -> (email, creds_json, fetched_at)
|
||||
_creds_cache: dict = {}
|
||||
_cache_ttl = 300 # 5 minutes
|
||||
|
||||
|
||||
def _get_sm():
|
||||
global _sm
|
||||
if _sm is None:
|
||||
with _sm_lock:
|
||||
if _sm is None:
|
||||
_sm = boto3.client('secretsmanager', region_name=os.environ.get('AWS_REGION', 'us-east-1'))
|
||||
return _sm
|
||||
|
||||
|
||||
def _actor_id_to_secret_name(actor_id: str) -> str:
|
||||
"""Convert actor_id to a valid Secrets Manager secret name."""
|
||||
# telegram:123456789 -> agent-claw/google-credentials/telegram-123456789
|
||||
safe = actor_id.replace(':', '-').replace('/', '-')
|
||||
return f'agent-claw/google-credentials/{safe}'
|
||||
|
||||
|
||||
def _fetch_credentials(actor_id: str) -> tuple[str, dict] | None:
|
||||
"""Fetch Google credentials for actor_id from Secrets Manager. Returns (email, creds_dict) or None."""
|
||||
now = time.time()
|
||||
cached = _creds_cache.get(actor_id)
|
||||
if cached and now - cached[2] < _cache_ttl:
|
||||
return cached[0], cached[1]
|
||||
|
||||
secret_name = _actor_id_to_secret_name(actor_id)
|
||||
try:
|
||||
secret = _get_sm().get_secret_value(SecretId=secret_name)['SecretString']
|
||||
creds = json.loads(secret)
|
||||
email = creds.get('client_email') or creds.get('email') or creds.get('user_email', '')
|
||||
_creds_cache[actor_id] = (email, creds, now)
|
||||
print(f'[proxy] Loaded credentials for actor={actor_id} email={email}', flush=True)
|
||||
return email, creds
|
||||
except Exception as e:
|
||||
print(f'[proxy] No credentials for actor={actor_id}: {e}', flush=True)
|
||||
return None
|
||||
|
||||
|
||||
def _write_credentials_file(email: str, creds: dict) -> str:
|
||||
"""Write credentials to /tmp/workspace_mcp_credentials/{email}.json. Returns path."""
|
||||
creds_dir = '/tmp/workspace_mcp_credentials'
|
||||
os.makedirs(creds_dir, exist_ok=True)
|
||||
path = f'{creds_dir}/{email}.json'
|
||||
with open(path, 'w') as f:
|
||||
json.dump(creds, f)
|
||||
return path
|
||||
|
||||
|
||||
class ProxyHandler(BaseHTTPRequestHandler):
|
||||
def log_message(self, format, *args):
|
||||
pass # suppress default access log
|
||||
|
||||
def _proxy(self):
|
||||
actor_id = self.headers.get('x-actor-id', '')
|
||||
|
||||
if actor_id:
|
||||
result = _fetch_credentials(actor_id)
|
||||
if result:
|
||||
email, creds = result
|
||||
_write_credentials_file(email, creds)
|
||||
os.environ['USER_GOOGLE_EMAIL'] = email
|
||||
os.environ['GOOGLE_WORKSPACE_MCP_CREDENTIALS_DIR'] = '/tmp/workspace_mcp_credentials'
|
||||
else:
|
||||
# No credentials found — proceed without setting email (workspace-mcp will use default or fail)
|
||||
print(f'[proxy] No Google credentials for actor={actor_id}, proceeding without', flush=True)
|
||||
|
||||
# Read request body
|
||||
content_length = int(self.headers.get('Content-Length', 0))
|
||||
body = self.rfile.read(content_length) if content_length > 0 else b''
|
||||
|
||||
# Build upstream request
|
||||
upstream_url = f'http://127.0.0.1:8081{self.path}'
|
||||
upstream_headers = {k: v for k, v in self.headers.items()
|
||||
if k.lower() not in ('host', 'content-length')}
|
||||
if body:
|
||||
upstream_headers['Content-Length'] = str(len(body))
|
||||
|
||||
req = urllib.request.Request(
|
||||
upstream_url,
|
||||
data=body or None,
|
||||
headers=upstream_headers,
|
||||
method=self.command,
|
||||
)
|
||||
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=60) as resp:
|
||||
self.send_response(resp.status)
|
||||
for k, v in resp.headers.items():
|
||||
if k.lower() not in ('transfer-encoding',):
|
||||
self.send_header(k, v)
|
||||
self.end_headers()
|
||||
self.wfile.write(resp.read())
|
||||
except urllib.error.HTTPError as e:
|
||||
self.send_response(e.code)
|
||||
for k, v in e.headers.items():
|
||||
if k.lower() not in ('transfer-encoding',):
|
||||
self.send_header(k, v)
|
||||
self.end_headers()
|
||||
self.wfile.write(e.read())
|
||||
except Exception as e:
|
||||
print(f'[proxy] Upstream error: {e}', flush=True)
|
||||
self.send_response(502)
|
||||
self.end_headers()
|
||||
self.wfile.write(b'Bad Gateway')
|
||||
|
||||
do_GET = _proxy
|
||||
do_POST = _proxy
|
||||
do_PUT = _proxy
|
||||
do_DELETE = _proxy
|
||||
do_OPTIONS = _proxy
|
||||
do_HEAD = _proxy
|
||||
do_PATCH = _proxy
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
port = int(os.environ.get('PROXY_PORT', 8080))
|
||||
server = HTTPServer(('0.0.0.0', port), ProxyHandler)
|
||||
print(f'[proxy] Listening on port {port}', flush=True)
|
||||
server.serve_forever()
|
||||
Reference in New Issue
Block a user