AppSync With Bedrock: GraphQL Resolvers That Call Tools Safely

AppSync With Bedrock: GraphQL Resolvers That Call Tools Safely

Exposing Bedrock through AppSync feels elegant: one GraphQL schema, Cognito auth, and a mutation that “just asks the assistant.” It is also how a single buggy client loops into a five-figure overnight bill. Production AppSync+Bedrock needs auth context propagation, strict input schemas, tool allowlists, and hard cost budgets — not a resolver that forwards raw strings to InvokeModel.

⚡ TL;DR: Keep GraphQL types narrow; put Bedrock calls behind Lambda resolvers that inherit Cognito sub/tenant claims; enforce per-tenant rate limits and token budgets in DynamoDB; validate tool arguments against JSON Schema before side effects. Stream via subscriptions only with abort and heartbeat. Illustrative SLO: p95 resolver setup < 300ms before first model token; hard deny when monthly tenant token budget exhausted.

Schema that cannot express unbounded spend

# 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
}

Pin modelId, temperature, and max tokens server-side. Clients choose product features, not foundation model SKUs.

Lambda resolver: auth → budget → Bedrock

// 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,
  };
};

IAM on this Lambda: bedrock:InvokeModel on the pinned model ARN, dynamodb:UpdateItem on the budget table, read-only access to the tenant’s KB if used — mirror Bedrock Agents guardrail posture.

Tools from GraphQL: never trust client tool JSON

If the assistant can call tools (create ticket, search code), define tools in the Lambda, not in the mutation input:

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

✅ Propagate tenant from JWT into every tool call.
❌ Let the model pass tenantId as a tool argument.

Combine with Knowledge Base metadata filters for multi-tenant RAG (companion post in this wave).

Rate limits, WAF, and abortable streams

  • API Gateway / AppSync: per-user throttling + AWS WAF rate rules
  • DynamoDB token buckets for soft product limits (see cost-controls companion)
  • For subscriptions/streaming: handle client disconnect; do not continue billing a full completion into the void
# 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

Sandbox any shell/code tools behind Lambda as in LLM Coding Agents sandboxes.

Closing checklist

✅ Dos
– ✅ Pin model ID and maxTokens server-side
– ✅ Pre-debit or atomically check per-tenant budgets
– ✅ Derive tenant from Cognito claims for tools and KB filters
– ✅ Validate tool args with JSON Schema additionalProperties: false
– ✅ Alarm on token burn and error spikes

❌ Don’ts
– ❌ Don’t expose raw InvokeModel inputs on the GraphQL surface
– ❌ Don’t trust client-supplied model IDs or tool names
– ❌ Don’t skip auth on “internal” AppSync APIs
– ❌ Don’t let resolvers hold * IAM on Bedrock and DynamoDB
– ❌ Don’t stream forever after the subscriber vanished

Related reading

Last updated on September 11, 2026


Discover more from CheatCoders

Subscribe to get the latest posts sent to your email.

Comments

No comments yet. Why don’t you start the discussion?

Leave a Reply