The Database for Agentic Workloads

// npm install digest-fetch@3.1.1
// Set TIDB_CLOUD_PUBLIC_KEY, TIDB_CLOUD_PRIVATE_KEY, and TIDB_CLUSTER_ID.
// Use API keys with permission to create branches on the selected cluster.
// Save as database-branching.mjs and run with Node.js.
import DigestClient from "digest-fetch";
import { randomUUID } from "node:crypto";
import { setTimeout as wait } from "node:timers/promises";

for (const key of [
  "TIDB_CLOUD_PUBLIC_KEY", "TIDB_CLOUD_PRIVATE_KEY", "TIDB_CLUSTER_ID",
]) {
  if (!process.env[key]) throw new Error(`Set ${key} first.`);
}
const client = new DigestClient(
  process.env.TIDB_CLOUD_PUBLIC_KEY,
  process.env.TIDB_CLOUD_PRIVATE_KEY
);
const base = "https://serverless.tidbapi.com/v1beta1";
const path = `/clusters/${encodeURIComponent(process.env.TIDB_CLUSTER_ID)}/branches`;

async function request(path, method = "GET", body) {
  const response = await client.fetch(`${base}${path}`, {
    method,
    headers: { "Content-Type": "application/json", Accept: "application/json" },
    body: body === undefined ? undefined : JSON.stringify(body),
    signal: AbortSignal.timeout(30000),
    redirect: "error",
  });
  if (!response.ok) throw new Error(`Branch API: HTTP ${response.status}`);
  return response.json();
}

let branch = await request(path, "POST", {
  displayName: `agent-test-${randomUUID().slice(0, 8)}`,
});
if (!branch.branchId) throw new Error("Branch API returned no branch ID.");
const branchId = branch.branchId;
console.log("Branch requested:", branchId);

const deadline = Date.now() + 5 * 60 * 1000;
while (branch.state !== "ACTIVE") {
  if (branch.state === "DELETED" || Date.now() >= deadline) {
    throw new Error(`Check branch ${branchId} in the console before retrying.`);
  }
  await wait(5000);
  branch = await request(`${path}/${encodeURIComponent(branchId)}`);
}
console.log("Branch ready:", branchId);
// Connect using this branch's credentials from the TiDB Cloud console.
// Delete this test branch in the console when your experiment is finished.

// npm install @tidbcloud/serverless@0.3.0
// Set TIDB_URL to a TiDB Cloud connection URL with a database selected.
// Save as persistent-agent-state.mjs and run with Node.js.
import { connect } from "@tidbcloud/serverless";
import { randomUUID } from "node:crypto";

if (!process.env.TIDB_URL) throw new Error("Set TIDB_URL first.");
const db = connect({ url: process.env.TIDB_URL });
await db.execute(`CREATE TABLE IF NOT EXISTS agent_memory_demo (
  agent_id VARCHAR(64) PRIMARY KEY,
  state JSON NOT NULL,
  updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
)`);

async function saveAgentState(agentId, state) {
  await db.execute(
    `INSERT INTO agent_memory_demo (agent_id, state, updated_at)
     VALUES (?, ?, NOW())
     ON DUPLICATE KEY UPDATE state = VALUES(state), updated_at = NOW()`,
    [agentId, JSON.stringify(state)]
  );
}

async function restoreAgentState(agentId) {
  const [row] = await db.execute(
    "SELECT state FROM agent_memory_demo WHERE agent_id = ?",
    [agentId]
  );
  return row ? row.state : null; // JSON columns are already decoded.
}

const agentId = randomUUID(); // Keep this ID to resume the same agent later.
const previousState = (await restoreAgentState(agentId)) ?? {};
await saveAgentState(agentId, {
  ...previousState,
  step: "summarize",
  outputs: ["draft-1.md"],
  tokensUsed: 4210,
});
console.log("Agent ID:", agentId);
console.log("Restored state:", await restoreAgentState(agentId));

// npm install @tidbcloud/serverless@0.3.0
// Set TIDB_URL to a TiDB Cloud connection URL with a database selected.
// Save as concurrent-agent-transactions.mjs and run with Node.js.
import { connect } from "@tidbcloud/serverless";
import { randomUUID } from "node:crypto";

if (!process.env.TIDB_URL) throw new Error("Set TIDB_URL first.");
const db = connect({ url: process.env.TIDB_URL });
await db.execute(`CREATE TABLE IF NOT EXISTS agent_tasks_demo (
  id VARCHAR(64) PRIMARY KEY,
  status VARCHAR(20) NOT NULL,
  agent_id VARCHAR(64)
)`);

async function assignTask(taskId, agentId) {
  const tx = await db.begin(); // Each worker gets its own transaction.
  try {
    const [task] = await tx.execute(
      "SELECT status FROM agent_tasks_demo WHERE id = ? FOR UPDATE",
      [taskId]
    );
    if (!task || task.status !== "pending") {
      await tx.rollback();
      return { agentId, claimed: false };
    }
    await tx.execute(
      "UPDATE agent_tasks_demo SET status = ?, agent_id = ? WHERE id = ?",
      ["running", agentId, taskId]
    );
    await tx.commit();
    return { agentId, claimed: true };
  } catch (error) {
    await tx.rollback().catch(() => {});
    throw error;
  }
}

const taskId = randomUUID(); // Use a fresh task on every run.
await db.execute(
  "INSERT INTO agent_tasks_demo (id, status) VALUES (?, ?)",
  [taskId, "pending"]
);
console.log(await Promise.all([
  assignTask(taskId, "agent-A"),
  assignTask(taskId, "agent-B"),
])); // Exactly one agent claims the task; the other returns claimed: false.