Here is an example code given by GLM-5.2, it is not perfect but will demonstrate what I want.
"""
A langgraph agent that waits for long *remote* jobs with a "partial pause",
keeps session memory across related TaskGroups, AND updates state
**incrementally per job** (not in one batch at the end).
Architecture
============
* Agent (FastAPI, port 8000)
GET / -> the web frontend (index.html)
POST /api/create -> create a new TaskGroup (optionally in an
existing session), kick off the graph
GET /api/stream/{tg_id} -> SSE stream of progress messages for one group
POST /cb/{job_id} -> callback the Job Manager POSTs results back to
* Job Manager (FastAPI, port 8001) -- a *separate* service, as the readme asks
POST /submit -> {job_id, rand_num, callback_url};
sleeps rand_num seconds, computes rand_num+1,
then POSTs {job_id, result} back to callback_url
Two-level memory (the readme's "maintain the memory")
=====================================================
Different thread_id per TaskGroup => partial pause (parked runs never block
other runs). A shared InMemoryStore keyed by session_id => related TaskGroups
(same session, different thread_ids) read each other's completed results.
Why per-job interrupt (Send fan-out) instead of one barrier interrupt
======================================================================
An earlier version parked the WHOLE TaskGroup on a single interrupt() and only
wrote state back when EVERY job was done, stashing partial results in an
in-memory JobRegistry meanwhile. Two things break that in a real program:
1. Partial results were NOT in the checkpoint -- they lived in a volatile
side-channel dict. An agent crash mid-wait lost them all (the raw files
survive on the FS, but the agent's "which jobs are done / their parsed
results" view was gone).
2. The barrier blocked per-job downstream: the readme wants "when step 2 of a
reaction is done, that reaction goes to step 3" -- impossible if nothing
proceeds until ALL step-2 jobs finish.
The LangChain team's fix (Forum_Disscussion.txt) is Send fan-out with ONE
interrupt() per job:
START -> node_plan --(Send per long job)--> node_wait_job --interrupt()
| (each branch parks
| on its OWN interrupt)
resume only job X's interrupt: v
Command(resume={interrupt_id: value}) node_collect (fan-in) -> END
When job X's /cb arrives we read the thread state, find X's pending interrupt
id, and resume ONLY that branch with Command(resume={id: value}). X's branch
wakes immediately, does X's post-processing, and its result is CHECKPOINTED
right then (durable). Other branches stay parked; staggered completion is
fine. When the last branch finishes, the fan-in node_collect runs and
publishes the group's final_data to the session store.
So: each job's result is durable the moment it lands; per-job downstream can
start immediately; the "partial pause" is real *within* a plan, not just
across plans. (For the demo's plain "sum at end" task this is more machinery
than strictly needed, but it is the shape the real DFT program requires.)
"""
from __future__ import annotations
import asyncio
import json
import operator
import random
import uuid
from typing import Annotated, List, Literal, Optional, TypedDict
import httpx
import uvicorn
from fastapi import FastAPI
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
from pydantic import BaseModel, Field
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import END, START, StateGraph
from langgraph.store.memory import InMemoryStore
from langgraph.types import Command, Send, interrupt
# --------------------------------------------------------------------------- #
# Constants #
# --------------------------------------------------------------------------- #
AGENT_HOST = "127.0.0.1"
AGENT_PORT = 8000
JOB_MANAGER_HOST = "127.0.0.1"
JOB_MANAGER_PORT = 8001
AGENT_URL = f"http://{AGENT_HOST}:{AGENT_PORT}"
JOB_MANAGER_URL = f"http://{JOB_MANAGER_HOST}:{JOB_MANAGER_PORT}"
HERE = __file__.rsplit("/", 1)[0]
STORE_KEY = "completed_groups"
# --------------------------------------------------------------------------- #
# Data models (Task/TaskGroup stay pydantic; graph state is a TypedDict so #
# that Send() branch args -- plain dicts -- flow into nodes naturally.) #
# --------------------------------------------------------------------------- #
class Task(BaseModel):
id: str
type: Literal["a", "b", "c"] = "a"
rand_num: float
job_result: float = -1.0
class TaskGroup(BaseModel):
id: str
final_data: float = -1.0
tasks: List[Task] = []
class MyState(TypedDict):
session_id: str
task_group: TaskGroup
short_results: List[dict] # node_plan only
job_results: Annotated[List[dict], operator.add] # one entry per branch
log: Annotated[List[str], operator.add]
# --------------------------------------------------------------------------- #
# Helpers #
# --------------------------------------------------------------------------- #
def is_long_task(t: Task) -> bool:
if t.type == "a":
return True
if t.type == "b":
return t.rand_num > 0
return False
def direct_result(t: Task) -> float:
return t.rand_num + 1
def expected_final_data(tg: TaskGroup) -> float:
return sum(t.rand_num + 1 for t in tg.tasks)
def get_taskgroup(task_group_id: str, ntask: int) -> TaskGroup:
ntask = ntask if ntask <= 5 else 5
task_type = ["a", "b", "c"]
tasks = []
for _ in range(ntask):
j = int(random.random() * 3)
if task_type[j] == "a":
r = random.random() + 1
elif task_type[j] == "b":
r = random.random() * 0.3 - 0.15
else:
r = random.random()
tasks.append(Task(id=str(uuid.uuid4()), type=task_type[j], rand_num=r))
return TaskGroup(id=task_group_id, tasks=tasks)
def _short(jid: str) -> str:
return jid[:8]
def _thread_id(session_id: str, tg_id: str) -> str:
return f"{session_id}/{tg_id}"
# --------------------------------------------------------------------------- #
# JobIndex: job_id -> (tg_id, session_id) so /cb can rebuild the thread_id. #
# Plus an `early_cache` for the race "a result arrives before its branch #
# reaches interrupt()" -- the branch pops it and skips the park entirely. #
# `resumed` guards against duplicate callbacks re-resolving an interrupt. #
# All sync => atomic w.r.t. the single asyncio event loop. #
# --------------------------------------------------------------------------- #
class JobIndex:
def __init__(self) -> None:
self._tg: dict[str, str] = {}
self._session: dict[str, str] = {}
self._early: dict[str, float] = {}
self._resumed: set[str] = set()
def register(self, job_id: str, tg_id: str, session_id: str) -> None:
self._tg[job_id] = tg_id
self._session[job_id] = session_id
def tg_of(self, job_id: str) -> Optional[str]:
return self._tg.get(job_id)
def session_of(self, job_id: str) -> Optional[str]:
return self._session.get(job_id)
def stash_early(self, job_id: str, result: float) -> None:
self._early[job_id] = result
def pop_early(self, job_id: str) -> Optional[float]:
return self._early.pop(job_id, None)
def mark_resumed(self, job_id: str) -> bool:
"""True if this job was NOT already resumed (so the caller should
proceed); False if it's a duplicate."""
if job_id in self._resumed:
return False
self._resumed.add(job_id)
return True
# --------------------------------------------------------------------------- #
# MessageBus: tg_id -> asyncio.Queue[str] #
# --------------------------------------------------------------------------- #
class MessageBus:
def __init__(self) -> None:
self._queues: dict[str, asyncio.Queue] = {}
def _queue(self, tg_id: str) -> asyncio.Queue:
if tg_id not in self._queues:
self._queues[tg_id] = asyncio.Queue()
return self._queues[tg_id]
def put(self, tg_id: str, msg: str) -> None:
self._queue(tg_id).put_nowait(msg)
def close(self, tg_id: str) -> None:
self._queue(tg_id).put_nowait(None)
def queue(self, tg_id: str) -> asyncio.Queue:
return self._queue(tg_id)
job_index = JobIndex()
bus = MessageBus()
session_store = InMemoryStore()
client: httpx.AsyncClient = httpx.AsyncClient(timeout=30.0)
def _read_completed(store, session_id: str) -> list:
item = store.get(session_id, STORE_KEY)
return list(item.value) if item and item.value else []
def _append_completed(store, session_id: str, entry: dict) -> None:
prev = _read_completed(store, session_id)
prev.append(entry)
store.put(session_id, STORE_KEY, prev)
def _find_interrupt_id(graph, thread_id: str, job_id: str) -> Optional[str]:
"""Read the thread's checkpoint, return the id of the PENDING interrupt
whose value carries `job_id`, or None if the branch hasn't parked yet.
Called only for jobs not already resumed (see JobIndex.mark_resumed), so a
match here is by construction the live, pending interrupt for this job."""
st = graph.get_state({"configurable": {"thread_id": thread_id}})
for t in st.tasks:
for intr in getattr(t, "interrupts", []):
v = intr.value
if isinstance(v, dict) and v.get("job_id") == job_id:
return intr.id
return None
# --------------------------------------------------------------------------- #
# The langgraph agent #
# --------------------------------------------------------------------------- #
async def _submit_to_manager(job_id: str, rand_num: float) -> None:
await client.post(
f"{JOB_MANAGER_URL}/submit",
json={
"job_id": job_id,
"rand_num": rand_num,
"callback_url": f"{AGENT_URL}/cb/{job_id}",
},
)
async def node_plan(state: MyState, *, store) -> dict:
"""Read session memory, compute short jobs directly, submit long jobs to
the remote manager. Long jobs are NOT waited for here -- they fan out to
one node_wait_job branch each via the conditional edge below."""
tg = state["task_group"]
session_id = state["session_id"]
msgs: List[str] = []
prior = _read_completed(store, session_id)
if prior:
summary = ", ".join(f"{_short(e['tg_id'])}={e['final_data']:.3f}" for e in prior)
msgs.append(f"Session {_short(session_id)}: {len(prior)} prior completed "
f"group(s) [{summary}] -- memory loaded.")
else:
msgs.append(f"Session {_short(session_id)}: no prior groups (new session).")
long_tasks = [t for t in tg.tasks if is_long_task(t)]
short_results: List[dict] = []
new_tasks: List[Task] = []
for t in tg.tasks:
if is_long_task(t):
job_index.register(t.id, tg.id, session_id) # sync, before any submit
new_tasks.append(t.model_copy())
else:
res = direct_result(t)
short_results.append({"job_id": t.id, "result": res})
new_tasks.append(t.model_copy(update={"job_result": res}))
msgs.append(f"Job {_short(t.id)} (type {t.type}, rand={t.rand_num:.3f}) "
f"is short, result={res:.3f}")
n_long = len(long_tasks)
n_short = len(tg.tasks) - n_long
expected = expected_final_data(tg)
msgs.append(f"There are {len(tg.tasks)} jobs ({n_long} long, {n_short} short). "
f"The expected final_data is {expected:.3f}")
# Submit long jobs concurrently; registration already happened (sync).
if long_tasks:
await asyncio.gather(*[_submit_to_manager(t.id, t.rand_num) for t in long_tasks])
for t in long_tasks:
msgs.append(f"Send job {_short(t.id)} (type {t.type}, "
f"rand={t.rand_num:.3f}) to remote...")
for m in msgs:
bus.put(tg.id, m)
return {"task_group": tg.model_copy(update={"tasks": new_tasks}),
"short_results": short_results, "log": msgs}
def fan_out(state: MyState):
"""One Send() per long job -> one node_wait_job branch each, parked on its
own interrupt(). No long jobs -> straight to the fan-in node_collect."""
tg = state["task_group"]
longs = [t for t in tg.tasks if is_long_task(t)]
if not longs:
return "node_collect"
return [
Send("node_wait_job",
{"job_id": t.id, "session_id": state["session_id"], "tg_id": tg.id})
for t in longs
]
async def node_wait_job(state, *, store) -> dict:
"""One branch per long job. Parks on its OWN interrupt(); the /cb for this
job resumes only this branch with Command(resume={id: value}). In a real
program this is where per-job post-processing lives (parse OUTCAR, validate,
publish artifact refs to the session store) -- and it runs the INSTANT this
job finishes, regardless of the other (still-parked) jobs.
Race: if this job's result arrived before the branch reached interrupt(),
JobIndex.pop_early returns it and we skip the park entirely.
"""
jid = state["job_id"]
cached = job_index.pop_early(jid)
if cached is not None:
val = cached
else:
val = interrupt({"job_id": jid}) # parks here until /cb resumes
return {"job_results": [{"job_id": jid, "result": float(val)}]}
async def node_collect(state: MyState, *, store) -> dict:
"""Fan-in: runs only after EVERY node_wait_job branch has completed (the
parallel-branch barrier). Merge short + long results, compute final_data,
publish the group summary to the shared session store."""
tg = state["task_group"]
results = {r["job_id"]: r["result"]
for r in list(state.get("short_results", [])) +
list(state.get("job_results", []))}
new_tasks = []
for t in tg.tasks:
res = results.get(t.id, t.job_result)
new_tasks.append(t.model_copy(update={"job_result": res}))
final_data = sum(t.job_result for t in new_tasks)
msg = f"The calculated final_data is {final_data:.3f}"
bus.put(tg.id, msg)
_append_completed(store, state["session_id"], {
"tg_id": tg.id, "final_data": final_data,
"n_tasks": len(tg.tasks),
})
return {"task_group": tg.model_copy(
update={"tasks": new_tasks, "final_data": final_data}),
"log": [msg]}
builder = StateGraph(MyState)
builder.add_node("node_plan", node_plan)
builder.add_node("node_wait_job", node_wait_job)
builder.add_node("node_collect", node_collect)
builder.add_edge(START, "node_plan")
builder.add_conditional_edges("node_plan", fan_out)
builder.add_edge("node_wait_job", "node_collect")
builder.add_edge("node_collect", END)
app_graph = builder.compile(checkpointer=InMemorySaver(), store=session_store)
async def run_task_group(tg_id: str, session_id: str, ntask: int) -> None:
"""First leg: invoke until the graph parks at the per-job interrupts (or, if
there are no long jobs, runs straight to END). Each /cb later resumes one
branch; the last resume closes the SSE stream."""
tg = get_taskgroup(tg_id, ntask)
result = await app_graph.ainvoke(
{"session_id": session_id, "task_group": tg,
"short_results": [], "job_results": [], "log": []},
config={"configurable": {"thread_id": _thread_id(session_id, tg_id)}},
)
if "__interrupt__" not in result:
bus.close(tg_id) # all-short or all-early -> already at END
async def _resume_one(tg_id: str, thread_id: str, iid: str, result: float) -> None:
"""Resume a single job's branch off the /cb request path, then close the
SSE stream if the whole group has reached END."""
await app_graph.ainvoke(
Command(resume={iid: result}),
config={"configurable": {"thread_id": thread_id}},
)
st = app_graph.get_state({"configurable": {"thread_id": thread_id}})
if not st.next: # () == no pending branches == END
bus.close(tg_id)
# --------------------------------------------------------------------------- #
# Job Manager (separate FastAPI service on JOB_MANAGER_PORT) #
# --------------------------------------------------------------------------- #
class SubmitJob(BaseModel):
job_id: str
rand_num: float
callback_url: str
job_manager_app = FastAPI(title="Remote Job Manager")
@job_manager_app.post("/submit")
async def submit_job(payload: SubmitJob) -> dict:
asyncio.create_task(_run_job(payload.job_id, payload.rand_num, payload.callback_url))
return {"ok": True, "job_id": payload.job_id}
async def _run_job(job_id: str, rand_num: float, callback_url: str) -> None:
await asyncio.sleep(rand_num)
result = rand_num + 1
try:
await client.post(callback_url, json={"job_id": job_id, "result": result})
except Exception as exc: # pragma: no cover
print(f"[job-manager] post-back failed job={_short(job_id)}: {exc}")
# --------------------------------------------------------------------------- #
# Agent web app (FastAPI on AGENT_PORT) #
# --------------------------------------------------------------------------- #
class CreateReq(BaseModel):
ntask: int = Field(default=3, ge=1, le=5)
session_id: Optional[str] = None
agent_app = FastAPI(title="TaskGroup Agent")
@agent_app.get("/")
async def index() -> FileResponse:
return FileResponse(f"{HERE}/index.html")
@agent_app.post("/api/create")
async def create(req: CreateReq) -> JSONResponse:
tg_id = str(uuid.uuid4())
session_id = req.session_id or str(uuid.uuid4())
asyncio.create_task(run_task_group(tg_id, session_id, req.ntask))
return JSONResponse({"task_group_id": tg_id, "session_id": session_id})
@agent_app.post("/cb/{job_id}")
async def job_callback(job_id: str, payload: dict) -> JSONResponse:
result = payload.get("result")
if result is None:
return JSONResponse({"ok": False, "job_id": job_id}, status_code=400)
result = float(result)
tg_id = job_index.tg_of(job_id)
if tg_id is None:
return JSONResponse({"ok": True, "job_id": job_id, "duplicate": True})
session_id = job_index.session_of(job_id)
thread_id = _thread_id(session_id, tg_id)
if not job_index.mark_resumed(job_id):
# already resumed this job -> idempotent no-op (duplicate callback)
return JSONResponse({"ok": True, "job_id": job_id, "duplicate": True})
bus.put(tg_id, f"job {_short(job_id)} is done, result is {result:.3f}")
iid = _find_interrupt_id(app_graph, thread_id, job_id)
if iid is None:
# branch hasn't parked yet (fast job beat the graph to interrupt());
# stash so node_wait_job pops it and skips the park.
job_index.stash_early(job_id, result)
else:
# resume ONLY this job's branch, off the request path.
asyncio.create_task(_resume_one(tg_id, thread_id, iid, result))
return JSONResponse({"ok": True, "job_id": job_id})
@agent_app.get("/api/stream/{tg_id}")
async def stream(tg_id: str) -> StreamingResponse:
queue = bus.queue(tg_id)
async def event_gen():
while True:
msg = await queue.get()
if msg is None:
yield f"data: {json.dumps({'done': True})}\n\n"
return
yield f"data: {json.dumps({'msg': msg})}\n\n"
return StreamingResponse(event_gen(), media_type="text/event-stream")
# --------------------------------------------------------------------------- #
# Entry point #
# --------------------------------------------------------------------------- #
async def main() -> None:
agent_cfg = uvicorn.Config(
agent_app, host=AGENT_HOST, port=AGENT_PORT, log_level="warning"
)
jm_cfg = uvicorn.Config(
job_manager_app, host=JOB_MANAGER_HOST, port=JOB_MANAGER_PORT, log_level="warning"
)
agent_server = uvicorn.Server(agent_cfg)
jm_server = uvicorn.Server(jm_cfg)
print(f"Agent: http://{AGENT_HOST}:{AGENT_PORT}")
print(f"Job Manager: http://{JOB_MANAGER_HOST}:{JOB_MANAGER_PORT}")
print("Open the Agent URL in a browser and click 'Create TaskGroup'.")
await asyncio.gather(agent_server.serve(), jm_server.serve())
if __name__ == "__main__":
try:
asyncio.run(main())
except KeyboardInterrupt:
pass
.