agent-claw/cdk/cdk.out/asset.c0db2a060be885d61722dfe6fbd3967e1c956826682078f6338123bf0c797e5b/handler.py

import json
import os
import time
import uuid
import boto3
import urllib.request
from typing import Any

# AWS clients
_ddb = None
_agentcore = None


def get_ddb():
    global _ddb
    if _ddb is None:
        _ddb = boto3.resource('dynamodb')
    return _ddb


def get_agentcore():
    global _agentcore
    if _agentcore is None:
        _agentcore = boto3.client('bedrock-agentcore', region_name='us-east-1')
    return _agentcore


def get_or_create_user(actor_id: str, from_info: dict) -> dict:
    """Look up user in registry, auto-registering on first contact."""
    table_name = os.environ.get('USERS_TABLE_NAME', '')
    if not table_name:
        return {'actor_id': actor_id, 'display_name': from_info.get('from_name', actor_id)}
    table = get_ddb().Table(table_name)
    response = table.get_item(Key={'actor_id': actor_id})
    item = response.get('Item')
    if item:
        return item
    now = int(time.time())
    item = {
        'actor_id': actor_id,
        'display_name': from_info.get('from_name') or actor_id,
        'telegram_username': from_info.get('from_username', ''),
        'created_at': str(now),
        'status': 'pending',
        'services': {},
    }
    table.put_item(Item=item)
    print(f'[agent-runner] Registered new user (pending): {actor_id}')
    return item


def update_user_status(actor_id: str, name: str, status: str) -> None:
    table_name = os.environ.get('USERS_TABLE_NAME', '')
    if not table_name:
        return
    table = get_ddb().Table(table_name)
    table.update_item(
        Key={'actor_id': actor_id},
        UpdateExpression='SET display_name = :n, #s = :s',
        ExpressionAttributeNames={'#s': 'status'},
        ExpressionAttributeValues={':n': name, ':s': status},
    )


def send_telegram_direct(chat_id: str, token: str, text: str) -> None:
    url = f'https://api.telegram.org/bot{token}/sendMessage'
    data = json.dumps({'chat_id': chat_id, 'text': text}).encode()
    req = urllib.request.Request(url, data=data, headers={'Content-Type': 'application/json'})
    urllib.request.urlopen(req, timeout=10)


def get_or_create_session(actor_id: str) -> str:
    """Look up active session for actor, or create a new one."""
    table = get_ddb().Table(os.environ['SESSION_TABLE_NAME'])

    response = table.get_item(Key={'actor_id': actor_id})
    item = response.get('Item')

    now = int(time.time())
    ttl_8hr = now + (8 * 3600)

    if item and item.get('ttl', 0) > now:
        # Active session exists — extend TTL
        table.update_item(
            Key={'actor_id': actor_id},
            UpdateExpression='SET #ttl = :ttl',
            ExpressionAttributeNames={'#ttl': 'ttl'},
            ExpressionAttributeValues={':ttl': ttl_8hr},
        )
        return item['session_id']

    # Create new session
    session_id = str(uuid.uuid4())
    table.put_item(Item={
        'actor_id': actor_id,
        'session_id': session_id,
        'created_at': str(now),
        'ttl': ttl_8hr,
    })
    return session_id


def handler(event, context):
    # ── Parse SQS records (FIFO — all from same actor) ───────────────────
    records = []
    for record in event.get('Records', []):
        try:
            records.append(json.loads(record['body']))
        except (json.JSONDecodeError, KeyError):
            continue

    if not records:
        return

    first = records[0]
    channel = first.get('channel', 'telegram')
    chat_id = first.get('chat_id', '')
    actor_id = f"{channel}:{chat_id}"

    # ── User registry ─────────────────────────────────────────────────────
    from_info = first.get('messages', [{}])[0]
    user_profile = get_or_create_user(actor_id, from_info)

    # ── Onboarding gate ─────────────────────────────────────────────────────
    table_name = os.environ.get('USERS_TABLE_NAME', '')
    if table_name and user_profile.get('status', 'active') == 'pending':
        raw_prompt = records[0]['messages'][0]['text'] if records else ''
        is_name_msg = bool(raw_prompt and len(raw_prompt.strip()) < 50 and '?' not in raw_prompt)
        if is_name_msg:
            name = raw_prompt.strip()
            update_user_status(actor_id, name=name, status='active')
            user_profile['display_name'] = name
            user_profile['status'] = 'active'
            prompt = f"[System: User just registered with name '{name}'. Welcome them warmly and ask how you can help.]"
        else:
            bot_token_secret_arn = os.environ.get('TELEGRAM_BOT_TOKEN_SECRET_ARN', '')
            bot_token = ''
            if bot_token_secret_arn:
                sm = boto3.client('secretsmanager', region_name='us-east-1')
                bot_token = sm.get_secret_value(SecretId=bot_token_secret_arn)['SecretString']
            send_telegram_direct(chat_id, bot_token, "Hi! I don't recognize you yet. What's your name?")
            return
    # ── Get or create AgentCore session ──────────────────────────────────
    session_id = get_or_create_session(actor_id)
    print(f"[agent-runner] actor={actor_id} session={session_id} user={user_profile.get('display_name', '')}")

    # ── Bundle messages ───────────────────────────────────────────────────
    if len(records) == 1:
        prompt = records[0]['messages'][0]['text']
    else:
        lines = [
            f"[{i+1}] {r['messages'][0]['text']}"
            for i, r in enumerate(records)
        ]
        prompt = f"You have {len(records)} queued messages:\n" + "\n".join(lines)

    # ── Build payload for AgentCore Runtime 1 ────────────────────────────
    payload: dict[str, Any] = {
        'prompt': prompt,
        'actor_id': actor_id,
        'session_id': session_id,
        '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', ''),
            'allowed': user_profile.get('allowed', True),
            'services': user_profile.get('services', {}),
        },
        'channel_adapter': {
            'type': channel,
            'target_id': str(chat_id),
            'bot_token_secret_arn': os.environ.get('TELEGRAM_BOT_TOKEN_SECRET_ARN', ''),
        },
    }

    # ── Invoke AgentCore Runtime 1 ────────────────────────────────────────
    runtime_arn = os.environ.get('RUNTIME_1_ARN', '')
    if not runtime_arn or runtime_arn == 'PLACEHOLDER_SET_AFTER_RUNTIME_DEPLOY':
        print(f"[agent-runner] RUNTIME_1_ARN not set — skipping AgentCore invoke")
        print(f"[agent-runner] Would have sent: {json.dumps(payload)[:200]}")
        return

    client = get_agentcore()
    response = client.invoke_agent_runtime(
        agentRuntimeArn=runtime_arn,
        runtimeSessionId=session_id,
        payload=json.dumps(payload).encode(),
    )

    # Consume streaming response (agent delivers to Telegram via send_message tool)
    for chunk in response.get('response', []):
        pass  # intentional no-op — agent handles delivery internally

    print(f"[agent-runner] Completed session={session_id} actor={actor_id}")