diff --git a/CHANGELOG.md b/CHANGELOG.md index e9dde1828..b7dcb4471 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/README.md b/README.md index 615d57c52..6811dcede 100644 --- a/README.md +++ b/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: diff --git a/cookbooks/durable-entrypoint/README.md b/cookbooks/durable-entrypoint/README.md new file mode 100644 index 000000000..d4b63785e --- /dev/null +++ b/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=`. | + +`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=' +``` + +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 +``` diff --git a/cookbooks/durable-entrypoint/agent.py b/cookbooks/durable-entrypoint/agent.py new file mode 100644 index 000000000..eef96c25c --- /dev/null +++ b/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")), + ) diff --git a/cookbooks/durable-entrypoint/app.yaml b/cookbooks/durable-entrypoint/app.yaml new file mode 100644 index 000000000..bf591d173 --- /dev/null +++ b/cookbooks/durable-entrypoint/app.yaml @@ -0,0 +1,3 @@ +command: + - python + - agent.py diff --git a/cookbooks/durable-entrypoint/openai_agent.py b/cookbooks/durable-entrypoint/openai_agent.py new file mode 100644 index 000000000..bee733c3c --- /dev/null +++ b/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} diff --git a/cookbooks/durable-entrypoint/requirements.txt b/cookbooks/durable-entrypoint/requirements.txt new file mode 100644 index 000000000..c8fb8a711 --- /dev/null +++ b/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 diff --git a/src/databricks_ai_bridge/durable_app/__init__.py b/src/databricks_ai_bridge/durable_app/__init__.py new file mode 100644 index 000000000..7af0fa2e7 --- /dev/null +++ b/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", +] diff --git a/src/databricks_ai_bridge/durable_app/app.py b/src/databricks_ai_bridge/durable_app/app.py new file mode 100644 index 000000000..285862867 --- /dev/null +++ b/src/databricks_ai_bridge/durable_app/app.py @@ -0,0 +1,227 @@ +"""ASGI application preserving the agent payload while adding durable metadata.""" + +from __future__ import annotations + +import asyncio +import copy +import json +from collections.abc import AsyncIterator, Awaitable, Callable +from contextlib import asynccontextmanager +from dataclasses import dataclass +from uuid import uuid4 + +from fastapi import FastAPI, HTTPException, Request, Response +from fastapi.responses import JSONResponse, StreamingResponse + +from databricks_ai_bridge.durable_runtime import ( + DatabricksDurableRuntime, + DurabilityStore, + DurableExecution, + DurableExecutionContext, + DurableExecutionStatus, + JsonObject, +) + +DurableAgentEntrypoint = Callable[[JsonObject, "DurableAgentContext"], Awaitable[JsonObject]] + + +@dataclass(frozen=True) +class DurableAgentContext: + """Stable run/session identifiers and attempt metadata for an agent.""" + + run_id: str + session_id: str + attempt: int + _execution_context: DurableExecutionContext + + @property + def is_recovery(self) -> bool: + return self.attempt > 1 + + async def emit(self, event: JsonObject) -> int: + return await self._execution_context.emit(event) + + +class DatabricksDurableApp: + """Host one JSON entrypoint without wrapping the application's request body.""" + + def __init__( + self, + *, + durability_store: DurabilityStore | None = None, + schema: str = "databricks_durable_app", + heartbeat_seconds: float = 3.0, + stale_seconds: float = 10.0, + scan_seconds: float = 3.0, + poll_seconds: float = 1.0, + path: str = "/invocations", + ) -> None: + self._entrypoint: DurableAgentEntrypoint | None = None + self._resume_entrypoint: DurableAgentEntrypoint | None = None + self._path = path.rstrip("/") or "/" + self.runtime = DatabricksDurableRuntime( + self._execute, + durability_store=durability_store, + schema=schema, + heartbeat_seconds=heartbeat_seconds, + stale_seconds=stale_seconds, + scan_seconds=scan_seconds, + poll_seconds=poll_seconds, + ) + + @asynccontextmanager + async def lifespan(_: FastAPI): + await self.runtime.start() + try: + yield + finally: + await self.runtime.stop() + + self.asgi_app = FastAPI(lifespan=lifespan) + self.asgi_app.add_api_route(self._path, self._submit, methods=["POST"]) + self.asgi_app.add_api_route(f"{self._path}/{{run_id}}", self._get, methods=["GET"]) + self.asgi_app.add_api_route( + f"{self._path}/{{run_id}}/events", + self._events, + methods=["GET"], + ) + self.asgi_app.add_api_route("/api/healthz", self._health, methods=["GET"]) + + def entrypoint(self, function: DurableAgentEntrypoint) -> DurableAgentEntrypoint: + """Register the single agent function invoked for every durable attempt.""" + if self._entrypoint is not None: + raise RuntimeError("DatabricksDurableApp supports one entrypoint") + self._entrypoint = function + return function + + def on_resume(self, function: DurableAgentEntrypoint) -> DurableAgentEntrypoint: + """Register the handler used after a stale attempt is reclaimed.""" + if self._resume_entrypoint is not None: + raise RuntimeError("DatabricksDurableApp supports one resume entrypoint") + self._resume_entrypoint = function + return function + + async def __call__(self, scope, receive, send) -> None: + await self.asgi_app(scope, receive, send) + + async def _execute( + self, + request: JsonObject, + execution_context: DurableExecutionContext, + ) -> JsonObject: + if self._entrypoint is None: + raise RuntimeError("register an @app.entrypoint before serving requests") + session_id = request.get("session_id") + payload = request.get("payload") + if not isinstance(session_id, str) or not isinstance(payload, dict): + raise TypeError("persisted request must contain session_id and payload") + context = DurableAgentContext( + run_id=execution_context.execution_id, + session_id=session_id, + attempt=execution_context.attempt, + _execution_context=execution_context, + ) + function = ( + self._resume_entrypoint + if context.is_recovery and self._resume_entrypoint is not None + else self._entrypoint + ) + return await function(copy.deepcopy(payload), context) + + async def _submit(self, http_request: Request) -> Response: + payload = await http_request.json() + if not isinstance(payload, dict): + raise HTTPException(400, "request body must be a JSON object") + + run_id = http_request.headers.get("Idempotency-Key") or str(uuid4()) + session_id = http_request.headers.get("Databricks-Agent-Session-Id") or str(uuid4()) + background = self._header_bool(http_request, "Databricks-Background", False) + stream = self._header_bool(http_request, "Databricks-Stream", False) + request: JsonObject = { + "session_id": session_id, + "payload": payload, + } + response_headers = { + "Databricks-Run-Id": run_id, + "Databricks-Agent-Session-Id": session_id, + } + + if stream: + await self.runtime.submit(run_id, request) + return StreamingResponse( + self._event_stream(run_id, 0), + media_type="text/event-stream", + headers=response_headers, + ) + + if background: + state = await self.runtime.submit(run_id, request) + status_code = 200 if state.is_terminal else 202 + response_headers["Location"] = f"{self._path}/{run_id}" + return JSONResponse( + self._state_payload(state), + status_code=status_code, + headers=response_headers, + ) + + result = await self.runtime.invoke(run_id, request) + return JSONResponse(result, headers=response_headers) + + async def _get(self, run_id: str) -> JSONResponse: + state = await self.runtime.get(run_id) + if state is None: + raise HTTPException(404, "run not found") + return JSONResponse(self._state_payload(state)) + + async def _events(self, run_id: str, after: int = 0) -> StreamingResponse: + if await self.runtime.get(run_id) is None: + raise HTTPException(404, "run not found") + return StreamingResponse( + self._event_stream(run_id, after), + media_type="text/event-stream", + ) + + async def _event_stream(self, run_id: str, after: int) -> AsyncIterator[str]: + cursor = after + while True: + events = await self.runtime.events(run_id, after_sequence=cursor) + for event in events: + cursor = event.sequence_number + yield ( + f"id: {cursor}\n" + f"event: {event.event.get('type', 'message')}\n" + f"data: {json.dumps(event.event)}\n\n" + ) + + state = await self.runtime.get(run_id) + if state is None or state.status in { + DurableExecutionStatus.COMPLETED, + DurableExecutionStatus.FAILED, + }: + return + await asyncio.sleep(0.25) + + @staticmethod + def _state_payload(state: DurableExecution) -> JsonObject: + return { + "run_id": state.execution_id, + "status": state.status.value, + "attempt": state.attempt, + "result": copy.deepcopy(state.response), + } + + @staticmethod + async def _health() -> JsonObject: + return {"status": "healthy"} + + @staticmethod + def _header_bool(request: Request, name: str, default: bool) -> bool: + value = request.headers.get(name) + if value is None: + return default + normalized = value.lower() + if normalized == "true": + return True + if normalized == "false": + return False + raise HTTPException(400, f"{name} must be true or false") diff --git a/src/databricks_ai_bridge/durable_runtime/README.md b/src/databricks_ai_bridge/durable_runtime/README.md new file mode 100644 index 000000000..82a839f85 --- /dev/null +++ b/src/databricks_ai_bridge/durable_runtime/README.md @@ -0,0 +1,136 @@ +# Databricks Durable Runtime + +`DatabricksDurableRuntime` is a transport-neutral durability layer for an +idempotent JSON request/response handler. It stores the request, response, and +heartbeat state in Lakebase and re-executes stale work after a process exits. + +It does not own agent session history. An OpenAI Agents SDK session, LangGraph +checkpointer, or other harness-managed store remains an executor concern. + +## Contract + +The caller provides: + +- a stable `execution_id`, used as the idempotency and recovery key; +- a JSON-object request; and +- an async executor that returns a JSON-object response. + +The executor receives `DurableExecutionContext`. On recovery, `attempt > 1` +and `is_recovery` is true. The runtime passes the exact persisted request on +every attempt and does not add a recovery message or reconstruct history. + +```python +from databricks_ai_bridge.durable_runtime import ( + DatabricksDurableRuntime, + DurableExecutionContext, +) + + +async def execute(request: dict, context: DurableExecutionContext) -> dict: + session = create_agent_session(request["session_id"]) + await context.emit({"type": "progress", "message": "agent started"}) + if context.is_recovery: + result = await resume_from_session(session) + else: + result = await start_from_request(request, session) + return {"output": result} + + +runtime = DatabricksDurableRuntime( + execute, + autoscaling_endpoint="projects/project/branches/branch/endpoints/primary", +) +``` + +`execution_id` identifies one durable request, not an SDK session. A response +ID or client idempotency key is normally the right value. A session ID can be +used only when the application intentionally allows one durable request per +session. For multi-turn agents, keep the harness session or conversation ID in +the persisted request and assign each invocation its own `execution_id`. + +Call `await runtime.start()` during process startup and `await runtime.stop()` +during shutdown. Shutdown cancels local tasks without marking them failed, so a +different process can claim them after the heartbeat becomes stale. + +## Input and Output Wiring + +- `submit(execution_id, request)` accepts background work and returns persisted + state immediately. +- `invoke(execution_id, request)` accepts work and waits for the persisted + response, including when another process owns the attempt. +- `get(execution_id)` returns status, attempt, heartbeat, request, and response. +- `wait(execution_id)` waits for a previously submitted response. +- `events(execution_id, after_sequence=N)` returns persisted events for replay. + +Executors may call `await context.emit(event)` to append an ordered JSON event. +Events are stored before delivery, survive worker replacement, and use the returned +sequence number as a replay cursor. HTTP and SSE adapters decide how those generic +events are presented to clients. + +Submitting the same ID and same JSON request is idempotent. If the response is +already complete, `invoke` returns the cached response. Reusing the ID with a +different request raises `DurableRequestConflictError`. + +An HTTP adapter remains small and transport-specific: + +```python +@app.post("/responses") +async def responses(request: RequestModel): + execution_id = stable_id_from(request) + payload = executor_payload(request) + if request.background: + state = await runtime.submit(execution_id, payload) + return status_response(state, status_code=202) + return await runtime.invoke(execution_id, payload) + + +@app.get("/responses/{execution_id}") +async def retrieve(execution_id: str): + state = await runtime.get(execution_id) + return status_or_response(state) +``` + +The adapter decides how IDs are supplied, which HTTP status shape to return, +and how JSON is converted to framework-specific request or response models. +`executor_payload` should remove transport-only fields such as `background`, +`stream`, polling cursors, or trace-return flags. Changing only the transport +mode must not create a request conflict for the same durable operation. + +For background execution, an adapter may generate the ID and return it in the +initial `202` response. For a blocking request that must survive a lost HTTP +connection, the client must supply a stable ID (for example an +`Idempotency-Key`, session ID, or response ID) so it can retry or retrieve the +same operation. The runtime does not hide this client/server recovery contract. + +Recommended HTTP error mappings are: + +- different request for an existing ID: `409 Conflict`; +- unknown ID on retrieval: `404 Not Found`; +- failed execution: a terminal failed-status payload; and +- blocking wait timeout: a gateway timeout while execution continues in the + background and remains retrievable by ID. + +## Durability Store + +The default schema is `databricks_durable_runtime`. It contains two tables: + +```text +executions + execution_id TEXT PRIMARY KEY + status TEXT + attempt INTEGER + heartbeat_at TIMESTAMPTZ + request JSONB + response JSONB + +execution_events + sequence_number BIGSERIAL PRIMARY KEY + execution_id TEXT + attempt INTEGER + event JSONB +``` + +Recovery is at-least-once. A process can exit after an external side effect but +before persisting its response, so executor tools must tolerate retries where +needed. A compare-and-swap claim on the stale row prevents multiple recovery +attempts from acquiring the same durability ownership. diff --git a/src/databricks_ai_bridge/durable_runtime/__init__.py b/src/databricks_ai_bridge/durable_runtime/__init__.py new file mode 100644 index 000000000..efba50012 --- /dev/null +++ b/src/databricks_ai_bridge/durable_runtime/__init__.py @@ -0,0 +1,49 @@ +"""Lakebase-backed durable execution for JSON request/response handlers. + +Install the memory extra before importing this module:: + + pip install databricks-ai-bridge[memory] +""" + +try: + import psycopg # noqa: F401 + import sqlalchemy # noqa: F401 +except ImportError as exc: + raise ImportError( + "DatabricksDurableRuntime requires databricks-ai-bridge[memory]. " + "Install it with: pip install databricks-ai-bridge[memory]" + ) from exc + +from databricks_ai_bridge.durable_runtime.runtime import DatabricksDurableRuntime +from databricks_ai_bridge.durable_runtime.store import ( + DEFAULT_DURABILITY_SCHEMA, + LakebaseDurabilityStore, +) +from databricks_ai_bridge.durable_runtime.types import ( + DurabilityStore, + DurableEvent, + DurableExecution, + DurableExecutionContext, + DurableExecutionFailedError, + DurableExecutionNotFoundError, + DurableExecutionStatus, + DurableExecutor, + DurableRequestConflictError, + JsonObject, +) + +__all__ = [ + "DEFAULT_DURABILITY_SCHEMA", + "DatabricksDurableRuntime", + "DurabilityStore", + "DurableEvent", + "DurableExecution", + "DurableExecutionContext", + "DurableExecutionFailedError", + "DurableExecutionNotFoundError", + "DurableExecutionStatus", + "DurableExecutor", + "DurableRequestConflictError", + "JsonObject", + "LakebaseDurabilityStore", +] diff --git a/src/databricks_ai_bridge/durable_runtime/runtime.py b/src/databricks_ai_bridge/durable_runtime/runtime.py new file mode 100644 index 000000000..a782a66ff --- /dev/null +++ b/src/databricks_ai_bridge/durable_runtime/runtime.py @@ -0,0 +1,334 @@ +"""Transport-neutral durable request runtime.""" + +from __future__ import annotations + +import asyncio +import copy +import json +import logging +from datetime import datetime, timezone +from typing import TYPE_CHECKING + +from databricks_ai_bridge.durable_runtime.store import LakebaseDurabilityStore +from databricks_ai_bridge.durable_runtime.types import ( + DurabilityStore, + DurableEvent, + DurableExecution, + DurableExecutionContext, + DurableExecutionFailedError, + DurableExecutionNotFoundError, + DurableExecutionStatus, + DurableExecutor, + JsonObject, +) + +if TYPE_CHECKING: + from databricks.sdk import WorkspaceClient + +logger = logging.getLogger(__name__) + + +def _copy_json_object(value: JsonObject, name: str) -> JsonObject: + if not isinstance(value, dict): + raise TypeError(f"{name} must be a JSON object, got {type(value).__name__}") + try: + return json.loads(json.dumps(value, allow_nan=False)) + except (TypeError, ValueError) as exc: + raise TypeError(f"{name} must be JSON serializable") from exc + + +class DatabricksDurableRuntime: + """Execute idempotent JSON requests with Lakebase-backed crash recovery. + + The runtime persists only request execution state. The executor owns agent + sessions, checkpoints, tools, and any recovery-specific prompt or behavior. + """ + + def __init__( + self, + executor: DurableExecutor | None = None, + *, + durability_store: DurabilityStore | None = None, + autoscaling_endpoint: str | None = None, + project: str | None = None, + branch: str | None = None, + workspace_client: WorkspaceClient | None = None, + schema: str = "databricks_durable_runtime", + heartbeat_seconds: float = 3.0, + stale_seconds: float = 10.0, + scan_seconds: float = 3.0, + poll_seconds: float = 1.0, + ) -> None: + if heartbeat_seconds <= 0: + raise ValueError("heartbeat_seconds must be positive") + if stale_seconds <= heartbeat_seconds: + raise ValueError("stale_seconds must be greater than heartbeat_seconds") + if scan_seconds <= 0: + raise ValueError("scan_seconds must be positive") + if poll_seconds <= 0: + raise ValueError("poll_seconds must be positive") + + self._executor = executor + self.durability_store = ( + durability_store + if durability_store is not None + else LakebaseDurabilityStore( + autoscaling_endpoint=autoscaling_endpoint, + project=project, + branch=branch, + workspace_client=workspace_client, + schema=schema, + ) + ) + self.heartbeat_seconds = heartbeat_seconds + self.stale_seconds = stale_seconds + self.scan_seconds = scan_seconds + self.poll_seconds = poll_seconds + self._tasks: dict[str, asyncio.Task[None]] = {} + self._scanner: asyncio.Task[None] | None = None + self._started = False + + async def execute( + self, + request: JsonObject, + context: DurableExecutionContext, + ) -> JsonObject: + """Run one attempt; subclasses may override this method.""" + if self._executor is None: + raise NotImplementedError("provide an executor or override execute()") + return await self._executor(request, context) + + async def start(self) -> None: + """Initialize storage and start proactive recovery scanning.""" + if self._started: + return + await self.durability_store.initialize() + self._started = True + self._scanner = asyncio.create_task( + self._scan_loop(), + name="databricks-durable-runtime-scanner", + ) + + async def stop(self) -> None: + """Stop local work, leaving active rows recoverable by another process.""" + if not self._started: + return + tasks = list(self._tasks.values()) + if self._scanner is not None: + self._scanner.cancel() + for task in tasks: + task.cancel() + await asyncio.gather( + *tasks, + *([self._scanner] if self._scanner is not None else []), + return_exceptions=True, + ) + self._tasks.clear() + self._scanner = None + self._started = False + await self.durability_store.close() + + async def submit(self, execution_id: str, request: JsonObject) -> DurableExecution: + """Accept an idempotent request and ensure recoverable work is scheduled.""" + self._require_started() + if not execution_id: + raise ValueError("execution_id must not be empty") + state = await self.durability_store.accept( + execution_id, + _copy_json_object(request, "request"), + ) + self._ensure_scheduled(state) + return state + + async def invoke( + self, + execution_id: str, + request: JsonObject, + *, + timeout: float | None = None, + ) -> JsonObject: + """Accept a request and wait for its persisted terminal response.""" + await self.submit(execution_id, request) + return await self.wait(execution_id, timeout=timeout) + + async def get(self, execution_id: str) -> DurableExecution | None: + """Return persisted state and schedule recovery if it is currently eligible.""" + self._require_started() + state = await self.durability_store.get(execution_id) + if state is not None: + self._ensure_scheduled(state) + return state + + async def wait( + self, + execution_id: str, + *, + timeout: float | None = None, + ) -> JsonObject: + """Wait for a completed response, including work owned by another process.""" + self._require_started() + + async def poll() -> JsonObject: + while True: + state = await self.get(execution_id) + if state is None: + raise DurableExecutionNotFoundError(execution_id) + if state.status == DurableExecutionStatus.COMPLETED: + if state.response is None: + raise RuntimeError( + f"execution {execution_id!r} completed without a response" + ) + return copy.deepcopy(state.response) + if state.status == DurableExecutionStatus.FAILED: + raise DurableExecutionFailedError(execution_id) + await asyncio.sleep(self.poll_seconds) + + if timeout is None: + return await poll() + return await asyncio.wait_for(poll(), timeout=timeout) + + async def events( + self, + execution_id: str, + *, + after_sequence: int | None = None, + ) -> list[DurableEvent]: + """Return persisted events after an optional replay cursor.""" + self._require_started() + return await self.durability_store.events(execution_id, after_sequence) + + def _ensure_scheduled(self, state: DurableExecution) -> None: + if not self._is_recoverable(state): + return + current = self._tasks.get(state.execution_id) + if current is not None and not current.done(): + return + task = asyncio.create_task( + self._execute_attempt(state.execution_id), + name=f"durable-execution-{state.execution_id}", + ) + self._tasks[state.execution_id] = task + task.add_done_callback(lambda completed: self._discard_task(state.execution_id, completed)) + + def _is_recoverable(self, state: DurableExecution) -> bool: + if state.status == DurableExecutionStatus.QUEUED: + return True + if state.status != DurableExecutionStatus.ACTIVE: + return False + if state.heartbeat_at is None: + return True + heartbeat_at = state.heartbeat_at + if heartbeat_at.tzinfo is None: + heartbeat_at = heartbeat_at.replace(tzinfo=timezone.utc) + age = (datetime.now(timezone.utc) - heartbeat_at).total_seconds() + return age >= self.stale_seconds + + async def _execute_attempt(self, execution_id: str) -> None: + try: + claimed = await self.durability_store.claim(execution_id, self.stale_seconds) + except Exception: + logger.exception("Failed to claim durable execution: %s", execution_id) + return + if claimed is None: + return + + heartbeat = asyncio.create_task( + self._heartbeat_loop(execution_id, claimed.attempt), + name=f"durable-heartbeat-{execution_id}-{claimed.attempt}", + ) + try: + + async def emit(event: JsonObject) -> int: + sequence_number = await self.durability_store.append_event( + execution_id, + claimed.attempt, + _copy_json_object(event, "event"), + ) + if sequence_number is None: + raise RuntimeError( + f"execution {execution_id!r} no longer owns attempt {claimed.attempt}" + ) + return sequence_number + + response = await self.execute( + copy.deepcopy(claimed.request), + DurableExecutionContext( + execution_id=execution_id, + attempt=claimed.attempt, + _emit=emit, + ), + ) + response = _copy_json_object(response, "executor response") + completed = await self.durability_store.complete( + execution_id, + claimed.attempt, + response, + ) + if not completed: + logger.info( + "Skipped completion after durability ownership changed: %s attempt=%d", + execution_id, + claimed.attempt, + ) + except asyncio.CancelledError: + raise + except Exception: + logger.exception( + "Durable execution failed: %s attempt=%d", + execution_id, + claimed.attempt, + ) + try: + await self.durability_store.fail(execution_id, claimed.attempt) + except Exception: + logger.exception( + "Failed to persist durable failure: %s attempt=%d", + execution_id, + claimed.attempt, + ) + finally: + heartbeat.cancel() + await asyncio.gather(heartbeat, return_exceptions=True) + + async def _heartbeat_loop(self, execution_id: str, attempt: int) -> None: + while True: + try: + owns_attempt = await self.durability_store.heartbeat(execution_id, attempt) + except Exception: + logger.warning( + "Durable heartbeat failed: %s attempt=%d", + execution_id, + attempt, + exc_info=True, + ) + await asyncio.sleep(self.heartbeat_seconds) + continue + if not owns_attempt: + return + await asyncio.sleep(self.heartbeat_seconds) + + async def _scan_loop(self) -> None: + while True: + try: + execution_ids = await self.durability_store.recoverable_execution_ids( + self.stale_seconds + ) + for execution_id in execution_ids: + state = await self.durability_store.get(execution_id) + if state is not None: + self._ensure_scheduled(state) + except asyncio.CancelledError: + raise + except Exception: + logger.exception("Databricks durable runtime recovery scan failed") + await asyncio.sleep(self.scan_seconds) + + def _discard_task(self, execution_id: str, completed: asyncio.Task[None]) -> None: + if self._tasks.get(execution_id) is completed: + self._tasks.pop(execution_id, None) + if not completed.cancelled(): + completed.exception() + + def _require_started(self) -> None: + if not self._started: + raise RuntimeError("DatabricksDurableRuntime.start() must be called first") diff --git a/src/databricks_ai_bridge/durable_runtime/store.py b/src/databricks_ai_bridge/durable_runtime/store.py new file mode 100644 index 000000000..5969904b9 --- /dev/null +++ b/src/databricks_ai_bridge/durable_runtime/store.py @@ -0,0 +1,394 @@ +"""Lakebase persistence for durable request execution.""" + +from __future__ import annotations + +import json +import os +import re +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any + +from sqlalchemy import text + +from databricks_ai_bridge.durable_runtime.types import ( + DurableEvent, + DurableExecution, + DurableExecutionStatus, + DurableRequestConflictError, + JsonObject, +) +from databricks_ai_bridge.lakebase import AsyncLakebaseSQLAlchemy + +if TYPE_CHECKING: + from databricks.sdk import WorkspaceClient + +DEFAULT_DURABILITY_SCHEMA = "databricks_durable_runtime" +_SCHEMA_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +def _serialize_json_object(value: JsonObject) -> str: + if not isinstance(value, dict): + raise TypeError(f"expected a JSON object, got {type(value).__name__}") + return json.dumps(value, allow_nan=False) + + +def _validate_execution_id(execution_id: str) -> None: + if not execution_id: + raise ValueError("execution_id must not be empty") + + +class LakebaseDurabilityStore: + """Store durable execution state in one Lakebase table.""" + + def __init__( + self, + *, + autoscaling_endpoint: str | None = None, + project: str | None = None, + branch: str | None = None, + workspace_client: WorkspaceClient | None = None, + schema: str = DEFAULT_DURABILITY_SCHEMA, + lakebase: AsyncLakebaseSQLAlchemy | None = None, + ) -> None: + if not _SCHEMA_NAME.fullmatch(schema): + raise ValueError(f"invalid durability schema name: {schema!r}") + + if lakebase is None: + autoscaling_endpoint = autoscaling_endpoint or os.getenv( + "LAKEBASE_AUTOSCALING_ENDPOINT" + ) + if autoscaling_endpoint is None: + project = project or os.getenv("LAKEBASE_AUTOSCALING_PROJECT") + branch = branch or os.getenv("LAKEBASE_AUTOSCALING_BRANCH") + lakebase = AsyncLakebaseSQLAlchemy( + autoscaling_endpoint=autoscaling_endpoint, + project=project, + branch=branch, + workspace_client=workspace_client, + schema=schema, + pool_pre_ping=True, + ) + + self._lakebase = lakebase + self._engine = lakebase.engine + self._table = f"{schema}.executions" + self._events_table = f"{schema}.execution_events" + + async def initialize(self) -> None: + await self._lakebase.create_schema() + async with self._engine.begin() as connection: + await connection.execute( + text( + f""" + CREATE TABLE IF NOT EXISTS {self._table} ( + execution_id TEXT PRIMARY KEY, + status TEXT NOT NULL, + attempt INTEGER NOT NULL DEFAULT 0, + heartbeat_at TIMESTAMPTZ, + request JSONB NOT NULL, + response JSONB, + CHECK (status IN ('QUEUED', 'ACTIVE', 'COMPLETED', 'FAILED')), + CHECK (jsonb_typeof(request) = 'object'), + CHECK (response IS NULL OR jsonb_typeof(response) = 'object') + ) + """ + ) + ) + await connection.execute( + text( + f""" + CREATE TABLE IF NOT EXISTS {self._events_table} ( + sequence_number BIGSERIAL PRIMARY KEY, + execution_id TEXT NOT NULL + REFERENCES {self._table}(execution_id) ON DELETE CASCADE, + attempt INTEGER NOT NULL, + event JSONB NOT NULL, + CHECK (jsonb_typeof(event) = 'object') + ) + """ + ) + ) + await connection.execute( + text( + f""" + CREATE INDEX IF NOT EXISTS execution_events_replay_idx + ON {self._events_table} (execution_id, sequence_number) + """ + ) + ) + await connection.execute( + text( + f""" + CREATE INDEX IF NOT EXISTS executions_recovery_idx + ON {self._table} (status, heartbeat_at) + WHERE status IN ('QUEUED', 'ACTIVE') + """ + ) + ) + + async def close(self) -> None: + await self._engine.dispose() + + async def accept(self, execution_id: str, request: JsonObject) -> DurableExecution: + _validate_execution_id(execution_id) + serialized_request = _serialize_json_object(request) + async with self._engine.begin() as connection: + await connection.execute( + text( + f""" + INSERT INTO {self._table} (execution_id, status, request) + VALUES (:execution_id, 'QUEUED', CAST(:request AS JSONB)) + ON CONFLICT (execution_id) DO NOTHING + """ + ), + {"execution_id": execution_id, "request": serialized_request}, + ) + row = ( + ( + await connection.execute( + text( + f""" + SELECT execution_id, status, attempt, heartbeat_at, + request::TEXT AS request_json, + response::TEXT AS response_json + FROM {self._table} + WHERE execution_id=:execution_id + """ + ), + {"execution_id": execution_id}, + ) + ) + .mappings() + .one() + ) + + state = self._to_execution(row) + if state.request != request: + raise DurableRequestConflictError( + f"execution {execution_id!r} was already accepted with a different request" + ) + return state + + async def get(self, execution_id: str) -> DurableExecution | None: + _validate_execution_id(execution_id) + async with self._engine.connect() as connection: + row = ( + ( + await connection.execute( + text( + f""" + SELECT execution_id, status, attempt, heartbeat_at, + request::TEXT AS request_json, + response::TEXT AS response_json + FROM {self._table} + WHERE execution_id=:execution_id + """ + ), + {"execution_id": execution_id}, + ) + ) + .mappings() + .one_or_none() + ) + return self._to_execution(row) if row is not None else None + + async def recoverable_execution_ids(self, stale_seconds: float) -> list[str]: + async with self._engine.connect() as connection: + rows = ( + ( + await connection.execute( + text( + f""" + SELECT execution_id + FROM {self._table} + WHERE status='QUEUED' + OR (status='ACTIVE' AND ( + heartbeat_at IS NULL + OR heartbeat_at < NOW() - (:stale * INTERVAL '1 second') + )) + ORDER BY heartbeat_at NULLS FIRST + """ + ), + {"stale": stale_seconds}, + ) + ) + .scalars() + .all() + ) + return list(rows) + + async def claim( + self, + execution_id: str, + stale_seconds: float, + ) -> DurableExecution | None: + _validate_execution_id(execution_id) + async with self._engine.begin() as connection: + row = ( + ( + await connection.execute( + text( + f""" + UPDATE {self._table} + SET status='ACTIVE', attempt=attempt+1, heartbeat_at=NOW() + WHERE execution_id=:execution_id + AND ( + status='QUEUED' + OR (status='ACTIVE' AND ( + heartbeat_at IS NULL + OR heartbeat_at < NOW() - (:stale * INTERVAL '1 second') + )) + ) + RETURNING execution_id, status, attempt, heartbeat_at, + request::TEXT AS request_json, + response::TEXT AS response_json + """ + ), + {"execution_id": execution_id, "stale": stale_seconds}, + ) + ) + .mappings() + .one_or_none() + ) + return self._to_execution(row) if row is not None else None + + async def heartbeat(self, execution_id: str, attempt: int) -> bool: + _validate_execution_id(execution_id) + async with self._engine.begin() as connection: + result = await connection.execute( + text( + f""" + UPDATE {self._table} + SET heartbeat_at=NOW() + WHERE execution_id=:execution_id + AND attempt=:attempt + AND status='ACTIVE' + """ + ), + {"execution_id": execution_id, "attempt": attempt}, + ) + return result.rowcount == 1 + + async def complete( + self, + execution_id: str, + attempt: int, + response: JsonObject, + ) -> bool: + _validate_execution_id(execution_id) + serialized_response = _serialize_json_object(response) + async with self._engine.begin() as connection: + result = await connection.execute( + text( + f""" + UPDATE {self._table} + SET status='COMPLETED', response=CAST(:response AS JSONB) + WHERE execution_id=:execution_id + AND attempt=:attempt + AND status='ACTIVE' + """ + ), + { + "execution_id": execution_id, + "attempt": attempt, + "response": serialized_response, + }, + ) + return result.rowcount == 1 + + async def fail(self, execution_id: str, attempt: int) -> bool: + _validate_execution_id(execution_id) + async with self._engine.begin() as connection: + result = await connection.execute( + text( + f""" + UPDATE {self._table} + SET status='FAILED' + WHERE execution_id=:execution_id + AND attempt=:attempt + AND status='ACTIVE' + """ + ), + {"execution_id": execution_id, "attempt": attempt}, + ) + return result.rowcount == 1 + + async def append_event( + self, + execution_id: str, + attempt: int, + event: JsonObject, + ) -> int | None: + """Append an event only while the caller owns the active attempt.""" + _validate_execution_id(execution_id) + serialized_event = _serialize_json_object(event) + async with self._engine.begin() as connection: + result = await connection.execute( + text( + f""" + INSERT INTO {self._events_table} (execution_id, attempt, event) + SELECT :execution_id, :attempt, CAST(:event AS JSONB) + WHERE EXISTS ( + SELECT 1 + FROM {self._table} + WHERE execution_id=:execution_id + AND attempt=:attempt + AND status='ACTIVE' + ) + RETURNING sequence_number + """ + ), + { + "execution_id": execution_id, + "attempt": attempt, + "event": serialized_event, + }, + ) + sequence_number = result.scalar_one_or_none() + return int(sequence_number) if sequence_number is not None else None + + async def events( + self, + execution_id: str, + after_sequence: int | None = None, + ) -> list[DurableEvent]: + """Return ordered events for one execution after an optional cursor.""" + _validate_execution_id(execution_id) + async with self._engine.connect() as connection: + result = await connection.execute( + text( + f""" + SELECT sequence_number, execution_id, attempt, + event::TEXT AS event_json + FROM {self._events_table} + WHERE execution_id=:execution_id + AND (:after_sequence IS NULL OR sequence_number > :after_sequence) + ORDER BY sequence_number + """ + ), + { + "execution_id": execution_id, + "after_sequence": after_sequence, + }, + ) + rows = result.mappings().all() + return [ + DurableEvent( + sequence_number=int(row["sequence_number"]), + execution_id=str(row["execution_id"]), + attempt=int(row["attempt"]), + event=json.loads(row["event_json"]), + ) + for row in rows + ] + + @staticmethod + def _to_execution(row: Mapping[str, Any]) -> DurableExecution: + return DurableExecution( + execution_id=str(row["execution_id"]), + status=DurableExecutionStatus(str(row["status"])), + attempt=int(row["attempt"]), + heartbeat_at=row["heartbeat_at"], + request=json.loads(row["request_json"]), + response=json.loads(row["response_json"]) if row["response_json"] else None, + ) diff --git a/src/databricks_ai_bridge/durable_runtime/types.py b/src/databricks_ai_bridge/durable_runtime/types.py new file mode 100644 index 000000000..a03b73f64 --- /dev/null +++ b/src/databricks_ai_bridge/durable_runtime/types.py @@ -0,0 +1,126 @@ +"""Public types for durable request execution.""" + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from datetime import datetime +from enum import Enum +from typing import Any, Protocol + +JsonObject = dict[str, Any] +DurableEventEmitter = Callable[[JsonObject], Awaitable[int]] + + +class DurableExecutionStatus(str, Enum): + """Lifecycle states stored by the durability layer.""" + + QUEUED = "QUEUED" + ACTIVE = "ACTIVE" + COMPLETED = "COMPLETED" + FAILED = "FAILED" + + +@dataclass(frozen=True) +class DurableExecution: + """Persisted state for one idempotent request.""" + + execution_id: str + status: DurableExecutionStatus + attempt: int + heartbeat_at: datetime | None + request: JsonObject + response: JsonObject | None + + @property + def is_terminal(self) -> bool: + return self.status in { + DurableExecutionStatus.COMPLETED, + DurableExecutionStatus.FAILED, + } + + +@dataclass(frozen=True) +class DurableEvent: + """One persisted event emitted by a durable execution attempt.""" + + sequence_number: int + execution_id: str + attempt: int + event: JsonObject + + +@dataclass(frozen=True) +class DurableExecutionContext: + """Attempt metadata passed to the caller-owned executor.""" + + execution_id: str + attempt: int + _emit: DurableEventEmitter | None = field(default=None, repr=False, compare=False) + + @property + def is_recovery(self) -> bool: + return self.attempt > 1 + + async def emit(self, event: JsonObject) -> int: + """Persist an ordered event and return its replay cursor.""" + if self._emit is None: + raise RuntimeError("event emission is not available for this execution context") + return await self._emit(event) + + +DurableExecutor = Callable[[JsonObject, DurableExecutionContext], Awaitable[JsonObject]] + + +class DurabilityStore(Protocol): + """Persistence contract used by :class:`DatabricksDurableRuntime`.""" + + async def initialize(self) -> None: ... + + async def close(self) -> None: ... + + async def accept(self, execution_id: str, request: JsonObject) -> DurableExecution: ... + + async def get(self, execution_id: str) -> DurableExecution | None: ... + + async def recoverable_execution_ids(self, stale_seconds: float) -> list[str]: ... + + async def claim( + self, + execution_id: str, + stale_seconds: float, + ) -> DurableExecution | None: ... + + async def heartbeat(self, execution_id: str, attempt: int) -> bool: ... + + async def complete( + self, + execution_id: str, + attempt: int, + response: JsonObject, + ) -> bool: ... + + async def fail(self, execution_id: str, attempt: int) -> bool: ... + + async def append_event( + self, + execution_id: str, + attempt: int, + event: JsonObject, + ) -> int | None: ... + + async def events( + self, + execution_id: str, + after_sequence: int | None = None, + ) -> list[DurableEvent]: ... + + +class DurableRequestConflictError(ValueError): + """Raised when an execution ID is reused with a different request.""" + + +class DurableExecutionNotFoundError(LookupError): + """Raised when waiting for an unknown execution ID.""" + + +class DurableExecutionFailedError(RuntimeError): + """Raised when a durable execution reaches the failed state.""" diff --git a/tests/databricks_ai_bridge/test_durable_app.py b/tests/databricks_ai_bridge/test_durable_app.py new file mode 100644 index 000000000..7bab77b94 --- /dev/null +++ b/tests/databricks_ai_bridge/test_durable_app.py @@ -0,0 +1,218 @@ +"""Tests for the generic durable entrypoint application.""" + +from unittest.mock import AsyncMock, patch + +import pytest + +pytest.importorskip("fastapi") +pytest.importorskip("psycopg") +pytest.importorskip("sqlalchemy") + +from fastapi.testclient import TestClient + +from databricks_ai_bridge.durable_app import DatabricksDurableApp +from databricks_ai_bridge.durable_runtime import ( + DurableExecution, + DurableExecutionContext, + DurableExecutionStatus, +) + + +@pytest.mark.asyncio +async def test_entrypoint_receives_stable_session_and_recovery_context(): + durable_app = DatabricksDurableApp(durability_store=AsyncMock()) + emitted: list[dict] = [] + + async def emit(event: dict) -> int: + emitted.append(event) + return 4 + + @durable_app.entrypoint + async def agent(payload, context): + cursor = await context.emit({"type": "progress"}) + return { + "payload": payload, + "run_id": context.run_id, + "session_id": context.session_id, + "attempt": context.attempt, + "is_recovery": context.is_recovery, + "cursor": cursor, + } + + result = await durable_app._execute( + {"session_id": "session-1", "payload": {"message": "hello"}}, + DurableExecutionContext("run-1", 2, _emit=emit), + ) + + assert result == { + "payload": {"message": "hello"}, + "run_id": "run-1", + "session_id": "session-1", + "attempt": 2, + "is_recovery": True, + "cursor": 4, + } + assert emitted == [{"type": "progress"}] + + +@pytest.mark.asyncio +async def test_resume_entrypoint_handles_recovered_attempt(): + durable_app = DatabricksDurableApp(durability_store=AsyncMock()) + + @durable_app.entrypoint + async def agent(payload, context): + return {"handler": "entrypoint"} + + @durable_app.on_resume + async def resume_agent(payload, context): + return { + "handler": "resume", + "payload": payload, + "session_id": context.session_id, + } + + result = await durable_app._execute( + {"session_id": "session-1", "payload": {"message": "hello"}}, + DurableExecutionContext("run-1", 2, _emit=AsyncMock()), + ) + + assert result == { + "handler": "resume", + "payload": {"message": "hello"}, + "session_id": "session-1", + } + + +def test_header_submission_preserves_application_payload(): + durable_app = DatabricksDurableApp(durability_store=AsyncMock()) + + @durable_app.entrypoint + async def agent(payload, context): + return payload + + submit = AsyncMock( + return_value=DurableExecution( + execution_id="run-1", + status=DurableExecutionStatus.QUEUED, + attempt=0, + heartbeat_at=None, + request={"session_id": "session-1", "payload": {"message": "hello"}}, + response=None, + ) + ) + + with ( + patch.object(durable_app.runtime, "start", new_callable=AsyncMock), + patch.object(durable_app.runtime, "stop", new_callable=AsyncMock), + patch.object(durable_app.runtime, "submit", submit), + TestClient(durable_app) as client, + ): + response = client.post( + "/invocations", + headers={ + "Idempotency-Key": "run-1", + "Databricks-Agent-Session-Id": "session-1", + "Databricks-Background": "true", + }, + json={"message": "hello"}, + ) + + assert response.status_code == 202 + assert response.json() == { + "run_id": "run-1", + "status": "QUEUED", + "attempt": 0, + "result": None, + } + assert response.headers["Databricks-Run-Id"] == "run-1" + assert response.headers["Databricks-Agent-Session-Id"] == "session-1" + submit.assert_awaited_once_with( + "run-1", + {"session_id": "session-1", "payload": {"message": "hello"}}, + ) + + +def test_foreground_submission_preserves_agent_response_body(): + durable_app = DatabricksDurableApp(durability_store=AsyncMock()) + + @durable_app.entrypoint + async def agent(payload, context): + return payload + + with ( + patch.object(durable_app.runtime, "start", new_callable=AsyncMock), + patch.object(durable_app.runtime, "stop", new_callable=AsyncMock), + patch.object( + durable_app.runtime, + "invoke", + new=AsyncMock(return_value={"answer": "unchanged"}), + ), + TestClient(durable_app) as client, + ): + response = client.post( + "/invocations", + headers={"Databricks-Agent-Session-Id": "session-1"}, + json={"message": "hello"}, + ) + + assert response.status_code == 200 + assert response.json() == {"answer": "unchanged"} + assert response.headers["Databricks-Agent-Session-Id"] == "session-1" + + +def test_header_background_stream_preserves_body_and_returns_run_id(): + durable_app = DatabricksDurableApp(durability_store=AsyncMock()) + + @durable_app.entrypoint + async def agent(payload, context): + return payload + + completed = DurableExecution( + execution_id="run-1", + status=DurableExecutionStatus.COMPLETED, + attempt=1, + heartbeat_at=None, + request={"session_id": "session-1", "payload": {"message": "hello"}}, + response={"message": "hello"}, + ) + submit = AsyncMock(return_value=completed) + + with ( + patch.object(durable_app.runtime, "start", new_callable=AsyncMock), + patch.object(durable_app.runtime, "stop", new_callable=AsyncMock), + patch.object(durable_app.runtime, "submit", submit), + patch.object(durable_app.runtime, "events", new=AsyncMock(return_value=[])), + patch.object(durable_app.runtime, "get", new=AsyncMock(return_value=completed)), + TestClient(durable_app) as client, + ): + response = client.post( + "/invocations", + headers={ + "Idempotency-Key": "run-1", + "Databricks-Agent-Session-Id": "session-1", + "Databricks-Background": "true", + "Databricks-Stream": "true", + }, + json={"message": "hello"}, + ) + + assert response.status_code == 200 + assert response.headers["Databricks-Run-Id"] == "run-1" + submit.assert_awaited_once_with( + "run-1", + {"session_id": "session-1", "payload": {"message": "hello"}}, + ) + + +def test_only_one_entrypoint_can_be_registered(): + durable_app = DatabricksDurableApp(durability_store=AsyncMock()) + + @durable_app.entrypoint + async def first(payload, context): + return payload + + with pytest.raises(RuntimeError, match="one entrypoint"): + + @durable_app.entrypoint + async def second(payload, context): + return payload diff --git a/tests/databricks_ai_bridge/test_durable_runtime.py b/tests/databricks_ai_bridge/test_durable_runtime.py new file mode 100644 index 000000000..debfe1614 --- /dev/null +++ b/tests/databricks_ai_bridge/test_durable_runtime.py @@ -0,0 +1,456 @@ +"""Tests for DatabricksDurableRuntime orchestration.""" + +import asyncio +import copy +from datetime import datetime, timedelta, timezone + +import pytest + +pytest.importorskip("psycopg") +pytest.importorskip("sqlalchemy") + +from databricks_ai_bridge.durable_runtime import ( + DatabricksDurableRuntime, + DurableEvent, + DurableExecution, + DurableExecutionContext, + DurableExecutionFailedError, + DurableExecutionStatus, + DurableRequestConflictError, +) + + +class MemoryDurabilityStore: + def __init__(self) -> None: + self.states: dict[str, DurableExecution] = {} + self.initialized = False + self.closed = False + self.heartbeats: list[tuple[str, int]] = [] + self.persisted_events: list[DurableEvent] = [] + + async def initialize(self) -> None: + self.initialized = True + + async def close(self) -> None: + self.closed = True + + async def accept(self, execution_id: str, request: dict) -> DurableExecution: + existing = self.states.get(execution_id) + if existing is not None: + if existing.request != request: + raise DurableRequestConflictError(execution_id) + return existing + state = DurableExecution( + execution_id=execution_id, + status=DurableExecutionStatus.QUEUED, + attempt=0, + heartbeat_at=None, + request=copy.deepcopy(request), + response=None, + ) + self.states[execution_id] = state + return state + + async def get(self, execution_id: str) -> DurableExecution | None: + return self.states.get(execution_id) + + async def recoverable_execution_ids(self, stale_seconds: float) -> list[str]: + return [ + execution_id + for execution_id, state in self.states.items() + if state.status == DurableExecutionStatus.QUEUED + or ( + state.status == DurableExecutionStatus.ACTIVE + and ( + state.heartbeat_at is None + or datetime.now(timezone.utc) - state.heartbeat_at + >= timedelta(seconds=stale_seconds) + ) + ) + ] + + async def claim( + self, + execution_id: str, + stale_seconds: float, + ) -> DurableExecution | None: + state = self.states[execution_id] + recoverable = state.status == DurableExecutionStatus.QUEUED or ( + state.status == DurableExecutionStatus.ACTIVE + and ( + state.heartbeat_at is None + or datetime.now(timezone.utc) - state.heartbeat_at + >= timedelta(seconds=stale_seconds) + ) + ) + if not recoverable: + return None + state = DurableExecution( + execution_id=state.execution_id, + status=DurableExecutionStatus.ACTIVE, + attempt=state.attempt + 1, + heartbeat_at=datetime.now(timezone.utc), + request=state.request, + response=None, + ) + self.states[execution_id] = state + return state + + async def heartbeat(self, execution_id: str, attempt: int) -> bool: + state = self.states[execution_id] + if state.status != DurableExecutionStatus.ACTIVE or state.attempt != attempt: + return False + self.heartbeats.append((execution_id, attempt)) + self.states[execution_id] = DurableExecution( + execution_id=state.execution_id, + status=state.status, + attempt=state.attempt, + heartbeat_at=datetime.now(timezone.utc), + request=state.request, + response=state.response, + ) + return True + + async def complete(self, execution_id: str, attempt: int, response: dict) -> bool: + state = self.states[execution_id] + if state.status != DurableExecutionStatus.ACTIVE or state.attempt != attempt: + return False + self.states[execution_id] = DurableExecution( + execution_id=state.execution_id, + status=DurableExecutionStatus.COMPLETED, + attempt=state.attempt, + heartbeat_at=state.heartbeat_at, + request=state.request, + response=copy.deepcopy(response), + ) + return True + + async def fail(self, execution_id: str, attempt: int) -> bool: + state = self.states[execution_id] + if state.status != DurableExecutionStatus.ACTIVE or state.attempt != attempt: + return False + self.states[execution_id] = DurableExecution( + execution_id=state.execution_id, + status=DurableExecutionStatus.FAILED, + attempt=state.attempt, + heartbeat_at=state.heartbeat_at, + request=state.request, + response=None, + ) + return True + + async def append_event( + self, + execution_id: str, + attempt: int, + event: dict, + ) -> int | None: + state = self.states[execution_id] + if state.status != DurableExecutionStatus.ACTIVE or state.attempt != attempt: + return None + persisted = DurableEvent( + sequence_number=len(self.persisted_events) + 1, + execution_id=execution_id, + attempt=attempt, + event=copy.deepcopy(event), + ) + self.persisted_events.append(persisted) + return persisted.sequence_number + + async def events( + self, + execution_id: str, + after_sequence: int | None = None, + ) -> list[DurableEvent]: + return [ + copy.deepcopy(event) + for event in self.persisted_events + if event.execution_id == execution_id + and (after_sequence is None or event.sequence_number > after_sequence) + ] + + +def make_runtime(executor, store=None, **kwargs): + return DatabricksDurableRuntime( + executor, + durability_store=store or MemoryDurabilityStore(), + heartbeat_seconds=0.01, + stale_seconds=0.05, + scan_seconds=0.01, + poll_seconds=0.005, + **kwargs, + ) + + +@pytest.mark.asyncio +async def test_invoke_persists_request_and_response(): + calls = [] + + async def execute(request: dict, context: DurableExecutionContext) -> dict: + calls.append((request, context)) + return {"output": request["input"]} + + store = MemoryDurabilityStore() + runtime = make_runtime(execute, store) + await runtime.start() + try: + response = await runtime.invoke("session-1", {"input": "hello"}) + state = await runtime.get("session-1") + finally: + await runtime.stop() + + assert response == {"output": "hello"} + assert state is not None + assert state.request == {"input": "hello"} + assert state.response == {"output": "hello"} + assert calls[0][1].attempt == 1 + assert calls[0][1].is_recovery is False + + +@pytest.mark.asyncio +async def test_completed_request_returns_cached_response_without_reexecution(): + call_count = 0 + + async def execute(request: dict, context: DurableExecutionContext) -> dict: + nonlocal call_count + call_count += 1 + return {"output": "done"} + + runtime = make_runtime(execute) + await runtime.start() + try: + first = await runtime.invoke("session-1", {"input": "hello"}) + second = await runtime.invoke("session-1", {"input": "hello"}) + finally: + await runtime.stop() + + assert first == second == {"output": "done"} + assert call_count == 1 + + +@pytest.mark.asyncio +async def test_same_id_with_different_request_is_rejected(): + async def execute(request: dict, context: DurableExecutionContext) -> dict: + return {"output": "done"} + + runtime = make_runtime(execute) + await runtime.start() + try: + await runtime.invoke("session-1", {"input": "first"}) + with pytest.raises(DurableRequestConflictError): + await runtime.submit("session-1", {"input": "second"}) + finally: + await runtime.stop() + + +@pytest.mark.asyncio +async def test_stale_attempt_reuses_request_and_marks_recovery(): + contexts = [] + + async def execute(request: dict, context: DurableExecutionContext) -> dict: + contexts.append((request, context)) + return {"output": "recovered"} + + store = MemoryDurabilityStore() + store.states["session-1"] = DurableExecution( + execution_id="session-1", + status=DurableExecutionStatus.ACTIVE, + attempt=1, + heartbeat_at=datetime.now(timezone.utc) - timedelta(minutes=1), + request={"input": "original"}, + response=None, + ) + runtime = make_runtime(execute, store) + await runtime.start() + try: + response = await runtime.wait("session-1") + finally: + await runtime.stop() + + assert response == {"output": "recovered"} + assert contexts == [ + ( + {"input": "original"}, + DurableExecutionContext(execution_id="session-1", attempt=2), + ) + ] + assert contexts[0][1].is_recovery is True + + +@pytest.mark.asyncio +async def test_executor_failure_is_persisted_as_terminal_state(): + async def execute(request: dict, context: DurableExecutionContext) -> dict: + raise RuntimeError("boom") + + runtime = make_runtime(execute) + await runtime.start() + try: + with pytest.raises(DurableExecutionFailedError): + await runtime.invoke("session-1", {"input": "hello"}) + state = await runtime.get("session-1") + finally: + await runtime.stop() + + assert state is not None + assert state.status == DurableExecutionStatus.FAILED + assert state.response is None + + +@pytest.mark.asyncio +async def test_submit_returns_before_background_execution_finishes(): + release = asyncio.Event() + + async def execute(request: dict, context: DurableExecutionContext) -> dict: + await release.wait() + return {"output": "done"} + + runtime = make_runtime(execute) + await runtime.start() + try: + state = await runtime.submit("session-1", {"input": "hello"}) + assert state.status == DurableExecutionStatus.QUEUED + release.set() + assert await runtime.wait("session-1") == {"output": "done"} + finally: + await runtime.stop() + + +@pytest.mark.asyncio +async def test_wait_observes_response_completed_by_another_process(): + executor_called = False + + async def execute(request: dict, context: DurableExecutionContext) -> dict: + nonlocal executor_called + executor_called = True + return {} + + store = MemoryDurabilityStore() + store.states["session-1"] = DurableExecution( + execution_id="session-1", + status=DurableExecutionStatus.ACTIVE, + attempt=1, + heartbeat_at=datetime.now(timezone.utc), + request={"input": "hello"}, + response=None, + ) + runtime = make_runtime(execute, store) + await runtime.start() + + async def complete_elsewhere(): + await asyncio.sleep(0.02) + await store.complete("session-1", 1, {"output": "remote"}) + + completion = asyncio.create_task(complete_elsewhere()) + try: + assert await runtime.wait("session-1") == {"output": "remote"} + finally: + await completion + await runtime.stop() + + assert executor_called is False + + +@pytest.mark.asyncio +async def test_blocking_timeout_does_not_cancel_execution(): + release = asyncio.Event() + + async def execute(request: dict, context: DurableExecutionContext) -> dict: + await release.wait() + return {"output": "done"} + + runtime = make_runtime(execute) + await runtime.start() + try: + with pytest.raises(TimeoutError): + await runtime.invoke("session-1", {"input": "hello"}, timeout=0.02) + state = await runtime.get("session-1") + assert state is not None + assert state.status == DurableExecutionStatus.ACTIVE + release.set() + assert await runtime.wait("session-1") == {"output": "done"} + finally: + await runtime.stop() + + +@pytest.mark.asyncio +async def test_executor_can_persist_replayable_events(): + async def execute(request: dict, context: DurableExecutionContext) -> dict: + sequence_number = await context.emit({"type": "progress", "step": 1}) + return {"last_sequence_number": sequence_number} + + runtime = make_runtime(execute) + await runtime.start() + try: + assert await runtime.invoke("session-1", {"input": "hello"}) == {"last_sequence_number": 1} + events = await runtime.events("session-1") + assert [(event.sequence_number, event.event) for event in events] == [ + (1, {"type": "progress", "step": 1}) + ] + assert await runtime.events("session-1", after_sequence=1) == [] + finally: + await runtime.stop() + + +@pytest.mark.asyncio +async def test_request_must_be_json_object(): + async def execute(request: dict, context: DurableExecutionContext) -> dict: + return {} + + runtime = make_runtime(execute) + await runtime.start() + try: + with pytest.raises(TypeError, match="request must be a JSON object"): + await runtime.submit("session-1", ["not", "an", "object"]) + finally: + await runtime.stop() + + +@pytest.mark.asyncio +async def test_subclass_can_own_execution_wiring(): + class Runtime(DatabricksDurableRuntime): + async def execute( + self, + request: dict, + context: DurableExecutionContext, + ) -> dict: + return {"attempt": context.attempt, "input": request["input"]} + + runtime = Runtime( + durability_store=MemoryDurabilityStore(), + heartbeat_seconds=0.01, + stale_seconds=0.05, + scan_seconds=0.01, + poll_seconds=0.005, + ) + await runtime.start() + try: + assert await runtime.invoke("session-1", {"input": "hello"}) == { + "attempt": 1, + "input": "hello", + } + finally: + await runtime.stop() + + +@pytest.mark.asyncio +async def test_start_and_stop_manage_store_lifecycle(): + async def execute(request: dict, context: DurableExecutionContext) -> dict: + return {} + + store = MemoryDurabilityStore() + runtime = make_runtime(execute, store) + + await runtime.start() + assert store.initialized is True + await runtime.stop() + assert store.closed is True + + +@pytest.mark.asyncio +async def test_runtime_requires_start_before_use(): + async def execute(request: dict, context: DurableExecutionContext) -> dict: + return {} + + runtime = make_runtime(execute) + with pytest.raises(RuntimeError, match=r"start\(\)"): + await runtime.submit("session-1", {}) diff --git a/tests/databricks_ai_bridge/test_durable_runtime_store.py b/tests/databricks_ai_bridge/test_durable_runtime_store.py new file mode 100644 index 000000000..174340e3a --- /dev/null +++ b/tests/databricks_ai_bridge/test_durable_runtime_store.py @@ -0,0 +1,207 @@ +"""Tests for the Lakebase durability store.""" + +from contextlib import asynccontextmanager +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock + +import pytest + +pytest.importorskip("psycopg") +pytest.importorskip("sqlalchemy") + +from databricks_ai_bridge.durable_runtime import ( + DurableExecutionStatus, + DurableRequestConflictError, + LakebaseDurabilityStore, +) + + +def mock_lakebase(): + connection = AsyncMock() + engine = MagicMock() + + @asynccontextmanager + async def begin(): + yield connection + + @asynccontextmanager + async def connect(): + yield connection + + engine.begin = begin + engine.connect = connect + engine.dispose = AsyncMock() + lakebase = MagicMock(engine=engine) + lakebase.create_schema = AsyncMock() + return lakebase, connection + + +def mapping_result(value): + result = MagicMock() + result.mappings.return_value.one.return_value = value + result.mappings.return_value.one_or_none.return_value = value + result.mappings.return_value.all.return_value = value if isinstance(value, list) else [value] + return result + + +def execution_row(**overrides): + row = { + "execution_id": "session-1", + "status": "QUEUED", + "attempt": 0, + "heartbeat_at": None, + "request_json": '{"input": "hello"}', + "response_json": None, + } + row.update(overrides) + return row + + +@pytest.mark.asyncio +async def test_initialize_creates_execution_and_event_tables(): + lakebase, connection = mock_lakebase() + store = LakebaseDurabilityStore(lakebase=lakebase) + + await store.initialize() + + sql = " ".join(str(call.args[0]) for call in connection.execute.await_args_list) + assert "databricks_durable_runtime.executions" in sql + assert "execution_id TEXT PRIMARY KEY" in sql + assert "request JSONB NOT NULL" in sql + assert "response JSONB" in sql + assert "databricks_durable_runtime.execution_events" in sql + assert "sequence_number BIGSERIAL PRIMARY KEY" in sql + lakebase.create_schema.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_accept_returns_existing_request_when_it_matches(): + lakebase, connection = mock_lakebase() + connection.execute.side_effect = [MagicMock(), mapping_result(execution_row())] + store = LakebaseDurabilityStore(lakebase=lakebase) + + state = await store.accept("session-1", {"input": "hello"}) + + assert state.execution_id == "session-1" + assert state.status == DurableExecutionStatus.QUEUED + assert state.request == {"input": "hello"} + + +@pytest.mark.asyncio +async def test_accept_rejects_same_id_with_different_request(): + lakebase, connection = mock_lakebase() + connection.execute.side_effect = [MagicMock(), mapping_result(execution_row())] + store = LakebaseDurabilityStore(lakebase=lakebase) + + with pytest.raises(DurableRequestConflictError): + await store.accept("session-1", {"input": "different"}) + + +@pytest.mark.asyncio +async def test_claim_returns_request_and_incremented_attempt(): + lakebase, connection = mock_lakebase() + heartbeat = datetime.now(timezone.utc) + connection.execute.return_value = mapping_result( + execution_row(status="ACTIVE", attempt=2, heartbeat_at=heartbeat) + ) + store = LakebaseDurabilityStore(lakebase=lakebase) + + state = await store.claim("session-1", 10) + + assert state is not None + assert state.attempt == 2 + assert state.heartbeat_at == heartbeat + assert state.request == {"input": "hello"} + + +@pytest.mark.asyncio +async def test_get_decodes_cached_response(): + lakebase, connection = mock_lakebase() + connection.execute.return_value = mapping_result( + execution_row( + status="COMPLETED", + attempt=1, + response_json='{"output": "done"}', + ) + ) + store = LakebaseDurabilityStore(lakebase=lakebase) + + state = await store.get("session-1") + + assert state is not None + assert state.status == DurableExecutionStatus.COMPLETED + assert state.response == {"output": "done"} + + +@pytest.mark.asyncio +async def test_complete_persists_response_for_owned_attempt(): + lakebase, connection = mock_lakebase() + connection.execute.return_value = MagicMock(rowcount=1) + store = LakebaseDurabilityStore(lakebase=lakebase) + + assert await store.complete("session-1", 2, {"output": "done"}) is True + + parameters = connection.execute.await_args.args[1] + assert parameters["execution_id"] == "session-1" + assert parameters["attempt"] == 2 + assert parameters["response"] == '{"output": "done"}' + + +@pytest.mark.asyncio +async def test_append_event_returns_replay_cursor_for_owned_attempt(): + lakebase, connection = mock_lakebase() + result = MagicMock() + result.scalar_one_or_none.return_value = 7 + connection.execute.return_value = result + store = LakebaseDurabilityStore(lakebase=lakebase) + + sequence_number = await store.append_event( + "session-1", + 2, + {"type": "progress", "step": 1}, + ) + + assert sequence_number == 7 + parameters = connection.execute.await_args.args[1] + assert parameters == { + "execution_id": "session-1", + "attempt": 2, + "event": '{"type": "progress", "step": 1}', + } + + +@pytest.mark.asyncio +async def test_events_returns_ordered_replay_data(): + lakebase, connection = mock_lakebase() + connection.execute.return_value = mapping_result( + [ + { + "sequence_number": 8, + "execution_id": "session-1", + "attempt": 2, + "event_json": '{"type": "progress", "step": 2}', + } + ] + ) + store = LakebaseDurabilityStore(lakebase=lakebase) + + events = await store.events("session-1", after_sequence=7) + + assert len(events) == 1 + assert events[0].sequence_number == 8 + assert events[0].attempt == 2 + assert events[0].event == {"type": "progress", "step": 2} + + +def test_schema_name_is_validated(): + lakebase, _ = mock_lakebase() + with pytest.raises(ValueError, match="invalid durability schema"): + LakebaseDurabilityStore(lakebase=lakebase, schema="bad-schema;drop") + + +@pytest.mark.asyncio +async def test_store_rejects_empty_execution_id(): + lakebase, _ = mock_lakebase() + store = LakebaseDurabilityStore(lakebase=lakebase) + with pytest.raises(ValueError, match="must not be empty"): + await store.accept("", {})