Claude Function Calling with API Gateway: Build a Secure Serverless LLM Endpoint

{
  "status": "success",
  "result": { "userId": "1234" }
}

// src/types.ts
import { z } from "zod";

/**
 * Claude sends a function call with a name and an arguments object.
 * We describe that shape with Zod so we can validate it at runtime.
 */
export const ClaudeAddUserSchema = z.object({
  name: z.string(),
  email: z.string().email(),
  age: z.number().int().positive().optional(),
});

/**
 * The full request body Claude will send.
 */
export const ClaudeRequestSchema = z.object({
  function: z.literal("addUser"),
  arguments: ClaudeAddUserSchema,
});

/**
 * Export TypeScript types derived from the schemas.
 */
export type ClaudeAddUser = z.infer<typeof ClaudeAddUserSchema>;
export type ClaudeRequest = z.infer<typeof ClaudeRequestSchema>;

// src/handler.ts
import { APIGatewayProxyEvent, APIGatewayProxyResult } from "aws-lambda";
import { DynamoDBClient, PutItemCommand } from "@aws-sdk/client-dynamodb";
import { ClaudeRequestSchema, ClaudeAddUser } from "./types";

/**
 * Create a DynamoDB client once so it can be reused across invocations.
 */
const ddb = new DynamoDBClient({});

/**
 * The Lambda entry point that API Gateway will invoke.
 */
export const handler = async (
  event: APIGatewayProxyEvent
): Promise<APIGatewayProxyResult> => {
  // -----------------------------------------------------------------
  // 1️⃣ Parse and validate Claude's JSON payload.
  // -----------------------------------------------------------------
  let payload: ClaudeAddUser;
  try {
    // The body arrives as a string; JSON.parse converts it to an object.
    const parsed = JSON.parse(event.body ?? "{}");
    // Zod validates shape and throws if something is missing or wrong.
    payload = ClaudeRequestSchema.parse(parsed).arguments;
  } catch (err) {
    // If validation fails, respond with Claude's error format.
    return {
      statusCode: 400,
      body: JSON.stringify({
        status: "error",
        message: "Invalid request payload",
        details: err instanceof Error ? err.message : String(err),
      }),
    };
  }

  // -----------------------------------------------------------------
  // 2️⃣ Build a type‑safe PutItem command for DynamoDB.
  // -----------------------------------------------------------------
  // The satisfies keyword makes sure the object conforms to PutItemCommandInput.
  const putCommand = {
    TableName: "Users",
    Item: {
      userId: { S: crypto.randomUUID() }, // primary key
      name: { S: payload.name },
      email: { S: payload.email },
      // age is optional; only add it if present.
      ...(payload.age && { age: { N: payload.age.toString() } }),
    },
  } satisfies PutItemCommand["input"]; // compile‑time check

  // Execute the command.
  await ddb.send(new PutItemCommand(putCommand));

  // -----------------------------------------------------------------
  // 3️⃣ Respond in Claude's expected schema.
  // -----------------------------------------------------------------
  return {
    statusCode: 200,
    body: JSON.stringify({
      status: "success",
      result: { userId: putCommand.Item.userId.S },
    }),
  };
};

// scripts/createApi.ts
import {
  APIGatewayClient,
  CreateRestApiCommand,
  GetResourcesCommand,
  CreateResourceCommand,
  PutMethodCommand,
  PutIntegrationCommand,
  CreateDeploymentCommand,
} from "@aws-sdk/client-api-gateway";

const client = new APIGatewayClient({});

/**
 * Helper to fetch the root resource ID ("/").
 */
async function getRootResourceId(apiId: string): Promise<string> {
  const resources = await client.send(
    new GetResourcesCommand({ restApiId: apiId })
  );
  const root = resources.items?.find((r) => r.path === "/");
  if (!root?.id) throw new Error("Root resource not found");
  return root.id;
}

/**
 * Main function that creates the API, the /addUser resource,
 * and wires it to the Lambda using proxy integration.
 */
export async function createClaudeApi(lambdaArn: string) {
  // 1️⃣ Create the API.
  const api = await client.send(
    new CreateRestApiCommand({
      name: "ClaudeFunctionCallingAPI",
      description: "Endpoint for Claude to invoke addUser",
      endpointConfiguration: { types: ["REGIONAL"] },
    })
  );
  const apiId = api.id!;
  console.log(`Created API ${apiId}`);

  // 2️⃣ Create /addUser resource.
  const rootId = await getRootResourceId(apiId);
  const addUserRes = await client.send(
    new CreateResourceCommand({
      restApiId: apiId,
      parentId: rootId,
      pathPart: "addUser",
    })
  );
  const addUserId = addUserRes.id!;

  // 3️⃣ Add POST method (Claude always POSTs JSON).
  await client.send(
    new PutMethodCommand({
      restApiId: apiId,
      resourceId: addUserId,
      httpMethod: "POST",
      authorizationType: "NONE", // we will add Cognito authorizer later
    })
  );

  // 4️⃣ Wire the method to Lambda using **proxy** integration.
  await client.send(
    new PutIntegrationCommand({
      restApiId: apiId,
      resourceId: addUserId,
      httpMethod: "POST",
      type: "AWS_PROXY", // critical! preserves nested JSON
      integrationHttpMethod: "POST",
      uri: `arn:aws:apigateway:${process.env.AWS_REGION}:lambda:path/2015-03-31/functions/${lambdaArn}/invocations`,
    })
  );

  // 5️⃣ Deploy the API to a stage called "prod".
  await client.send(
    new CreateDeploymentCommand({
      restApiId: apiId,
      stageName: "prod",
      description: "Initial deployment for Claude function calls",
    })
  );

  console.log(`API deployed to https://${apiId}.execute-api.${process.env.AWS_REGION}.amazonaws.com/prod/addUser`);
}

// scripts/allowApiGateway.ts
import { LambdaClient, AddPermissionCommand } from "@aws-sdk/client-lambda";

const client = new LambdaClient({});

export async function grantApiInvoke(lambdaArn: string, apiArn: string) {
  await client.send(
    new AddPermissionCommand({
      FunctionName: lambdaArn,
      StatementId: "APIGatewayInvoke",
      Action: "lambda:InvokeFunction",
      Principal: "apigateway.amazonaws.com",
      SourceArn: `${apiArn}/*/*`,
    })
  );
  console.log("Permission granted for API Gateway to invoke Lambda");
}

// scripts/addCognitoAuthorizer.ts
import {
  APIGatewayClient,
  CreateAuthorizerCommand,
  UpdateMethodCommand,
} from "@aws-sdk/client-api-gateway";

const client = new APIGatewayClient({});

export async function attachCognitoAuthorizer(
  apiId: string,
  resourceId: string,
  userPoolArn: string
) {
  // 1️⃣ Create the authorizer object.
  const authorizer = await client.send(
    new CreateAuthorizerCommand({
      restApiId: apiId,
      name: "CognitoAuthorizer",
      type: "COGNITO_USER_POOLS",
      providerARNs: [userPoolArn],
      identitySource: "method.request.header.Authorization", // JWT header
    })
  );

  // 2️⃣ Update the POST method to require the authorizer.
  await client.send(
    new UpdateMethodCommand({
      restApiId: apiId,
      resourceId,
      httpMethod: "POST",
      patchOperations: [
        {
          op: "replace",
          path: "/authorizationType",
          value: "COGNITO_USER_POOLS",
        },
        {
          op: "replace",
          path: "/authorizerId",
          value: authorizer.id!,
        },
      ],
    })
  );

  console.log("Cognito authorizer attached to /addUser POST");
}