Prototype header-based durable agent entrypoint

CHANGELOG.md

@@ -9,6 +9,8 @@
 
 ### Improvements
 - databricks-openai: `DatabricksOpenAI` and `AsyncDatabricksOpenAI` clients now follow HTTP redirects by default, configurable via the new `follow_redirects` parameter (#445)
+- databricks-ai-bridge: Add a transport-neutral `DatabricksDurableRuntime` with Lakebase request, response, event persistence, and stale-attempt recovery
+- databricks-ai-bridge: Add an AgentCore-style `DatabricksDurableApp` entrypoint prototype
 
 ### Bug Fixes
 - databricks-ai-bridge: Genie now returns the full answer text aggregated from all text attachments (#432)

README.md

@@ -24,6 +24,23 @@ For frameworks without dedicated integration packages:
 pip install databricks-ai-bridge
 ```
 
+## Durable Runtime
+
+[`DatabricksDurableRuntime`](./src/databricks_ai_bridge/durable_runtime/README.md)
+adds Lakebase-backed request, result, and event persistence with heartbeat-based
+stale-attempt recovery around a caller-owned async handler. The handler remains
+responsible for agent sessions and checkpoints.
+
+```sh
+pip install 'databricks-ai-bridge[memory]'
+```
+
+The [`DatabricksDurableApp`](./src/databricks_ai_bridge/durable_app/app.py)
+prototype supplies a complete ASGI application around one `@app.entrypoint`.
+Its header-based protocol preserves the application's JSON request and
+foreground response bodies. See the
+[minimal cookbook](./cookbooks/durable-entrypoint/README.md).
+
 ### Install from source
 
 With https:

cookbooks/durable-entrypoint/README.md

@@ -0,0 +1,148 @@
+# Durable Header Entrypoint
+
+This cookbook shows the decorator-style app with the application's JSON body
+left unchanged. Durable submission metadata is supplied through headers.
+
+```python
+@app.entrypoint
+async def agent(payload, context): ...
+
+@app.on_resume
+async def resume_agent(payload, context): ...
+```
+
+The server validates the headers, persists their normalized values with the
+original body, and constructs `context` for each attempt. Recovery does not
+replay HTTP headers; it rebuilds `context` from the durable run record.
+
+## Background, streaming, and client impact
+
+| Capability | Developer contract | Client contract |
+| --- | --- | --- |
+| Background | Return the application's normal JSON result. The app stores status and the final result. | Keep the existing body; add `Idempotency-Key`, `Databricks-Agent-Session-Id`, and `Databricks-Background: true`. A non-streaming request returns `202`; poll `GET /invocations/{run_id}`. |
+| Durable streaming | Convert SDK events to JSON and call `await context.emit(event)`. | Add `Databricks-Stream: true`. Save SSE `id` values and reconnect through `GET /invocations/{run_id}/events?after=<id>`. |
+
+`background=true, stream=true` starts one durable run and immediately opens its
+event stream. The run ID comes back in `Databricks-Run-Id`; disconnecting does
+not cancel the run.
+
+### OpenAI Agents SDK: before and after
+
+The OpenAI Agents SDK is in-process, so a deployed client previously used
+whatever route the developer created around `Runner.run_streamed()`:
+
+```python
+payload = {"message": "hello"}
+async with http.stream("POST", "/invocations", json=payload): ...
+```
+
+The request body and route can remain unchanged. The client only adds runtime
+headers, then uses the new polling/replay routes after the POST:
+
+```python
+async with http.stream(
+    "POST",
+    "/invocations",
+    json=payload,
+    headers={
+        "Idempotency-Key": "run-1",
+        "Databricks-Agent-Session-Id": "conversation-1",
+        "Databricks-Background": "true",
+        "Databricks-Stream": "true",
+    },
+) as response:
+    run_id = response.headers["Databricks-Run-Id"]
+    async for line in response.aiter_lines():
+        last_event_id = remember_sse_id(line, last_event_id)
+
+final = (await http.get(f"/invocations/{run_id}")).json()
+```
+
+### LangGraph SDK: before and after
+
+A native LangGraph client uses several framework routes, not one invocation
+route:
+
+```python
+thread = await langgraph.threads.create()
+run = await langgraph.runs.create(
+    thread["thread_id"], "agent", input={"messages": messages}
+)
+result = await langgraph.runs.join(thread["thread_id"], run["run_id"])
+
+async for event in langgraph.runs.stream(
+    thread["thread_id"], "agent", input={"messages": messages}
+):
+    consume(event)
+```
+
+Headers preserve the body of a single existing endpoint; they do not make the
+LangGraph threads/runs protocol compatible. Without a LangGraph adapter, the
+client changes to the `httpx` invocation call above and passes `thread_id` in
+`Databricks-Agent-Session-Id`.
+
+References: [OpenAI Agents SDK streaming](https://openai.github.io/openai-agents-python/streaming/),
+[LangGraph background runs](https://docs.langchain.com/langsmith/runs), and
+[LangGraph resumable streaming](https://docs.langchain.com/langsmith/streaming).
+
+## Durable HITL flow
+
+HITL is modeled as two durable runs. The runtime does not keep a worker alive
+while waiting for a person.
+
+1. A background streamed proposal completes with `requires_action`.
+2. The client reviews the persisted result.
+3. The client submits approval as another background streamed run using the
+   same `session_id`.
+
+Start the proposal and watch its persisted event stream:
+
+```bash
+curl -N -X POST localhost:8000/invocations \
+  -H 'content-type: application/json' \
+  -H 'idempotency-key: proposal-1' \
+  -H 'databricks-agent-session-id: approval-session-1' \
+  -H 'databricks-background: true' \
+  -H 'databricks-stream: true' \
+  -d '{
+    "action": "publish the release notes"
+  }'
+```
+
+The request body is the application payload, not a runtime envelope. Poll
+`GET /invocations/proposal-1`. Its persisted result contains
+`result.status=requires_action`. Then approve it:
+
+```bash
+curl -N -X POST localhost:8000/invocations \
+  -H 'content-type: application/json' \
+  -H 'idempotency-key: approval-1' \
+  -H 'databricks-agent-session-id: approval-session-1' \
+  -H 'databricks-background: true' \
+  -H 'databricks-stream: true' \
+  -d '{
+    "action": "publish the release notes",
+    "decision": "approve",
+    "wait_seconds": 60
+  }'
+```
+
+Stop the process while the approved action is waiting. A new process reclaims
+the stale run, calls `@app.on_resume` with the original payload and same
+`session_id`, and appends events to the existing durable stream. Reconnect with:
+
+```bash
+curl -N 'localhost:8000/invocations/approval-1/events?after=<last-event-id>'
+```
+
+Poll `GET /invocations/approval-1` for the authoritative final result. External side
+effects remain at-least-once and must be idempotent.
+
+## Run
+
+```bash
+pip install -r requirements.txt
+export OPENAI_API_KEY=...
+export LAKEBASE_AUTOSCALING_ENDPOINT=projects/.../endpoints/...
+python agent.py
+```

cookbooks/durable-entrypoint/agent.py

@@ -0,0 +1,48 @@
+"""OpenAI Agents SDK loop hosted by the SDK-provided durable entrypoint."""
+
+import os
+
+import uvicorn
+from openai_agent import run_openai_agent
+
+from databricks_ai_bridge.durable_app import DatabricksDurableApp, DurableAgentContext
+
+app = DatabricksDurableApp()
+
+
+@app.entrypoint
+async def agent(payload: dict, context: DurableAgentContext) -> dict:
+    result = await run_openai_agent(
+        payload=payload,
+        session_id=context.session_id,
+        emit=context.emit,
+    )
+
+    return {
+        "result": result,
+        "session_id": context.session_id,
+        "attempt": context.attempt,
+    }
+
+
+@app.on_resume
+async def resume_agent(payload: dict, context: DurableAgentContext) -> dict:
+    result = await run_openai_agent(
+        payload=payload,
+        session_id=context.session_id,
+        emit=context.emit,
+        is_recovery=True,
+    )
+    return {
+        "result": result,
+        "session_id": context.session_id,
+        "attempt": context.attempt,
+    }
+
+
+if __name__ == "__main__":
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=int(os.getenv("DATABRICKS_APP_PORT", "8000")),
+    )

cookbooks/durable-entrypoint/app.yaml

@@ -0,0 +1,3 @@
+command:
+  - python
+  - agent.py

cookbooks/durable-entrypoint/openai_agent.py

@@ -0,0 +1,85 @@
+"""OpenAI Agents SDK loop for durable streaming and human approval."""
+
+import asyncio
+import os
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from agents import Agent, Runner, function_tool
+from databricks_openai.agents import AsyncDatabricksSession
+
+EventEmitter = Callable[[dict[str, Any]], Awaitable[int]]
+
+
+@function_tool
+async def complete_approved_action(action: str, wait_seconds: int) -> str:
+    """Simulate an approved side effect after a bounded delay."""
+    if wait_seconds < 0 or wait_seconds > 300:
+        raise ValueError("wait_seconds must be between 0 and 300")
+    await asyncio.sleep(wait_seconds)
+    return f"Completed approved action: {action}"
+
+
+def _create_agent() -> Agent:
+    return Agent(
+        name="Durable approval assistant",
+        model=os.getenv("OPENAI_MODEL", "gpt-4.1-mini"),
+        instructions=(
+            "For a PROPOSAL request, produce a short plan and do not call tools. "
+            "For an APPROVED request, always call complete_approved_action. "
+            "For a REJECTED request, acknowledge the rejection without calling tools."
+        ),
+        tools=[complete_approved_action],
+    )
+
+
+async def run_openai_agent(
+    payload: dict[str, Any],
+    session_id: str,
+    emit: EventEmitter,
+    *,
+    is_recovery: bool = False,
+) -> dict[str, Any]:
+    session = AsyncDatabricksSession(
+        session_id=session_id,
+        autoscaling_endpoint=os.environ["LAKEBASE_AUTOSCALING_ENDPOINT"],
+        schema=os.getenv(
+            "OPENAI_AGENT_SESSION_SCHEMA", "durable_openai_agent_sessions"
+        ),
+    )
+    action = str(payload["action"])
+    decision = payload.get("decision")
+    wait_seconds = int(payload.get("wait_seconds", 0))
+    if decision is None:
+        instruction = f"PROPOSAL: Draft a plan for this action: {action}"
+    elif decision == "approve":
+        instruction = (
+            f"APPROVED: Complete this action: {action}. "
+            f"Call complete_approved_action with wait_seconds={wait_seconds}."
+        )
+    elif decision == "reject":
+        instruction = f"REJECTED: Do not perform this action: {action}"
+    else:
+        raise ValueError("decision must be approve, reject, or omitted")
+
+    if is_recovery:
+        instruction = (
+            "RECOVERY: Continue the interrupted durable run using the existing session. "
+            + instruction
+        )
+
+    result = Runner.run_streamed(_create_agent(), input=instruction, session=session)
+
+    async for event in result.stream_events():
+        if event.type == "raw_response_event":
+            await emit(event.data.model_dump(mode="json"))
+
+    output = str(result.final_output or "")
+    if decision is None:
+        response = {"status": "requires_action", "action": action, "proposal": output}
+        await emit({"type": "agent.approval_required", "action": action})
+        return response
+
+    status = "completed" if decision == "approve" else "rejected"
+    await emit({"type": f"agent.{status}", "action": action})
+    return {"status": status, "action": action, "output": output}

cookbooks/durable-entrypoint/requirements.txt

@@ -0,0 +1,3 @@
+databricks-ai-bridge[agent-server]
+databricks-openai[memory]>=0.17.0
+openai-agents>=0.19.4,<0.20

src/databricks_ai_bridge/durable_app/__init__.py

@@ -0,0 +1,21 @@
+"""AgentCore-style durable entrypoint hosted by Databricks AI Bridge."""
+
+try:
+    import fastapi  # noqa: F401
+except ImportError as exc:
+    raise ImportError(
+        "DatabricksDurableApp requires databricks-ai-bridge[agent-server]. "
+        "Install it with: pip install databricks-ai-bridge[agent-server]"
+    ) from exc
+
+from databricks_ai_bridge.durable_app.app import (
+    DatabricksDurableApp,
+    DurableAgentContext,
+    DurableAgentEntrypoint,
+)
+
+__all__ = [
+    "DatabricksDurableApp",
+    "DurableAgentContext",
+    "DurableAgentEntrypoint",
+]