Streaming Structured Working Memory

This example demonstrates how to create an agent that maintains a todo list using structured working memory (schema-based), even with minimal context. For a Markdown-based version, see the advanced working memory example.

Setup

Let’s break down how to create an agent with structured working memory. We’ll build a todo list manager that remembers tasks as a JSON object.

1. Setting up Structured Memory

We’ll configure the memory system with a short context window and provide a Zod schema for the todo list. For persistence, you can configure a storage provider like @mastra/libsql:


import { Memory } from "@mastra/memory";
import { LibSQLStore } from "@mastra/libsql";
import { z } from "zod";
 
// Define the schema for the todo list
const todoListSchema = z.object({
  items: z.array(
    z.object({
      description: z.string(),
      due: z.string().optional(),
      started: z.string().optional(),
      status: z.enum(["active", "completed"]).default("active"),
    })
  ),
});
 
const memory = new Memory({
  options: {
    lastMessages: 1,
    workingMemory: {
      enabled: true,
      schema: todoListSchema,
    },
  },
  storage: new LibSQLStore({
    url: "file:../mastra.db",
  }),
});

2. Creating the Todo List Agent

We’ll create an agent that uses this structured memory. The instructions should guide the agent to maintain the todo list as a JSON object matching the schema.


import { openai } from "@ai-sdk/openai";
import { Agent } from "@mastra/core/agent";
 
const todoAgent = new Agent({
  name: "TODO Agent",
  instructions:
    "You are a helpful todolist AI agent. Maintain the user's todo list as a JSON object according to the provided schema. Each item should include a description, status (active/completed), and optional due/started dates. When the user adds, completes, or updates a task, update the JSON accordingly. Always show the current todo list at the start of the chat.",
  model: openai("gpt-4o-mini"),
  memory,
});

Note: When workingMemory.schema is set, a default system prompt is automatically injected to instruct the agent on how to use structured working memory. You can override or extend this prompt as needed.

Usage Example

The agent’s responses will contain <working_memory>{...}</working_memory> tags with the updated JSON. We’ll look at two ways to handle this:

Basic Usage

For simple cases, you can use maskStreamTags to hide the working memory updates from users:


import { randomUUID } from "crypto";
import { maskStreamTags } from "@mastra/core/utils";
 
// Start a conversation
const threadId = randomUUID();
const resourceId = "SOME_USER_ID";
 
// Add a new todo item
const response = await todoAgent.stream(
  "Add a task: Build a new feature for our app. It should take about 2 hours and needs to be done by next Friday.",
  {
    threadId,
    resourceId,
  },
);
 
// Process the stream, hiding working memory updates
for await (const chunk of maskStreamTags(
  response.textStream,
  "working_memory",
)) {
  process.stdout.write(chunk);
}

Advanced Usage with UI Feedback

For a better user experience, you can show loading states while working memory is being updated:


// Add lifecycle hooks to provide UI feedback
const maskedStream = maskStreamTags(response.textStream, "working_memory", {
  // Called when a working_memory tag starts
  onStart: () => showLoadingSpinner("Updating todo list..."),
  // Called when a working_memory tag ends
  onEnd: () => hideLoadingSpinner(),
  // Called with the content that was masked
  onMask: (chunk) => console.debug("Updated todo list:", chunk),
});
 
// Process the masked stream
for await (const chunk of maskedStream) {
  process.stdout.write(chunk);
}

Example: Structured Working Memory State

After several interactions, the working memory might look like:


{
  "items": [
    {
      "description": "Buy groceries",
      "due": "2025-07-01",
      "started": "2025-06-24",
      "status": "active"
    },
    {
      "description": "Finish project report",
      "due": "2025-07-05",
      "status": "completed"
    }
  ]
}

The agent can refer to or update this list directly in JSON, ensuring structured, type-safe memory retention across conversations.

This example demonstrates:

Setting up a memory system with structured working memory using a Zod schema
Creating an agent that manages a todo list as a JSON object
Using maskStreamTags to hide memory updates from users
Providing UI loading states during memory updates with lifecycle hooks
Maintaining persistent, structured memory even with a short context window

Even with only one message in context (lastMessages: 1), the agent maintains the complete todo list in structured working memory. Each time the agent responds, it updates the JSON object, ensuring persistence across interactions.

To learn more about agent memory, including other memory types and storage options, see the Memory documentation.