Durable execution for agent-openai-advanced + agent-langgraph-advanced

agent-langgraph-advanced/agent_server/start_server.py

@@ -41,6 +41,16 @@ def transform_stream_event(self, event, response_id):
     poll_interval_seconds=float(os.getenv("POLL_INTERVAL_SECONDS", "1.0")),
 )
 
+log_level = os.getenv("LOG_LEVEL", "INFO")
+_lvl = getattr(logging, log_level.upper(), logging.INFO)
+logging.getLogger("agent_server").setLevel(_lvl)
+# Surface [durable] lifecycle logs from LongRunningAgentServer into apps logs.
+logging.getLogger("databricks_ai_bridge").setLevel(_lvl)
+if not logging.getLogger().handlers:
+    logging.basicConfig(level=_lvl, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
+else:
+    logging.getLogger().setLevel(_lvl)
+
 # Define the app as a module level variable to enable multiple workers
 app = agent_server.app  # noqa: F841
 setup_mlflow_git_based_version_tracking()

agent-langgraph-advanced/pyproject.toml

@@ -39,6 +39,12 @@ setup = [
 [tool.uv]
 default-groups = ["dev", "setup"]
 
+# TEMPORARY: point at the open PR branch while ML-64230 durable-resume
+# changes are in review. Revert to the registry release once merged.
+[tool.uv.sources]
+databricks-ai-bridge = { git = "https://github.com/databricks/databricks-ai-bridge", branch = "dhruv0811/durable-execution-resume" }
+databricks-langchain = { git = "https://github.com/databricks/databricks-ai-bridge", branch = "dhruv0811/durable-execution-resume", subdirectory = "integrations/langchain" }
+
 
 [tool.pytest.ini_options]
 base_url = "http://localhost:8000"

agent-openai-advanced/agent_server/start_server.py

@@ -56,7 +56,18 @@ def transform_stream_event(self, event, response_id):
 )
 
 log_level = os.getenv("LOG_LEVEL", "INFO")
-logging.getLogger("agent_server").setLevel(getattr(logging, log_level.upper(), logging.INFO))
+_lvl = getattr(logging, log_level.upper(), logging.INFO)
+logging.getLogger("agent_server").setLevel(_lvl)
+# Surface [durable] lifecycle logs from LongRunningAgentServer into apps logs.
+# These are INFO-level in databricks_ai_bridge but the library logger defaults
+# to WARNING unless the host process sets it explicitly.
+logging.getLogger("databricks_ai_bridge").setLevel(_lvl)
+# Ensure the root handler actually emits at this level too. uvicorn sets up
+# its own handlers for 'uvicorn.*' but leaves root untouched.
+if not logging.getLogger().handlers:
+    logging.basicConfig(level=_lvl, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
+else:
+    logging.getLogger().setLevel(_lvl)
 
 # Wrap the existing lifespan to ensure session tables are created before serving requests
 _original_lifespan = agent_server.app.router.lifespan_context

agent-openai-advanced/pyproject.toml

@@ -43,6 +43,12 @@ setup = [
 [tool.uv]
 default-groups = ["dev", "setup"]
 
+# TEMPORARY: point at the open PR branch while ML-64230 durable-resume
+# changes are in review. Revert to the registry release once merged.
+[tool.uv.sources]
+databricks-ai-bridge = { git = "https://github.com/databricks/databricks-ai-bridge", branch = "dhruv0811/durable-execution-resume" }
+databricks-openai = { git = "https://github.com/databricks/databricks-ai-bridge", branch = "dhruv0811/durable-execution-resume", subdirectory = "integrations/openai" }
+
 [tool.pytest.ini_options]
 base_url = "http://localhost:8000"
 addopts = "-n 7"

e2e-chatbot-app-next/packages/ai-sdk-providers/src/providers-server.ts

@@ -10,7 +10,10 @@ import {
 } from '@chat-template/auth';
 import { createDatabricksProvider } from '@databricks/ai-sdk-provider';
 import { extractReasoningMiddleware, wrapLanguageModel } from 'ai';
-import { shouldInjectContextForEndpoint } from './request-context';
+import {
+  getApiProxyUrl,
+  shouldInjectContextForEndpoint,
+} from './request-context';
 
 // Header keys for passing context through streamText headers
 export const CONTEXT_HEADER_CONVERSATION_ID = 'x-databricks-conversation-id';
@@ -71,7 +74,7 @@ export async function getWorkspaceHostname(): Promise<string> {
 // Environment variable to enable SSE logging
 const LOG_SSE_EVENTS = process.env.LOG_SSE_EVENTS === 'true';
 
-const API_PROXY = process.env.API_PROXY;
+const API_PROXY = getApiProxyUrl();
 
 // Cache for endpoint details to check task type and OBO scopes
 const endpointDetailsCache = new Map<

e2e-chatbot-app-next/packages/ai-sdk-providers/src/request-context.ts

@@ -2,11 +2,28 @@
  * Utility functions for request context handling.
  */
 
+/**
+ * Resolve the URL the AI SDK provider should POST to, or ``undefined`` when
+ * the chatbot should talk directly to a Databricks serving endpoint.
+ *
+ * Resolution order:
+ *   1. Explicit ``API_PROXY`` env var — caller knows best.
+ *   2. ``DATABRICKS_SERVING_ENDPOINT`` set → direct-endpoint mode; no proxy.
+ *   3. Advanced-template default → route via this Node server's own
+ *      ``/invocations`` proxy.
+ */
+export function getApiProxyUrl(): string | undefined {
+  if (process.env.API_PROXY) return process.env.API_PROXY;
+  if (process.env.DATABRICKS_SERVING_ENDPOINT) return undefined;
+  const port = process.env.CHAT_APP_PORT || process.env.PORT || '3000';
+  return `http://localhost:${port}/invocations`;
+}
+
 /**
  * Determines whether context should be injected based on endpoint type.
  *
  * Context is injected when:
- * 1. Using API_PROXY environment variable, OR
+ * 1. The Express /invocations proxy is in play (explicit or inferred), OR
  * 2. Endpoint task type is 'agent/v2/chat' or 'agent/v1/responses'
  *
  * @param endpointTask - The task type of the serving endpoint (optional)
@@ -15,9 +32,7 @@
 export function shouldInjectContextForEndpoint(
   endpointTask: string | undefined,
 ): boolean {
-  const API_PROXY = process.env.API_PROXY;
-
-  if (API_PROXY) {
+  if (getApiProxyUrl()) {
     return true;
   }