From 571ae057ad4027f0308d4eb51db5fd0cc462ffd4 Mon Sep 17 00:00:00 2001
From: Antigravity <bot@antigravity.ai>
Date: Fri, 29 May 2026 14:05:19 -0400
Subject: [PATCH] fix: lazy-import AsyncPostgresSaver to fix ops-worker/gumbo
 startup crash (ModuleNotFoundError langgraph.checkpoint.postgres)

---
 agents/gumbo/gumbo/graph.py         | 6 +++++-
 agents/ops_agent/ops_agent/graph.py | 6 +++++-
 agents/ops_agent/requirements.txt   | 1 +
 3 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/agents/gumbo/gumbo/graph.py b/agents/gumbo/gumbo/graph.py
index b61d8c2..affc551 100644
--- a/agents/gumbo/gumbo/graph.py
+++ b/agents/gumbo/gumbo/graph.py
@@ -7,8 +7,8 @@ from typing import TypedDict
 
 from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_openai import ChatOpenAI
-from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
 from langgraph.graph import END, StateGraph
+# AsyncPostgresSaver imported lazily in run_gumbo() to avoid startup crash
 
 from gumbo.mcp_fs import fetch_text_via_mcp
 
@@ -53,6 +53,10 @@ def build_graph_builder() -> StateGraph:
 
 
 async def run_gumbo(object_key: str, thread_id: str, conn_string: str) -> GumboState:
+    try:
+        from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+    except ImportError:
+        from langgraph_checkpoint_postgres.aio import AsyncPostgresSaver
     builder = build_graph_builder()
     async with AsyncPostgresSaver.from_conn_string(conn_string) as checkpointer:
         await checkpointer.setup()
diff --git a/agents/ops_agent/ops_agent/graph.py b/agents/ops_agent/ops_agent/graph.py
index 8e41815..6935764 100644
--- a/agents/ops_agent/ops_agent/graph.py
+++ b/agents/ops_agent/ops_agent/graph.py
@@ -4,7 +4,7 @@ from typing import TypedDict
 from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_openai import ChatOpenAI
 from langgraph.graph import StateGraph, END
-from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+# AsyncPostgresSaver imported lazily below (only when conn_string is provided)
 
 from ops_agent.tools import (
     k8s_get_status, k8s_get_logs, k8s_restart_pod, k8s_scale_deployment,
@@ -214,6 +214,10 @@ async def run_ops_graph(target_type: str, target_id: str, is_approved: bool = Fa
     }
     
     if conn_string:
+        try:
+            from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+        except ImportError:
+            from langgraph_checkpoint_postgres.aio import AsyncPostgresSaver
         async with AsyncPostgresSaver.from_conn_string(conn_string) as checkpointer:
             await checkpointer.setup()
             graph = builder.compile(checkpointer=checkpointer)
diff --git a/agents/ops_agent/requirements.txt b/agents/ops_agent/requirements.txt
index 145a3a1..414363f 100644
--- a/agents/ops_agent/requirements.txt
+++ b/agents/ops_agent/requirements.txt
@@ -1,4 +1,5 @@
 langgraph>=0.2.0
+langgraph-checkpoint-postgres>=2.0.0
 langchain-core>=0.3.0
 langchain-openai>=0.2.0
 temporalio>=1.8.0