Templates: Durable prose-recovery + always-rotate (companion to ai-bridge#425)

agent-langgraph-advanced/agent_server/agent.py

@@ -40,7 +40,7 @@
 logging.getLogger("mlflow.utils.autologging_utils").setLevel(logging.ERROR)
 sp_workspace_client = WorkspaceClient()
 
-LLM_ENDPOINT_NAME = "databricks-claude-sonnet-4-5"
+LLM_ENDPOINT_NAME = "databricks-gpt-5-2"
 LAKEBASE_CONFIG = init_lakebase_config()
 
 
@@ -123,6 +123,20 @@ async def stream_handler(
             # For on-behalf-of user authentication, pass get_user_workspace_client() to init_agent.
             agent = await init_agent(store=store, checkpointer=checkpointer)
 
+            # When the checkpointer already has prior turns for this thread,
+            # the chat client's full-history echo is redundant — `add_messages`
+            # would append duplicates (it dedupes by `id`, but MLflow's
+            # `responses_to_cc` doesn't preserve IDs, so dedup never fires).
+            # Forward only the latest user message; the checkpointer prepends
+            # the rest.
+            state = await agent.aget_state(config)
+            if state and state.values.get("messages") and input_state["messages"]:
+                last_user = next(
+                    (m for m in reversed(input_state["messages"]) if m.get("role") == "user"),
+                    None,
+                )
+                input_state["messages"] = [last_user] if last_user else []
+
             async for event in process_agent_astream_events(
                 agent.astream(input_state, config, stream_mode=["updates", "messages"])
             ):

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-prose-recovery" }
+databricks-langchain = { git = "https://github.com/databricks/databricks-ai-bridge", branch = "dhruv0811/durable-execution-prose-recovery", 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/agent_server/utils.py

@@ -29,6 +29,7 @@ class LakebaseConfig:
     autoscaling_endpoint: Optional[str]
     autoscaling_project: Optional[str]
     autoscaling_branch: Optional[str]
+    memory_schema: Optional[str] = None
 
     @property
     def description(self) -> str:
@@ -103,13 +104,15 @@ def init_lakebase_config() -> LakebaseConfig:
             "  Option 3 (provisioned): LAKEBASE_INSTANCE_NAME=<your-instance-name>\n"
         )
 
+    memory_schema = os.getenv("LAKEBASE_AGENT_MEMORY_SCHEMA") or None
+
     # Priority: endpoint > project+branch > instance_name (mutually exclusive in the library)
     if endpoint:
-        return LakebaseConfig(instance_name=None, autoscaling_endpoint=endpoint, autoscaling_project=None, autoscaling_branch=None)
+        return LakebaseConfig(instance_name=None, autoscaling_endpoint=endpoint, autoscaling_project=None, autoscaling_branch=None, memory_schema=memory_schema)
     elif has_autoscaling:
-        return LakebaseConfig(instance_name=None, autoscaling_endpoint=None, autoscaling_project=project, autoscaling_branch=branch)
+        return LakebaseConfig(instance_name=None, autoscaling_endpoint=None, autoscaling_project=project, autoscaling_branch=branch, memory_schema=memory_schema)
     else:
-        return LakebaseConfig(instance_name=resolve_lakebase_instance_name(raw_name), autoscaling_endpoint=None, autoscaling_project=None, autoscaling_branch=None)
+        return LakebaseConfig(instance_name=resolve_lakebase_instance_name(raw_name), autoscaling_endpoint=None, autoscaling_project=None, autoscaling_branch=None, memory_schema=memory_schema)
 
 
 def get_lakebase_access_error_message(lakebase_description: str) -> str:
@@ -187,9 +190,12 @@ async def deduplicate_input(request: ResponsesAgentRequest, session: AsyncDatabr
     """Return the input messages to pass to the Runner, avoiding duplication with session history.
 
     When a client sends the full conversation history AND the session already has
-    that history persisted, passing everything through would duplicate messages.
-    If the session already covers the prior turns, only the latest message is needed
-    since the session will prepend the full history automatically.
+    that history persisted, passing everything through would duplicate messages
+    in the LLM call (Runner combines session items + input items, and the OpenAI
+    SDK's `_dedupe_key` doesn't dedupe role-bearing items — see
+    `agents/run_internal/items.py:224-250`). If the session already has any
+    items, the prior turns are persisted there and we only need to forward the
+    latest user message.
     """
     messages = [i.model_dump() for i in request.input]
     # Normalize assistant message content from string to structured list format.
@@ -204,7 +210,10 @@ async def deduplicate_input(request: ResponsesAgentRequest, session: AsyncDatabr
         ):
             msg["content"] = [{"type": "output_text", "text": msg["content"], "annotations": []}]
     session_items = await session.get_items()
-    if len(session_items) >= len(messages) - 1:
+    # Trust the session as authoritative for prior turns. Forward only the
+    # latest message (the new user turn). The Runner will prepend session
+    # history on the LLM call automatically.
+    if session_items and len(messages) > 1:
         return [messages[-1]]
     return messages
 

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-prose-recovery" }
+databricks-openai = { git = "https://github.com/databricks/databricks-ai-bridge", branch = "dhruv0811/durable-execution-prose-recovery", subdirectory = "integrations/openai" }
+
 [tool.pytest.ini_options]
 base_url = "http://localhost:8000"
 addopts = "-n 7"