Hardening Agent Tool Calls: Prompt Injection Defense and Supply Chain Trust in Multi-Tenant Agentic Systems

import { z } from 'zod';

// Define strict per-tool input schemas — no 'any', no passthrough
const QueryToolSchema = z.object({
  query: z.string()
    .max(500)
    .regex(/^[a-zA-Z0-9\s_\-\.]+$/, 'No special characters allowed'),
  tenantId: z.string().uuid(),
});

type ToolCallContext = { callerTenantId: string };

function validateToolCall<T>(
  schema: z.ZodSchema<T>,
  rawArgs: unknown,
  ctx: ToolCallContext
): T {
  // Parse and throw on violation — never coerce silently
  const parsed = schema.parse(rawArgs);

  // Enforce tenant boundary at the validation layer, not inside handler logic
  if ('tenantId' in (parsed as object)) {
    const { tenantId } = parsed as { tenantId: string };
    if (tenantId !== ctx.callerTenantId) {
      throw new Error(`Tenant boundary violation: caller ${ctx.callerTenantId} attempted access to ${tenantId}`);
    }
  }

  return parsed;
}

// Usage inside your tool dispatch middleware
export function dispatchToolCall(
  toolName: string,
  rawArgs: unknown,
  ctx: ToolCallContext
) {
  const schemas: Record<string, z.ZodSchema> = {
    query_tool: QueryToolSchema,
  };

  const schema = schemas[toolName];
  if (!schema) throw new Error(`Unknown tool: ${toolName}`);

  const validatedArgs = validateToolCall(schema, rawArgs, ctx);
  return executeToolHandler(toolName, validatedArgs);
}

// Consequence levels assigned at tool registration — not runtime
enum ConsequenceLevel {
  READ = 0,
  WRITE = 1,
  FINANCIAL = 2,
  IRREVERSIBLE = 3,
}

interface ToolRegistration {
  name: string;
  schema: z.ZodSchema;
  consequenceLevel: ConsequenceLevel;
  spendLimitCents?: number;
}

async function policyGatedDispatch(
  tool: ToolRegistration,
  validatedArgs: unknown,
  ctx: ToolCallContext
): Promise<unknown> {
  if (tool.consequenceLevel >= ConsequenceLevel.FINANCIAL) {
    const { amount } = validatedArgs as { amount: number };
    const withinLimit = await policyEngine.checkSpendLimit(
      ctx.callerTenantId,
      amount,
      tool.spendLimitCents ?? 0
    );

    if (!withinLimit) {
      // Route to human-in-the-loop approval, not inline execution
      return await humanApprovalQueue.submit({
        toolName: tool.name,
        args: validatedArgs,
        tenantId: ctx.callerTenantId,
        requestedAt: new Date().toISOString(),
      });
    }
  }

  return executeToolHandler(tool.name, validatedArgs);
}

# Verify SLSA attestation before allowing image into agent deployment pipeline
# Using cosign as the verification tool

import subprocess
import sys

def verify_image_provenance(
    image_ref: str,
    expected_signer: str,
    slsa_level: int = 2
) -> bool:
    """
    Gate on SLSA attestation before deploying any agent container image.
    Returns True only if attestation is valid and meets minimum SLSA level.
    """
    result = subprocess.run(
        [
            "cosign", "verify-attestation",
            "--type", "slsaprovenance",
            "--certificate-identity", expected_signer,
            "--certificate-oidc-issuer", "https://token.actions.githubusercontent.com",
            image_ref
        ],
        capture_output=True,
        text=True
    )

    if result.returncode != 0:
        print(f"Attestation verification FAILED for {image_ref}: {result.stderr}")
        return False

    # Parse attestation payload and check SLSA build level
    # (abbreviated — real implementation parses JSON predicate)
    print(f"Attestation verified for {image_ref} at SLSA level {slsa_level}")
    return True

# Called from deployment pipeline — fail fast, never deploy unverified images
if not verify_image_provenance(
    image_ref="my-registry/agent-worker:sha256-abc123",
    expected_signer="ci-system@myorg.com"
):
    sys.exit(1)

// Mandatory security annotation block on every tool definition
// Enforced by CI lint check — not optional documentation

/**
 * @tool query_customer_records
 * @security-annotation
 *   ACCEPTS: tenantId (UUID), searchTerm (alphanumeric max 100 chars)
 *   AFFECTS: read-only access to tenant-scoped customer table
 *   INJECTION SURFACE: searchTerm is passed to PostgreSQL ILIKE — 
 *     validated against /^[a-zA-Z0-9\s]+$/ before use
 *   ADVERSARIAL INPUTS: path traversal N/A, SQL injection mitigated by
 *     parameterized query + regex gate, tenant crossing mitigated by
 *     row-level security policy on tenant_id column
 *   REVIEWED BY: @sam — 2026-08-19
 */
export const queryCustomerRecordsTool = {
  name: 'query_customer_records',
  schema: z.object({
    tenantId: z.string().uuid(),
    searchTerm: z.string().max(100).regex(/^[a-zA-Z0-9\s]+$/),
  }),
  consequenceLevel: ConsequenceLevel.READ,
  handler: async ({ tenantId, searchTerm }: { tenantId: string; searchTerm: string }) => {
    // parameterized — no string interpolation into query
    return db.query(
      'SELECT id, name FROM customers WHERE tenant_id = $1 AND name ILIKE $2',
      [tenantId, `%${searchTerm}%`]
    );
  },
};

// Prompt template fragment for Claude Code / Copilot when generating tool handlers
// Embed security constraints as generation context, not post-review checklist

const TOOL_GENERATION_CONTEXT = `
When generating a tool handler, you MUST:
1. Define a Zod schema with explicit .max() and .regex() on every string field
2. Never interpolate validated input directly into SQL, shell commands, or file paths
3. Include a @security-annotation JSDoc block listing: ACCEPTS, AFFECTS, INJECTION SURFACE
4. Assign a ConsequenceLevel enum value at the tool registration level
5. Inject tenantId from the call context object, never from LLM-provided arguments

Example constraint pattern:
  z.string().max(200).regex(/^[a-zA-Z0-9_\-\.]+$/, 'Alphanumeric and safe punctuation only')

Never generate:
  db.query('SELECT * FROM ' + tableName)  // injection surface
  fs.readFile(userProvidedPath)            // path traversal surface
  args.tenantId                           // must come from ctx, not LLM output
`;

// Inject this as a system-level constraint in your AI coding assistant config
// so it applies to every tool generation session, not just when you remember to ask