Database Sessions & Persistence

import { createTool } from "@iqai/adk";
import z from "zod";

export const counterTool = createTool({
	name: "increment_counter",
	description: "Increment a named counter",
	schema: z.object({
		counterName: z.string().describe("Name of the counter"),
		amount: z.number().default(1).describe("Amount to increment"),
	}),
	fn: ({ counterName, amount }, context) => {
		const counters = context.state.get("counters", {});
		const oldValue = counters[counterName] || 0;
		const newValue = oldValue + amount;
		const newCounters = { ...counters, [counterName]: newValue };
		context.state.set("counters", newCounters);
		return {
			counterName,
			oldValue,
			newValue,
			increment: amount,
		};
	},
});

export const saveCounterReportTool = createTool({
	name: "save_counter_report",
	description: "Save a report of all counter values to an artifact file",
	schema: z.object({
		filename: z.string().describe("Name of the file to save the report to"),
	}),
	fn: async ({ filename }, context) => {
		const counters = context.state.get("counters", {});
		const report = Object.entries(counters)
			.map(([name, value]) => `${name}: ${value}`)
			.join("\n");

		await context.saveArtifact(filename, {
			text: report || "No counters found",
		});

		return {
			success: true,
			filename,
			countersCount: Object.keys(counters).length,
		};
	},
});

import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import {
	AgentBuilder,
	createDatabaseSessionService,
	InMemoryArtifactService,
} from "@iqai/adk";
import { counterTool, saveCounterReportTool } from "./tools";

export async function getRootAgent() {
	const dbDir = path.join(os.tmpdir(), "adk-examples");
	if (!fs.existsSync(dbDir)) {
		fs.mkdirSync(dbDir, { recursive: true });
	}

	const sessionService = createDatabaseSessionService(
		getSqliteConnectionString("sessions"),
	);
	const artifactService = new InMemoryArtifactService();

	return await AgentBuilder.withModel(
		process.env.LLM_MODEL || "gemini-3-flash-preview",
	)
		.withTools(counterTool, saveCounterReportTool)
		.withSessionService(sessionService)
		.withArtifactService(artifactService)
		.withEventsCompaction({
			compactionInterval: 3,
			overlapSize: 1,
		})
		.build();
}

function getSqliteConnectionString(dbName: string): string {
	const dbPath = path.join(os.tmpdir(), "adk-examples", `${dbName}.db`);
	return `sqlite://${dbPath}`;
}

import { BaseSessionService, Session } from "@iqai/adk";
import { ask } from "../utils";
import { getRootAgent } from "./agents/agent";

async function main() {
	const { runner, sessionService, session } = await getRootAgent();

	// 1. Persistence: Counter stored in SQLite database
	console.log("\n📊 Persistence: Counter stored in DB...");
	await ask(runner, "Increment 'visits' counter by 1 and show its value");

	// 2. Artifacts: Save counter report to file
	console.log("\n📁 Artifacts: Saving counter report...");
	await ask(runner, 'Save current counter values to "counter-report.txt"');

	// 3. Rewind: Undo the last counter increment
	console.log("\n🔄 Rewind: Demonstrating time-travel...");

	// Show current state
	let currentSession = await sessionService.getSession(
		session.appName,
		session.userId,
		session.id,
	);
	console.log("📍 State before big increment:", currentSession.state);

	// Make a big increment that we'll want to undo
	await ask(runner, "Increment 'visits' counter by 100");

	// Get the invocationId of the increment operation
	const invocationId = await getInvocationIdWithStateDelta(
		sessionService,
		session,
	);

	// Show the state after the big increment
	currentSession = await sessionService.getSession(
		session.appName,
		session.userId,
		session.id,
	);
	console.log("📍 State after +100 increment:", currentSession.state);

	// Rewind to before the +100 increment - undoing that change!
	if (invocationId) {
		await runner.rewind({
			userId: session.userId,
			sessionId: session.id,
			rewindBeforeInvocationId: invocationId,
		});

		// Show the state after rewind - it should be back to what it was before
		currentSession = await sessionService.getSession(
			session.appName,
			session.userId,
			session.id,
		);
		console.log("⏪ State after rewind:", currentSession.state);
		console.log("✨ Successfully rewound - the +100 never happened!");

		// Verify with the agent (it will use the restored state)
		await ask(runner, "Increment 'visits' by 1 and show the new value");
	}

	// 4. Compaction: Multiple counter ops trigger auto-summarization
	console.log("\n📦 Compaction: Multiple counter operations...");
	const counters = ["logins", "clicks", "views", "shares"];
	for (const counter of counters) {
		await runner.ask(`Increment '${counter}' counter by 1`);
		await logCompactions(sessionService, session);
	}
	await ask(runner, "Show all counters");

	console.log("\n✅ All features demonstrated with counters!");
}

/**
 * Gets the invocationId of the last event that has a state delta.
 */
async function getInvocationIdWithStateDelta(
	sessionService: BaseSessionService,
	session: Session,
): Promise<string | undefined> {
	const currentSession = await sessionService.getSession(
		session.appName,
		session.userId,
		session.id,
	);

	// Search backwards through events to find the last state-changing event
	for (let i = currentSession.events.length - 1; i >= 0; i--) {
		const event = currentSession.events[i];
		if (
			event.actions?.stateDelta &&
			Object.keys(event.actions.stateDelta).length > 0
		) {
			return event.invocationId;
		}
	}

	return undefined;
}

async function logCompactions(
	sessionService: BaseSessionService,
	session: Session,
) {
	const updatedSession = await sessionService.getSession(
		session.appName,
		session.userId,
		session.id,
	);

	const compactions = updatedSession.events
		.filter((e) => e.actions?.compaction)
		.map((e) => e.actions.compaction);

	if (compactions.length === 0) return;

	for (const [i, c] of compactions.entries()) {
		const parts = c.compactedContent?.parts ?? [];
		const text = parts.map((p: any) => p.text).join("\n");
		console.log(`📦 Compaction ${i + 1}: ${text.substring(0, 100)}...\n`);
	}
}

main().catch(console.error);

node index.ts