Appsync Bedrock: GraphQL Resolvers That Call Tools Safely

# schema.graphql — illustrative
type AskResult {
  requestId: ID!
  answer: String!
  citations: [String!]!
  inputTokens: Int!
  outputTokens: Int!
}

input AskInput {
  question: String! @length(max: 4000)
  repoId: ID!
  # ❌ Never: freeform modelId from the client
  # ❌ Never: raw toolName + JSON blob
}

type Mutation {
  askCodeAssistant(input: AskInput!): AskResult!
    @aws_cognito_user_pools
    @rateLimit(limit: 30, duration: 60) # illustrative directive / WAF pairing
}

// resolvers/askCodeAssistant.ts
import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
import { DynamoDBDocumentClient, UpdateCommand } from "@aws-sdk/lib-dynamodb";
import {
  BedrockRuntimeClient,
  ConverseCommand,
} from "@aws-sdk/client-bedrock-runtime";

const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({}));
const bedrock = new BedrockRuntimeClient({});
const BUDGET_TABLE = process.env.BUDGET_TABLE!;
const MODEL_ID = process.env.MODEL_ID!; // pinned
const MAX_OUT = 1024;

export const handler = async (event: any) => {
  const tenantId = event.identity?.claims?.["custom:tenant_id"];
  const sub = event.identity?.sub;
  if (!tenantId || !sub) throw new Error("Unauthorized");

  const question: string = event.arguments.input.question;
  const repoId: string = event.arguments.input.repoId;
  if (!/^[\w-]{1,64}$/.test(repoId)) throw new Error("Invalid repoId");

  // Atomic budget consume — fail closed
  try {
    await ddb.send(
      new UpdateCommand({
        TableName: BUDGET_TABLE,
        Key: { pk: `TENANT#${tenantId}`, sk: "BUDGET#MONTH" },
        UpdateExpression: "ADD tokensUsed :t SET updatedAt = :u",
        ConditionExpression: "tokensUsed < maxTokens",
        ExpressionAttributeValues: {
          ":t": 1500, // conservative pre-debit; reconcile with actual later
          ":u": new Date().toISOString(),
        },
      })
    );
  } catch {
    throw new Error("Token budget exceeded");
  }

  const resp = await bedrock.send(
    new ConverseCommand({
      modelId: MODEL_ID,
      system: [
        {
          text: `Tenant ${tenantId}. Answer only about repo ${repoId}. Refuse secrets.`,
        },
      ],
      messages: [{ role: "user", content: [{ text: question }] }],
      inferenceConfig: { maxTokens: MAX_OUT, temperature: 0.2 },
      // toolConfig only if tools are allowlisted below
    })
  );

  const text =
    resp.output?.message?.content?.map((c) => ("text" in c ? c.text : "")).join("") ??
    "";
  return {
    requestId: event.request?.requestId,
    answer: text,
    citations: [],
    inputTokens: resp.usage?.inputTokens ?? 0,
    outputTokens: resp.usage?.outputTokens ?? 0,
  };
};

const TOOLS = [
  {
    toolSpec: {
      name: "search_repo",
      description: "Search symbols in the caller's tenant repo",
      inputSchema: {
        json: {
          type: "object",
          properties: { query: { type: "string", maxLength: 200 } },
          required: ["query"],
          additionalProperties: false,
        },
      },
    },
  },
];

// When executing search_repo, force tenantId from Cognito claims — never from model args

# CloudWatch alarm — illustrative
aws cloudwatch put-metric-alarm \
  --alarm-name appsync-bedrock-tokens-high \
  --metric-name InputTokenSum \
  --namespace CheatCoders/AppSyncBedrock \
  --threshold 5000000 \
  --comparison-operator GreaterThanThreshold \
  --evaluation-periods 1 \
  --period 3600 \
  --statistic Sum \
  --alarm-actions arn:aws:sns:us-east-1:123456789012:ai-spend