AG-UI Chat History

The chat-history layer turns an AG-UI agent’s event stream into durable, typed chat history on top of the Memory Service ORM. It gives you three things:

Typed chat models — Chat (a DRSession) and Message (a DREvent), with nested ToolCall / Reasoning models living inside a single message event body. One Memory Service event per logical message.
A transparent storage wrapper — AGUIStorageAgent forwards an inner agent’s events verbatim (in real time) while a background task folds the same stream into persisted history: inbound user messages are stored, prior history is replayed into the run, and every text / tool-call / reasoning delta is captured.
Disconnect survival and cancellation — StreamPersistenceManager runs the storage agent behind an unbounded queue so persistence completes even if the client stops reading, and exposes cancellation that marks in-flight records interrupted.

Everything is imported from a single package:

from datarobot.application_utils.chat_history import (
    AGUIAgent,
    AGUIStorageAgent,
    ChatRepository,
    ChatSessionRegistry,
    MessageRepository,
    RunHandle,
    StreamPersistenceManager,
)

The chat-history package depends on both persistence and ag_ui; the persistence sub-package never imports ag_ui, so the ORM stays transport-agnostic.

Quick-start

Wire the repositories, wrap your inner agent, and run it through the manager. StreamPersistenceManager.run returns a RunHandle you iterate for the client-facing stream, cancel, or await to completion.

DRMemoryServiceClient() needs a configured endpoint and token — set DATAROBOT_ENDPOINT / DATAROBOT_API_TOKEN or pass them as constructor arguments (see Environment variables).

import asyncio
from uuid import uuid4

from ag_ui.core import RunAgentInput, UserMessage

from datarobot.application_utils.chat_history import (
    AGUIStorageAgent,
    ChatRepository,
    ChatSessionRegistry,
    MessageRepository,
    StreamPersistenceManager,
)
from datarobot.application_utils.persistence import DRMemoryServiceClient, DRMemorySpace


async def main(inner_agent) -> None:
    user_id = uuid4()

    async with DRMemoryServiceClient() as client:
        space = await DRMemorySpace.post(client, deduplication_key="my-chat-app-v1")

        # A registry shares the chat_uuid -> session_id cache across repos.
        registry = ChatSessionRegistry(space)
        chat_repo = ChatRepository(space, registry)
        message_repo = MessageRepository(space, registry)

        # The factory builds a fresh storage agent per run; the manager owns the
        # in-flight-run registry keyed by (thread_id, run_id).
        manager = StreamPersistenceManager(
            lambda: AGUIStorageAgent("assistant", user_id, chat_repo, message_repo, inner_agent)
        )

        input = RunAgentInput(
            thread_id="thread-1",
            run_id="run-1",
            messages=[UserMessage(id="m1", role="user", content="Hello!")],
            tools=[],
            context=[],
            state={},
            forwarded_props={},
        )

        handle = await manager.run(input)

        # Stream events to the client. Stopping early (a disconnect) does NOT
        # abort persistence — the producer keeps draining behind the scenes.
        async for event in handle.events():
            ...  # forward `event` to your UI / SSE response

        # Wait for the background persistence to finish, then read history back.
        await handle.wait()
        history = await message_repo.get_chat_messages(
            (await chat_repo.get_chat_by_thread_id(user_id, "thread-1")).chat_uuid
        )
        for message in history:
            print(message.role, message.content, message.tool_calls, message.reasonings)


# `inner_agent` is any AGUIAgent: `run(input) -> AsyncGenerator[BaseEvent]`.
asyncio.run(main(inner_agent=...))

Cancel a run through the handle (or manager.cancel(thread_id, run_id)):

handle = await manager.run(input)
...
handle.cancel()  # -> True if a live run was found, False otherwise
await handle.wait()  # in-flight records are finalized as `interrupted`

Data model

Model	Base	Stored as
`Chat`	`DRSession`	One session per thread. `thread_id` is a range key (indexed `//thread/{thread_id}/` lookup); `dedup_key = sha256("chat"\0user_uuid\0thread_id)` makes create idempotent. Metadata: `name`, `chat_uuid`, `user_uuid`.
`Message`	`DREvent` (`event_type="message"`)	One event per logical message. Body carries `v`, `role`, `status`, `in_progress`, and the nested `tool_calls: list[ToolCall]` / `reasonings: list[Reasoning]`.
`ToolCall`	`BaseModel`	Nested in a message body — `tool_call_id`, `name`, `arguments`, `content`, `status`.
`Reasoning`	`BaseModel`	Nested in a message body — `name`, `content`, `status`.

The nested models round-trip through the Memory Service ORM’s typed serialization (see Memory Service ORM); empty content / arguments are transparently encoded to a zero-width placeholder to satisfy the service’s min_length=1.

Role values: developer, system, assistant, user, tool, reasoning. MessageStatus values: active, complete, interrupted, errored.

ChatRepository and MessageRepository are the CRUD facades (create/adopt a chat idempotently, append messages, mutate nested tool-calls / reasonings, read history sorted by sequence id). ChatSessionRegistry resolves chat_uuid → session_id via a bounded in-process cache, then an indexed chat:<uuid> EntityLocator point lookup — never a full-space scan.

Secondary index and consistency

The Memory Service has no cross-session event query, so resolving an entity by its application UUID — chat_uuid → session, message_uuid → chat, or a tool-call / reasoning uuid → parent message — could otherwise only be done by scanning every session in the space. Instead the layer keeps a dedup-keyed secondary index: one EntityLocator (itself a DRSession) per locatable entity, looked up by the exact key "<kind>:<uuid>" (kind is one of chat, message, tool_call, reasoning). Every cold lookup is a single indexed point get — there are no full-space scans.

Locators live in the same Memory Space as the chats; their //loc/ description prefix keeps them out of every Chat.list result (and vice-versa), so no second space is needed.

Best-effort writes and the consistency trade-off

Each create writes the entity first (the event or session — the source of truth), then writes its locator best-effort. A locator write that fails is logged and swallowed; it never fails the operation. The deliberate consequences:

The worker that created the entity is unaffected — its in-process cache is warm.
If a locator write is lost, a different replica (or a lookup after a restart) that addresses the entity by uuid — get_message, update_message, a tool-call / reasoning update — may report “not found” for an entity that actually exists.
The data is not lost: get_chat_messages lists the chat’s events directly and still returns the message.

This backend therefore offers no strong cross-replica read-after-write consistency for uuid-addressed lookups. A caller that cannot tolerate that needs a transactional backend, not this one.

Orphans

delete_chat soft-deletes the session (and its events) in one shot; there is no per-message delete, and it does not walk and delete the chat’s locators. Those locators are left as orphans — cheap, bounded by the same soft-delete TTL as the sessions they point at, and harmless: a stale locator resolves to a dead session and is treated as not-found, and because UUIDs never collide it can never resurrect or mis-route a new entity.

Run-outcome status semantics

Every message / tool-call / reasoning record carries a status. How a run ends determines the terminal status of the records that were still in_progress:

How the run ends	Terminal status	Mechanism
Normal completion	`complete`	The inner agent emits `RunFinishedEvent`; `handle_run_lifecycle` flips the active message (and its tool-calls / reasonings) to `in_progress=False`, `status=complete`.
Explicit cancel	`interrupted`	`RunHandle.cancel()` / `StreamPersistenceManager.cancel()` cancels the producer task; the `CancelledError` propagates into `AGUIStorageAgent.run`, whose `finally` calls `_finalize_interrupted` to flip every still-`in_progress` record to `status=interrupted`.
Inner-agent raw crash	`errored`	The inner agent raises (rather than emitting a terminal `RunErrorEvent`). `AGUIStorageAgent.run` records the exception, its `finally` calls `_finalize_errored` to flip still-active records to `status=errored`, then re-raises — the manager’s producer catches it and synthesizes a client-facing `RunErrorEvent(code=INTERNAL_ERROR)` so the client never hangs.
Client disconnect	`complete`	The client stops reading `RunHandle.events()`, but the producer drains the inner agent into its unbounded queue regardless, so the run finishes normally and persists as `complete`. `await handle.wait()` blocks until that drain finishes.

An in-band RunErrorEvent emitted by the inner agent is a normal terminal event: handle_run_lifecycle records those records as errored directly. _finalize_errored is specifically the safety net for a raw exception that escaped the inner agent without a terminal event.

Extensibility

AGUIStorageAgent is an open state machine: every seam a consumer might want to customize is a method you can override, and it depends only on repository Protocols. There are four extension points.

1. Event-dispatch registry

event_handlers() maps each AG-UI event type to a handler-method name. Extend it (walking the MRO, so subclasses of known events resolve automatically) to persist brand-new or custom event types; unrecognized events fall through to handle_unknown_event (a no-op by default).

class MyStorageAgent(AGUIStorageAgent):
    @classmethod
    def event_handlers(cls):
        return {**super().event_handlers(), MyCustomEvent: "handle_my_custom_event"}

    async def handle_my_custom_event(self, state, chat, event): ...  # persist the custom event

2. Category handlers

handle_text_message, handle_tool_call, handle_reasoning, handle_run_lifecycle and handle_step own each event family. Override one to change how a whole category is stored — e.g., to extract structured content instead of appending raw text deltas.

class StructuredTextAgent(AGUIStorageAgent):
    async def handle_text_message(self, state, chat, event):
        # e.g., parse `event` into structured fields before buffering / flushing
        await super().handle_text_message(state, chat, event)

3. Message-build hooks

build_message_create, build_tool_call_create, build_reasoning_create and message_update_fields construct the DTOs the repository persists. Override them to populate extra fields declared on your Message / ToolCall / Reasoning subclasses.

class TaggedStorageAgent(AGUIStorageAgent):
    def build_message_create(self, state, chat, agui_id, role):
        base = super().build_message_create(state, chat, agui_id, role)
        return MyMessageCreate(**base.model_dump(), team="research")

    def message_update_fields(self, state, event):
        return {"finished_at": state.current_event_timestamp}

4. Repository coupling (Protocols)

The agent references only ChatRepositoryLike / MessageRepositoryLike (runtime_checkable typing.Protocols), so any conforming backend — a SQL store, an in-memory fake for tests — is a drop-in replacement with no imports of the concrete classes.

class InMemoryChatRepo:  # structurally satisfies ChatRepositoryLike
    async def get_chat_by_thread_id(self, user_uuid, thread_id): ...
    async def create_chat(self, chat_data): ...

    # ...remaining protocol methods...


agent = AGUIStorageAgent("assistant", user_id, InMemoryChatRepo(), my_message_repo, inner)

Public API

Models

class datarobot.application_utils.chat_history.Chat

A chat thread, persisted as one Memory Service session.

Field mapping

thread_idAnnotated[str, DRRangeKey]: Indexed description segment (//thread/{thread_id}/); powers the fast get_chat_by_thread_id lookup.
dedup_keyAnnotated[str, DRDeduplicationKey]: Idempotent-create key (see chat_deduplication_key()).
name, chat_uuid, user_uuid: Session metadata.

The session’s single participants entry is the user’s participant ID (see participant_id()); the agent is not a session participant.

class datarobot.application_utils.chat_history.Message

A chat message, persisted as one Memory Service event under a Chat.

Tool calls and reasoning steps live in tool_calls / reasonings inside this event’s body (not as separate events). v is the body payload schema version and gates reads.

See the module docstring for why the application timestamp is named timestamp rather than created_at.

content carries the same zero-width placeholder codec as the nested ToolCall / Reasoning content fields, so an empty message round-trips the Memory Service min_length=1 constraint transparently on raw ORM reads (list() / get()) as well as through MessageRepository.

class datarobot.application_utils.chat_history.ToolCall

A tool call nested inside a Message body.

arguments and content carry the zero-width placeholder codec so an empty string round-trips through the Memory Service min_length=1 constraint transparently.

class datarobot.application_utils.chat_history.Reasoning

A reasoning step nested inside a Message body.

content carries the zero-width placeholder codec so an empty string round-trips through the Memory Service min_length=1 constraint transparently.

class datarobot.application_utils.chat_history.Role

Message source role, mirroring the AG-UI message roles.

See https://docs.ag-ui.com/concepts/messages.

class datarobot.application_utils.chat_history.MessageStatus: Lifecycle status of a message, tool call or reasoning step.

Repositories

class datarobot.application_utils.chat_history.ChatRepositoryLike

Structural interface for chat persistence.

Any object exposing this method set (a SQL store, an in-memory fake) is an accepted chat repository; consumers depend on this Protocol rather than on ChatRepository.

async create_chat(chat_data)

Create a chat (idempotent by (user, thread)) or return the existing one.

Return type:: Chat

async get_chat_by_thread_id(user_uuid, thread_id)

Return the chat for a (user, thread_id) pair, or None.

Return type:: Chat | None

async get_all_chats(user_uuid)

Return every chat, optionally scoped to a single user.

Return type:: Sequence[Chat]

async update_chat_name(chat_uuid, name)

Rename a chat, returning the updated chat, or None when it is unknown.

Return type:: Chat | None

async delete_chat(chat_uuid)

Delete a chat, returning the deleted chat, or None when it is unknown.

Return type:: Chat | None

class datarobot.application_utils.chat_history.MessageRepositoryLike

Structural interface for message persistence (one event per message).

transaction()

Return an async context manager scoping a batch of writes.

Return type:: AbstractAsyncContextManager[None]

async create_message(message_data)

Persist a new message as one event and return it.

Return type:: Message

async update_message(message_uuid, update)

Patch a message in place, or return None when it is unknown.

Return type:: Message | None

async create_message_tool_call(data)

Append a tool call to its parent message’s body.

Return type:: ToolCall

async update_message_tool_call(uuid, update)

Patch a nested tool call, or return None when it is unknown.

Return type:: ToolCall | None

async create_message_reasoning(data)

Append a reasoning step to its parent message’s body.

Return type:: Reasoning

async update_message_reasoning(uuid, update)

Patch a nested reasoning step, or return None when it is unknown.

Return type:: Reasoning | None

async get_message(message_uuid)

Return a message by its application UUID, or None.

Return type:: Message | None

async get_message_by_agui_id(chat_uuid, agui_id)

Return a message by its AG-UI id within a chat, or None.

Return type:: Message | None

async get_tool_call_by_agui_id(message_uuid, agui_id)

Return a tool call by its AG-UI id within a message, or None.

Return type:: ToolCall | None

async get_chat_messages(chat_uuid)

Return every message in a chat, ordered oldest first.

Return type:: Sequence[Message]

async get_last_messages(chat_uuids)

Return the most recent message for each of the given chats.

Return type:: dict[UUID, Message]

class datarobot.application_utils.chat_history.ChatSessionRegistry

Map an app chat UUID to a Memory Service session id.

A bounded in-process cache covers hot paths. Because a Chat’s indexed description is keyed by thread_id (not by chat_uuid), a cold-cache resolve — e.g. on a replica that did not create the chat, or after a process restart — reads the dedup-keyed chat:<uuid> EntityLocator (an indexed O(1) point lookup), not a full-space scan. Fast, indexed (user, thread_id) lookups live on ChatRepository.get_chat_by_thread_id() instead.

property locators: LocatorIndex: The shared LocatorIndex for uuid → location lookups.

register(chat_uuid, session_id)

Cache the chat_uuid → session_id mapping.

Return type:: None

unregister(chat_uuid)

Drop a cached mapping (e.g. after the chat is deleted).

Return type:: None

get_session_id(chat_uuid)

Return the cached session id for a chat, without hitting the service.

Return type:: str | None

async resolve(chat_uuid)

Resolve a chat UUID to its session id via the locator index on a cache miss.

Parameters:: chat_uuid (UUID) – The application chat identifier.
Returns:: The Memory Service session id, or None when no chat:<uuid> locator exists (e.g. a lost best-effort index write, or an unknown chat).
Return type:: str | None

class datarobot.application_utils.chat_history.ChatRepository

Chat persistence backed by Memory Service sessions.

async create_chat(chat_data)

Create a chat, short-circuiting to the existing one for a known (user, thread).

Parameters:: chat_data (ChatCreate) – Must carry both user_uuid and thread_id; they derive the deduplication key and the single session participant.
Returns:: The created (or adopted) chat.
Return type:: Chat
Raises:: ValueError – If user_uuid or thread_id is missing.

async get_chat_by_thread_id(user_uuid, thread_id)

Return the chat for a (user, thread_id) pair, or None.

Tries the indexed description filter first (participant + thread id), then falls back to a participant-scoped scan for robustness.

Parameters:

user_uuid (UUID) – Owning user.
thread_id (str) – AG-UI thread identifier.

Return type:

Chat | None

async get_all_chats(user_uuid)

Return every chat, optionally scoped to a single user.

Parameters:: user_uuid (UUID | None) – When given, only chats participant-scoped to this user are returned.
Return type:: Sequence[Chat]

async update_chat_name(chat_uuid, name)

Rename a chat, retrying on a version conflict.

Parameters:

chat_uuid (UUID) – The chat to rename.
name (str) – The new display name.

Returns:

The updated chat, or None when no session carries the chat UUID.

Return type:

Chat | None

async delete_chat(chat_uuid)

Delete a chat and drop its registry entry.

Parameters:: chat_uuid (UUID) – The chat to delete.
Returns:: The deleted chat, or None when no session carries the chat UUID.
Return type:: Chat | None

class datarobot.application_utils.chat_history.MessageRepository

Message persistence backed by session events — one event per logical message.

Tool calls and reasoning steps are stored as typed nested models inside the parent message’s event body; a mutation re-serializes and patches the whole event body. Bounded in-process caches short-circuit the uuid → chat and child → parent lookups; on a cold cache these resolve via the dedup-keyed EntityLocator index (an O(1) point lookup), never a full-space scan.

transaction()

No-op batching scope; the Memory Service has no cross-document transaction.

Return type:: AsyncGenerator[None, None]

async create_message(message_data)

Persist a new message as a single event.

Parameters:: message_data (MessageCreate) – Message fields; chat_uuid is required.
Returns:: The persisted message (base content decoded).
Return type:: Message
Raises:: ValueError – If chat_uuid is missing.

async update_message(message_uuid, update)

Patch a message’s own fields in place.

Parameters:

message_uuid (UUID) – Application UUID of the message.
update (MessageUpdate) – Only the explicitly-set, non-None fields are applied.

Returns:

The updated message, or None when it does not exist.

Return type:

Message | None

async create_message_tool_call(data)

Append a tool call to its parent message’s body.

Parameters:: data (MessageToolCallCreate) – Tool-call fields; message_uuid names the parent message.
Returns:: The newly appended tool call.
Return type:: ToolCall
Raises:: ValueError – If the parent message does not exist.

async update_message_tool_call(uuid, update)

Patch a nested tool call in place.

Parameters:

uuid (UUID) – The tool call UUID.
update (MessageToolCallUpdate) – Only explicitly-set, non-None fields are applied.

Returns:

The updated tool call, or None when it does not exist.

Return type:

ToolCall | None

async create_message_reasoning(data)

Append a reasoning step to its parent message’s body.

Parameters:: data (MessageReasoningCreate) – Reasoning fields; message_uuid names the parent message.
Returns:: The newly appended reasoning step.
Return type:: Reasoning
Raises:: ValueError – If the parent message does not exist.

async update_message_reasoning(uuid, update)

Patch a nested reasoning step in place.

Parameters:

uuid (UUID) – The reasoning UUID.
update (MessageReasoningUpdate) – Only explicitly-set, non-None fields are applied.

Returns:

The updated reasoning step, or None when it does not exist.

Return type:

Reasoning | None

async get_message(message_uuid)

Return a message by its application UUID, or None.

Return type:: Message | None

async get_message_by_agui_id(chat_uuid, agui_id)

Return a message by its AG-UI ID within a chat, or None.

Return type:: Message | None

async get_tool_call_by_agui_id(message_uuid, agui_id)

Return a tool call by its AG-UI ID within a message, or None.

Return type:: ToolCall | None

async get_chat_messages(chat_uuid)

Return every message in a chat, ordered oldest first (by sequence ID).

Return type:: Sequence[Message]

async get_last_messages(chat_uuids)

Return the most recent message for each of the given chats.

Parameters:: chat_uuids (list[UUID]) – Chats to fetch the tail message for.
Returns:: Maps each chat UUID that has at least one message to its latest one.
Return type:: dict[UUID, Message]

AG-UI storage

class datarobot.application_utils.chat_history.AGUIAgent

Minimal AG-UI agent contract: a named object exposing an event stream.

abstractmethod run(input)

Yield the agent’s AG-UI BaseEvent stream.

Return type:: AsyncGenerator[BaseEvent, None]

class datarobot.application_utils.chat_history.AGUIStorageAgent

Wrap an inner AG-UI agent, persisting its event stream as chat history.

The wrapper is transparent: run() yields the inner agent’s events unchanged and in real time. Persistence happens on a separate background task fed from an internal queue, so a slow or failing store never stalls the outgoing stream, and consumer disconnection does not abort persistence.

classmethod event_handlers()

Return the AG-UI-event-type → handler-method-name dispatch table.

Override (typically {**super().event_handlers(), CustomEvent: "..."}) to register a handler for a new event type.

Returns:: Maps each handled event class to the name of the instance method that processes it.
Return type:: dict[type[ag_ui.core.BaseEvent], str]

async handle_unknown_event(state, chat, event)

Handle an event with no registered handler.

The default implementation ignores the event. Override to persist custom event types.

Return type:: None

translate(messages)

Translate stored messages into AG-UI history messages.

Delegates to the injected translate callable (default translate_messages). Override to customize the replayed history shape.

Return type:: list[ExtendedBaseMessage]

build_message_create(state, chat, agui_id, role)

Build the DTO used to create a new (agent) message.

Override to populate extra fields declared on a Message subclass (returning a matching MessageCreate subclass).

Parameters:

state (StorageState) – Current machine state (for active_step and the event timestamp).
chat (Chat) – The chat the message belongs to.
agui_id (str | None) – The AG-UI message ID, when known.
role (str | None) – The message role, defaulting to assistant.

Returns:

The DTO passed to MessageRepositoryLike.create_message().

Return type:

MessageCreate

build_tool_call_create(state, tool_call_id, tool_call_name)

Build the DTO used to append a tool call to the active message.

Override to populate extra fields declared on a ToolCall subclass.

Return type:: MessageToolCallCreate

build_reasoning_create(state, name)

Build the DTO used to append a reasoning step to the active message.

Override to populate extra fields declared on a Reasoning subclass.

Return type:: MessageReasoningCreate

message_update_fields(state, event)

Return extra fields to merge into a terminal message update.

The default is empty. Override to persist extra fields (declared on a Message subclass) when a message completes. Unknown keys are ignored by the base MessageUpdate.

Return type:: dict[str, Any]

async run(input)

Persist inbound user messages, replay history, then stream the inner agent.

An inbound message not already persisted in this chat must be a user message; a new non-user message yields a terminal RunErrorEvent with the ErrorCodes.INVALID_INPUT code and stops the run. “Already persisted” spans the tool calls and reasoning steps nested in a stored message, so a client that echoes its full message list back — tool results and reasoning steps included, as AG-UI clients normally do — replays cleanly instead of tripping that guard on a record that is in fact already stored. The inner agent’s stream is yielded verbatim while a background task persists it; on cancellation, still-active records are flipped to interrupted, and when the inner agent crashes with a raw exception they are flipped to errored and the exception is re-raised.

Parameters:: input (RunAgentInput) – The AG-UI run input. input.messages is replaced in place with the full translated chat history before the inner agent runs.
Yields:: ag_ui.core.BaseEvent – The inner agent’s events (or a terminal error event).
Return type:: AsyncGenerator[BaseEvent, None]

async handle_run_lifecycle(state, chat, event)

Reset state on run start; flush and finalize records on finish / error.

Return type:: None

async handle_step(state, chat, event)

Track the active step name across StepStarted / StepFinished.

Return type:: None

async handle_text_message(state, chat, event)

Fold text-message events onto the active message’s content.

Override to store structured content instead of raw appended deltas.

Return type:: None

async handle_tool_call(state, chat, event)

Fold tool-call events onto a tool call nested in the active message.

Return type:: None

async handle_reasoning(state, chat, event)

Fold reasoning events onto a reasoning step nested in the active message.

Handles both the current AG-UI Reasoning* events and the deprecated Thinking* events with identical persistence semantics; the two families map one-to-one:

Deprecated `Thinking*`	Current `Reasoning*`
`ThinkingStartEvent`	`ReasoningStartEvent`
`ThinkingEndEvent`	`ReasoningEndEvent`
`ThinkingTextMessageStartEvent`	`ReasoningMessageStartEvent`
`ThinkingTextMessageContentEvent`	`ReasoningMessageContentEvent`
`ThinkingTextMessageEndEvent`	`ReasoningMessageEndEvent`
(none)	`ReasoningMessageChunkEvent`
(none)	`ReasoningEncryptedValueEvent`

The Reasoning* events additionally carry a message_id which is persisted as the reasoning step’s agui_id for correlation; the Thinking* events carry only an optional title (persisted as the step’s name).

Override to store structured content instead of raw appended deltas.

Return type:: None

async flush_message_buffer(state)

Persist and clear buffered message content, if any.

Return type:: None

async flush_tool_call_buffer(state)

Persist and clear buffered tool-call arguments, if any.

Return type:: None

async flush_reasoning_buffer(state)

Persist and clear buffered reasoning content, if any.

Return type:: None

Stream manager

class datarobot.application_utils.chat_history.StreamPersistenceManager

Run AG-UI storage agents so their output survives client disconnects.

The manager builds a fresh AGUIStorageAgent per run from an injected factory, then spawns a background producer task that drains the agent’s stream into an unbounded queue. It owns the instance-scoped registry of in-flight runs, keyed by (thread_id, run_id), used by cancel().

async run(input, *args, **kwargs)

Start a run and return a RunHandle.

Spawns a producer task that builds the storage agent (via the factory, with args / kwargs), iterates its stream, and drains every event into an unbounded queue. The run is registered under (thread_id, run_id) for the lifetime of the producer; the producer unregisters itself when it finishes.

Parameters:

input (ag_ui.core.RunAgentInput) – The AG-UI run input; its thread_id / run_id key the registry.
*args (ParamSpecArgs) – Forwarded verbatim to the agent factory.
**kwargs (ParamSpecKwargs) – Forwarded verbatim to the agent factory.

Returns:

A handle exposing the run’s event stream, cancellation and a completion await.

Return type:

RunHandle

cancel(thread_id, run_id)

Cancel the run keyed by (thread_id, run_id).

Cancels the producer task, propagating asyncio.CancelledError into the storage agent’s run so its interrupt finalization marks still-active records interrupted.

Parameters:

thread_id (str) – The run’s AG-UI thread id.
run_id (str) – The run’s AG-UI run id.

Returns:

True if a live run was found and cancellation requested; False when no matching run exists (already finished or unknown).

Return type:

bool

class datarobot.application_utils.chat_history.RunHandle

A handle to one in-flight run started by StreamPersistenceManager.

Variables:

thread_id (str) – The AG-UI thread ID of the run.
run_id (str) – The AG-UI run ID of the run.

async events()

Yield the run’s events until the terminating sentinel.

Reads the unbounded producer queue, sleeping briefly when it is empty. The generator ends when it dequeues NoMoreEvents; because the producer always enqueues that sentinel, this can never hang. A consumer may stop iterating at any time (a client disconnect) — the producer keeps draining and persisting regardless.

Yields:: ag_ui.core.BaseEvent – Each event the producer forwarded, in order.
Return type:: AsyncGenerator[BaseEvent, None]

cancel()

Cancel this run.

Returns:: True if a live run was found and cancellation requested; False if the run had already finished.
Return type:: bool

async wait()

Wait until the producer finishes (the stream is fully drained and persisted).

Returns after the storage agent’s run — including its guaranteed final flush and, on cancellation, interrupt finalization — has completed. This never raises for a cancelled or failed run: the producer captures those outcomes internally (a failure is surfaced as a synthesized RunErrorEvent).

Return type:: None

Identifier helpers

The deterministic key derivations the models rely on (all importable from datarobot.application_utils.chat_history):

datarobot.application_utils.chat_history.constants.chat_deduplication_key(user_uuid, thread_id)

Return the deduplication key for a chat, keyed by user and AG-UI thread id.

Parameters:

user_uuid (UUID) – Owning user’s UUID.
thread_id (str) – AG-UI thread identifier.

Returns:

The SHA-256 of "chat", the user UUID and the thread id (NUL-separated), truncated to DEDUPLICATION_KEY_LENGTH characters. Idempotent: a retried create for the same (user, thread) adopts the existing session.

Return type:

str

datarobot.application_utils.chat_history.constants.session_deduplication_key(namespace, *parts)

Build a stable, namespaced deduplication key for idempotent session create.

Parameters:

namespace (str) – Logical document namespace (e.g. "chat").
*parts (str) – Ordered key components; combined with NUL separators before hashing.

Returns:

The lowercase hex SHA-256 digest, truncated to DEDUPLICATION_KEY_LENGTH characters.

Return type:

str

datarobot.application_utils.chat_history.constants.participant_id(user_uuid, *, override=None)

Return a stable 24-hex participant id for a user.

Unlike the agent-application helper this is transport-agnostic: it takes an explicit override argument instead of reading request/middleware context.

Parameters:

user_uuid (UUID) – The user’s UUID; hashed to derive a deterministic ObjectId-shaped id.
override (str | None) – An explicit participant id (e.g. a DataRobot user id). When it normalizes to a valid 24-hex value it is used verbatim; otherwise the derived value is returned.

Returns:

A 24-character lowercase hex participant id.

Return type:

str

datarobot.application_utils.chat_history.constants.normalize_participant_id(raw)

Normalize a caller-supplied participant id to 24-char lowercase hex, or None.

Parameters:: raw (str | None) – A candidate participant id (e.g. a X-DataRobot-User-Id header value).
Returns:: The normalized 24-hex id, or None when raw is missing or not a valid 24-character hexadecimal string.
Return type:: str | None

datarobot.application_utils.chat_history.constants.DEDUPLICATION_KEY_LENGTH: int = 64: Memory Service deduplication keys may be up to 72 characters; we truncate the hex digest to a stable 64.

Running acceptance tests

The chat-history acceptance suite drives a scripted inner agent against a live Memory Service. It is skipped by default and requires credentials:

export DATAROBOT_ENDPOINT="https://app.datarobot.com/api/v2"
export DATAROBOT_API_TOKEN="<your-token>"

pytest tests/application_utils/chat_history/acceptance -m integration -vv