Skip to content

Generalize long-running agent recovery and replay - #467

Open
shivam5 wants to merge 16 commits into
databricks:mainfrom
shivam5:durable-server
Open

Generalize long-running agent recovery and replay#467
shivam5 wants to merge 16 commits into
databricks:mainfrom
shivam5:durable-server

Conversation

@shivam5

@shivam5 shivam5 commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

What did you change, and why?

Change: Generalize LongRunningAgentServer around two explicit crash-recovery strategies while keeping durable background execution and SSE replay in the server.

  • ResumeStrategy.EVENT_LOG reconstructs recovery context from the immediately previous attempt’s durable events and rotates the agent session to <session-key>::attempt-N.
  • ResumeStrategy.AGENT_SESSION keeps the same agent session key and sends a fixed recovery prompt; the agent SDK/harness restores its own transcript.
  • Heartbeats, stale-attempt detection, atomic claiming, request persistence, terminal-response persistence, and ordered SSE event persistence/replay are server-owned behavior for both strategies.
  • An optional @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.
  • Agent-session recovery accepts context.conversation_id, custom_inputs.session_id, or custom_inputs.thread_id as the harness session anchor. If all are absent, the server warns and injects the generated durable response_id as context.conversation_id.
  • Event-log recovery requires @stream(), because the streamed events are its recovery input. Agent-session recovery preserves the original invoke/stream handler mode.
  • Completed and failed terminal Responses payloads are stored directly so polling does not reconstruct authoritative results from stream fragments.
  • Public OpenAI Agents SDK and LangGraph cookbooks use one shared, restart-safe App entry point with separate bundle targets for the two recovery strategies.
  • Lakebase connections honor the Databricks Apps PGDATABASE resource binding instead of always connecting to databricks_postgres.
  • SSE retrieval advances past sequence 0 correctly 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

Strategy Input passed to the resumed agent loop Session behavior Event-log use
EVENT_LOG Original request plus [RECOVERY] prose containing the previous attempt’s durable events Rotates to <session-key>::attempt-N SSE replay and recovery context
AGENT_SESSION Fixed [RECOVERY] prompt Reuses the same harness session so the SDK reloads its transcript SSE replay only

Recovery 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.

Store Owner Physical tables Purpose
Runtime durability LongRunningAgentServer agent_server.responses Durable response_id, original request, terminal response, status, heartbeat, attempt number, and original invoke/stream mode
Durable event log LongRunningAgentServer agent_server.messages Ordered stream events/output items used for starting_after SSE replay; also read by EVENT_LOG recovery
Agent session store Agent SDK/harness Example: <session schema>.agent_sessions and .agent_messages Logical conversation transcript used by AGENT_SESSION recovery; the server neither creates nor interprets this transcript

response_id identifies 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 persisted original_request; it is not a separate durability-table column.

Public cookbooks

The PR contains the same small wait-tool agent in two harnesses:

  1. cookbooks/openai-sdk-agent
  2. cookbooks/langgraph-agent

Each cookbook has event_log and agent_session bundle targets. A root
app.yaml and shared App entry point keep the command and resource bindings
available after databricks apps stop / start; the stable App name selects the
same 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

uv run pytest -q \
  tests/databricks_ai_bridge/test_long_running_db.py \
  tests/databricks_ai_bridge/test_long_running_server.py \
  tests/databricks_ai_bridge/test_lakebase.py
# 238 passed

uv run pytest -q tests/databricks_ai_bridge \
  --ignore=tests/databricks_ai_bridge/test_model_serving_obo_credential_strategy.py
# 329 passed, 120 skipped
# The ignored OBO credential-strategy test is unrelated to this change.

uv run ruff format --check .
uv run ruff check .
uv run ty check
# All exit 0; ty reports only three pre-existing unused-ignore warnings.

App build logs also confirm dependencies install successfully through the Databricks PyPI proxy.

Deployed environments

Both latest deployments are RUNNING, their active deployments are SUCCEEDED, and their dedicated Lakebase branches are READY.

Strategy Databricks App Runtime URL Lakebase tables
Event-log recovery App overview https://openai-agent-framework-sse-1653573648247579.staging.aws.databricksapps.com durable-framework-sse
Agent-session recovery App overview https://openai-agent-session-sse-1653573648247579.staging.aws.databricksapps.com durable-session-sse

Platform state was checked with:

env -u DATABRICKS_HOST -u DATABRICKS_TOKEN -u DATABRICKS_CONFIG_PROFILE \
  databricks apps get <app> --profile ml-inference-staging -o json

env -u DATABRICKS_HOST -u DATABRICKS_TOKEN -u DATABRICKS_CONFIG_PROFILE \
  databricks postgres get-branch \
  projects/shivam-openai-agent-on-apps/branches/<branch> \
  --profile ml-inference-staging -o json

Client protocol exercised

All live tests used the public POST /responses, GET /responses/{response_id}, and SSE replay contract:

curl -N -X POST "$APP_URL/responses" \
  -H "Authorization: Bearer $APP_TOKEN" \
  -H "Content-Type: application/json" \
  --data '{
    "background": true,
    "stream": true,
    "input": [{"role": "user", "content": "<long tool-call request>"}],
    "custom_inputs": {"session_id": "<manual test session>"}
  }'

curl -N \
  "$APP_URL/responses/$RESPONSE_ID?stream=true&starting_after=$LAST_SEQUENCE" \
  -H "Authorization: Bearer $APP_TOKEN"

curl "$APP_URL/responses/$RESPONSE_ID" \
  -H "Authorization: Bearer $APP_TOKEN"

Happy path and disconnect/replay

  • Event-log App: response resp_b6e5d83e6b1f44549a8fae9b completed on attempt 1. For response resp_bba2c7517912463abf54b2ed, the client disconnected after sequence 1; replay returned sequences 2–35, response.completed, and [DONE].
  • Agent-session App: response resp_e2da1b63155b4280902b6409 completed on attempt 1. For response resp_8c1da65aafc14e269fd4c87b, the client disconnected after sequence 10; replay returned sequences 11–36, response.completed, and [DONE].
  • Polling the completed durable response_id returned the stored terminal Responses payload in both Apps.

Real App stop/start: event-log recovery

Input session: pr467-event-crash-20260821T184027Z
Durable response: resp_1b2c026380fb4fb2a7c1f3d8

  1. The client started a 30-second tool call, captured events through sequence 10, and intentionally disconnected after six seconds (curl: (28), HTTP 200 with 8,579 bytes already received).
  2. databricks apps stop openai-agent-framework-sse --profile ml-inference-staging stopped compute while the tool was still running.
  3. Lakebase after the stop showed in_progress, attempt 1, no terminal response, 11 event rows (sequences 0–10), and one message in the original SDK session.
  4. After databricks apps start ..., stale-heartbeat claiming advanced the row to attempt 2.
  5. Reconnect from cursor 10 returned sequences 11–35. The first event was response.resumed; the stream ended with response.completed and [DONE].
  6. Final Lakebase state is completed, attempt 2. Attempt 1 has 11 event rows; attempt 2 has 25. The SDK session rotated to pr467-event-crash-20260821T184027Z::attempt-2.
  7. The resumed loop received the original user input plus an 8,444-character [RECOVERY] input containing serialized prior event/tool-call evidence and produced PR467_EVENT_CRASH_RECOVERED_20260821T184027Z.

Real App stop/start: agent-session recovery

Input session: pr467-session-crash-20260821T184221Z
Durable response: resp_516816b5e26f49a4a98248f1

  1. The same 30-second tool-call scenario was stopped with attempt 1 still in_progress; Lakebase had 11 event rows (sequences 0–10), one SDK user message, and no terminal response.
  2. Restart claimed attempt 2 without rotating the session ID.
  3. Correct cursor replay:
    curl -N \
      "$APP_URL/responses/resp_516816b5e26f49a4a98248f1?stream=true&starting_after=10" \
      -H "Authorization: Bearer $APP_TOKEN"
    returned HTTP 200, 37 JSON events (sequences 11–47), then [DONE]. The first event was response.resumed; the last was response.completed.
  4. Final agent_server.responses state is completed, attempt 2, with a persisted terminal response. Event rows are 11 for attempt 1 and 37 for attempt 2.
  5. The same SDK session contains five messages: the original user input, the fixed 255-character [RECOVERY] prompt, the resumed tool call, its output, and the final assistant message. No ::attempt-2 session exists.
  6. The client received PR467_SESSION_CRASH_RECOVERED_20260821T184221Z, demonstrating that harness-owned history was restored and the resumed loop completed.

The Lakebase checks queried agent_server.responses, grouped agent_server.messages by attempt/sequence, and inspected the SDK-owned agent_messages rows 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-staging using the
exact local PR wheel and separate empty databases under
projects/shivam-debug/branches/pr-cuj-harness:

Strategy App Database
Event log langgraph-event-recovery pr467-langgraph-event
Agent session langgraph-session-recovery pr467-langgraph-session

The test for each App used background=true, stream=true, a 60-second
wait_for_completion tool call, an actual databricks apps stop while the tool
was pending, databricks apps start, SSE reconnection, final polling, and a
follow-up conversational turn.

LangGraph EVENT_LOG recovery

Original conversation: pr467-langgraph-event_log-20260825T184118Z
Durable response: resp_8d13bf3284084485a85bffe7

  1. Attempt 1 persisted one event: sequence 0, response.output_item.done with a function_call. The App was stopped before a tool output existed.
  2. Restart claimed attempt 2. Sequence 1 was response.resumed with rotated conversation ID pr467-langgraph-event_log-20260825T184118Z::attempt-2.
  3. Attempt 2 persisted four rows (sequences 1–4): response.resumed, a new function_call, function_call_output, and the final assistant message. This shows prose recovery started a fresh LangGraph thread and re-invoked the interrupted tool.
  4. Final polling returned completed, attempt 2, a persisted terminal response, and PR467_LANGGRAPH_EVENT_LOG_20260825T184118Z.
  5. The original LangGraph thread has 3 checkpoints; the rotated attempt-2 thread has 8 checkpoints.
  6. The client then sent a new turn using the rotated ID from response.resumed. The response was FOLLOWUP_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.
  7. The live initial SSE stream emitted sequence 0 only once. This specifically verifies the sequence-zero cursor regression fixed in this update.

LangGraph AGENT_SESSION recovery

Conversation: pr467-langgraph-agent_session-20260825T184625Z
Durable response: resp_387945bf600244e9ae579223

  1. Attempt 1 persisted one event: sequence 0, the pending function_call.
  2. Restart claimed attempt 2 and emitted response.resumed at sequence 1 with the unchanged conversation ID.
  3. Attempt 2 persisted only response.resumed, function_call_output, and the final assistant message (sequences 1–3). There was no second function-call event, showing astream(None) resumed the pending LangGraph checkpoint instead of rebuilding the turn from prose.
  4. Final polling returned completed, attempt 2, a persisted terminal response, and PR467_LANGGRAPH_AGENT_SESSION_20260825T184625Z.
  5. Lakebase contains one LangGraph thread with 8 checkpoints; no ::attempt-2 thread was created.
  6. A follow-up using the same conversation ID returned 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.yaml resource bindings, and Lakebase clients ignoring the
Apps-provided PGDATABASE value.

Durable HITL cookbook flow

Both the OpenAI Agents SDK and LangGraph cookbooks now show the same agent-managed HITL pattern:

  1. The client submits a proposal with background=true and stream=true.
  2. The first durable Response completes with APPROVAL_REQUIRED; no worker or lease is held while waiting for a person.
  3. The client submits approval as a second background streamed Response using the same SDK session or LangGraph thread.
  4. The second run includes a 60-second tool call, so the App can be stopped and recovered while approved work is active.
  5. Both terminal Responses and both ordered event streams are persisted. The harness session/checkpointer preserves conversational state between the proposal and approval turns.

This keeps responsibilities explicit: LongRunningAgentServer makes 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

Concern Contract in this PR
Background background=true creates a durable Responses record. The client receives an in_progress response ID and polls GET /responses/{id}.
Streaming @stream() yields Responses events; the server persists them and reconnects with ?stream=true&starting_after=<sequence>.
OpenAI Agents SDK client The SDK loop remains server-side. A custom deployment client adopts Responses; an existing Responses-compatible client keeps its shape and changes the base URL.
LangGraph SDK client This cookbook runs LangGraph behind Responses, not LangGraph Agent Server. Native threads/runs calls change to Responses unless a LangGraph protocol adapter is added.

Both public cookbook READMEs now include concrete before/after client snippets and the exact background=true, stream=true, polling, and reconnect flow.

@shivam5 shivam5 changed the title Add standalone durable agent server Generalize long-running agent recovery and replay Aug 20, 2026
Comment thread src/databricks_ai_bridge/long_running/AGENTS.md Outdated
@shivam5
shivam5 marked this pull request as ready for review August 20, 2026 20:34
@shivam5
shivam5 marked this pull request as draft August 20, 2026 20:34
@shivam5

shivam5 commented Aug 20, 2026

Copy link
Copy Markdown
Contributor Author

@smurching - Should we keep some examples in this repo as well?
Or just keep everything in https://github.com/databricks/app-templates?

@shivam5
shivam5 marked this pull request as ready for review August 20, 2026 22:25
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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

some of the mermaid diagrams aren't formatted correctly, I don't think it's due to any changes in this pr, but can we fix it if it's an issue?

Image

Comment thread src/databricks_ai_bridge/long_running/server.py Outdated

@jamesbxwu jamesbxwu left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Went through first pass and here are a few structural comments:

  1. Can we get rid of the sse_replay knob and keep the existing behavior of always writing to the event log and support streaming replay?
  2. Should the examples folder be defined at the top level vs inside of src/long_running? It doesn't seem that examples as they are today covers everything inside of the databricks-ai-bridge repo
  3. Consider renaming examples to cookbooks
  4. 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 examples folder that explains the different configurations for the long running agent server and have a single summary table comparing them?
  5. 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

Comment thread src/databricks_ai_bridge/long_running/server.py Outdated
Comment thread examples/openai-sdk-agent/agent_recovery_polling/handlers.py Outdated
Comment thread examples/openai-sdk-agent/agent_recovery_polling/TEST_RESULTS.md Outdated
Comment thread examples/openai-sdk-agent/framework_recovery_sse/TEST_RESULTS.md Outdated
Comment on lines +27 to +28
| 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 |

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is it still possible to rename these to be more accurate such as agent_server.checkpoints and ``agent_server.events`?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what is this used for?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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."""

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we improve the documentation here to cover

  1. what are previous events
  2. what does _default_request include and why should this seemingly generic concept of a ResumeContext need to differentiate default vs not?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

added documentation. The default request is use the server default. If it still sounds confusing I can take a stab to simplify this.

Comment thread src/databricks_ai_bridge/long_running/server.py Outdated
Comment thread src/databricks_ai_bridge/long_running/server.py Outdated
Comment thread src/databricks_ai_bridge/long_running/server.py Outdated
Comment thread src/databricks_ai_bridge/long_running/server.py

@jamesbxwu jamesbxwu left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

had a few more questions on the combination of using ResumeStrategy.AGENT_SESSION and on_resume() together, let's chat offline

Comment on lines +17 to +22
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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this method needs a bit more explanation

_on_resume_function: Callable[..., Any] | None = None


class ResumeStrategy(str, Enum):

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
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",

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

agent responses can be quite long and verbose, any concerns on db row/column size limits if we store this?

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +206 to +211
def _agent_session_recovery_message() -> dict[str, Any]:
return {
"type": "message",
"role": "user",
"content": AGENT_SESSION_RECOVERY_PROMPT,
}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
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:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is this redundant with line 1014 above?

Suggested change
if return_trace_id:

Comment on lines +44 to +47
@on_resume()
async def resume(request, context: ResumeContext):
resumed = await context.default_resume_request(request)
return resumed

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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]]:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's not immediately clear that on_resume() is used to reconstruct an incoming request during crash recovery.

  1. if that is indeed the hook we are exposing, then we should rename this to build_recovery_request
  2. right now the decorator here does not restrict the input and output, should we do that?

Comment on lines +55 to +56
- `ResumeStrategy.AGENT_SESSION` keeps the original session key and replaces
the request input with a recovery prompt. The agent SDK supplies the history.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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()?

@shivam5

shivam5 commented Aug 26, 2026

Copy link
Copy Markdown
Contributor Author

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants