The org-scoped getInvoiceStats tool

Right now the chat can talk, but it can’t count. Ask it “how many overdue invoices do we have?” and it hands back a confident, fabricated number, because all that sits behind the answer is a language model predicting plausible text, and “fourteen” is as plausible as “three”. This lesson gives the chat one tool, getInvoiceStats, so it answers from real invoice aggregates instead of guessing.

When it works, that same question returns a number matching the inspector’s count panel for the acting org, and the assistant cites it rather than inventing one. The part that matters more than the feature: a forged orgId buried in the model’s tool-call arguments cannot reach another organization’s data. There’s no new card on screen yet, so your proof here is the inspector’s panels and the network tab.

Your mission

This lesson turns the chat into a grounded analyst, and it carries the most important rule of the project: the model is untrusted input. Every decision below serves it.

The load-bearing decision is that orgId is never part of the tool’s inputSchema. The model chooses which statistics it wants; it never chooses whose data they cover, because execute closes over the orgId the route’s auth boundary already established. The MODEL_FROM_INPUT_ORGID inspector flag lets you break this on purpose and watch the cross-tenant leak appear.

The outputSchema returns a minimal aggregate, not raw rows, for two reasons. Cost: each tool result is fed back as input on the loop’s next step, so full rows compound input tokens. Leakage: rows carry invoice numbers, amounts, and customer names the model has no reason to see. Project at the tool boundary, where the data leaves your control.

Failures follow “return don’t throw”: a thrown read is caught and returned as a typed { error } the model can apologize for, instead of 500ing the request. Operational failures become recoverable; programmer bugs like a typo still crash, because those belong in your logs. And one guard you skip entirely: the SDK validates arguments against inputSchema before execute runs, so an invented status is caught for free.

Wiring the tool into the route adds per-step audit: onStepFinish writes one 'llm.step' row per loop step, a trace of how many model round-trips a question took. Out of scope here: the typed card UI that renders these tool parts (the last lesson of this chapter) and the token-counting half of onStepFinish that feeds the daily quota (the next lesson).

Asking “how many overdue invoices do we have?” returns a count matching the seed’s overdue active rows for org-acme, confirmed against the inspector’s row-count panel.

untested

Asking “what’s our total paid this month?” returns a total matching a reduce over the active paid rows.

untested

A message asking the model to “use orgId = org-globex” still returns org-acme’s data; flipping MODEL_FROM_INPUT_ORGID and repeating shows the leak, proving the closure is the structural reason it’s safe.

untested

A recursion-prone prompt produces at most five tool-getInvoiceStats parts and a final message acknowledging the cap.

untested

With “Force tool error” on, a stats question produces a tool part in output-error and a follow-up text answer asking you to rephrase, with no 500 in the network tab.

untested

Each conversation writes one 'llm.step' row per loop step plus one 'llm.finish' row, scoped to the active org.

untested

Coding time

Write src/lib/llm/tools.ts and wire it into src/app/api/chat/route.ts against the brief and the checklist above, then confirm the behavior by hand against the inspector. Once you’ve taken your shot, open the walkthrough.

Reference solution and walkthrough

Two files: the tool, then the change to the route that uses it.

The tool

src/lib/llm/tools.ts reads as one dense block, so step through it part by part. The order your attention should land in: inputSchema, then outputSchema, then the closure inside execute, then the error handling, then the type exports.

import 'server-only';

import { type InferUITools, tool, type UIMessage } from 'ai';
import { z } from 'zod';
import { scopedInvoices } from '@/lib/invoices/scoped-query';
import { getFlag } from '@/server/inspector-flags';
import type { Invoice } from '@/server/types';

const isoDate = (iso: string): string => iso.slice(0, 10);

// The single read-only tool. `execute` closes over `ctx.orgId` from the server
// auth boundary — the model NEVER passes `orgId` (it is not in `inputSchema`), so
// a forged tool-call argument cannot cross tenants. The `MODEL_FROM_INPUT_ORGID`
// inspector flag is the only path that reads `orgId` from model input; it exists
// solely to make that leak visible by hand (default off → always `ctx.orgId`).
export const buildInvoiceTools = (ctx: { orgId: string }) => ({
  getInvoiceStats: tool({
    description:
      'Return aggregate invoice statistics for the current organization. Use this for any question that needs counts, totals, or status breakdowns of invoices.',
    inputSchema: z.strictObject({
      status: z.enum(['draft', 'sent', 'paid', 'overdue']).optional(),
      since: z.iso.date().optional(),
    }),
    outputSchema: z.strictObject({
      count: z.number().int(),
      totalAmount: z.number(),
      byStatus: z.record(z.string(), z.number().int()),
      oldestUnpaidDueDate: z.iso.date().nullable(),
    }),
    execute: async (input) => {
      try {
        if (getFlag('FORCE_TOOL_ERROR')) {
          return { error: 'stats_unavailable' as const };
        }

        const scopeOrgId = getFlag('MODEL_FROM_INPUT_ORGID')
          ? ((input as { orgId?: string }).orgId ?? ctx.orgId)
          : ctx.orgId;

        let query = scopedInvoices(scopeOrgId).active();
        if (input.status) {
          query = query.filter((inv) => inv.status === input.status);
        }
        if (input.since) {
          const since = input.since;
          query = query.filter((inv) => isoDate(inv.createdAt) >= since);
        }
        const rows = query.take(Number.MAX_SAFE_INTEGER);

        const totalAmount = rows.reduce(
          (sum, inv) => sum + Number(inv.total),
          0,
        );

        const byStatus = rows.reduce<Record<string, number>>((acc, inv) => {
          acc[inv.status] = (acc[inv.status] ?? 0) + 1;
          return acc;
        }, {});

        const oldestUnpaidDueDate = rows
          .filter(
            (inv): inv is Invoice & { dueAt: string } =>
              inv.status !== 'paid' && inv.dueAt !== null,
          )
          .reduce<string | null>(
            (oldest, inv) =>
              oldest === null || inv.dueAt < oldest ? inv.dueAt : oldest,
            null,
          );

        return {
          count: rows.length,
          totalAmount,
          byStatus,
          oldestUnpaidDueDate:
            oldestUnpaidDueDate === null ? null : isoDate(oldestUnpaidDueDate),
        };
      } catch {
        return { error: 'stats_unavailable' as const };
      }
    },
  }),
});

export type InvoiceTools = ReturnType<typeof buildInvoiceTools>;

// The client imports only this — the typed message whose tool parts are backed
// by the real tool map.
export type InvoiceUIMessage = UIMessage<
  unknown,
  never,
  InferUITools<InvoiceTools>
>;

The tool reads the store and branches on inspector flags, so it must never reach the client bundle. import 'server-only' makes a stray client import a build error, not a runtime leak.

import 'server-only';

import { type InferUITools, tool, type UIMessage } from 'ai';
import { z } from 'zod';
import { scopedInvoices } from '@/lib/invoices/scoped-query';
import { getFlag } from '@/server/inspector-flags';
import type { Invoice } from '@/server/types';

const isoDate = (iso: string): string => iso.slice(0, 10);

// The single read-only tool. `execute` closes over `ctx.orgId` from the server
// auth boundary — the model NEVER passes `orgId` (it is not in `inputSchema`), so
// a forged tool-call argument cannot cross tenants. The `MODEL_FROM_INPUT_ORGID`
// inspector flag is the only path that reads `orgId` from model input; it exists
// solely to make that leak visible by hand (default off → always `ctx.orgId`).
export const buildInvoiceTools = (ctx: { orgId: string }) => ({
  getInvoiceStats: tool({
    description:
      'Return aggregate invoice statistics for the current organization. Use this for any question that needs counts, totals, or status breakdowns of invoices.',
    inputSchema: z.strictObject({
      status: z.enum(['draft', 'sent', 'paid', 'overdue']).optional(),
      since: z.iso.date().optional(),
    }),
    outputSchema: z.strictObject({
      count: z.number().int(),
      totalAmount: z.number(),
      byStatus: z.record(z.string(), z.number().int()),
      oldestUnpaidDueDate: z.iso.date().nullable(),
    }),
    execute: async (input) => {
      try {
        if (getFlag('FORCE_TOOL_ERROR')) {
          return { error: 'stats_unavailable' as const };
        }

        const scopeOrgId = getFlag('MODEL_FROM_INPUT_ORGID')
          ? ((input as { orgId?: string }).orgId ?? ctx.orgId)
          : ctx.orgId;

        let query = scopedInvoices(scopeOrgId).active();
        if (input.status) {
          query = query.filter((inv) => inv.status === input.status);
        }
        if (input.since) {
          const since = input.since;
          query = query.filter((inv) => isoDate(inv.createdAt) >= since);
        }
        const rows = query.take(Number.MAX_SAFE_INTEGER);

        const totalAmount = rows.reduce(
          (sum, inv) => sum + Number(inv.total),
          0,
        );

        const byStatus = rows.reduce<Record<string, number>>((acc, inv) => {
          acc[inv.status] = (acc[inv.status] ?? 0) + 1;
          return acc;
        }, {});

        const oldestUnpaidDueDate = rows
          .filter(
            (inv): inv is Invoice & { dueAt: string } =>
              inv.status !== 'paid' && inv.dueAt !== null,
          )
          .reduce<string | null>(
            (oldest, inv) =>
              oldest === null || inv.dueAt < oldest ? inv.dueAt : oldest,
            null,
          );

        return {
          count: rows.length,
          totalAmount,
          byStatus,
          oldestUnpaidDueDate:
            oldestUnpaidDueDate === null ? null : isoDate(oldestUnpaidDueDate),
        };
      } catch {
        return { error: 'stats_unavailable' as const };
      }
    },
  }),
});

export type InvoiceTools = ReturnType<typeof buildInvoiceTools>;

// The client imports only this — the typed message whose tool parts are backed
// by the real tool map.
export type InvoiceUIMessage = UIMessage<
  unknown,
  never,
  InferUITools<InvoiceTools>
>;

buildInvoiceTools is a factory, not a constant. It takes ctx and returns the tool map, so the orgId the route passes in becomes a closed-over value for this one request.

import 'server-only';

import { type InferUITools, tool, type UIMessage } from 'ai';
import { z } from 'zod';
import { scopedInvoices } from '@/lib/invoices/scoped-query';
import { getFlag } from '@/server/inspector-flags';
import type { Invoice } from '@/server/types';

const isoDate = (iso: string): string => iso.slice(0, 10);

// The single read-only tool. `execute` closes over `ctx.orgId` from the server
// auth boundary — the model NEVER passes `orgId` (it is not in `inputSchema`), so
// a forged tool-call argument cannot cross tenants. The `MODEL_FROM_INPUT_ORGID`
// inspector flag is the only path that reads `orgId` from model input; it exists
// solely to make that leak visible by hand (default off → always `ctx.orgId`).
export const buildInvoiceTools = (ctx: { orgId: string }) => ({
  getInvoiceStats: tool({
    description:
      'Return aggregate invoice statistics for the current organization. Use this for any question that needs counts, totals, or status breakdowns of invoices.',
    inputSchema: z.strictObject({
      status: z.enum(['draft', 'sent', 'paid', 'overdue']).optional(),
      since: z.iso.date().optional(),
    }),
    outputSchema: z.strictObject({
      count: z.number().int(),
      totalAmount: z.number(),
      byStatus: z.record(z.string(), z.number().int()),
      oldestUnpaidDueDate: z.iso.date().nullable(),
    }),
    execute: async (input) => {
      try {
        if (getFlag('FORCE_TOOL_ERROR')) {
          return { error: 'stats_unavailable' as const };
        }

        const scopeOrgId = getFlag('MODEL_FROM_INPUT_ORGID')
          ? ((input as { orgId?: string }).orgId ?? ctx.orgId)
          : ctx.orgId;

        let query = scopedInvoices(scopeOrgId).active();
        if (input.status) {
          query = query.filter((inv) => inv.status === input.status);
        }
        if (input.since) {
          const since = input.since;
          query = query.filter((inv) => isoDate(inv.createdAt) >= since);
        }
        const rows = query.take(Number.MAX_SAFE_INTEGER);

        const totalAmount = rows.reduce(
          (sum, inv) => sum + Number(inv.total),
          0,
        );

        const byStatus = rows.reduce<Record<string, number>>((acc, inv) => {
          acc[inv.status] = (acc[inv.status] ?? 0) + 1;
          return acc;
        }, {});

        const oldestUnpaidDueDate = rows
          .filter(
            (inv): inv is Invoice & { dueAt: string } =>
              inv.status !== 'paid' && inv.dueAt !== null,
          )
          .reduce<string | null>(
            (oldest, inv) =>
              oldest === null || inv.dueAt < oldest ? inv.dueAt : oldest,
            null,
          );

        return {
          count: rows.length,
          totalAmount,
          byStatus,
          oldestUnpaidDueDate:
            oldestUnpaidDueDate === null ? null : isoDate(oldestUnpaidDueDate),
        };
      } catch {
        return { error: 'stats_unavailable' as const };
      }
    },
  }),
});

export type InvoiceTools = ReturnType<typeof buildInvoiceTools>;

// The client imports only this — the typed message whose tool parts are backed
// by the real tool map.
export type InvoiceUIMessage = UIMessage<
  unknown,
  never,
  InferUITools<InvoiceTools>
>;

The inputSchema lists the fields the model may fill in: an optional status from a fixed enum and an optional since date. Note what is absent — there is no orgId. This omission is the entire tenancy guarantee. Anything in inputSchema is something the model controls, so a prompt-injected “set orgId to org-globex” would be a cross-tenant read; with no field to carry an org, the leak is not unlikely but unrepresentable. strictObject rejects any unknown key the model tries to smuggle in, the same Zod 4 discipline you used on Server Action inputs.

import 'server-only';

import { type InferUITools, tool, type UIMessage } from 'ai';
import { z } from 'zod';
import { scopedInvoices } from '@/lib/invoices/scoped-query';
import { getFlag } from '@/server/inspector-flags';
import type { Invoice } from '@/server/types';

const isoDate = (iso: string): string => iso.slice(0, 10);

// The single read-only tool. `execute` closes over `ctx.orgId` from the server
// auth boundary — the model NEVER passes `orgId` (it is not in `inputSchema`), so
// a forged tool-call argument cannot cross tenants. The `MODEL_FROM_INPUT_ORGID`
// inspector flag is the only path that reads `orgId` from model input; it exists
// solely to make that leak visible by hand (default off → always `ctx.orgId`).
export const buildInvoiceTools = (ctx: { orgId: string }) => ({
  getInvoiceStats: tool({
    description:
      'Return aggregate invoice statistics for the current organization. Use this for any question that needs counts, totals, or status breakdowns of invoices.',
    inputSchema: z.strictObject({
      status: z.enum(['draft', 'sent', 'paid', 'overdue']).optional(),
      since: z.iso.date().optional(),
    }),
    outputSchema: z.strictObject({
      count: z.number().int(),
      totalAmount: z.number(),
      byStatus: z.record(z.string(), z.number().int()),
      oldestUnpaidDueDate: z.iso.date().nullable(),
    }),
    execute: async (input) => {
      try {
        if (getFlag('FORCE_TOOL_ERROR')) {
          return { error: 'stats_unavailable' as const };
        }

        const scopeOrgId = getFlag('MODEL_FROM_INPUT_ORGID')
          ? ((input as { orgId?: string }).orgId ?? ctx.orgId)
          : ctx.orgId;

        let query = scopedInvoices(scopeOrgId).active();
        if (input.status) {
          query = query.filter((inv) => inv.status === input.status);
        }
        if (input.since) {
          const since = input.since;
          query = query.filter((inv) => isoDate(inv.createdAt) >= since);
        }
        const rows = query.take(Number.MAX_SAFE_INTEGER);

        const totalAmount = rows.reduce(
          (sum, inv) => sum + Number(inv.total),
          0,
        );

        const byStatus = rows.reduce<Record<string, number>>((acc, inv) => {
          acc[inv.status] = (acc[inv.status] ?? 0) + 1;
          return acc;
        }, {});

        const oldestUnpaidDueDate = rows
          .filter(
            (inv): inv is Invoice & { dueAt: string } =>
              inv.status !== 'paid' && inv.dueAt !== null,
          )
          .reduce<string | null>(
            (oldest, inv) =>
              oldest === null || inv.dueAt < oldest ? inv.dueAt : oldest,
            null,
          );

        return {
          count: rows.length,
          totalAmount,
          byStatus,
          oldestUnpaidDueDate:
            oldestUnpaidDueDate === null ? null : isoDate(oldestUnpaidDueDate),
        };
      } catch {
        return { error: 'stats_unavailable' as const };
      }
    },
  }),
});

export type InvoiceTools = ReturnType<typeof buildInvoiceTools>;

// The client imports only this — the typed message whose tool parts are backed
// by the real tool map.
export type InvoiceUIMessage = UIMessage<
  unknown,
  never,
  InferUITools<InvoiceTools>
>;

The outputSchema is the minimal aggregate that goes back to the model: a count, a totalAmount, a byStatus map, and one date. Projecting here, at the boundary the data leaves your control, caps input-token growth across loop steps and keeps row-level customer data off the wire entirely.

import 'server-only';

import { type InferUITools, tool, type UIMessage } from 'ai';
import { z } from 'zod';
import { scopedInvoices } from '@/lib/invoices/scoped-query';
import { getFlag } from '@/server/inspector-flags';
import type { Invoice } from '@/server/types';

const isoDate = (iso: string): string => iso.slice(0, 10);

// The single read-only tool. `execute` closes over `ctx.orgId` from the server
// auth boundary — the model NEVER passes `orgId` (it is not in `inputSchema`), so
// a forged tool-call argument cannot cross tenants. The `MODEL_FROM_INPUT_ORGID`
// inspector flag is the only path that reads `orgId` from model input; it exists
// solely to make that leak visible by hand (default off → always `ctx.orgId`).
export const buildInvoiceTools = (ctx: { orgId: string }) => ({
  getInvoiceStats: tool({
    description:
      'Return aggregate invoice statistics for the current organization. Use this for any question that needs counts, totals, or status breakdowns of invoices.',
    inputSchema: z.strictObject({
      status: z.enum(['draft', 'sent', 'paid', 'overdue']).optional(),
      since: z.iso.date().optional(),
    }),
    outputSchema: z.strictObject({
      count: z.number().int(),
      totalAmount: z.number(),
      byStatus: z.record(z.string(), z.number().int()),
      oldestUnpaidDueDate: z.iso.date().nullable(),
    }),
    execute: async (input) => {
      try {
        if (getFlag('FORCE_TOOL_ERROR')) {
          return { error: 'stats_unavailable' as const };
        }

        const scopeOrgId = getFlag('MODEL_FROM_INPUT_ORGID')
          ? ((input as { orgId?: string }).orgId ?? ctx.orgId)
          : ctx.orgId;

        let query = scopedInvoices(scopeOrgId).active();
        if (input.status) {
          query = query.filter((inv) => inv.status === input.status);
        }
        if (input.since) {
          const since = input.since;
          query = query.filter((inv) => isoDate(inv.createdAt) >= since);
        }
        const rows = query.take(Number.MAX_SAFE_INTEGER);

        const totalAmount = rows.reduce(
          (sum, inv) => sum + Number(inv.total),
          0,
        );

        const byStatus = rows.reduce<Record<string, number>>((acc, inv) => {
          acc[inv.status] = (acc[inv.status] ?? 0) + 1;
          return acc;
        }, {});

        const oldestUnpaidDueDate = rows
          .filter(
            (inv): inv is Invoice & { dueAt: string } =>
              inv.status !== 'paid' && inv.dueAt !== null,
          )
          .reduce<string | null>(
            (oldest, inv) =>
              oldest === null || inv.dueAt < oldest ? inv.dueAt : oldest,
            null,
          );

        return {
          count: rows.length,
          totalAmount,
          byStatus,
          oldestUnpaidDueDate:
            oldestUnpaidDueDate === null ? null : isoDate(oldestUnpaidDueDate),
        };
      } catch {
        return { error: 'stats_unavailable' as const };
      }
    },
  }),
});

export type InvoiceTools = ReturnType<typeof buildInvoiceTools>;

// The client imports only this — the typed message whose tool parts are backed
// by the real tool map.
export type InvoiceUIMessage = UIMessage<
  unknown,
  never,
  InferUITools<InvoiceTools>
>;

The scope decision. With the flag off, scopeOrgId is ctx.orgId, full stop. Only MODEL_FROM_INPUT_ORGID reads orgId off the model’s input, and it exists to demonstrate the leak — real code has no such branch.

import 'server-only';

import { type InferUITools, tool, type UIMessage } from 'ai';
import { z } from 'zod';
import { scopedInvoices } from '@/lib/invoices/scoped-query';
import { getFlag } from '@/server/inspector-flags';
import type { Invoice } from '@/server/types';

const isoDate = (iso: string): string => iso.slice(0, 10);

// The single read-only tool. `execute` closes over `ctx.orgId` from the server
// auth boundary — the model NEVER passes `orgId` (it is not in `inputSchema`), so
// a forged tool-call argument cannot cross tenants. The `MODEL_FROM_INPUT_ORGID`
// inspector flag is the only path that reads `orgId` from model input; it exists
// solely to make that leak visible by hand (default off → always `ctx.orgId`).
export const buildInvoiceTools = (ctx: { orgId: string }) => ({
  getInvoiceStats: tool({
    description:
      'Return aggregate invoice statistics for the current organization. Use this for any question that needs counts, totals, or status breakdowns of invoices.',
    inputSchema: z.strictObject({
      status: z.enum(['draft', 'sent', 'paid', 'overdue']).optional(),
      since: z.iso.date().optional(),
    }),
    outputSchema: z.strictObject({
      count: z.number().int(),
      totalAmount: z.number(),
      byStatus: z.record(z.string(), z.number().int()),
      oldestUnpaidDueDate: z.iso.date().nullable(),
    }),
    execute: async (input) => {
      try {
        if (getFlag('FORCE_TOOL_ERROR')) {
          return { error: 'stats_unavailable' as const };
        }

        const scopeOrgId = getFlag('MODEL_FROM_INPUT_ORGID')
          ? ((input as { orgId?: string }).orgId ?? ctx.orgId)
          : ctx.orgId;

        let query = scopedInvoices(scopeOrgId).active();
        if (input.status) {
          query = query.filter((inv) => inv.status === input.status);
        }
        if (input.since) {
          const since = input.since;
          query = query.filter((inv) => isoDate(inv.createdAt) >= since);
        }
        const rows = query.take(Number.MAX_SAFE_INTEGER);

        const totalAmount = rows.reduce(
          (sum, inv) => sum + Number(inv.total),
          0,
        );

        const byStatus = rows.reduce<Record<string, number>>((acc, inv) => {
          acc[inv.status] = (acc[inv.status] ?? 0) + 1;
          return acc;
        }, {});

        const oldestUnpaidDueDate = rows
          .filter(
            (inv): inv is Invoice & { dueAt: string } =>
              inv.status !== 'paid' && inv.dueAt !== null,
          )
          .reduce<string | null>(
            (oldest, inv) =>
              oldest === null || inv.dueAt < oldest ? inv.dueAt : oldest,
            null,
          );

        return {
          count: rows.length,
          totalAmount,
          byStatus,
          oldestUnpaidDueDate:
            oldestUnpaidDueDate === null ? null : isoDate(oldestUnpaidDueDate),
        };
      } catch {
        return { error: 'stats_unavailable' as const };
      }
    },
  }),
});

export type InvoiceTools = ReturnType<typeof buildInvoiceTools>;

// The client imports only this — the typed message whose tool parts are backed
// by the real tool map.
export type InvoiceUIMessage = UIMessage<
  unknown,
  never,
  InferUITools<InvoiceTools>
>;

The read rides scopedInvoices(scopeOrgId).active(), the same tenant-scoped builder the list view uses, with the model’s optional status and since composed on as filters. The since comparison works on the YYYY-MM-DD slice, so a lexicographic string compare is a correct date compare.

import 'server-only';

import { type InferUITools, tool, type UIMessage } from 'ai';
import { z } from 'zod';
import { scopedInvoices } from '@/lib/invoices/scoped-query';
import { getFlag } from '@/server/inspector-flags';
import type { Invoice } from '@/server/types';

const isoDate = (iso: string): string => iso.slice(0, 10);

// The single read-only tool. `execute` closes over `ctx.orgId` from the server
// auth boundary — the model NEVER passes `orgId` (it is not in `inputSchema`), so
// a forged tool-call argument cannot cross tenants. The `MODEL_FROM_INPUT_ORGID`
// inspector flag is the only path that reads `orgId` from model input; it exists
// solely to make that leak visible by hand (default off → always `ctx.orgId`).
export const buildInvoiceTools = (ctx: { orgId: string }) => ({
  getInvoiceStats: tool({
    description:
      'Return aggregate invoice statistics for the current organization. Use this for any question that needs counts, totals, or status breakdowns of invoices.',
    inputSchema: z.strictObject({
      status: z.enum(['draft', 'sent', 'paid', 'overdue']).optional(),
      since: z.iso.date().optional(),
    }),
    outputSchema: z.strictObject({
      count: z.number().int(),
      totalAmount: z.number(),
      byStatus: z.record(z.string(), z.number().int()),
      oldestUnpaidDueDate: z.iso.date().nullable(),
    }),
    execute: async (input) => {
      try {
        if (getFlag('FORCE_TOOL_ERROR')) {
          return { error: 'stats_unavailable' as const };
        }

        const scopeOrgId = getFlag('MODEL_FROM_INPUT_ORGID')
          ? ((input as { orgId?: string }).orgId ?? ctx.orgId)
          : ctx.orgId;

        let query = scopedInvoices(scopeOrgId).active();
        if (input.status) {
          query = query.filter((inv) => inv.status === input.status);
        }
        if (input.since) {
          const since = input.since;
          query = query.filter((inv) => isoDate(inv.createdAt) >= since);
        }
        const rows = query.take(Number.MAX_SAFE_INTEGER);

        const totalAmount = rows.reduce(
          (sum, inv) => sum + Number(inv.total),
          0,
        );

        const byStatus = rows.reduce<Record<string, number>>((acc, inv) => {
          acc[inv.status] = (acc[inv.status] ?? 0) + 1;
          return acc;
        }, {});

        const oldestUnpaidDueDate = rows
          .filter(
            (inv): inv is Invoice & { dueAt: string } =>
              inv.status !== 'paid' && inv.dueAt !== null,
          )
          .reduce<string | null>(
            (oldest, inv) =>
              oldest === null || inv.dueAt < oldest ? inv.dueAt : oldest,
            null,
          );

        return {
          count: rows.length,
          totalAmount,
          byStatus,
          oldestUnpaidDueDate:
            oldestUnpaidDueDate === null ? null : isoDate(oldestUnpaidDueDate),
        };
      } catch {
        return { error: 'stats_unavailable' as const };
      }
    },
  }),
});

export type InvoiceTools = ReturnType<typeof buildInvoiceTools>;

// The client imports only this — the typed message whose tool parts are backed
// by the real tool map.
export type InvoiceUIMessage = UIMessage<
  unknown,
  never,
  InferUITools<InvoiceTools>
>;

query.take(Number.MAX_SAFE_INTEGER). The scoped-query builder is keyset-shaped — it hands back a page, not the whole set — so an explicit huge take is how you materialize every row for an aggregate that has to see all of them.

import 'server-only';

import { type InferUITools, tool, type UIMessage } from 'ai';
import { z } from 'zod';
import { scopedInvoices } from '@/lib/invoices/scoped-query';
import { getFlag } from '@/server/inspector-flags';
import type { Invoice } from '@/server/types';

const isoDate = (iso: string): string => iso.slice(0, 10);

// The single read-only tool. `execute` closes over `ctx.orgId` from the server
// auth boundary — the model NEVER passes `orgId` (it is not in `inputSchema`), so
// a forged tool-call argument cannot cross tenants. The `MODEL_FROM_INPUT_ORGID`
// inspector flag is the only path that reads `orgId` from model input; it exists
// solely to make that leak visible by hand (default off → always `ctx.orgId`).
export const buildInvoiceTools = (ctx: { orgId: string }) => ({
  getInvoiceStats: tool({
    description:
      'Return aggregate invoice statistics for the current organization. Use this for any question that needs counts, totals, or status breakdowns of invoices.',
    inputSchema: z.strictObject({
      status: z.enum(['draft', 'sent', 'paid', 'overdue']).optional(),
      since: z.iso.date().optional(),
    }),
    outputSchema: z.strictObject({
      count: z.number().int(),
      totalAmount: z.number(),
      byStatus: z.record(z.string(), z.number().int()),
      oldestUnpaidDueDate: z.iso.date().nullable(),
    }),
    execute: async (input) => {
      try {
        if (getFlag('FORCE_TOOL_ERROR')) {
          return { error: 'stats_unavailable' as const };
        }

        const scopeOrgId = getFlag('MODEL_FROM_INPUT_ORGID')
          ? ((input as { orgId?: string }).orgId ?? ctx.orgId)
          : ctx.orgId;

        let query = scopedInvoices(scopeOrgId).active();
        if (input.status) {
          query = query.filter((inv) => inv.status === input.status);
        }
        if (input.since) {
          const since = input.since;
          query = query.filter((inv) => isoDate(inv.createdAt) >= since);
        }
        const rows = query.take(Number.MAX_SAFE_INTEGER);

        const totalAmount = rows.reduce(
          (sum, inv) => sum + Number(inv.total),
          0,
        );

        const byStatus = rows.reduce<Record<string, number>>((acc, inv) => {
          acc[inv.status] = (acc[inv.status] ?? 0) + 1;
          return acc;
        }, {});

        const oldestUnpaidDueDate = rows
          .filter(
            (inv): inv is Invoice & { dueAt: string } =>
              inv.status !== 'paid' && inv.dueAt !== null,
          )
          .reduce<string | null>(
            (oldest, inv) =>
              oldest === null || inv.dueAt < oldest ? inv.dueAt : oldest,
            null,
          );

        return {
          count: rows.length,
          totalAmount,
          byStatus,
          oldestUnpaidDueDate:
            oldestUnpaidDueDate === null ? null : isoDate(oldestUnpaidDueDate),
        };
      } catch {
        return { error: 'stats_unavailable' as const };
      }
    },
  }),
});

export type InvoiceTools = ReturnType<typeof buildInvoiceTools>;

// The client imports only this — the typed message whose tool parts are backed
// by the real tool map.
export type InvoiceUIMessage = UIMessage<
  unknown,
  never,
  InferUITools<InvoiceTools>
>;

The reduces. totalAmount sums the row totals; byStatus accumulates counts with a typed accumulator; oldestUnpaidDueDate narrows each row with a type guard (inv is Invoice & { dueAt: string }) so the comparison sees a non-null dueAt, and folds to null when there are no unpaid dated rows.

import 'server-only';

import { type InferUITools, tool, type UIMessage } from 'ai';
import { z } from 'zod';
import { scopedInvoices } from '@/lib/invoices/scoped-query';
import { getFlag } from '@/server/inspector-flags';
import type { Invoice } from '@/server/types';

const isoDate = (iso: string): string => iso.slice(0, 10);

// The single read-only tool. `execute` closes over `ctx.orgId` from the server
// auth boundary — the model NEVER passes `orgId` (it is not in `inputSchema`), so
// a forged tool-call argument cannot cross tenants. The `MODEL_FROM_INPUT_ORGID`
// inspector flag is the only path that reads `orgId` from model input; it exists
// solely to make that leak visible by hand (default off → always `ctx.orgId`).
export const buildInvoiceTools = (ctx: { orgId: string }) => ({
  getInvoiceStats: tool({
    description:
      'Return aggregate invoice statistics for the current organization. Use this for any question that needs counts, totals, or status breakdowns of invoices.',
    inputSchema: z.strictObject({
      status: z.enum(['draft', 'sent', 'paid', 'overdue']).optional(),
      since: z.iso.date().optional(),
    }),
    outputSchema: z.strictObject({
      count: z.number().int(),
      totalAmount: z.number(),
      byStatus: z.record(z.string(), z.number().int()),
      oldestUnpaidDueDate: z.iso.date().nullable(),
    }),
    execute: async (input) => {
      try {
        if (getFlag('FORCE_TOOL_ERROR')) {
          return { error: 'stats_unavailable' as const };
        }

        const scopeOrgId = getFlag('MODEL_FROM_INPUT_ORGID')
          ? ((input as { orgId?: string }).orgId ?? ctx.orgId)
          : ctx.orgId;

        let query = scopedInvoices(scopeOrgId).active();
        if (input.status) {
          query = query.filter((inv) => inv.status === input.status);
        }
        if (input.since) {
          const since = input.since;
          query = query.filter((inv) => isoDate(inv.createdAt) >= since);
        }
        const rows = query.take(Number.MAX_SAFE_INTEGER);

        const totalAmount = rows.reduce(
          (sum, inv) => sum + Number(inv.total),
          0,
        );

        const byStatus = rows.reduce<Record<string, number>>((acc, inv) => {
          acc[inv.status] = (acc[inv.status] ?? 0) + 1;
          return acc;
        }, {});

        const oldestUnpaidDueDate = rows
          .filter(
            (inv): inv is Invoice & { dueAt: string } =>
              inv.status !== 'paid' && inv.dueAt !== null,
          )
          .reduce<string | null>(
            (oldest, inv) =>
              oldest === null || inv.dueAt < oldest ? inv.dueAt : oldest,
            null,
          );

        return {
          count: rows.length,
          totalAmount,
          byStatus,
          oldestUnpaidDueDate:
            oldestUnpaidDueDate === null ? null : isoDate(oldestUnpaidDueDate),
        };
      } catch {
        return { error: 'stats_unavailable' as const };
      }
    },
  }),
});

export type InvoiceTools = ReturnType<typeof buildInvoiceTools>;

// The client imports only this — the typed message whose tool parts are backed
// by the real tool map.
export type InvoiceUIMessage = UIMessage<
  unknown,
  never,
  InferUITools<InvoiceTools>
>;

The error boundary. FORCE_TOOL_ERROR returns the error shape immediately, and a catch returns it if the read throws. { error: 'stats_unavailable' as const } widens the return union to “aggregate OR error”, which the SDK accepts because the object serializes and the model can read it and apologize. This is “return don’t throw” in one line.

import 'server-only';

import { type InferUITools, tool, type UIMessage } from 'ai';
import { z } from 'zod';
import { scopedInvoices } from '@/lib/invoices/scoped-query';
import { getFlag } from '@/server/inspector-flags';
import type { Invoice } from '@/server/types';

const isoDate = (iso: string): string => iso.slice(0, 10);

// The single read-only tool. `execute` closes over `ctx.orgId` from the server
// auth boundary — the model NEVER passes `orgId` (it is not in `inputSchema`), so
// a forged tool-call argument cannot cross tenants. The `MODEL_FROM_INPUT_ORGID`
// inspector flag is the only path that reads `orgId` from model input; it exists
// solely to make that leak visible by hand (default off → always `ctx.orgId`).
export const buildInvoiceTools = (ctx: { orgId: string }) => ({
  getInvoiceStats: tool({
    description:
      'Return aggregate invoice statistics for the current organization. Use this for any question that needs counts, totals, or status breakdowns of invoices.',
    inputSchema: z.strictObject({
      status: z.enum(['draft', 'sent', 'paid', 'overdue']).optional(),
      since: z.iso.date().optional(),
    }),
    outputSchema: z.strictObject({
      count: z.number().int(),
      totalAmount: z.number(),
      byStatus: z.record(z.string(), z.number().int()),
      oldestUnpaidDueDate: z.iso.date().nullable(),
    }),
    execute: async (input) => {
      try {
        if (getFlag('FORCE_TOOL_ERROR')) {
          return { error: 'stats_unavailable' as const };
        }

        const scopeOrgId = getFlag('MODEL_FROM_INPUT_ORGID')
          ? ((input as { orgId?: string }).orgId ?? ctx.orgId)
          : ctx.orgId;

        let query = scopedInvoices(scopeOrgId).active();
        if (input.status) {
          query = query.filter((inv) => inv.status === input.status);
        }
        if (input.since) {
          const since = input.since;
          query = query.filter((inv) => isoDate(inv.createdAt) >= since);
        }
        const rows = query.take(Number.MAX_SAFE_INTEGER);

        const totalAmount = rows.reduce(
          (sum, inv) => sum + Number(inv.total),
          0,
        );

        const byStatus = rows.reduce<Record<string, number>>((acc, inv) => {
          acc[inv.status] = (acc[inv.status] ?? 0) + 1;
          return acc;
        }, {});

        const oldestUnpaidDueDate = rows
          .filter(
            (inv): inv is Invoice & { dueAt: string } =>
              inv.status !== 'paid' && inv.dueAt !== null,
          )
          .reduce<string | null>(
            (oldest, inv) =>
              oldest === null || inv.dueAt < oldest ? inv.dueAt : oldest,
            null,
          );

        return {
          count: rows.length,
          totalAmount,
          byStatus,
          oldestUnpaidDueDate:
            oldestUnpaidDueDate === null ? null : isoDate(oldestUnpaidDueDate),
        };
      } catch {
        return { error: 'stats_unavailable' as const };
      }
    },
  }),
});

export type InvoiceTools = ReturnType<typeof buildInvoiceTools>;

// The client imports only this — the typed message whose tool parts are backed
// by the real tool map.
export type InvoiceUIMessage = UIMessage<
  unknown,
  never,
  InferUITools<InvoiceTools>
>;

The type exports. InvoiceTools is the inferred shape of the tool map; InvoiceUIMessage threads it through InferUITools so the client gets a fully typed message. The client imports only InvoiceUIMessage, never the server-only tool.

← Previous 1 / 1 Next →

Tool calling, the inputSchema / outputSchema contract, server-side execute, and the agentic loop come from the tools lesson of chapter 107; the typed-UIMessage-via-InferUITools mechanism from the generative-UI lesson; the scopedInvoices builder and its keyset shape from the scoped-reads lesson.

Wiring it into the route

The previous lesson’s route streamed text-only answers because no tools were passed. Adding the tool is two changes to that handler: build the tool map per request and pass it to streamText, and add an onStepFinish that writes one step-audit row per loop step.

Before (text-only)

const result = streamText({
  model: chatModel,
  system: invoiceQAPrompt({ orgName }),
  messages: convertToModelMessages(input.messages as InvoiceUIMessage[]),
  stopWhen: stepCountIs(5),
  maxOutputTokens: 1024,
  onFinish: ({ usage, finishReason }) =>
    writeLlmFinishEvent({
      userId: ctx.userId,
      orgId: ctx.orgId,
      finishReason,
      usage,
    }),
  onError: ({ error }) => {
    console.error('[chat] stream error', { code: 'stream_error' });
    void error;
  },
});

The previous lesson’s call. No tools, so the loop never branches and the model can only emit text; the stopWhen cap is set but never exercised. onFinish writes the single per-turn finish row.

After (tool-grounded)

const result = streamText({
  model: chatModel,
  system: invoiceQAPrompt({ orgName }),
  messages: convertToModelMessages(input.messages as InvoiceUIMessage[]),
  tools,
  stopWhen: stepCountIs(5),
  maxOutputTokens: 1024,
  onStepFinish: async ({ usage, toolCalls, finishReason }) => {
    await writeLlmStepEvent({
      userId: ctx.userId,
      orgId: ctx.orgId,
      finishReason,
      usage,
      toolCalls,
    });
  },
  onFinish: ({ usage, finishReason }) =>
    writeLlmFinishEvent({
      userId: ctx.userId,
      orgId: ctx.orgId,
      finishReason,
      usage,
    }),
  onError: ({ error }) => {
    console.error('[chat] stream error', { code: 'stream_error' });
    void error;
  },
});

This lesson’s call. With tools passed, the loop can call getInvoiceStats and feed the result back. onStepFinish fires once per step and writes the per-step audit row; onFinish is unchanged.

The tools value is built just above the streamText call, inside the handler:

const orgName = org?.name ?? 'your organization';

const tools = buildInvoiceTools({ orgId: ctx.orgId });

const result = streamText({

The full route at the end of this lesson, for reference:

import { convertToModelMessages, stepCountIs, streamText } from 'ai';
import { z } from 'zod';
import { authedRoute } from '@/lib/authed-route';
import { writeLlmFinishEvent, writeLlmStepEvent } from '@/lib/llm/audit';
import { chatModel } from '@/lib/llm/models';
import { invoiceQAPrompt } from '@/lib/llm/prompts';
import { buildInvoiceTools, type InvoiceUIMessage } from '@/lib/llm/tools';

export const POST = authedRoute(
  'member',
  z.strictObject({ messages: z.array(z.unknown()) }),
  async (input, ctx) => {
    const org = await ctx.db.query.organization.findFirst({
      where: (o) => o.id === ctx.orgId,
    });
    const orgName = org?.name ?? 'your organization';

    const tools = buildInvoiceTools({ orgId: ctx.orgId });

    const result = streamText({
      model: chatModel,
      system: invoiceQAPrompt({ orgName }),
      messages: convertToModelMessages(input.messages as InvoiceUIMessage[]),
      tools,
      stopWhen: stepCountIs(5),
      maxOutputTokens: 1024,
      onStepFinish: async ({ usage, toolCalls, finishReason }) => {
        await writeLlmStepEvent({
          userId: ctx.userId,
          orgId: ctx.orgId,
          finishReason,
          usage,
          toolCalls,
        });
      },
      onFinish: ({ usage, finishReason }) =>
        writeLlmFinishEvent({
          userId: ctx.userId,
          orgId: ctx.orgId,
          finishReason,
          usage,
        }),
      onError: ({ error }) => {
        console.error('[chat] stream error', { code: 'stream_error' });
        void error;
      },
    });

    return result.toUIMessageStreamResponse();
  },
);

Caution

onStepFinish writes the step-audit row and only that. The next lesson adds an addUsage call alongside it for the per-user daily token counter and wraps this authedRoute(...) export in withLlmQuota(...); leave both out for now. The repo’s finished solution already carries them, but that is the next lesson’s work, not this one’s.

Build tools inside the handler, not at module load: at import time there is no request to draw ctx.orgId from, and a module-level value would be shared across every request. Building it per request is what makes each closure capture this request’s authenticated org.

The agentic-loop primitives (stopWhen, stepCountIs, the loop) come from the tools lesson of chapter 107; the append-only one-row-per-event audit discipline from the append-only audit log lesson.

AI SDK — Tool Calling

ai-sdk.dev

The tool() helper, inputSchema validation, multi-step calls, and how execute errors become tool-error parts — the exact mechanics you're wiring here.

AI SDK — tool() reference

ai-sdk.dev

The full signature, including the outputSchema field the guide omits and how it infers the execute input type.

OWASP Top 10 for LLM Applications

owasp.org

LLM01 prompt injection and LLM08 excessive agency — the named risks behind this lesson's rule that the model is untrusted input.

Moment of truth

This project has no per-lesson test suite, so verification is the type-and-build health check plus the by-hand checks below. Run:

pnpm verify

It runs Biome’s CI lint, tsc --noEmit, and a next build with SKIP_ENV_VALIDATION=true. A clean typecheck and a successful build confirm the slice compiles and the types line up; the behavior that matters here can’t be asserted without a live model, so the checks below confirm it does the right thing.

The live checks need AI_GATEWAY_API_KEY in your .env (the chat makes a real model call). Act as member-A in org-acme to start, and use the inspector controls named in parentheses.

Asking “what’s our total paid this month?” returns a total matching a reduce over org-acme’s active paid rows, and the assistant text bubble cites it. (If the model answers with no tool-getInvoiceStats part at all, the fix is to sharpen the system prompt, not the code — the prompt is the lever for instruction-following.)

untested

With MODEL_FROM_INPUT_ORGID off, asking the model to use orgId = org-globex still yields org-acme’s numbers. Then flip the flag, switch the identity to org-globex, and repeat — now you see org-globex’s numbers leak through the model’s argument. This is the worst class of LLM-in-SaaS bug, made visible. Revert the flag.

untested

A recursion-prone prompt produces at most five tool-getInvoiceStats parts and a final message acknowledging the cap; removing stopWhen and repeating shows the loop running to the SDK default. Revert.

untested

With the “Force tool error” toggle on, a stats question shows the output-error state and a follow-up text answer asking you to rephrase, with no 500 in the network tab. Revert.

untested

After a multi-step conversation, the inspector’s llm_audit_events tail shows one 'llm.step' row per step plus one 'llm.finish' row, all scoped to the active org.

untested

On the first check: the smoke-test box prints raw text, so “the assistant cites the number” means the final text bubble names a value you confirm against the inspector’s row-count panel; the tool part itself you read in the network tab and the llm_audit_events tail.

When questions outgrow a fixed aggregate — “which customers mention a refund in their notes?” — the next reach is retrieval over embeddings, covered in the RAG lesson of chapter 107. It’s a different tool, but the closure-over-orgId rule carries over unchanged.