Skip to content

[Feature]: Message-level branching (fork) within a session #317

@GenerQAQ

Description

@GenerQAQ

Feature Description

Enable true message-level branching within a single session by allowing store_message to accept an optional parent_id. This lets users fork the conversation at any point — creating a new branch in the message tree without creating a new session.

Currently, parent_id is auto-assigned to the latest message by created_at, making every session a linear chain. With this change, users can point a new message at any ancestor, forming a real tree — just like pi-mono's /fork and /tree behavior.

Use cases:

  • An agent conversation went off track at turn 5; send a new message with parent_id pointing to turn 4's message, branching off from there.
  • Explore multiple strategies in parallel within the same session — each branch is a different attempt.
  • "Checkpoint and retry" workflow: continue from a known-good message without duplicating anything.

Example:

Session message tree (before fork):
  M1(user) → M2(asst) → M3(user) → M4(asst) → M5(user) → M6(asst)

User sends a new message with parent_id = M2:
  M1(user) → M2(asst) → M3(user) → M4(asst) → M5(user) → M6(asst)
                 ↘ M7(user)   ← new branch starts here

Then assistant responds to M7:
  M1(user) → M2(asst) → M3(user) → M4(asst) → M5(user) → M6(asst)
                 ↘ M7(user) → M8(asst)

No new session is created. M7 and M8 live in the same session, sharing M1 and M2 with the original branch.

Proposed Changes

1. StoreMessage — accept optional parent_id (small change)

API:

POST /api/v1/session/:session_id/messages

// Existing body fields stay the same, add optional parent_id:
{
    "blob": { ... },
    "format": "openai",
    "meta": {},
    "parent_id": "uuid-of-message-to-branch-from"   // NEW, optional
}
  • If parent_id is provided: validate it belongs to the same session, use it as the parent.
  • If parent_id is omitted: keep current behavior (auto-assign to the latest message by created_at).

Response: Same as current store_message response (the created message object).

Error cases:

  • parent_id is not a valid UUID → 400
  • parent_id message does not exist or belongs to a different session → 404

Python SDK:

# Normal (no fork, backward compatible)
client.sessions.store_message(session_id, blob=..., format="openai")

# Fork from a specific message
client.sessions.store_message(
    session_id,
    blob=...,
    format="openai",
    parent_id="message-uuid-to-fork-from"
)

TypeScript SDK:

// Normal (no fork, backward compatible)
await client.sessions.storeMessage(sessionId, blob, { format: "openai" });

// Fork from a specific message
await client.sessions.storeMessage(sessionId, blob, {
    format: "openai",
    parentId: "message-uuid-to-fork-from",
});

2. GetMessages — support branch-aware retrieval (large change)

Once a session has branches, the current GET /messages (which returns all messages in created_at order) will mix messages from different branches. We need a way to retrieve a specific branch path.

API — add optional leaf_id parameter:

GET /api/v1/session/:session_id/messages?leaf_id=<message-uuid>
  • If leaf_id is provided: walk the parent_id chain from that message to the root, return only messages on that path (in chronological order).
  • If leaf_id is omitted: default to the latest leaf (most recent message with no children), preserving backward compatibility for linear sessions.

Python SDK:

# Current behavior (all messages or latest branch)
msgs = client.sessions.get_messages(session_id, format="acontext")

# Get a specific branch
msgs = client.sessions.get_messages(session_id, format="acontext", leaf_id="leaf-message-uuid")

TypeScript SDK:

// Current behavior
const msgs = await client.sessions.getMessages(sessionId, { format: "acontext" });

// Get a specific branch
const msgs = await client.sessions.getMessages(sessionId, {
    format: "acontext",
    leafId: "leaf-message-uuid",
});

3. List branches / leaves (new endpoint, medium change)

To navigate the tree, users need to discover what branches exist.

API:

GET /api/v1/session/:session_id/branches

Response:

{
    "code": 0,
    "data": {
        "leaves": [
            {
                "message_id": "M6-uuid",
                "depth": 6,
                "created_at": "2025-01-01T00:00:06Z"
            },
            {
                "message_id": "M8-uuid",
                "depth": 4,
                "created_at": "2025-01-01T00:00:08Z"
            }
        ]
    },
    "msg": ""
}

Each leaf represents the tip of a branch. Users can then call GET /messages?leaf_id=<leaf> to retrieve that branch's conversation.

4. CORE — tree-aware session buffer (large change)

The CORE module currently reads messages linearly by created_at for task extraction and LLM context building. With branching, CORE needs to:

  • When processing a newly stored message, walk its parent_id chain to build the correct branch context (not mix in messages from sibling branches).
  • Task extraction should be scoped to the branch the new message belongs to.

This is the most architecturally significant change and may need a phased approach.

Impact Component/Area

  • Client SDK (Python)
  • Client SDK (TypeScript)
  • Core Service (Python)
  • API Server (Go)
  • UI/Dashboard (Next.js)
  • CLI Tool
  • Documentation
  • Other (please specify)

Additional Context

Current state of parent_id:

  • The parent_id column and foreign key already exist on the messages table in both Go (GORM) and Python (SQLAlchemy) ORMs.
  • The Parent / Children relationships are defined in the model.
  • However, parent_id is currently auto-assigned (always points to the latest message) and never accepted from client input.
  • CORE never queries by parent_id — all processing is ORDER BY created_at.
  • parent_id is only exposed in the format=acontext response; other formats (OpenAI, Anthropic, Gemini) don't include it.

Suggested implementation order:

  1. StoreMessage accepts parent_id (small, backward compatible, unblocks everything)
  2. GetMessages supports leaf_id (enables branch reading)
  3. List branches endpoint (enables branch discovery)
  4. CORE tree-aware processing (enables correct task extraction on branches)

Metadata

Metadata

Assignees

Labels

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions