diff --git a/CHANGELOG.md b/CHANGELOG.md index d8b336d7e..5d68d32c2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### 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 persistence and stale-attempt recovery ### 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..063b74cbf 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,20 @@ 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/response persistence, heartbeat detection, and +stale-attempt recovery around a caller-owned async handler. The handler remains +responsible for agent sessions and checkpoints. + +See the [OpenAI Agents SDK App](./examples/openai-sdk-agent/README.md) for a +complete FastAPI, background polling, SDK session, and Databricks Apps example. + +```sh +pip install 'databricks-ai-bridge[memory]' +``` + ### Install from source With https: diff --git a/examples/openai-sdk-agent/OBSERVATIONS.md b/examples/openai-sdk-agent/OBSERVATIONS.md new file mode 100644 index 000000000..c93fef07e --- /dev/null +++ b/examples/openai-sdk-agent/OBSERVATIONS.md @@ -0,0 +1,155 @@ +# Live Test Observations + +Tests ran on 2026-08-19 against the example in this PR, using an isolated App +name and schemas so they did not affect the earlier experiment. + +- App: `open-ai-sdk-runtime` +- URL: `https://open-ai-sdk-runtime-1653573648247579.staging.aws.databricksapps.com` +- Lakebase branch: `projects/shivam-openai-agent-on-apps/branches/agent-app` +- Database: `databricks_postgres` +- Runtime schema: `openai_sdk_runtime_durability` +- SDK session schema: `openai_sdk_runtime_sessions` + +The PR packages were built as local wheels and included in the test deployment +because `DatabricksDurableRuntime` was not released yet. All remaining packages +installed through the workspace package repository, and the App finished in +`ACTIVE / RUNNING`; no package-proxy failure occurred. + +## Common request and database query + +Each test supplied a stable ID in `custom_inputs.session_id`: + +```bash +curl -sS --max-time 900 \ + -o response.json \ + -w 'http_code=%{http_code}\ntime_total=%{time_total}\n' \ + -X POST "$APP_URL/responses" \ + -H "Authorization: Bearer $APP_TOKEN" \ + -H 'Content-Type: application/json' \ + --data @request.json +``` + +The following query joined runtime state with SDK history: + +```sql +SELECT e.status, e.attempt, e.heartbeat_at, + e.request, e.response, count(m.id) AS sdk_messages +FROM openai_sdk_runtime_durability.executions e +LEFT JOIN openai_sdk_runtime_sessions.agent_messages m + ON m.session_id = e.execution_id +WHERE e.execution_id = :execution_id +GROUP BY e.execution_id; +``` + +## Test 1: blocking happy path + +Execution: `runtime-happy-20260819T221342Z` + +```text +http_code=200 +time_total=74.347104 +curl_exit=0 +``` + +Lakebase after completion: + +```text +status=COMPLETED attempt=1 sdk_messages=30 +request=persisted response=persisted +``` + +The runtime row contained the normalized Responses request and final response. +The SDK table independently contained the model and tool history. + +## Test 2: cache and conflict + +Posting the exact Test 1 request and ID again returned in 0.258 seconds. The +original and cached response files had the same SHA-256: + +```text +726b28b5a378e0b6e176a9894deeceb07f28ed8c7865ba21b81fd3324a98a23e +``` + +Changing only the input while retaining the ID returned: + +```text +http_code=409 +execution 'runtime-happy-20260819T221342Z' was already accepted with a different request +``` + +This verifies exact-request idempotency rather than ID-only response reuse. + +## Test 3: blocking client disconnect + +Execution: `runtime-disconnect-20260819T221458Z` + +The client process was terminated after Lakebase showed `ACTIVE`, attempt `1`, +and one SDK message: + +```text +curl_exit=143 +``` + +The runtime task continued without the HTTP client. Its final state was: + +```text +status=COMPLETED attempt=1 sdk_messages=36 +request=persisted response=persisted +``` + +`GET /responses/runtime-disconnect-20260819T221458Z` then returned the completed +response with HTTP `200`. Unlike the earlier custom supervisor experiment, a +disconnected client can retrieve the persisted result by its stable ID. + +## Test 4: background stop/start recovery + +Execution: `runtime-crash-20260819T222154Z` + +The background request returned immediately: + +```text +http_code=202 +time_total=0.223404 +status=in_progress +``` + +The App was stopped after the first attempt had persisted SDK history: + +```text +before stop: status=ACTIVE attempt=1 sdk_messages=5 response=NULL +after stop: status=ACTIVE attempt=1 sdk_messages=7 response=NULL +``` + +While compute was stopped, retrieval returned HTTP `503`. The runtime row and +SDK history remained in Lakebase. After `databricks apps start`, the scanner +claimed the stale row: + +```text +after restart: status=ACTIVE attempt=2 sdk_messages=12 response=NULL +final: status=COMPLETED attempt=2 sdk_messages=31 response=persisted +``` + +SDK message `112` was the fixed recovery note. It contained neither the PR URL +nor the old temporary workspace. The following recovered tool calls nevertheless +used `/tmp/openai-sdk-agent-0oxbysef/repo` and checked out PR 459. Those values +were present only in messages `105` through `111`, proving that attempt `2` +reopened and used the persisted SDK session history. The old pod-local directory +was gone, so the agent recreated it. + +After completion: + +- `GET /responses/runtime-crash-20260819T222154Z` returned HTTP `200`. +- Reposting the exact background request returned the cached response in 0.221 seconds. +- Retrieved and cached responses had the same SHA-256: + `4929b07e0931a5a487efac00f4d13910b4ccf2d15e3edb470f670775eb34f5bc`. + +## Result + +The live tests verify the intended separation: + +- `DatabricksDurableRuntime` persists request, response, status, attempt, and heartbeat. +- `AsyncDatabricksSession` persists replayable agent and tool history. +- Recovery replays the persisted request to the executor, while this executor + intentionally resumes the agent with only the SDK session and recovery note. +- A stable execution ID lets a client poll after a disconnect or App restart. +- Pod-local files and in-flight tool processes are not durable. diff --git a/examples/openai-sdk-agent/README.md b/examples/openai-sdk-agent/README.md new file mode 100644 index 000000000..49377759e --- /dev/null +++ b/examples/openai-sdk-agent/README.md @@ -0,0 +1,114 @@ +# Durable OpenAI Agents SDK App + +This example is the PR-review agent from the custom-runtime experiment, with +its application-owned durability package replaced by `DatabricksDurableRuntime`. +The agent loop in `review_agent.py` and the OpenAI Agents SDK session in +`sessions.py` remain application concerns. + +See [Live Test Observations](./OBSERVATIONS.md) for blocking, cache/conflict, +client-disconnect, and real App stop/start recovery results. + +## Responsibilities + +```text +client + -> FastAPI adapter (app.py) + -> DatabricksDurableRuntime + -> Lakebase: openai_sdk_agent_durability.executions + -> executor (execute_durable_review) + -> OpenAI Agents SDK + tools + -> Lakebase: openai_sdk_agent_sessions.agent_messages +``` + +`DatabricksDurableRuntime` owns request/response persistence, exact-request +idempotency, heartbeats, stale-attempt claims, and process-start recovery. The +executor owns the SDK session and recovery behavior. On attempt 1 it starts the +review from the request. On attempt 2 or later it reopens the same SDK session +and supplies only the fixed recovery note; it does not reconstruct an agent +prompt from the durability request. + +This example intentionally allows one durable request per SDK session, so +`custom_inputs.session_id` is also the runtime `execution_id`. A multi-turn +application should use a separate execution ID for each invocation and keep its +conversation or session ID in the persisted request. + +## HTTP contract + +- `POST /responses` and `POST /invocations` run in blocking mode by default. +- Set `background: true` to receive `202` with an ID and poll + `GET /responses/{execution_id}`. +- Repeating the same normalized request and ID returns the cached response. +- Reusing an ID with a different request returns `409 Conflict`. +- `background` and `stream` are transport fields and are not persisted. + Streaming is rejected because this example does not implement it. + +Clients should supply a stable `custom_inputs.session_id`. A generated ID can +be returned to a connected client, but a blocking client that loses its +connection before receiving that ID cannot later identify the execution. + +Example background request: + +```bash +SESSION_ID="review-$(date -u +%Y%m%dT%H%M%SZ)" + +curl -X POST "$APP_URL/responses" \ + -H "Authorization: Bearer $APP_TOKEN" \ + -H 'Content-Type: application/json' \ + -d "{ + \"background\": true, + \"input\": [{\"role\": \"user\", \"content\": \"Execute the complete PR CUJ.\"}], + \"custom_inputs\": { + \"session_id\": \"$SESSION_ID\", + \"pr_url\": \"https://github.com/databricks/databricks-ai-bridge/pull/459\", + \"minimum_minutes\": 0 + } + }" + +curl "$APP_URL/responses/$SESSION_ID" \ + -H "Authorization: Bearer $APP_TOKEN" +``` + +## Lakebase state + +Both stores use the App's `postgres` resource but separate schemas: + +| Owner | Schema and table | Persisted state | +| --- | --- | --- | +| Runtime | `openai_sdk_agent_durability.executions` | execution ID, status, attempt, heartbeat, normalized request, final response | +| OpenAI Agents SDK | `openai_sdk_agent_sessions.agent_messages` | replayable user, assistant, tool-call, and tool-output items | + +The runtime provides at-least-once recovery. Pod-local files and in-flight tool +processes do not survive a crash, and tools must tolerate retries. + +## Deploy + +Install from the repository checkout while developing this unreleased runtime: + +```bash +uv venv +uv pip install -e '../..[memory]' -e '../../integrations/openai[memory]' +uv pip install 'openai-agents>=0.19.4,<0.20' 'mcp>=1.29.0,<2' \ + 'mlflow>=3.10.1' 'fastapi>=0.129.0' 'uvicorn>=0.41.0' +``` + +For Databricks Apps, configure one Lakebase branch/database and one secret, then +deploy with an explicitly selected profile: + +```bash +databricks bundle deploy -t dev --profile \ + --var="lakebase_branch=projects//branches/" \ + --var="lakebase_database=projects//branches//databases/" \ + --var="openai_secret_scope=" \ + --var="openai_secret_key=" + +databricks bundle run open_ai_sdk_agent -t dev --profile \ + --var="lakebase_branch=projects//branches/" \ + --var="lakebase_database=projects//branches//databases/" \ + --var="openai_secret_scope=" \ + --var="openai_secret_key=" +``` + +After the runtime is released, the App build installs `requirements.txt` +directly. When deploying this PR before release, replace the +`databricks-ai-bridge` requirement with an installable wheel or Git ref that +contains `DatabricksDurableRuntime`. diff --git a/examples/openai-sdk-agent/app.py b/examples/openai-sdk-agent/app.py new file mode 100644 index 000000000..63047251d --- /dev/null +++ b/examples/openai-sdk-agent/app.py @@ -0,0 +1,186 @@ +"""OpenAI Agents SDK app backed by DatabricksDurableRuntime.""" + +import os +import re +from contextlib import asynccontextmanager +from uuid import uuid4 + +import uvicorn +from fastapi import FastAPI, HTTPException, Request, Response +from mlflow.types.responses import ResponsesAgentRequest, ResponsesAgentResponse +from review_agent import execute_review, resume_review +from sessions import create_session, initialize_sessions + +from databricks_ai_bridge.durable_runtime import ( + DatabricksDurableRuntime, + DurableExecution, + DurableExecutionContext, + DurableExecutionFailedError, + DurableExecutionStatus, + DurableRequestConflictError, + JsonObject, +) + +PR_URL = re.compile(r"^https://github\.com/[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+/pull/[0-9]+/?$") +RECOVERY_NOTE = ( + "This is a crash-recovery attempt. Continue the interrupted task using only " + "the persisted SDK session history. Inspect the shell workspace and safely " + "repeat any interrupted tool." +) +RESPONSE_STATUS = { + DurableExecutionStatus.QUEUED: "in_progress", + DurableExecutionStatus.ACTIVE: "in_progress", + DurableExecutionStatus.COMPLETED: "completed", + DurableExecutionStatus.FAILED: "failed", +} + + +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 getattr(request.context, "conversation_id", None): + return str(request.context.conversation_id) + return str(uuid4()) + + +def _review_inputs(request: ResponsesAgentRequest) -> tuple[str, float]: + custom_inputs = dict(request.custom_inputs or {}) + pr_url = str(custom_inputs.get("pr_url") or "") + if not PR_URL.fullmatch(pr_url): + raise HTTPException(400, "custom_inputs.pr_url must be a public GitHub pull-request URL") + try: + minimum_minutes = float(custom_inputs.get("minimum_minutes", 30)) + except (TypeError, ValueError) as exc: + raise HTTPException(400, "custom_inputs.minimum_minutes must be a number") from exc + if not 0 <= minimum_minutes <= 60: + raise HTTPException(400, "custom_inputs.minimum_minutes must be between 0 and 60") + return pr_url, minimum_minutes + + +def _durable_request(request: ResponsesAgentRequest, session_id: str) -> JsonObject: + payload = request.model_dump(mode="json", exclude_none=True) + payload.pop("background", None) + payload.pop("stream", None) + custom_inputs = dict(payload.get("custom_inputs") or {}) + custom_inputs["session_id"] = session_id + payload["custom_inputs"] = custom_inputs + return payload + + +def _message(text: str) -> dict[str, object]: + return { + "id": str(uuid4()), + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": text}], + } + + +async def execute_durable_review( + request: JsonObject, + context: DurableExecutionContext, +) -> JsonObject: + agent_request = ResponsesAgentRequest.model_validate(request) + session_id = _session_id(agent_request) + pr_url, minimum_minutes = _review_inputs(agent_request) + session = create_session(session_id) + if context.is_recovery: + report = await resume_review(session, RECOVERY_NOTE) + else: + report = await execute_review(pr_url, minimum_minutes, session) + response = ResponsesAgentResponse.model_validate( + { + "id": context.execution_id, + "status": "completed", + "output": [_message(report)], + "custom_outputs": { + "execution_id": context.execution_id, + "session_id": session_id, + "attempt": context.attempt, + }, + } + ) + return response.model_dump(mode="json", exclude_none=True) + + +@asynccontextmanager +async def lifespan(application: FastAPI): + await initialize_sessions() + durable_runtime = DatabricksDurableRuntime( + execute_durable_review, + schema=os.getenv("LAKEBASE_DURABILITY_SCHEMA", "openai_sdk_agent_durability"), + ) + await durable_runtime.start() + application.state.durable_runtime = durable_runtime + try: + yield + finally: + await durable_runtime.stop() + + +app = FastAPI(title="Durable OpenAI SDK PR review agent", lifespan=lifespan) + + +def _status_response(state: DurableExecution) -> ResponsesAgentResponse: + if state.status == DurableExecutionStatus.COMPLETED: + if state.response is None: + raise RuntimeError(f"execution {state.execution_id!r} completed without a response") + return ResponsesAgentResponse.model_validate(state.response) + custom_inputs = dict(state.request.get("custom_inputs") or {}) + return ResponsesAgentResponse( + id=state.execution_id, + status=RESPONSE_STATUS[state.status], + output=[], + custom_outputs={ + "execution_id": state.execution_id, + "session_id": custom_inputs.get("session_id", state.execution_id), + "attempt": state.attempt, + }, + ) + + +@app.get("/health") +@app.get("/api/healthz") +async def health() -> dict: + return {"ok": True, "model": "gpt-5.6-luna"} + + +@app.post("/responses") +@app.post("/invocations") +async def invoke( + request: ResponsesAgentRequest, + response: Response, + http_request: Request, +) -> ResponsesAgentResponse: + if request.stream: + raise HTTPException(400, "streaming is not implemented by this example") + _review_inputs(request) + execution_id = _session_id(request) + payload = _durable_request(request, execution_id) + durable_runtime: DatabricksDurableRuntime = http_request.app.state.durable_runtime + try: + if bool(getattr(request, "background", False)): + state = await durable_runtime.submit(execution_id, payload) + if not state.is_terminal: + response.status_code = 202 + return _status_response(state) + result = await durable_runtime.invoke(execution_id, payload) + return ResponsesAgentResponse.model_validate(result) + except DurableRequestConflictError as exc: + raise HTTPException(409, str(exc)) from exc + except DurableExecutionFailedError as exc: + raise HTTPException(500, str(exc)) from exc + + +@app.get("/responses/{execution_id}") +async def retrieve(execution_id: str, request: Request) -> ResponsesAgentResponse: + durable_runtime: DatabricksDurableRuntime = request.app.state.durable_runtime + state = await durable_runtime.get(execution_id) + if state is None: + raise HTTPException(404, f"execution {execution_id!r} was not found") + return _status_response(state) + + +if __name__ == "__main__": + uvicorn.run("app:app", host="0.0.0.0", port=int(os.getenv("DATABRICKS_APP_PORT", "8000"))) diff --git a/examples/openai-sdk-agent/app.yaml b/examples/openai-sdk-agent/app.yaml new file mode 100644 index 000000000..6a1dbb283 --- /dev/null +++ b/examples/openai-sdk-agent/app.yaml @@ -0,0 +1,12 @@ +command: ["python", "app.py"] +env: + - name: OPENAI_API_KEY + valueFrom: openai-api-key + - name: OPENAI_MODEL + value: gpt-5.6-luna + - name: LAKEBASE_AUTOSCALING_ENDPOINT + valueFrom: postgres + - name: LAKEBASE_SESSION_SCHEMA + value: openai_sdk_agent_sessions + - name: LAKEBASE_DURABILITY_SCHEMA + value: openai_sdk_agent_durability diff --git a/examples/openai-sdk-agent/databricks.yml b/examples/openai-sdk-agent/databricks.yml new file mode 100644 index 000000000..ad4b4f888 --- /dev/null +++ b/examples/openai-sdk-agent/databricks.yml @@ -0,0 +1,48 @@ +bundle: + name: open_ai_sdk_agent + +variables: + lakebase_branch: + description: Full Lakebase branch resource name + lakebase_database: + description: Full Lakebase database resource name + openai_secret_scope: + description: Databricks secret scope containing the OpenAI API key + openai_secret_key: + description: Databricks secret key containing the OpenAI API key + +resources: + apps: + open_ai_sdk_agent: + name: open-ai-sdk-agent + description: OpenAI Agents SDK example using DatabricksDurableRuntime + source_code_path: ./ + config: + command: ["python", "app.py"] + env: + - name: OPENAI_API_KEY + value_from: openai-api-key + - name: OPENAI_MODEL + value: gpt-5.6-luna + - name: LAKEBASE_AUTOSCALING_ENDPOINT + value_from: postgres + - name: LAKEBASE_SESSION_SCHEMA + value: openai_sdk_agent_sessions + - name: LAKEBASE_DURABILITY_SCHEMA + value: openai_sdk_agent_durability + resources: + - name: openai-api-key + secret: + scope: ${var.openai_secret_scope} + key: ${var.openai_secret_key} + permission: READ + - name: postgres + postgres: + branch: ${var.lakebase_branch} + database: ${var.lakebase_database} + permission: CAN_CONNECT_AND_CREATE + +targets: + dev: + mode: development + default: true diff --git a/examples/openai-sdk-agent/requirements.txt b/examples/openai-sdk-agent/requirements.txt new file mode 100644 index 000000000..9f30a15e0 --- /dev/null +++ b/examples/openai-sdk-agent/requirements.txt @@ -0,0 +1,7 @@ +openai-agents>=0.19.4,<0.20 +mcp>=1.29.0,<2 +databricks-openai[memory]>=0.17.0 +databricks-ai-bridge[memory]>=0.21.0 +mlflow>=3.10.1 +fastapi>=0.129.0 +uvicorn>=0.41.0 diff --git a/examples/openai-sdk-agent/review_agent.py b/examples/openai-sdk-agent/review_agent.py new file mode 100644 index 000000000..a7232658a --- /dev/null +++ b/examples/openai-sdk-agent/review_agent.py @@ -0,0 +1,146 @@ +"""OpenAI Agents SDK loop for the PR-review example.""" + +import asyncio +import os +import time +from pathlib import Path +from tempfile import TemporaryDirectory + +from agents import Agent, Runner, function_tool + +MODEL = "gpt-5.6-luna" +SHELL_TIMEOUT_SECONDS = 15 * 60 +SHELL_OUTPUT_LIMIT = 60_000 +SENSITIVE_ENVIRONMENT_MARKERS = ( + "CREDENTIAL", + "DATABRICKS", + "KEY", + "LAKEBASE", + "PASSWORD", + "SECRET", + "TOKEN", +) +CUJS = ( + ("quality", "discover and run the repository's formatting, lint, and type checks"), + ("package-tests", "run all PR-relevant package test suites and report each separately"), + ("build-install", "build distributions, install them cleanly, and verify imports"), + ("agent-e2e", "launch an affected example agent, invoke it, and retain useful logs"), +) + + +def _model() -> str: + configured = os.getenv("OPENAI_MODEL", MODEL) + if configured != MODEL: + raise RuntimeError(f"This cost-controlled example only allows {MODEL}; got {configured}") + return configured + + +def _shell_environment(workspace: Path) -> dict[str, str]: + environment = { + name: value + for name, value in os.environ.items() + if not any(marker in name.upper() for marker in SENSITIVE_ENVIRONMENT_MARKERS) + } + environment["HOME"] = str(workspace) + return environment + + +def _limit_output(output: str) -> str: + if len(output) <= SHELL_OUTPUT_LIMIT: + return output + half_limit = SHELL_OUTPUT_LIMIT // 2 + return f"{output[:half_limit]}\n... output truncated ...\n{output[-half_limit:]}" + + +def create_agent(workspace: Path) -> Agent: + """Create an SDK agent with a shell backed by the Databricks App process.""" + + @function_tool + async def run_shell(command: str) -> str: + """Run one shell command in the review workspace and return its status and output.""" + process = await asyncio.create_subprocess_shell( + command, + cwd=workspace, + env=_shell_environment(workspace), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + ) + try: + stdout, _ = await asyncio.wait_for(process.communicate(), timeout=SHELL_TIMEOUT_SECONDS) + except TimeoutError: + process.kill() + stdout, _ = await process.communicate() + output = stdout.decode(errors="replace") + return _limit_output( + f"Command timed out after {SHELL_TIMEOUT_SECONDS} seconds.\n{output}" + ) + except asyncio.CancelledError: + process.kill() + await process.communicate() + raise + output = stdout.decode(errors="replace") + return _limit_output(f"Exit status: {process.returncode}\n{output}") + + return Agent( + name="Sequential PR CUJ reviewer", + model=_model(), + instructions=( + "Review the supplied public GitHub PR by actually executing every CUJ in order. " + "Use run_shell for checkout, package installation, tests, and app execution. " + "Do not claim a check passed without command output. Keep changes inside the review " + "environment and finish with a concise evidence-based Markdown report." + ), + tools=[run_shell], + ) + + +def build_prompt(pr_url: str, iteration: int, workspace: Path, recovery_note: str = "") -> str: + steps = "\n".join(f"{index}. {name}: {goal}" for index, (name, goal) in enumerate(CUJS, 1)) + return f"""PR: {pr_url} +Iteration: {iteration} +Workspace: {workspace} + +First clone the repository into {workspace}/repo and check out the PR head. Each run_shell call +starts in {workspace}, so use an explicit `cd {workspace}/repo` for repository commands. Then +execute these CUJs sequentially in the same environment: +{steps} + +If a CUJ fails, investigate the cause and continue with the remaining CUJs. Include exact commands, +exit status, and important output in the final report. {recovery_note} +""" + + +async def execute_review( + pr_url: str, minimum_minutes: float, session, recovery_note: str = "" +) -> str: + """Run complete CUJ iterations sequentially until the minimum wall time is met.""" + started = time.monotonic() + iteration = 0 + reports: list[str] = [] + with TemporaryDirectory(prefix="openai-sdk-agent-") as temporary_directory: + workspace = Path(temporary_directory) + while iteration == 0 or time.monotonic() - started < minimum_minutes * 60: + iteration += 1 + result = await Runner.run( + create_agent(workspace), + build_prompt(pr_url, iteration, workspace, recovery_note if iteration == 1 else ""), + session=session, + max_turns=100, + ) + reports.append(f"## Iteration {iteration}\n\n{result.final_output}") + if minimum_minutes <= 0: + break + return "\n\n".join(reports) + + +async def resume_review(session, recovery_note: str) -> str: + """Resume using only the SDK session history and a fixed recovery note.""" + with TemporaryDirectory(prefix="openai-sdk-agent-") as temporary_directory: + workspace = Path(temporary_directory) + result = await Runner.run( + create_agent(workspace), + recovery_note, + session=session, + max_turns=100, + ) + return str(result.final_output) diff --git a/examples/openai-sdk-agent/sessions.py b/examples/openai-sdk-agent/sessions.py new file mode 100644 index 000000000..60c0ead02 --- /dev/null +++ b/examples/openai-sdk-agent/sessions.py @@ -0,0 +1,18 @@ +"""OpenAI Agents SDK session history stored in the App's Lakebase resource.""" + +import os + +from databricks_openai.agents import AsyncDatabricksSession + + +def create_session(session_id: str) -> AsyncDatabricksSession: + return AsyncDatabricksSession( + session_id=session_id, + autoscaling_endpoint=os.environ["LAKEBASE_AUTOSCALING_ENDPOINT"], + schema=os.getenv("LAKEBASE_SESSION_SCHEMA", "openai_sdk_agent_sessions"), + ) + + +async def initialize_sessions() -> None: + session = create_session("__startup__") + await session._ensure_tables() 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..0007f65ad --- /dev/null +++ b/src/databricks_ai_bridge/durable_runtime/README.md @@ -0,0 +1,123 @@ +# 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"]) + 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. + +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 one table: + +```text +executions + execution_id TEXT PRIMARY KEY + status TEXT + attempt INTEGER + heartbeat_at TIMESTAMPTZ + request JSONB + response 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..ce0202e48 --- /dev/null +++ b/src/databricks_ai_bridge/durable_runtime/__init__.py @@ -0,0 +1,47 @@ +"""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, + DurableExecution, + DurableExecutionContext, + DurableExecutionFailedError, + DurableExecutionNotFoundError, + DurableExecutionStatus, + DurableExecutor, + DurableRequestConflictError, + JsonObject, +) + +__all__ = [ + "DEFAULT_DURABILITY_SCHEMA", + "DatabricksDurableRuntime", + "DurabilityStore", + "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..81597ba65 --- /dev/null +++ b/src/databricks_ai_bridge/durable_runtime/runtime.py @@ -0,0 +1,309 @@ +"""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, + 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) + + 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: + response = await self.execute( + copy.deepcopy(claimed.request), + DurableExecutionContext( + execution_id=execution_id, + attempt=claimed.attempt, + ), + ) + 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..0e9aecd37 --- /dev/null +++ b/src/databricks_ai_bridge/durable_runtime/store.py @@ -0,0 +1,301 @@ +"""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 ( + 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" + + 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 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 + + @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..bbadce43b --- /dev/null +++ b/src/databricks_ai_bridge/durable_runtime/types.py @@ -0,0 +1,99 @@ +"""Public types for durable request execution.""" + +from dataclasses import dataclass +from datetime import datetime +from enum import Enum +from typing import Any, Protocol + +JsonObject = dict[str, Any] + + +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 DurableExecutionContext: + """Attempt metadata passed to the caller-owned executor.""" + + execution_id: str + attempt: int + + @property + def is_recovery(self) -> bool: + return self.attempt > 1 + + +class DurableExecutor(Protocol): + async def __call__( + self, + request: JsonObject, + context: DurableExecutionContext, + ) -> 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: ... + + +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_runtime.py b/tests/databricks_ai_bridge/test_durable_runtime.py new file mode 100644 index 000000000..4a3658de8 --- /dev/null +++ b/tests/databricks_ai_bridge/test_durable_runtime.py @@ -0,0 +1,405 @@ +"""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, + 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]] = [] + + 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 + + +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_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..6d9197fec --- /dev/null +++ b/tests/databricks_ai_bridge/test_durable_runtime_store.py @@ -0,0 +1,159 @@ +"""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 + 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_single_request_response_table(): + 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 "messages" not 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"}' + + +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("", {})