agent-claw/cdk/lib/agent-claw-stack.ts

import * as cdk from 'aws-cdk-lib';
import * as s3 from 'aws-cdk-lib/aws-s3';
import * as s3deploy from 'aws-cdk-lib/aws-s3-deployment';
import * as dynamodb from 'aws-cdk-lib/aws-dynamodb';
import * as sqs from 'aws-cdk-lib/aws-sqs';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as apigatewayv2 from 'aws-cdk-lib/aws-apigatewayv2';
import * as apigatewayv2integrations from 'aws-cdk-lib/aws-apigatewayv2-integrations';
import * as iam from 'aws-cdk-lib/aws-iam';
import * as secretsmanager from 'aws-cdk-lib/aws-secretsmanager';
import { SqsEventSource } from 'aws-cdk-lib/aws-lambda-event-sources';
import { Construct } from 'constructs';
import * as path from 'path';
import { execSync } from 'child_process';

export class AgentClawStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
    super(scope, id, props);

    // ── Context parameters ─────────────────────────────────────────────────
    const telegramBotTokenSecretArn = this.node.tryGetContext('telegramBotTokenSecretArn') as string | undefined;
    const braveApiKeySecretArn = this.node.tryGetContext('braveApiKeySecretArn') as string | undefined;
    const existingWorkspaceBucketName = this.node.tryGetContext('workspaceBucketName') as string | undefined;
    const runtime1Arn = this.node.tryGetContext('runtime1Arn') as string | undefined;

    if (!telegramBotTokenSecretArn) {
      throw new Error('Context param required: telegramBotTokenSecretArn');
    }
    if (!braveApiKeySecretArn) {
      throw new Error('Context param required: braveApiKeySecretArn');
    }

    // ── Secrets (reference existing) ───────────────────────────────────────
    const botTokenSecret = secretsmanager.Secret.fromSecretCompleteArn(
      this, 'TelegramBotToken', telegramBotTokenSecretArn
    );
    const braveApiKeySecret = secretsmanager.Secret.fromSecretCompleteArn(
      this, 'BraveApiKey', braveApiKeySecretArn
    );

    // ── S3 workspace bucket ────────────────────────────────────────────────
    const workspaceBucket = existingWorkspaceBucketName
      ? s3.Bucket.fromBucketName(this, 'WorkspaceBucket', existingWorkspaceBucketName)
      : new s3.Bucket(this, 'WorkspaceBucket', {
          bucketName: `agent-claw-workspace-${this.account}`,
          removalPolicy: cdk.RemovalPolicy.RETAIN,
          versioned: false,
          encryption: s3.BucketEncryption.S3_MANAGED,
        });

    // Seed workspace files on deploy (only if bucket was created by us)
    if (!existingWorkspaceBucketName) {
      new s3deploy.BucketDeployment(this, 'WorkspaceFiles', {
        sources: [s3deploy.Source.asset(path.join(__dirname, '../../workspace'))],
        destinationBucket: workspaceBucket as s3.Bucket,
      });
    }

    // ── DynamoDB session store ─────────────────────────────────────────────
    const sessionTable = new dynamodb.Table(this, 'SessionStore', {
      tableName: 'agent-claw-sessions',
      partitionKey: { name: 'actor_id', type: dynamodb.AttributeType.STRING },
      billingMode: dynamodb.BillingMode.PAY_PER_REQUEST,
      timeToLiveAttribute: 'ttl',
      removalPolicy: cdk.RemovalPolicy.RETAIN,
    });

    // ── DynamoDB user registry ─────────────────────────────────────────────
    const usersTable = new dynamodb.Table(this, 'UsersTable', {
      tableName: 'agent-claw-users',
      partitionKey: { name: 'actor_id', type: dynamodb.AttributeType.STRING },
      billingMode: dynamodb.BillingMode.PAY_PER_REQUEST,
      removalPolicy: cdk.RemovalPolicy.RETAIN,
    });

    // ── SQS FIFO message queue ─────────────────────────────────────────────
    const messageQueue = new sqs.Queue(this, 'MessageQueue', {
      queueName: 'agent-claw-messages.fifo',
      fifo: true,
      contentBasedDeduplication: false,
      visibilityTimeout: cdk.Duration.seconds(900),
      receiveMessageWaitTime: cdk.Duration.seconds(20),
    });

    // ── Lambda: tg-ingest ─────────────────────────────────────────────────
    const tgIngestFn = new lambda.Function(this, 'TgIngest', {
      functionName: 'agent-claw-tg-ingest',
      runtime: lambda.Runtime.PYTHON_3_12,
      handler: 'handler.handler',
      code: lambda.Code.fromAsset(path.join(__dirname, '../../src/lambdas/tg-ingest')),
      timeout: cdk.Duration.seconds(10),
      memorySize: 128,
      environment: {
        MESSAGE_QUEUE_URL: messageQueue.queueUrl,
        TELEGRAM_BOT_TOKEN_SECRET_ARN: telegramBotTokenSecretArn,
        TELEGRAM_WEBHOOK_SECRET: '', // set via SSM or direct env after deploy
      },
    });
    messageQueue.grantSendMessages(tgIngestFn);
    botTokenSecret.grantRead(tgIngestFn);

    // ── Lambda: agent-runner ───────────────────────────────────────────────
    const agentRunnerFn = new lambda.Function(this, 'AgentRunner', {
      functionName: 'agent-claw-agent-runner',
      runtime: lambda.Runtime.PYTHON_3_12,
      handler: 'handler.handler',
      code: lambda.Code.fromAsset(path.join(__dirname, '../../src/lambdas/agent-runner')),
      timeout: cdk.Duration.seconds(900),
      memorySize: 256,
      environment: {
        SESSION_TABLE_NAME: sessionTable.tableName,
        WORKSPACE_BUCKET_NAME: workspaceBucket.bucketName,
        TELEGRAM_BOT_TOKEN_SECRET_ARN: telegramBotTokenSecretArn,
        BRAVE_API_KEY_SECRET_ARN: braveApiKeySecretArn,
        RUNTIME_1_ARN: runtime1Arn ?? 'PLACEHOLDER_SET_AFTER_RUNTIME_DEPLOY',
        AWS_REGION_NAME: 'us-east-1',
      },
    });

    sessionTable.grantReadWriteData(agentRunnerFn);
    usersTable.grantReadWriteData(agentRunnerFn);
    agentRunnerFn.addEnvironment('USERS_TABLE_NAME', usersTable.tableName);
    workspaceBucket.grantRead(agentRunnerFn);
    botTokenSecret.grantRead(agentRunnerFn);
    braveApiKeySecret.grantRead(agentRunnerFn);
    messageQueue.grantConsumeMessages(agentRunnerFn);

    // AgentCore invoke permission
    agentRunnerFn.addToRolePolicy(new iam.PolicyStatement({
      actions: ['bedrock-agentcore:InvokeAgentRuntime'],
      resources: ['*'],
    }));

    // SQS event source
    agentRunnerFn.addEventSource(new SqsEventSource(messageQueue, {
      batchSize: 10,
      enabled: true,
    }));

    // ── API Gateway HTTP ───────────────────────────────────────────────────
    const httpApi = new apigatewayv2.HttpApi(this, 'WebhookApi', {
      apiName: 'agent-claw-webhook',
    });

    httpApi.addRoutes({
      path: '/telegram',
      methods: [apigatewayv2.HttpMethod.POST],
      integration: new apigatewayv2integrations.HttpLambdaIntegration(
        'TgIngestIntegration', tgIngestFn
      ),
    });

    // ── AgentCore Runtime 1 ────────────────────────────────────────────────
    // NOTE: AgentCore CDK L2 constructs are in preview. Using CfnResource.
    // The runtime1Arn output below needs to be fed back as context param
    // on subsequent deploys so agent-runner can invoke it.

    // IAM execution role for Runtime 1
    const runtime1Role = new iam.Role(this, 'Runtime1Role', {
      assumedBy: new iam.ServicePrincipal('bedrock-agentcore.amazonaws.com'),
      description: 'Execution role for agent-claw Runtime 1 (main assistant)',
    });
    runtime1Role.addToPolicy(new iam.PolicyStatement({
      actions: [
        'bedrock:InvokeModel',
        'bedrock:InvokeModelWithResponseStream',
      ],
      resources: ['*'],
    }));
    workspaceBucket.grantRead(runtime1Role);
    botTokenSecret.grantRead(runtime1Role);
    braveApiKeySecret.grantRead(runtime1Role);
    usersTable.grantReadWriteData(runtime1Role);
    // Google secret grants added after workspace_mcp section below
    runtime1Role.addToPolicy(new iam.PolicyStatement({
      actions: [
        'bedrock-agentcore:CreateEvent',
        'bedrock-agentcore:ListEvents',
        'bedrock-agentcore:RetrieveMemoryRecords',
      ],
      resources: ['*'],
    }));

    // AgentCore Runtime 1 resource (CfnResource — L2 in preview)
    // CodeZip: packages src/runtime-1/ as a zip and uploads to S3
    // TODO: Replace with L2 construct when aws-cdk-lib includes it stable
    // For now, the runtime ARN must be created manually or via CLI after first synth
    // and fed back as context param runtime1Arn.

    // ── Google Workspace MCP ──────────────────────────────────────────────
    const googleOAuthClientSecret = secretsmanager.Secret.fromSecretNameV2(
      this, 'GoogleOAuthClient', 'agent-claw/google-oauth-client'
    );

    // workspace-mcp Lambda execution role (import existing — created during initial setup)
    const _workspaceMcpRole = iam.Role.fromRoleName(
      this, 'WorkspaceMcpRole', 'agent-claw-workspace-mcp-role'
    );
    googleOAuthClientSecret.grantRead(_workspaceMcpRole);
    // Grant workspace-mcp role read access to all per-user Google credential secrets
    (_workspaceMcpRole as iam.Role).addToPrincipalPolicy?.(new iam.PolicyStatement({
      sid: 'PerUserGoogleCredentialsRead',
      actions: ['secretsmanager:GetSecretValue'],
      resources: [`arn:aws:secretsmanager:${this.region}:${this.account}:secret:agent-claw/google-credentials/*`],
    }));

    // workspace-mcp Lambda — import existing (created with zip + layer, no Docker)
    const workspaceMcpFn = lambda.Function.fromFunctionName(
      this, 'WorkspaceMcp', 'agent-claw-workspace-mcp'
    );

    // Function URL — AWS_IAM auth (already created, reference for policy attachment)
    const workspaceMcpFunctionUrl = 'https://25hugrzw4uwtueeg77jsmft6lq0wunmd.lambda-url.us-east-1.on.aws';
    const workspaceMcpMcpUrl = workspaceMcpFunctionUrl + '/mcp';

    // AgentCore execution role — grant InvokeFunctionUrl identity policy
    runtime1Role.addToPolicy(new iam.PolicyStatement({
      sid: 'WorkspaceMcpInvoke',
      actions: ['lambda:InvokeFunctionUrl'],
      resources: [workspaceMcpFn.functionArn],
      conditions: { StringEquals: { 'lambda:FunctionUrlAuthType': 'AWS_IAM' } },
    }));

    // Grant AgentCore execution role read access to Google OAuth client + per-user credentials
    googleOAuthClientSecret.grantRead(runtime1Role);
    runtime1Role.addToPolicy(new iam.PolicyStatement({
      sid: 'PerUserGoogleCredentialsReadRuntime',
      actions: ['secretsmanager:GetSecretValue'],
      resources: [`arn:aws:secretsmanager:${this.region}:${this.account}:secret:agent-claw/google-credentials/*`],
    }));

    // Pass workspace_mcp MCP URL to agent-runner (informational)
    agentRunnerFn.addEnvironment('WORKSPACE_MCP_URL', workspaceMcpMcpUrl);

    // ── OAuth Handler Lambda ───────────────────────────────────────────────
    const oauthHandlerFn = new lambda.Function(this, 'OAuthHandler', {
      functionName: 'agent-claw-oauth-handler',
      runtime: lambda.Runtime.PYTHON_3_12,
      handler: 'handler.handler',
      code: lambda.Code.fromAsset(path.join(__dirname, '../../src/lambdas/oauth-handler')),
      timeout: cdk.Duration.seconds(30),
      memorySize: 128,
      environment: {
        GOOGLE_OAUTH_CLIENT_SECRET_ARN: googleOAuthClientSecret.secretArn,
        USERS_TABLE_NAME: usersTable.tableName,
        // OAUTH_REDIRECT_URI set after API GW URL is known — injected via addEnvironment below
        OAUTH_REDIRECT_URI: 'PLACEHOLDER',
      },
    });
    googleOAuthClientSecret.grantRead(oauthHandlerFn);
    usersTable.grantReadWriteData(oauthHandlerFn);
    // Grant OAuth handler write access to per-user credential secrets
    oauthHandlerFn.addToRolePolicy(new iam.PolicyStatement({
      sid: 'PerUserGoogleCredentialsWrite',
      actions: [
        'secretsmanager:CreateSecret',
        'secretsmanager:PutSecretValue',
        'secretsmanager:GetSecretValue',
      ],
      resources: [`arn:aws:secretsmanager:${this.region}:${this.account}:secret:agent-claw/google-credentials/*`],
    }));

    // Add OAuth routes to existing HTTP API
    httpApi.addRoutes({
      path: '/oauth/start',
      methods: [apigatewayv2.HttpMethod.GET],
      integration: new apigatewayv2integrations.HttpLambdaIntegration(
        'OAuthStartIntegration', oauthHandlerFn
      ),
    });
    httpApi.addRoutes({
      path: '/oauth/callback',
      methods: [apigatewayv2.HttpMethod.GET],
      integration: new apigatewayv2integrations.HttpLambdaIntegration(
        'OAuthCallbackIntegration', oauthHandlerFn
      ),
    });

    // Set OAUTH_REDIRECT_URI now that we have the API URL
    const oauthRedirectUri = `${httpApi.url}oauth/callback`;
    oauthHandlerFn.addEnvironment('OAUTH_REDIRECT_URI', oauthRedirectUri);

    // Pass OAuth start URL to AgentCore runtime so agent can generate connect links
    const oauthStartUrl = `${httpApi.url}oauth/start`;
    // NOTE: AgentCore runtime env vars are set in agentcore.json / agentcore deploy, not CDK.
    // Output the URL so it can be manually set in agentcore.json OAUTH_START_URL.

    // ── Outputs ────────────────────────────────────────────────────────────

    new cdk.CfnOutput(this, 'WorkspaceMcpFunctionUrl', {
      value: workspaceMcpFunctionUrl,
      description: 'workspace-mcp Lambda Function URL (MCP endpoint for Gmail/Calendar)',
    });
    new cdk.CfnOutput(this, 'OAuthStartUrl', {
      value: oauthStartUrl,
      description: 'Google OAuth start URL — set as OAUTH_START_URL in agentcore.json',
    });
    new cdk.CfnOutput(this, 'OAuthRedirectUri', {
      value: oauthRedirectUri,
      description: 'Google OAuth redirect URI — register in Google Cloud Console',
    });

    new cdk.CfnOutput(this, 'WebhookUrl', {
      value: `${httpApi.url}telegram`,
      description: 'Register this URL with Telegram BotFather as webhook endpoint',
    });

    new cdk.CfnOutput(this, 'WorkspaceBucketName', {
      value: workspaceBucket.bucketName,
      description: 'S3 bucket containing agent workspace files',
    });

    new cdk.CfnOutput(this, 'SessionTableName', {
      value: sessionTable.tableName,
      description: 'DynamoDB table for session mapping',
    });

    new cdk.CfnOutput(this, 'UsersTableName', {
      value: usersTable.tableName,
      description: 'DynamoDB user registry table',
    });

    new cdk.CfnOutput(this, 'MessageQueueUrl', {
      value: messageQueue.queueUrl,
      description: 'SQS FIFO queue for incoming messages',
    });

    new cdk.CfnOutput(this, 'Runtime1RoleArn', {
      value: runtime1Role.roleArn,
      description: 'IAM execution role ARN for AgentCore Runtime 1',
    });
  }
}