fix(watsonx): persist the orchestrate-issued thread_id for conversation continuity

[guidovizoso]

Problem

watsonx Orchestrate's /chat/completions endpoint is stateful: conversation history is stored server-side, keyed by a thread_id that watsonx issues and returns in the SSE stream. It ignores client-invented thread IDs (creating a fresh thread instead) and only acts on the last user message in the payload.

Both watsonx adapters (TypeScript and Python) sent the AG-UI threadId as the X-IBM-THREAD-ID header and discarded the thread_id returned in the stream. As a result, every turn started a brand-new orchestrate conversation and the agent lost all prior context — multi-turn memory was completely broken.

Reported by a user of @ag-ui/watsonx@0.0.1 with CopilotKit chat, who confirmed that extracting the thread_id from the RAW SSE events and replaying it as X-IBM-THREAD-ID restores conversation memory.

Fix

Both adapters now natively support the watsonx Orchestrate conversation model:

• Omit X-IBM-THREAD-ID on the first turn of a thread so orchestrate creates the thread (sending a client-invented ID just produced orphan threads).
• Capture the watsonx-issued thread_id from SSE chunks — including chunks that carry no choices, which is where it can first appear. The TypeScript adapter persists it in the stream's finally so a mid-stream error doesn't lose it; Python persists eagerly as chunks arrive.
• Store the AG-UI threadId → watsonx thread_id mapping in a thread-ID store on the agent: in-memory by default, pluggable via threadIdStore (TS) / thread_id_store (Python) with sync-or-async get/set for server deployments where agent instances are created per-request. clone() shares the store so clones continue existing threads.
• Reuse the stored watsonx ID as X-IBM-THREAD-ID on subsequent turns.
• Trim the payload when continuing a known thread: only messages from the last user message onward are sent (trailing assistant/tool messages are kept so tool results still reach watsonx), since earlier context lives server-side and orchestrate ignores the rest.

Existing workarounds that consume the RAW events keep working — RAW emission is unchanged.

Testing

• TypeScript: 65 tests pass, including 9 new ones in thread-persistence.test.ts covering first-turn header omission, capture/reuse, per-thread tracking, choice-less chunks, payload trimming, custom/async stores, and clone(). tsc --noEmit and tsdown build clean.
• Python: 43 tests pass, including 8 new ones in TestThreadPersistence mirroring the TypeScript coverage.
• One pre-existing test per language asserted the old header behavior and was updated to assert the new contract.

[mxmzb]

Nice find on the root cause — capturing the orchestrate-issued thread_id, omitting the header on the first turn, and sharing the store through clone() is the right shape, and I like that MESSAGES_SNAPSHOT still uses the untrimmed input and RAW emission is genuinely untouched.

I checked out the branch and ran both suites: 65 TS / 43 Python pass, tsc --noEmit clean. Claims all hold up.

Four things I'd want addressed before merge. All four I confirmed by running code against the branch, not by reading the diff.

---

1. A dead watsonx thread_id is never invalidated, so the conversation wedges permanently

Once a bad ID lands in the store there's no path back out. I seeded the store with wx-dead, made the chat call return 404, and after the failed run the store still held wx-dead. Every subsequent turn on that AG-UI thread re-sends it and fails identically until the process restarts.

Worth calling out that this is a new failure mode: before this PR the header was always the AG-UI threadId, so a thread watsonx didn't recognise just meant a fresh thread. Now an expired or deleted orchestrate thread is terminal.

Suggestion: when the chat call returns 4xx and we sent a stored ID, clear the mapping — ideally retry once with the header omitted so the turn still succeeds.

2. Python: the store read sits outside the try, so the stream dies with no terminal event

agent.py:186 does the get, but the try that produces RUN_ERROR doesn't start until line 219. A store get that raises escapes the async generator entirely. What I actually observed: RUN_STARTED, then nothing — no RUN_ERROR, no RUN_FINISHED. The exception surfaces inside FastAPI's event_generator and the SSE response just dies mid-flight.

The TypeScript side handles the identical case correctly (RUN_STARTED → RUN_ERROR), so this is also a cross-language divergence. Moving line 186 inside the try should be all it takes.

3. TypeScript: a threadIdStore.set() failure in the finally breaks the event lifecycle

index.ts:336 awaits the store write inside the finally. With a store whose set throws, I got:

RUN_STARTED, STEP_STARTED, RAW, TEXT_MESSAGE_START, TEXT_MESSAGE_CONTENT, STEP_FINISHED, RUN_ERROR

No TEXT_MESSAGE_END, no MESSAGES_SNAPSHOT, no RUN_FINISHED. The answer streamed through fine, and the client is left holding an unterminated text message plus an error.

Second-order problem with the same line: because it throws from a finally, it also replaces any genuine stream exception that was already propagating, so the real cause gets swallowed. A try/catch-and-log around the write fixes both.

4. Python: the store write is ahead of the content handling and awaited inline

agent.py:284-292 runs before the choices block, on the very first chunk. With a throwing store the entire answer is lost — I measured zero TEXT_MESSAGE_CONTENT deltas, where TS at least delivers the text before erroring.

There's a quieter version of this too: a merely slow store — which is exactly the database-backed case the README recommends — applies backpressure to SSE consumption on every single chunk, not just once. Persisting after the content handling (and guarding it) covers both.

---

One design question rather than a defect

On a tool-result turn the message list is [user, assistant(tool_calls), tool] with no new user message, and messagesSinceLastUser re-sends that older user message alongside the continuation header. If orchestrate really does only act on the last user message in the payload, it may re-answer the previous question instead of consuming the tool result.

Not a regression — the old code sent the full history every turn — but the header is new, so the combination is untested territory. Worth confirming against a real orchestrate agent that calls a tool before we rely on it.

Minor

Two concurrent first turns on the same threadId both omit the header, so you get two watsonx threads and one orphan, last write wins. Probably fine to note and move on.

Happy to push the fixes for 1–4 if it's easier than doing them yourself.

[mxmzb]

Blocking on the four items in my comment above. Short version of what needs to change:

1. Clear the stored thread_id when watsonx rejects it. Right now a dead orchestrate thread wedges the AG-UI thread permanently — I seeded wx-dead, returned 404, and the store still held wx-dead afterwards. This is a new failure mode; pre-PR a stale thread just meant a fresh one.
2. Python: move the thread_id_store.get() at agent.py:186 inside the try. A raising get escapes the generator, so the client sees RUN_STARTED and then nothing at all — no RUN_ERROR, no RUN_FINISHED. TS already handles this case, so it's a divergence too.
3. TS: guard the threadIdStore.set() at index.ts:336. Throwing from that finally drops TEXT_MESSAGE_END / MESSAGES_SNAPSHOT / RUN_FINISHED and masks any real stream exception that was already propagating.
4. Python: move the store write at agent.py:284-292 after the content handling and guard it. As placed, it fires on the first chunk, so a failing store loses the whole answer, and a slow one (the database-backed case the README recommends) applies backpressure on every chunk.

The design question about re-sending the last user message on tool-result turns isn't blocking, but I'd like an answer before this goes out — that one wants a real orchestrate agent with a tool to settle.

Approach and test coverage are good; this is all failure-path handling around the new store. Happy to push the fixes if that's faster for you.

[ranst91]

I'm "requesting changes" for a different reason than @mxmzb

The idea of persisting a map between watsonx thread ID and client sent thread ID would cause all sorts of mismatches and confusion, to be visible with observability tools. I'd like to explore different ways. I know that watsonx would allow you to strictly define the ID yourself, so that's a direction to explore.

Dismissing as to not block this PR