Resume a codegen thread from its checkpoint instead of re-entering it

composer/io/context.py

@@ -43,6 +43,7 @@
 from composer.diagnostics.jsonl_sink import emit as _emit_jsonl
 
 from langgraph._internal._typing import StateLike
+from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.graph.state import CompiledStateGraph
 
 from langchain_core.runnables import RunnableConfig
@@ -323,6 +324,23 @@ async def run_to_completion[I: StateLike, S: StateLike, C: StateLike | None](
         retry=retry,
     )
 
+async def latest_checkpoint_of(
+    checkpointer: BaseCheckpointSaver, thread_id: str
+) -> str | None:
+    """The id of the most recent checkpoint on ``thread_id``, or ``None`` when
+    the thread has none yet.
+
+    Pass the result as ``run_to_completion``'s ``checkpoint_id`` to continue an
+    interrupted run instead of re-entering it: a run config that names a
+    checkpoint makes :func:`composer.io.graph_runner.run_graph` drop the input,
+    so the graph picks up the checkpoint's pending tasks rather than applying
+    the input afresh from ``START``. Re-applying it would re-run the entry node
+    over a message history that already has an initial prompt in it, which the
+    provider rejects (a second system message, no longer at position 0).
+    """
+    tup = await checkpointer.aget_tuple({"configurable": {"thread_id": thread_id}})
+    return None if tup is None else tup.checkpoint["id"]
+
 class RetryPolicy(ABC):
     """The transient-failure ("floor") retry contract: which exceptions are
     worth re-running from the last checkpoint, and how long to back off

composer/workflow/executor.py

@@ -40,7 +40,7 @@
 from composer.tools.search import cvl_manual_tools
 from composer.templates.loader import load_jinja_template
 from composer.io.protocol import CodeGenIOHandler, WorkflowPurpose
-from composer.io.context import with_handler, run_to_completion
+from composer.io.context import with_handler, run_to_completion, latest_checkpoint_of
 from composer.diagnostics.timing import set_current_task_id
 from composer.ui.codegen_events import CodeGenEventHandler
 from composer.core.state import AIComposerInput, AIComposerExtra
@@ -289,10 +289,18 @@ async def _run_codegen(
 
     thread_id = workflow_options.thread_id
 
+    # An explicitly named thread may already hold a checkpoint (an earlier run
+    # of this same command that crashed, or one the caller stopped). Continue
+    # from it unless the caller named a checkpoint of their own.
+    resume_from = workflow_options.checkpoint_id
     if thread_id is None:
         thread_id = "crypto_session_" + str(uuid.uuid1())
         await handler.log_workflow_thread(WorkflowPurpose.CODEGEN, thread_id)
         logger.info(f"Selected thread id: {thread_id}")
+    elif resume_from is None:
+        resume_from = await latest_checkpoint_of(checkpointer, thread_id)
+        if resume_from is not None:
+            logger.info(f"Resuming thread {thread_id} from checkpoint {resume_from}")
 
     mem_root = memory_namespace or thread_id
 
@@ -481,7 +489,7 @@ async def _run_codegen(
                     flow_input,
                     thread_id=thread_id,
                     context=work_context,
-                    checkpoint_id=workflow_options.checkpoint_id,
+                    checkpoint_id=resume_from,
                     recursion_limit=workflow_options.recursion_limit,
                     description="Code generation",
                 )

tests/test_graph_resume.py

@@ -0,0 +1,167 @@
+"""Re-running an interrupted run on its own thread must continue it, not restart it.
+
+Codegen takes ``--thread-id``, so the way to pick up a crashed run is to invoke
+the same command again against the same thread. That only works if the second
+invocation resumes the thread's last checkpoint: langgraph treats a non-``None``
+input as an update from ``START``, so passing the original input again re-enters
+the entry node on top of a message history that already holds an initial prompt.
+The provider then rejects the request outright ("multiple non-consecutive system
+messages"), which is how the bug surfaced in the wild rather than as a quiet
+duplicate.
+
+The graphs here are two-node pregel loops over an in-memory checkpointer, with
+an entry node that injects a system prompt the way graphcore's does. No LLM, no
+Postgres.
+"""
+import operator
+import uuid
+from typing import Annotated, Any, TypedDict
+
+import pytest
+
+from langgraph.checkpoint.memory import InMemorySaver
+from langgraph.graph import StateGraph, START, END
+from langgraph.graph.state import CompiledStateGraph
+
+from composer.io.context import (
+    latest_checkpoint_of,
+    run_to_completion,
+    with_handler,
+)
+from composer.io.event_handler import NullEventHandler
+
+pytestmark = pytest.mark.asyncio
+
+
+class ResumeState(TypedDict):
+    messages: Annotated[list[str], operator.add]
+
+
+class _NullIOHandler:
+    async def log_checkpoint_id(self, *, path: list[str], checkpoint_id: str) -> None:
+        pass
+
+    async def log_state_update(self, path: list[str], st: dict) -> None:
+        pass
+
+    async def log_start(self, *, path: list[str], description: str, tool_id: str | None) -> None:
+        pass
+
+    async def log_end(self, path: list[str]) -> None:
+        pass
+
+    async def human_interaction(self, ty: Any, debug_thunk: Any) -> str:
+        raise AssertionError("no HITL interaction expected in resume tests")
+
+
+class _WedgedOnce(Exception):
+    """The crash that ends the first invocation."""
+
+
+def _build_graph(
+    saver: InMemorySaver, node_runs: list[str], fail_on_first_work: bool = True
+) -> CompiledStateGraph[ResumeState, None, ResumeState, ResumeState]:
+    """entry -> work. ``entry`` injects the system prompt (as graphcore's entry
+    node does); ``work`` raises the first time it is reached."""
+    work_calls = 0
+
+    async def entry(state: ResumeState) -> dict:
+        node_runs.append("entry")
+        return {"messages": ["system: you author CVL"]}
+
+    async def work(state: ResumeState) -> dict:
+        nonlocal work_calls
+        work_calls += 1
+        node_runs.append("work")
+        if fail_on_first_work and work_calls == 1:
+            raise _WedgedOnce("wedged")
+        return {"messages": ["assistant: done"]}
+
+    builder = StateGraph(ResumeState)
+    builder.add_node("entry", entry)
+    builder.add_node("work", work)
+    builder.add_edge(START, "entry")
+    builder.add_edge("entry", "work")
+    builder.add_edge("work", END)
+    return builder.compile(checkpointer=saver)
+
+
+async def _run(
+    graph: CompiledStateGraph[ResumeState, None, ResumeState, ResumeState],
+    *,
+    thread_id: str,
+    checkpoint_id: str | None = None,
+) -> ResumeState:
+    async with with_handler(_NullIOHandler(), NullEventHandler()):
+        return await run_to_completion(
+            graph,
+            {"messages": ["user: implement the contract"]},
+            thread_id=thread_id,
+            context=None,
+            recursion_limit=25,
+            description="resume test",
+            checkpoint_id=checkpoint_id,
+        )
+
+
+async def test_a_fresh_thread_has_no_checkpoint_to_resume():
+    saver = InMemorySaver()
+    assert await latest_checkpoint_of(saver, uuid.uuid1().hex) is None
+
+
+async def test_rerunning_the_same_thread_resumes_instead_of_re_entering():
+    saver = InMemorySaver()
+    node_runs: list[str] = []
+    graph = _build_graph(saver, node_runs)
+    tid = uuid.uuid1().hex
+
+    with pytest.raises(_WedgedOnce):
+        await _run(graph, thread_id=tid)
+    assert node_runs == ["entry", "work"]
+
+    resume_from = await latest_checkpoint_of(saver, tid)
+    assert resume_from is not None, "the crashed attempt should have left a checkpoint"
+
+    final = await _run(graph, thread_id=tid, checkpoint_id=resume_from)
+
+    # The completed node stays completed; only the failed one re-runs.
+    assert node_runs == ["entry", "work", "work"]
+    assert final["messages"].count("system: you author CVL") == 1
+    assert final["messages"].count("user: implement the contract") == 1
+    assert final["messages"][-1] == "assistant: done"
+
+
+async def test_without_the_resume_point_the_entry_node_runs_twice():
+    """The regression this guards against: the same second invocation, minus the
+    resume point, re-enters ``entry`` and duplicates the initial prompt."""
+    saver = InMemorySaver()
+    node_runs: list[str] = []
+    graph = _build_graph(saver, node_runs)
+    tid = uuid.uuid1().hex
+
+    with pytest.raises(_WedgedOnce):
+        await _run(graph, thread_id=tid)
+
+    final = await _run(graph, thread_id=tid)
+
+    assert node_runs == ["entry", "work", "entry", "work"]
+    assert final["messages"].count("system: you author CVL") == 2
+
+
+async def test_resuming_a_finished_thread_is_a_no_op():
+    """A re-run of a thread that already completed returns its final state
+    rather than authoring a second time."""
+    saver = InMemorySaver()
+    node_runs: list[str] = []
+    graph = _build_graph(saver, node_runs, fail_on_first_work=False)
+    tid = uuid.uuid1().hex
+
+    first = await _run(graph, thread_id=tid)
+    assert node_runs == ["entry", "work"]
+
+    resume_from = await latest_checkpoint_of(saver, tid)
+    assert resume_from is not None
+    again = await _run(graph, thread_id=tid, checkpoint_id=resume_from)
+
+    assert node_runs == ["entry", "work"]
+    assert again["messages"] == first["messages"]