# How to Use the Memory System in Ax for Context Management

> Master Ax memory for context management. Learn to track interactions, rewind with tags, and integrate persistence using AxMemory and MCP servers.

- Repository: [Ax/ax](https://github.com/ax-llm/ax)
- Tags: how-to-guide
- Published: 2026-02-25

---

**Ax provides a built-in, session-aware memory layer through the `AxMemory` class that automatically tracks interactions, supports tag-based rewinding, and integrates with external persistence via MCP servers.**

The memory system in Ax for context management enables AI agents to maintain state across multi-turn conversations, recall function results, and selectively prune irrelevant history. Implemented in the `ax-llm/ax` repository, this system centers on the **`AxMemory`** class in [`src/ax/mem/memory.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/mem/memory.ts), which orchestrates per-session storage and retrieval operations.

## Core Architecture of Ax's Memory System

Understanding the memory system requires familiarity with three primary components that handle data persistence, session isolation, and external integration.

### AxMemory and MemoryImpl

The **`AxMemory`** class serves as the public API that agents and applications interact with directly. It maintains a default memory instance and a map of per-session **`MemoryImpl`** instances.

Key methods in [`src/ax/mem/memory.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/mem/memory.ts) include:

- **`addRequest`** – Stores user input with role and content
- **`addResponse`** – Records assistant replies as indexed chat entries
- **`addFunctionResults`** – Persists tool call outputs for later reference
- **`updateResult`** – Merges streaming partial responses (thought blocks) into complete entries
- **`addTag`**, **`rewindToTag`**, **`removeByTag`** – Tag-based history manipulation
- **`history`**, **`getLast`**, **`reset`** – Retrieval and clearing operations

### Memory Data Model

`MemoryImpl` stores an array of **`AxMemoryData`** objects, where each entry contains:

- **`role`** – `"user"`, `"assistant"`, or `"function"`
- **`chat`** – An array of `{ index, value }` objects holding original payloads
- **`tags`** – Optional string identifiers for advanced pruning

When processing streaming responses, `updateResult` intelligently merges partial **`thoughtBlocks`** until a signature appears, ensuring the final stored representation accurately reflects the complete assistant output.

## Session Management and Isolation

The memory system in Ax for context management automatically isolates conversations by **`sessionId`**, preventing data leakage between users or chat rooms.

### Per-Session Memory Instances

When you call memory methods with a `sessionId` parameter, `AxMemory` routes the operation to a dedicated `MemoryImpl` instance for that session. If no `sessionId` is provided, the framework uses a singleton default memory.

This architecture enables:

- **Multi-tenancy** – Multiple users sharing the same agent process without context collision
- **Selective resets** – Clearing one user's history without affecting others
- **Persistent sessions** – Reattaching to existing session data across agent restarts (when paired with external storage)

## Tag-Based Context Pruning

Tags provide fine-grained control over conversation history, allowing agents to bookmark specific turns and later rewind or delete from that point.

### Adding and Using Tags

Attach a tag to the most recent memory entry using **`addTag(name)`**. Once tagged, you can:

- **`rewindToTag(name)`** – Returns the tagged entry and all subsequent history, removing it from active memory (useful for "undo" functionality)
- **`removeByTag(name)`** – Deletes all entries containing the specified tag, effectively pruning stale context

Both methods throw descriptive errors if the requested tag does not exist, preventing accidental data loss from typos.

## Integrating External Memory via MCP

For persistence across process restarts or shared memory between distributed agents, Ax supports the Model Context Protocol (MCP) through **`AxMCPClient`**.

### MCP Server-Memory Setup

The **`AxDBMemory`** class in [`src/ax/db/memory.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/db/memory.ts) provides an in-memory database interface used by the MCP server-memory plugin. To integrate external persistence:

1. **Initialize the MCP client** with a transport connecting to `@modelcontextprotocol/server-memory`
2. **Convert to functions** using `mcp.toFunction()` to expose memory operations as agent-callable tools
3. **Attach to agent** via the `functions` parameter in `agent()`

This pattern, demonstrated in [`src/examples/mcp-client-memory.ts`](https://github.com/ax-llm/ax/blob/main/src/examples/mcp-client-memory.ts), allows agents to store and retrieve memories via a remote database while maintaining the local `AxMemory` cache for fast context access.

## Practical Implementation Examples

### Basic Memory Operations

```typescript
import { AxMemory } from '@ax-llm/ax';

// Initialize memory (default session)
const mem = new AxMemory();

// Store user query
mem.addRequest([{ role: 'user', content: 'Explain quantum computing' }]);

// Store assistant response
mem.addResponse([{ index: 0, content: 'Quantum computing uses qubits...' }]);

// Retrieve full conversation
const history = mem.history(0);
console.log(history);

```

### Session Isolation

```typescript
const sessionA = 'user-alice';
const sessionB = 'user-bob';

// Alice's conversation
mem.addRequest([{ role: 'user', content: 'My favorite color is blue' }], sessionA);
mem.addResponse([{ index: 0, content: 'Noted: Alice likes blue' }], sessionA);

// Bob's conversation (isolated)
mem.addRequest([{ role: 'user', content: 'I prefer red' }], sessionB);

// Alice's history remains unaffected by Bob's messages
const aliceHistory = mem.history(0, sessionA);

```

### Tag-Based Rewinding

```typescript
// Mark a decision point
mem.addTag('decision:tool-selection');

// Continue conversation...
mem.addRequest([{ role: 'user', content: 'Use the weather tool' }]);
mem.addFunctionResults([{ name: 'getWeather', result: 'Sunny, 72°F' }]);

// Oops, wrong tool selected - rewind to decision point
const prunedHistory = mem.rewindToTag('decision:tool-selection');
// prunedHistory contains everything from the tag onward, removed from active memory

```

### MCP External Memory Integration

```typescript
import { AxMCPClient, axCreateMCPStdioTransport } from '@ax-llm/ax-tools';
import { agent, AxAI } from '@ax-llm/ax';

// Connect to MCP server-memory
const transport = axCreateMCPStdioTransport({
  command: 'npx',
  args: ['-y', '@modelcontextprotocol/server-memory'],
});
const mcp = new AxMCPClient(transport);
await mcp.init();

// Expose as agent functions
const memoryFns = mcp.toFunction();

// Create agent with external memory capabilities
const memoryAgent = agent(
  'userMessage:string, userId:string -> reply:string',
  {
    functions: { local: memoryFns },
  }
);

// Agent can now store/retrieve persistent memories via the MCP service

```

## Summary

The memory system in Ax for context management provides a robust foundation for building stateful AI agents through the **`AxMemory`** class in [`src/ax/mem/memory.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/mem/memory.ts). Key capabilities include:

- **Automatic session isolation** – Per-session `MemoryImpl` instances prevent context leakage between users
- **Streaming-safe storage** – The `updateResult` method merges partial thoughts without data loss
- **Tag-based manipulation** – `addTag`, `rewindToTag`, and `removeByTag` enable precise context pruning
- **External persistence** – MCP client integration via `AxMCPClient` supports durable storage across process restarts

By combining these features, developers can implement sophisticated context management strategies that balance immediate retrieval performance with long-term memory persistence.

## Frequently Asked Questions

### How does Ax handle memory for streaming responses?

Ax processes streaming responses through the **`updateResult`** method in [`src/ax/mem/memory.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/mem/memory.ts), which intelligently merges partial **`thoughtBlocks`** until a signature appears. This ensures that intermediate streaming chunks are accumulated into complete, coherent memory entries rather than fragmented pieces, preserving the full context of assistant reasoning.

### Can multiple users share the same Ax agent without leaking conversation history?

Yes, the memory system in Ax for context management automatically isolates sessions using the **`sessionId`** parameter. When provided, `AxMemory` routes operations to dedicated `MemoryImpl` instances stored in an internal map. This architecture ensures that User A's history remains completely separate from User B's, even when both interact with the same agent process simultaneously.

### What is the difference between `rewindToTag` and `removeByTag`?

**`rewindToTag`** returns the tagged entry and all subsequent history while removing it from active memory, effectively creating an "undo" point that extracts context for potential reuse. In contrast, **`removeByTag`** permanently deletes all entries containing the specified tag without returning them, serving as a pruning mechanism to discard stale or irrelevant context. Both methods throw errors if the tag does not exist, preventing accidental operations on undefined markers.

### How do I persist Ax memory across server restarts?

To achieve durable persistence, integrate the **MCP server-memory** plugin using **`AxMCPClient`** from `@ax-llm/ax-tools`. This client connects to `@modelcontextprotocol/server-memory` via stdio transport, converting memory operations into agent-callable functions. While `AxMemory` maintains fast in-process caching, the MCP client ensures that critical memories survive process termination by storing them in an external database service.