agent-claw/cdk/cdk.out/asset.7053cd1618f5f520a7aac409588128f920d8fe76791c1dbcc65610454d1a5387/handler.py

import json
import os
import time
import uuid
import boto3
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),
        'allowed': True,
    }
    table.put_item(Item=item)
    print(f'[agent-runner] Registered new user: {actor_id}')
    return item


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)

    # ── 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', ''),
            'allowed': user_profile.get('allowed', True),
        },
        '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}")