LongRunningAgentServer: Durable prose-recovery + always-rotate

[dhruv0811]

Summary

Adds durable, crash-resumable execution to MLflow ResponsesAgent handlers via LongRunningAgentServer. A pod that dies mid-stream is seamlessly recovered by another pod via prose recovery — the resumed attempt receives a single user message containing the prior attempt's stream events as JSON plus a directive to figure out what's done and continue.

The bridge is minimal: heartbeat, proactive stale scan, atomic CAS claim, conversation_id rotation, prose-recovery message construction. Echo dedup (when the chat client echoes the full conversation history) is the agent handler's responsibility, since the agent owns its SDK session/checkpointer state.

Full design documentation in src/databricks_ai_bridge/long_running/AGENTS.md including 4 customer journeys, sequence diagrams, ER diagram, and the planned TaskFlow migration mapping.

High-level flow

flowchart TD
    A[Client: POST /responses<br/>background:true] --> B[Bridge: store FULL original_request,<br/>spawn handler, heartbeat every 3s]
    B --> C{Run completes?}
    C -->|yes| D[status=completed<br/>events streamable via GET]
    C -->|pod crashes| E[heartbeat stops]
    E -->|stale > 10s| F1{stale detected by:}
    F1 --> F2[client GET /responses/&#123;id&#125;]
    F1 --> F3[scanner loop ~30s ± 50% jitter]
    F2 --> G[CAS claim attempt+=1<br/>build prose recovery msg<br/>rotate conv_id ::attempt-N]
    F3 --> G
    G --> H[Re-invoke handler on<br/>fresh rotated SDK session]
    H --> C
    G -.->|response.resumed sentinel<br/>carries rotated conv_id| I[Chatbot captures rotation,<br/>uses rotated id on next turn<br/>→ lands on clean rotated session]

Loading

What the bridge does

Heartbeat + atomic CAS claim. The owning pod writes heartbeat_at every 3s. The heartbeat write is CAS-checked on attempt_number — if another pod has bumped the row to attempt_number = N+1, the prior pod's heartbeat at attempt N returns rowcount=0 and stops. Ownership is implicit: the pod that last successfully heartbeats at the current attempt is the de facto owner. There is no owner_pod_id column.

When GET /responses/{id} (or the proactive scanner) sees the row's heartbeat is older than 10s, the new pod atomically claims via:

UPDATE responses
SET attempt_number = attempt_number + 1, heartbeat_at = now()
WHERE response_id = :rid AND attempt_number = :current AND heartbeat stale
RETURNING attempt_number

Postgres row locking ensures exactly one pod wins.

Prose recovery. On claim, the new pod filters the message log to the prior attempt's stream events, json.dumps-serializes them, and wraps them in a single user message:

[RECOVERY] The previous attempt of this agent task crashed mid-execution.
Below is the raw stream-event log from that attempt as JSON. Some tool calls
may have completed and some may have been interrupted before returning a
result. Inspect the events, figure out what is already done versus
in-progress / not completed, and continue the task from where it left off.
If a tool call was interrupted, you may re-invoke it if its result is still
needed.

Events: <json.dumps>

The model reads this and decides what to redo, what to skip. SDK-agnostic — no provider-specific tool-pair structure, no synthetic tool events, no per-SDK adapter code.

Always-rotate conversation_id. On resume, context.conversation_id is rotated to {base}::attempt-N so the SDK opens a fresh thread/session for the resumed attempt. The original session, which may carry orphan tool_use / partial checkpoint state from the crashed attempt, is never read again.

The rotated value is surfaced in the SSE response.resumed sentinel; cooperating chat clients (see app-templates#204) capture it and use the rotated session for subsequent turns.

Proactive stale-scan loop. Each pod runs a background scanner every ~30s (± 50% jitter) that queries for in-progress responses with stale heartbeats and tries to claim+resume them. Crashed responses get recovered even when no client polls. CAS-claim semantics ensure only one pod wins on contention.

Stream resume. GET /responses/{id}?stream=true&starting_after=N replays the message log past the cursor; a response.resumed sentinel marks attempt boundaries.

What is NOT in the bridge

• No SDK adapter wrappers. integrations/openai/.../session.py and integrations/langchain/.../checkpoint.py are byte-identical to main. No aget_tuple / get_items overrides.
• No echo dedup. The bridge does not manipulate request.input shape. When the chat client echoes the full prior conversation, the agent handler is responsible for deduping its input against the SDK's session/checkpointer state. See companion PR app-templates#204 for the per-template ~10-line implementations.

Files

Component	File
LongRunningAgentServer (heartbeat, claim, retrieve, rotation, scanner)	long_running/server.py
Prose recovery message builder	long_running/server.py::_build_prose_recovery_message
Async repository (CAS, heartbeat, append, retrieve, find-stale)	long_running/repository.py
Durability columns + idempotent migrations	long_running/db.py, models.py
Settings (heartbeat tuning, timeouts, scan jitter)	long_running/settings.py
Design doc with mermaid diagrams + CUJs	long_running/AGENTS.md
Tests	tests/databricks_ai_bridge/test_long_running_*.py

Settings

Setting	Default	Rationale
heartbeat_interval_seconds	3.0	Frequent enough that GC pauses don't trip stale detection
heartbeat_stale_threshold_seconds	10.0	Three missed heartbeats = unambiguously dead
task_timeout_seconds	3600	Hard ceiling per attempt
poll_interval_seconds	1.0	Stream-retrieve poll rate
stale_scan_interval_seconds	30.0	Scanner cadence
stale_scan_jitter_fraction	0.5	±50% jitter so multi-pod scans don't synchronize

Test plan

• 104 unit tests pass (tests/databricks_ai_bridge/test_long_running_server.py, test_long_running_db.py)
• ruff check, ruff format --check, ty check clean
• UI testing on dhruv-lg-adv-durable-prose (LangGraph + Claude Sonnet 4.5) — multi-tool kill mid-deep_research + multi-turn ✅
• UI testing on dhruv-oai-adv-durable-prose (OpenAI Agents SDK + GPT-5) — same matrix ✅
• UI testing on dhruv-oai-cl-durable-prose (OpenAI Agents SDK + Claude Sonnet 4.5) — same matrix ✅

Companion PR

app-templates#204 — chatbot proxy alias capture (cooperation for cross-turn rotation) + per-template UI-echo dedup.

Known follow-ups (non-blocking)

• DELETE /responses/{id} cancellation (currently 501).
• LISTEN/NOTIFY replacing the 1s poll loop in stream retrieve.
• Metrics: heartbeat_miss_count, claim_latency_seconds, attempts_per_response, stale_scan_iterations.

See AGENTS.md §5 for the planned TaskFlow migration — LongRunningAgentServer's HTTP surface and @invoke()/@stream() contract stay the same when the engine internals swap.

Text.Only.-.Prose.mov

Tool.Calling.Multiturn.-.Prose.mov

[bbqiu]

okay i left a few comments, but i think we should revisit the overall appraoch

it seems like there are lots of little hacks/cases that we added for compatibility with each of hte specific langgraph / openai agents examples that we have. taking a step back, i think the server should be responsible for:

1. heartbeat per response that is running
2. loop that scans for stale heartbeats (should have some random delay + jitter to decrease multiplle readers of the db at once)
3. resumption logic that rotates the conversation ID + summarizes existing things
   a. for followup messages after rotation, the dedupe should be in agent logic. this case of history from client + history from internal agent history should be handled

[bbqiu]

note for myself to come back later

[bbqiu]

this deduping should already happen within the agent right? this doesn't belong here

this is not very diff from client submitting more info than what is already in the session thread

[dhruv0811]

Did some analysis with claude on why this is needed. I agree that this might not be the best place for this, happy to discuss if we want to either move this into the Openai, Langgraph wrappers or the app templates themselves.

agent-openai-advanced/agent_server/utils.py:189 does have a deduplicate_input function — but its heuristic (len(session_items) >= len(messages) - 1) breaks under always-rotate. The OpenAI SDK's own deduplicate_input_items_preferring_latest (agents/run_internal/items.py:280) does dedup tool calls by call_id, but _dedupe_key (line 224) explicitly returns None for any item with a role field — so user/assistant message echo isn't covered. LangGraph's add_messages dedupes by id, but MLflow's responses_to_cc doesn't preserve IDs (mlflow/types/responses.py:315-383), so cross-boundary dedup never fires there either.

[dhruv0811]

After discussion on Slack, we decided to move the dedup logic into the app-templates instead!

[bbqiu]

i think we can simplify this to create a new UUID for the convo thread. are tehre benefits to keeping the naming scheme?

[dhruv0811]

We have some form of debug readability with this. Say the user looks into their conv store and sees the json dump + our injected system prompt, they might be confused as to where this came from. The id having a attempt-N part to it helps provide some visibility into this.

[bbqiu]

(nit for an edge case) whenever we write terminal status to the db, we should check that the current attempt number == the attempt number in the db (meaning this pod still owns the current response)

[bbqiu]

claude's comment w/ more details:

## Guard terminal status writes on `attempt_number`

Every place we write a terminal status (`completed` / `failed`) should verify this pod still owns the run by gating on the `attempt_number` the task was started with. Without this guard, a pod whose heartbeat stalled and lost ownership via `claim_stale_response` can clobber the row's status while another pod is mid-resume — pollers and streamers stop early on a stale `completed`/`failed`.

**Sites that need the guard:**
- `_do_background_stream` final `update_response_status(response_id, "completed")` (`server.py:937`)
- `_do_background_invoke` final `update_response_status(response_id, "completed")` (`server.py:1018`)
- `_task_scope` error cleanup `update_response_status(response_id, "failed")` (`server.py:821`)
- `_deferred_mark_failed` `update_response_status(response_id, "failed")` (`server.py:94`)
- Pre-flight `update_response_status(response_id, "failed")` when no stream/invoke fn is registered (`server.py:871`, `:978`)

**Shape:**
```python
async def update_response_status(
    response_id: str,
    status: str,
    *,
    expected_attempt_number: int | None = None,
    expected_current_status: str | None = None,
) -> bool:
    stmt = update(Response).where(Response.response_id == response_id)
    if expected_attempt_number is not None:
        stmt = stmt.where(Response.attempt_number == expected_attempt_number)
    if expected_current_status is not None:
        stmt = stmt.where(Response.status == expected_current_status)
    ...

Thread attempt_number through to every callsite. It's already in scope inside _run_background_stream / _run_background_invoke; _task_scope and _deferred_mark_failed need it passed in (currently they only get response_id).

A False return is the caller's signal "you no longer own this run" — log it and exit cleanly, don't escalate to an error.

[bbqiu]

lgtm after fixing the last nit

[bbqiu]

claude's comment w/ more details:

## Guard terminal status writes on `attempt_number`

Every place we write a terminal status (`completed` / `failed`) should verify this pod still owns the run by gating on the `attempt_number` the task was started with. Without this guard, a pod whose heartbeat stalled and lost ownership via `claim_stale_response` can clobber the row's status while another pod is mid-resume — pollers and streamers stop early on a stale `completed`/`failed`.

**Sites that need the guard:**
- `_do_background_stream` final `update_response_status(response_id, "completed")` (`server.py:937`)
- `_do_background_invoke` final `update_response_status(response_id, "completed")` (`server.py:1018`)
- `_task_scope` error cleanup `update_response_status(response_id, "failed")` (`server.py:821`)
- `_deferred_mark_failed` `update_response_status(response_id, "failed")` (`server.py:94`)
- Pre-flight `update_response_status(response_id, "failed")` when no stream/invoke fn is registered (`server.py:871`, `:978`)

**Shape:**
```python
async def update_response_status(
    response_id: str,
    status: str,
    *,
    expected_attempt_number: int | None = None,
    expected_current_status: str | None = None,
) -> bool:
    stmt = update(Response).where(Response.response_id == response_id)
    if expected_attempt_number is not None:
        stmt = stmt.where(Response.attempt_number == expected_attempt_number)
    if expected_current_status is not None:
        stmt = stmt.where(Response.status == expected_current_status)
    ...

Thread attempt_number through to every callsite. It's already in scope inside _run_background_stream / _run_background_invoke; _task_scope and _deferred_mark_failed need it passed in (currently they only get response_id).

A False return is the caller's signal "you no longer own this run" — log it and exit cleanly, don't escalate to an error.

[bbqiu]

should we only do this with a matching attempt number?