Generalize long-running agent recovery and replay - #467
Conversation
|
@smurching - Should we keep some examples in this repo as well? |
| For framework-managed recovery's rotated-session flow to work cross-turn, a cooperating chat UI needs to: | ||
|
|
||
| 1. **Capture the rotated `conversation_id` from the SSE `response.resumed` event** when one is emitted during a streaming retrieve. | ||
| 2. **Use the rotated value as `context.conversation_id` on subsequent requests** for the same chat. |
jamesbxwu
left a comment
There was a problem hiding this comment.
Went through first pass and here are a few structural comments:
- Can we get rid of the
sse_replayknob and keep the existing behavior of always writing to the event log and support streaming replay? - Should the
examplesfolder be defined at the top level vs inside ofsrc/long_running? It doesn't seem thatexamplesas they are today covers everything inside of thedatabricks-ai-bridgerepo - Consider renaming
examplestocookbooks - I found the examples to be incredibly hard to follow and their corresponding README.md to be difficult to compare. Can we have a top level README.md in the
examplesfolder that explains the different configurations for the long running agent server and have a single summary table comparing them? - I feel the review agent is too complicated to understand, let's keep the initial examples/cookbook very simple and easy to follow. At this time, we are focused on demonstrating top level wiring, lifecycle hooks and etc and less about showing a complicated review agent
| | Runtime durability | `agent_server.responses` | Original request, terminal response, status, heartbeat, attempt, handler mode | | ||
| | Durable event log | `agent_server.messages` | No rows; the table exists only for schema compatibility | |
There was a problem hiding this comment.
is it still possible to rename these to be more accurate such as agent_server.checkpoints and ``agent_server.events`?
There was a problem hiding this comment.
We can, although there would be backward incompatibility. Should we think about this in a later PR and keep the names as is for now?
| Integer, nullable=False, server_default="1", default=1 | ||
| ) | ||
| original_request: Mapped[str | None] = mapped_column(Text, nullable=True) | ||
| response: Mapped[str | None] = mapped_column(Text, nullable=True) |
There was a problem hiding this comment.
In cases, where we will not store the event log, we need to store the final response to return the respone in background mode.
If it is a background non-streaming agent with agent managed recovery, we won't be storing the event logs, but need to store the final result.
|
|
||
| @dataclass(frozen=True) | ||
| class ResumeContext: | ||
| """Metadata and default behavior available to an ``@on_resume`` handler.""" |
There was a problem hiding this comment.
can we improve the documentation here to cover
- what are previous events
- what does _default_request include and why should this seemingly generic concept of a
ResumeContextneed to differentiate default vs not?
There was a problem hiding this comment.
added documentation. The default request is use the server default. If it still sounds confusing I can take a stab to simplify this.
jamesbxwu
left a comment
There was a problem hiding this comment.
had a few more questions on the combination of using ResumeStrategy.AGENT_SESSION and on_resume() together, let's chat offline
| def _session_id(request: ResponsesAgentRequest) -> str: | ||
| custom_inputs = dict(request.custom_inputs or {}) | ||
| if custom_inputs.get("session_id"): | ||
| return str(custom_inputs["session_id"]) | ||
| if request.context and request.context.conversation_id: | ||
| return request.context.conversation_id |
There was a problem hiding this comment.
this method needs a bit more explanation
| _on_resume_function: Callable[..., Any] | None = None | ||
|
|
||
|
|
||
| class ResumeStrategy(str, Enum): |
There was a problem hiding this comment.
let's rename this to "recovery" so that it is differentiated from planned pause/resume such as HITL? Let's update all references from resume -> recovery in that case.
| class ResumeStrategy(str, Enum): | |
| class RecoveryStrategy(str, Enum): |
| f"ALTER TABLE {AGENT_DB_SCHEMA}.responses " | ||
| "ADD COLUMN IF NOT EXISTS attempt_number INTEGER NOT NULL DEFAULT 1", | ||
| f"ALTER TABLE {AGENT_DB_SCHEMA}.responses ADD COLUMN IF NOT EXISTS original_request TEXT", | ||
| f"ALTER TABLE {AGENT_DB_SCHEMA}.responses ADD COLUMN IF NOT EXISTS terminal_response TEXT", |
There was a problem hiding this comment.
agent responses can be quite long and verbose, any concerns on db row/column size limits if we store this?
There was a problem hiding this comment.
this server file is becoming quite long (1500+ lines) and combines invoke entry points, heartbeat + recovery. I personally had a hard time tracking all the changes. Can we refactor them into such logical parts for better modularity? I recognize this may be out of scope for this PR but let's discuss whether this should be deferred.
| def _agent_session_recovery_message() -> dict[str, Any]: | ||
| return { | ||
| "type": "message", | ||
| "role": "user", | ||
| "content": AGENT_SESSION_RECOVERY_PROMPT, | ||
| } |
There was a problem hiding this comment.
instead of defining a method, why not just make this a constant similar to AGENT_SESSION_RECOVERY_PROMPT?
| emitted tool calls / outputs / narrative, and an ``[INTERRUPTED]`` synthetic | ||
| output paired with any tool call that didn't finish. Completed work is | ||
| preserved; only the interrupted step re-runs. | ||
| When a heartbeat becomes stale, another pod atomically claims the execution |
There was a problem hiding this comment.
| When a heartbeat becomes stale, another pod atomically claims the execution | |
| When a heartbeat becomes stale, another worker atomically claims the execution |
| stream_event={"trace_id": span.trace_id}, | ||
| attempt_number=attempt_number, | ||
| ) | ||
| if return_trace_id: |
There was a problem hiding this comment.
is this redundant with line 1014 above?
| if return_trace_id: |
| @on_resume() | ||
| async def resume(request, context: ResumeContext): | ||
| resumed = await context.default_resume_request(request) | ||
| return resumed |
There was a problem hiding this comment.
i didn't see an example of this in the cookbook, did I miss it?
| return _on_resume_function | ||
|
|
||
|
|
||
| def on_resume() -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: |
There was a problem hiding this comment.
it's not immediately clear that on_resume() is used to reconstruct an incoming request during crash recovery.
- if that is indeed the hook we are exposing, then we should rename this to
build_recovery_request - right now the decorator here does not restrict the input and output, should we do that?
| - `ResumeStrategy.AGENT_SESSION` keeps the original session key and replaces | ||
| the request input with a recovery prompt. The agent SDK supplies the history. |
There was a problem hiding this comment.
how should the developer reason about using AGENT_SESSION and implementing on_resume? My understanding is that setting ResumeStrategy.AGENT_SESSION means I get automatic recovery if my harness/sdk has a session store so when should I implement on_resume()?
|
Interface comparison set:
The first three use the same tiny progress agent and the same run/heartbeat/event durability semantics so the developer-facing differences are directly comparable. |

What did you change, and why?
Change: Generalize
LongRunningAgentServeraround two explicit crash-recovery strategies while keeping durable background execution and SSE replay in the server.ResumeStrategy.EVENT_LOGreconstructs recovery context from the immediately previous attempt’s durable events and rotates the agent session to<session-key>::attempt-N.ResumeStrategy.AGENT_SESSIONkeeps the same agent session key and sends a fixed recovery prompt; the agent SDK/harness restores its own transcript.@on_resume()request transformer runs only after this process wins a stale-attempt claim. Without an override,ResumeContext.default_resume_request()applies the selected strategy. The resulting request then goes through the same@invoke()or@stream()handler mode as the original attempt; normal handlers do not contain recovery branches.context.conversation_id,custom_inputs.session_id, orcustom_inputs.thread_idas the harness session anchor. If all are absent, the server warns and injects the generated durableresponse_idascontext.conversation_id.@stream(), because the streamed events are its recovery input. Agent-session recovery preserves the original invoke/stream handler mode.PGDATABASEresource binding instead of always connecting todatabricks_postgres.0correctly instead of replaying the first event on every polling iteration.Why: The earlier implementation mixed runtime durability, SSE replay, and agent conversation recovery. This makes the ownership boundaries explicit and lets an agent harness use its native session persistence instead of forcing every recovery through generated event-log prose.
Recovery contract
EVENT_LOG[RECOVERY]prose containing the previous attempt’s durable events<session-key>::attempt-NAGENT_SESSION[RECOVERY]promptRecovery restarts a handler; it does not resume a suspended Python coroutine. Interrupted side effects are therefore at-least-once and tools still need idempotency.
Three separate stores
The implementation has three logically independent stores. They may share one Lakebase database, but they have different owners and purposes.
LongRunningAgentServeragent_server.responsesresponse_id, original request, terminal response, status, heartbeat, attempt number, and original invoke/stream modeLongRunningAgentServeragent_server.messagesstarting_afterSSE replay; also read byEVENT_LOGrecovery<session schema>.agent_sessionsand.agent_messagesAGENT_SESSIONrecovery; the server neither creates nor interprets this transcriptresponse_ididentifies the durable HTTP operation. The harness session ID identifies conversational history. They are not inherently the same. For agent-session recovery, the selected session anchor remains inside the persistedoriginal_request; it is not a separate durability-table column.Public cookbooks
The PR contains the same small wait-tool agent in two harnesses:
cookbooks/openai-sdk-agentcookbooks/langgraph-agentEach cookbook has
event_logandagent_sessionbundle targets. A rootapp.yamland shared App entry point keep the command and resource bindingsavailable after
databricks apps stop/start; the stable App name selects thesame recovery strategy and SDK-owned schema on restart. Handler and agent logic
remain shared, while the small strategy-specific entry points show the direct
wiring. Deployment helpers and test-result artifacts are not part of the PR.
How do you know it works?
Automated validation
App build logs also confirm dependencies install successfully through the Databricks PyPI proxy.
Deployed environments
Both latest deployments are
RUNNING, their active deployments areSUCCEEDED, and their dedicated Lakebase branches areREADY.durable-framework-ssedurable-session-ssePlatform state was checked with:
Client protocol exercised
All live tests used the public
POST /responses,GET /responses/{response_id}, and SSE replay contract:Happy path and disconnect/replay
resp_b6e5d83e6b1f44549a8fae9bcompleted on attempt 1. For responseresp_bba2c7517912463abf54b2ed, the client disconnected after sequence 1; replay returned sequences 2–35,response.completed, and[DONE].resp_e2da1b63155b4280902b6409completed on attempt 1. For responseresp_8c1da65aafc14e269fd4c87b, the client disconnected after sequence 10; replay returned sequences 11–36,response.completed, and[DONE].response_idreturned the stored terminal Responses payload in both Apps.Real App stop/start: event-log recovery
Input session:
pr467-event-crash-20260821T184027ZDurable response:
resp_1b2c026380fb4fb2a7c1f3d8curl: (28), HTTP 200 with 8,579 bytes already received).databricks apps stop openai-agent-framework-sse --profile ml-inference-stagingstopped compute while the tool was still running.in_progress, attempt 1, no terminal response, 11 event rows (sequences 0–10), and one message in the original SDK session.databricks apps start ..., stale-heartbeat claiming advanced the row to attempt 2.response.resumed; the stream ended withresponse.completedand[DONE].completed, attempt 2. Attempt 1 has 11 event rows; attempt 2 has 25. The SDK session rotated topr467-event-crash-20260821T184027Z::attempt-2.[RECOVERY]input containing serialized prior event/tool-call evidence and producedPR467_EVENT_CRASH_RECOVERED_20260821T184027Z.Real App stop/start: agent-session recovery
Input session:
pr467-session-crash-20260821T184221ZDurable response:
resp_516816b5e26f49a4a98248f1in_progress; Lakebase had 11 event rows (sequences 0–10), one SDK user message, and no terminal response.[DONE]. The first event wasresponse.resumed; the last wasresponse.completed.agent_server.responsesstate iscompleted, attempt 2, with a persisted terminal response. Event rows are 11 for attempt 1 and 37 for attempt 2.[RECOVERY]prompt, the resumed tool call, its output, and the final assistant message. No::attempt-2session exists.PR467_SESSION_CRASH_RECOVERED_20260821T184221Z, demonstrating that harness-owned history was restored and the resumed loop completed.The Lakebase checks queried
agent_server.responses, groupedagent_server.messagesby attempt/sequence, and inspected the SDK-ownedagent_messagesrows for the exact session IDs. No test-result files are added to the open-source diff; the evidence is recorded here.LangGraph durability validation — August 25, 2026
Both LangGraph cookbook targets were deployed in
ml-inference-stagingusing theexact local PR wheel and separate empty databases under
projects/shivam-debug/branches/pr-cuj-harness:langgraph-event-recoverypr467-langgraph-eventlanggraph-session-recoverypr467-langgraph-sessionThe test for each App used
background=true,stream=true, a 60-secondwait_for_completiontool call, an actualdatabricks apps stopwhile the toolwas pending,
databricks apps start, SSE reconnection, final polling, and afollow-up conversational turn.
LangGraph EVENT_LOG recovery
Original conversation:
pr467-langgraph-event_log-20260825T184118ZDurable response:
resp_8d13bf3284084485a85bffe7response.output_item.donewith afunction_call. The App was stopped before a tool output existed.response.resumedwith rotated conversation IDpr467-langgraph-event_log-20260825T184118Z::attempt-2.response.resumed, a newfunction_call,function_call_output, and the final assistantmessage. This shows prose recovery started a fresh LangGraph thread and re-invoked the interrupted tool.completed, attempt 2, a persisted terminal response, andPR467_LANGGRAPH_EVENT_LOG_20260825T184118Z.response.resumed. The response wasFOLLOWUP_REMEMBERS:PR467_LANGGRAPH_EVENT_LOG_20260825T184118Z, proving that the client must adopt the emitted rotated conversation ID for later turns. The old ID points at the abandoned attempt-1 checkpoint.LangGraph AGENT_SESSION recovery
Conversation:
pr467-langgraph-agent_session-20260825T184625ZDurable response:
resp_387945bf600244e9ae579223function_call.response.resumedat sequence 1 with the unchanged conversation ID.response.resumed,function_call_output, and the final assistantmessage(sequences 1–3). There was no second function-call event, showingastream(None)resumed the pending LangGraph checkpoint instead of rebuilding the turn from prose.completed, attempt 2, a persisted terminal response, andPR467_LANGGRAPH_AGENT_SESSION_20260825T184625Z.::attempt-2thread was created.FOLLOWUP_REMEMBERS:PR467_LANGGRAPH_AGENT_SESSION_20260825T184625Z.The live deployment also caught and fixed three cookbook/runtime issues before
recording these passing runs: duplicate bundle App source paths, missing
restart-time
app.yamlresource bindings, and Lakebase clients ignoring theApps-provided
PGDATABASEvalue.Durable HITL cookbook flow
Both the OpenAI Agents SDK and LangGraph cookbooks now show the same agent-managed HITL pattern:
background=trueandstream=true.APPROVAL_REQUIRED; no worker or lease is held while waiting for a person.This keeps responsibilities explicit:
LongRunningAgentServermakes each Response durable, while the agent/harness defines the HITL protocol and stores the state needed by the next turn. Tool side effects remain at-least-once and require idempotency.Background and streaming client contract
background=truecreates a durable Responses record. The client receives anin_progressresponse ID and pollsGET /responses/{id}.@stream()yields Responses events; the server persists them and reconnects with?stream=true&starting_after=<sequence>.Both public cookbook READMEs now include concrete before/after client snippets and the exact
background=true, stream=true, polling, and reconnect flow.