diff --git a/composer/io/context.py b/composer/io/context.py index e15162d0..eba24fae 100644 --- a/composer/io/context.py +++ b/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 diff --git a/composer/workflow/executor.py b/composer/workflow/executor.py index 5e1a920b..448f06a6 100644 --- a/composer/workflow/executor.py +++ b/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", ) diff --git a/tests/test_graph_resume.py b/tests/test_graph_resume.py new file mode 100644 index 00000000..4f53f25b --- /dev/null +++ b/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"]