docs + refactor: Session API documentation and AI SDK decoupling

docs/index.md

@@ -44,6 +44,7 @@
 - [Server-Driven Messages](./server-driven-messages.md) - Autonomous agent workflows: scheduled follow-ups, queue processing, webhooks, chained reasoning
 - TODO: [Using AI Models](./using-ai-models.md) - OpenAI, Anthropic, Workers AI, and other providers
 - TODO: [RAG (Retrieval Augmented Generation)](./rag.md) - Vector search with Vectorize
+- [Sessions (Experimental)](./sessions.md) - Persistent conversation storage with tree-structured messages, context blocks, compaction, and search
 - [Workspace (Experimental)](./workspace.md) - Durable virtual filesystem backed by SQLite + R2
 - [Codemode (Experimental)](./codemode.md) - LLM-generated executable code for tool orchestration
 - [Client Tools Continuation](./client-tools-continuation.md) - Handling tool calls across client/server

experimental/session-memory/src/server.ts

@@ -85,7 +85,7 @@ export class ChatAgent extends Agent<Env> {
     const result = streamText({
       model: this.getAI(),
       system: await this.session.freezeSystemPrompt(),
-      messages: await convertToModelMessages(truncated),
+      messages: await convertToModelMessages(truncated as UIMessage[]),
       tools: await this.session.tools(),
       stopWhen: stepCountIs(5)
     });
@@ -140,7 +140,7 @@ export class ChatAgent extends Agent<Env> {
 
   @callable()
   getMessages(): UIMessage[] {
-    return this.session.getHistory();
+    return this.session.getHistory() as UIMessage[];
   }
 
   @callable()

experimental/session-multichat/README.md

@@ -72,7 +72,6 @@ SessionManager.create(agent)
   .onCompaction(fn)                  // register compaction function
   .compactAfter(tokenThreshold)      // auto-compact threshold
   .withCachedPrompt(provider?)       // cache frozen system prompt
-  .maxContextMessages(count)         // limit context messages per session
 ```
 
 ## Session Lifecycle

experimental/session-multichat/src/server.ts

@@ -106,7 +106,7 @@ export class MultiSessionAgent extends Agent<Env> {
     const result = streamText({
       model: this.getAI(),
       system: await session.freezeSystemPrompt(),
-      messages: await convertToModelMessages(truncated),
+      messages: await convertToModelMessages(truncated as UIMessage[]),
       tools: { ...(await session.tools()), ...this.manager.tools() },
       stopWhen: stepCountIs(5)
     });
@@ -149,7 +149,7 @@ export class MultiSessionAgent extends Agent<Env> {
 
   @callable()
   getHistory(chatId: string): UIMessage[] {
-    return this.manager.getSession(chatId).getHistory();
+    return this.manager.getSession(chatId).getHistory() as UIMessage[];
   }
 
   @callable()

experimental/session-search/src/server.ts

@@ -92,7 +92,7 @@ export class SearchAgent extends Agent<Env> {
     const result = streamText({
       model: this.getAI(),
       system: await this.session.freezeSystemPrompt(),
-      messages: await convertToModelMessages(truncated),
+      messages: await convertToModelMessages(truncated as UIMessage[]),
       tools: await this.session.tools(),
       stopWhen: stepCountIs(5)
     });
@@ -135,7 +135,7 @@ export class SearchAgent extends Agent<Env> {
 
   @callable()
   getMessages(): UIMessage[] {
-    return this.session.getHistory();
+    return this.session.getHistory() as UIMessage[];
   }
 
   @callable()

packages/agents/src/experimental/memory/index.ts

@@ -8,7 +8,6 @@ export {
   Session,
   AgentSessionProvider,
   AgentContextProvider,
-  type MessageQueryOptions,
   type SessionProvider,
   type SearchResult,
   type StoredCompaction,

packages/agents/src/experimental/memory/session/context.ts

@@ -15,7 +15,8 @@
  * - SearchProvider (get+search+set?)  → searchable via search_context tool
  */
 
-import { jsonSchema, type ToolSet } from "ai";
+import type { ToolSet } from "ai";
+import { z } from "zod";
 import { estimateStringTokens } from "../utils/tokens";
 import { isSearchProvider, type SearchProvider } from "./search";
 import { isSkillProvider, type SkillProvider } from "./skills";
@@ -651,9 +652,9 @@ export class ContextBlocks {
 
       toolSet.set_context = {
         description: `Write to a context block. Available blocks:\n${blockDescriptions.join("\n")}\n\nWrites are durable and persist across sessions.`,
-        inputSchema: jsonSchema({
+        inputSchema: z.fromJSONSchema({
           type: "object" as const,
-          properties: properties as Record<string, object>,
+          properties: properties as Record<string, Record<string, unknown>>,
           required
         }),
         execute: async ({
@@ -713,7 +714,7 @@ export class ContextBlocks {
           "Available skill blocks: " +
           skillLabels.map((l) => `"${l}"`).join(", ") +
           ". Check the system prompt for available keys.",
-        inputSchema: jsonSchema({
+        inputSchema: z.fromJSONSchema({
           type: "object" as const,
           properties: {
             label: {
@@ -746,7 +747,7 @@ export class ContextBlocks {
           (loadedList.length > 0
             ? " Currently loaded: " + loadedList.join(", ") + "."
             : " No skills currently loaded."),
-        inputSchema: jsonSchema({
+        inputSchema: z.fromJSONSchema({
           type: "object" as const,
           properties: {
             label: {
@@ -782,7 +783,7 @@ export class ContextBlocks {
           "Available searchable blocks: " +
           searchLabels.map((l) => `"${l}"`).join(", ") +
           ".",
-        inputSchema: jsonSchema({
+        inputSchema: z.fromJSONSchema({
           type: "object" as const,
           properties: {
             label: {

packages/agents/src/experimental/memory/session/index.ts

@@ -49,4 +49,8 @@ export type { SearchProvider } from "./search";
 export { AgentSearchProvider, isSearchProvider } from "./search";
 export type { SkillProvider } from "./skills";
 export { isSkillProvider, R2SkillProvider } from "./skills";
-export type { MessageQueryOptions, SessionOptions } from "./types";
+export type {
+  SessionMessage,
+  SessionMessagePart,
+  SessionOptions
+} from "./types";

packages/agents/src/experimental/memory/session/manager.ts

@@ -6,14 +6,15 @@
  * Cross-session search and tools.
  */
 
-import type { UIMessage } from "ai";
-import { jsonSchema, type ToolSet } from "ai";
+import type { ToolSet } from "ai";
+import { z } from "zod";
 import type { CompactResult } from "../utils/compaction-helpers";
 import type { WritableContextProvider } from "./context";
 import type { StoredCompaction } from "./provider";
 import type { SqlProvider } from "./providers/agent";
 import type { SearchProvider } from "./search";
 import { Session, type SessionContextOptions } from "./session";
+import type { SessionMessage } from "./types";
 
 export interface SessionInfo {
   id: string;
@@ -35,27 +36,23 @@ interface PendingManagerContext {
   options: SessionContextOptions;
 }
 
-export interface SessionManagerOptions {
-  maxContextMessages?: number;
-}
+export interface SessionManagerOptions {}
 
 export class SessionManager {
   private agent!: SqlProvider;
-  private _maxContextMessages = 100;
   private _pending: PendingManagerContext[] = [];
   private _cachedPrompt?: WritableContextProvider | true;
   private _compactionFn?:
-    | ((messages: UIMessage[]) => Promise<CompactResult | null>)
+    | ((messages: SessionMessage[]) => Promise<CompactResult | null>)
     | null;
   private _tokenThreshold?: number;
   private _sessions = new Map<string, Session>();
   private _historyLabel?: string;
   private _tableReady = false;
   private _ready = false;
 
-  constructor(agent: SqlProvider, options: SessionManagerOptions = {}) {
+  constructor(agent: SqlProvider, _options: SessionManagerOptions = {}) {
     this.agent = agent;
-    this._maxContextMessages = options.maxContextMessages ?? 100;
     this._ready = true;
     this._ensureTable();
   }
@@ -69,7 +66,7 @@ export class SessionManager {
    *   .withContext("soul", { provider: { get: async () => "You are helpful." } })
    *   .withContext("memory", { description: "Learned facts", maxTokens: 1100 })
    *   .withCachedPrompt()
-   *   .maxContextMessages(50);
+   *   .compactAfter(100_000);
    *
    * // Each getSession(id) auto-creates namespaced providers:
    * //   memory key: "memory_<sessionId>"
@@ -80,7 +77,6 @@ export class SessionManager {
   static create(agent: SqlProvider): SessionManager {
     const mgr: SessionManager = Object.create(SessionManager.prototype);
     mgr.agent = agent;
-    mgr._maxContextMessages = 100;
     mgr._pending = [];
     mgr._compactionFn = null;
     mgr._tokenThreshold = undefined;
@@ -102,17 +98,12 @@ export class SessionManager {
     return this;
   }
 
-  maxContextMessages(count: number): this {
-    this._maxContextMessages = count;
-    return this;
-  }
-
   /**
    * Register a compaction function propagated to all sessions.
    * Called by `Session.compact()` to compress message history.
    */
   onCompaction(
-    fn: (messages: UIMessage[]) => Promise<CompactResult | null>
+    fn: (messages: SessionMessage[]) => Promise<CompactResult | null>
   ): this {
     this._compactionFn = fn;
     return this;
@@ -278,7 +269,7 @@ export class SessionManager {
 
   async append(
     sessionId: string,
-    message: UIMessage,
+    message: SessionMessage,
     parentId?: string
   ): Promise<string> {
     await this.getSession(sessionId).appendMessage(message, parentId);
@@ -288,7 +279,7 @@ export class SessionManager {
 
   async upsert(
     sessionId: string,
-    message: UIMessage,
+    message: SessionMessage,
     parentId?: string
   ): Promise<string> {
     const session = this.getSession(sessionId);
@@ -304,7 +295,7 @@ export class SessionManager {
 
   async appendAll(
     sessionId: string,
-    messages: UIMessage[],
+    messages: SessionMessage[],
     parentId?: string
   ): Promise<string | null> {
     const session = this.getSession(sessionId);
@@ -317,7 +308,7 @@ export class SessionManager {
     return lastParent;
   }
 
-  getHistory(sessionId: string, leafId?: string): UIMessage[] {
+  getHistory(sessionId: string, leafId?: string): SessionMessage[] {
     return this.getSession(sessionId).getHistory(leafId);
   }
 
@@ -337,7 +328,7 @@ export class SessionManager {
 
   // ── Branching ──────────────────────────────────────────────────
 
-  getBranches(sessionId: string, messageId: string): UIMessage[] {
+  getBranches(sessionId: string, messageId: string): SessionMessage[] {
     return this.getSession(sessionId).getBranches(messageId);
   }
 
@@ -357,7 +348,7 @@ export class SessionManager {
     let parentId: string | null = null;
     for (const msg of history) {
       const newId = crypto.randomUUID();
-      const copy: UIMessage = { ...msg, id: newId };
+      const copy: SessionMessage = { ...msg, id: newId };
       await newSession.appendMessage(copy, parentId);
       parentId = newId;
     }
@@ -368,12 +359,6 @@ export class SessionManager {
 
   // ── Compaction ────────────────────────────────────────────────
 
-  needsCompaction(sessionId: string): boolean {
-    return (
-      this.getSession(sessionId).getPathLength() > this._maxContextMessages
-    );
-  }
-
   addCompaction(
     sessionId: string,
     summary: string,
@@ -470,7 +455,7 @@ export class SessionManager {
       session_search: {
         description:
           "Search past conversations for relevant context. Searches across all sessions.",
-        inputSchema: jsonSchema({
+        inputSchema: z.fromJSONSchema({
           type: "object" as const,
           properties: {
             query: { type: "string" as const, description: "Search query" }

packages/agents/src/experimental/memory/session/provider.ts

@@ -4,7 +4,7 @@
  * Pure storage for tree-structured messages with compaction overlays and search.
  */
 
-import type { UIMessage } from "ai";
+import type { SessionMessage } from "./types";
 
 export interface SearchResult {
   id: string;
@@ -29,17 +29,17 @@ export interface StoredCompaction {
 export interface SessionProvider {
   // ── Read ────────────────────────────────────────────────────────
 
-  getMessage(id: string): UIMessage | null;
+  getMessage(id: string): SessionMessage | null;
 
   /**
    * Get conversation as a path from root to leaf.
    * Applies compaction overlays. If leafId is null, uses the latest leaf.
    */
-  getHistory(leafId?: string | null): UIMessage[];
+  getHistory(leafId?: string | null): SessionMessage[];
 
-  getLatestLeaf(): UIMessage | null;
+  getLatestLeaf(): SessionMessage | null;
 
-  getBranches(messageId: string): UIMessage[];
+  getBranches(messageId: string): SessionMessage[];
 
   getPathLength(leafId?: string | null): number;
 
@@ -49,9 +49,9 @@ export interface SessionProvider {
    * Append a message. Parented to the latest leaf unless parentId is provided.
    * Idempotent — same message.id twice is a no-op.
    */
-  appendMessage(message: UIMessage, parentId?: string | null): void;
+  appendMessage(message: SessionMessage, parentId?: string | null): void;
 
-  updateMessage(message: UIMessage): void;
+  updateMessage(message: SessionMessage): void;
 
   deleteMessages(messageIds: string[]): void;