Durable Execution과 Workflow 복구

import os
import time
import psycopg
from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.postgres import PostgresSaver

# DB 커넥션 스트링 정의 (PostgreSQL 16+ 기준)
DB_URI = "postgresql://postgres:password@localhost:5432/agent_store?sslmode=disable"

# 1. 에이전트가 공유할 스키마 상태 정의
class DurableState(TypedDict):
    task: str
    tool_status: str
    idempotency_key: str
    retry_count: int

# 2. 가상의 안전 가드 레이어가 내장된 비장애/장애 유발 Tool 함수
def execute_payment_api(idempotency_key: str, amount: int) -> bool:
    """
    실제 프로덕션 환경에서는 영속 저장소(Redis/DB)에 idempotency_key가 
    존재하는지 선제 체크 쿼리를 날리는 방어선이 구축되어야 합니다.
    """
    print(f"[Tool Server] 멱등키 [{idempotency_key}] 검증 중...")
    # 가상의 네트워크 처리 실행
    time.sleep(0.5)
    return True

# 3. 그래프 노드 정의
def payment_node(state: DurableState):
    current_key = state["idempotency_key"]
    
    print(f"\n[Worker] >>> payment_node 실행 시작 (멱등키: {current_key})")
    
    # 툴 호출 가드 레이어 통과
    success = execute_payment_api(current_key, 50000)
    
    # [의도적 장애 주입 시나리오]
    # 최초 실행 시, 툴 호출은 성공했으나 상태를 저장하기 직전에 Worker 프로세스가 강제 종료(OOM)된 상황을 시뮬레이션
    if state["retry_count"] == 0:
        print("[CRITICAL ERROR] 상태를 DB 체크포인트에 쓰기 직전 Worker 프로세스가 사망했습니다 (OOM) !!!")
        # 실제 상용 환경에서는 프로세스가 킬당하는 상황이므로 os._exit()로 프로세스를 강제 중단
        os._exit(1)
        
    return {
        "tool_status": "PAYMENT_SUCCESS",
        "retry_count": state["retry_count"] + 1
    }

# 4. 엔지니어링 파이프라인 컴파일 및 인프라 영속 바인딩
def build_durable_agent():
    # PostgreSQL 연결 및 체크포인트 테이블 초기화
    # LangGraph가 내부적으로 필요 테이블을 가동 시 자동 인덱싱 및 생성합니다.
    conn = psycopg.connect(DB_URI, autocommit=True)
    checkpointer = PostgresSaver(conn)
    # 메모리 압축이 필요한 대규모 클러스터인 경우 checkpointer.setup() 수행
    
    # 그래프 빌드
    workflow = StateGraph(DurableState)
    workflow.add_node("payment_processor", payment_node)
    
    workflow.set_entry_point("payment_processor")
    workflow.add_edge("payment_processor", END)
    
    # checkpointer 객체를 주입하며 컴파일하여 Durable Runtime 레이어 구축
    app = workflow.compile(checkpointer=checkpointer)
    return app

# 외부 트리거 가동 메인 제어 루프
if __name__ == "__main__":
    # 주의: 이 스크립트는 최초 가동 시 1번 인덱스에서 의도적으로 다운됩니다.
    # 이후 다른 분산 세션 워커가 동일 thread_id로 상태를 로드하면 이어서 복구 가동됩니다.
    agent_app = build_durable_agent()
    
    # 특정 단일 사용자 유저 스레드 세션 지정
    config = {"configurable": {"thread_id": "session_user_99a8"}}
    
    initial_state = {
        "task": "50000원 대금 결제 요청",
        "tool_status": "INIT",
        "idempotency_key": "tx_req_uuid_00192f", # 고정된 멱등 키 주입
        "retry_count": 0
    }
    
    try:
        print("첫 번째 워커 인스턴스가 에이전트를 구동합니다.")
        agent_app.invoke(initial_state, config)
    except SystemExit:
        print("\n[인프라 모니터링] 워커 1호기 사망 감지. 즉시 새 컨테이너 스케줄링 및 복구를 시도합니다...\n")
        
        # 새롭게 띄워진 2호기 워커가 동일한 thread_id 컨텍스트로 이어서 진입하는 시나리오
        recovered_state = {
            "retry_count": 1 # 오답 노트 전이 변수 명시
        }
        
        print("두 번째 새 가용 워커가 바통을 이어받아 동일 세션을 복구 재개합니다.")
        # 내부적으로 PostgresSaver가 기존 세션의 'payment_processor' 진입 전 스냅샷을 자동 인출
        final_result = agent_app.invoke(recovered_state, config)
        print(f" 최종 복구 실행 완료 결과: {final_result}")