diff --git a/test/custode/early_session_recovery_test.exs b/test/custode/early_session_recovery_test.exs index 7b4b595..a76a475 100644 --- a/test/custode/early_session_recovery_test.exs +++ b/test/custode/early_session_recovery_test.exs @@ -4,7 +4,16 @@ defmodule Custode.EarlySessionRecoveryTest do import Custode.TestHelpers import Ecto.Query, only: [from: 2] - alias Custode.{Agents, ConversationArcs, ExecutionFacts, Repo} + alias Custode.{ + AgentHandoff, + Agents, + ConversationArcs, + ExecutionFacts, + OperatorMessages, + Repo, + Routine + } + alias Custode.Operator.Actions setup do @@ -26,12 +35,15 @@ defmodule Custode.EarlySessionRecoveryTest do :ok end - for provider <- [:claude, :codex], failure <- [:timeout, :worker_killed] do + for provider <- [:claude, :codex], + failure <- [:timeout, :worker_killed], + ordering <- [:ordinary, :replay_first] do @provider provider @failure failure + @ordering ordering @tag provider: provider - test "#{provider} resumes its early handle after #{failure} and an agent restart" do + test "#{provider} resumes its early handle after #{failure} with #{ordering} admission" do provider = @provider workspace = tmp_workspace!() routine = routine_fixture!(workspace, %{provider: provider, cron: :manual}) @@ -40,8 +52,8 @@ defmodule Custode.EarlySessionRecoveryTest do native_id = uid("native") write_cli!(binary, provider, native_id, :block) - assert {:ok, :started} = Actions.message(routine.id, "first turn") - first = latest_job(routine.id) + assert {:ok, initial, :created} = Actions.message_with_receipt(routine.id, "first turn") + first = correlated_job(initial) timeout = if @failure == :timeout, do: 1_500, else: 30_000 first = prepare_job!(first, binary, timeout) worker = worker(provider) @@ -59,24 +71,49 @@ defmodule Custode.EarlySessionRecoveryTest do end) assert Task.yield(task, 0) == nil + [{old_pid, _state}] = Registry.lookup(registry(provider), routine.id) + monitor = Process.monitor(old_pid) stop_first_turn(task, @failure) mark_terminal!(first, "discarded") + assert Repo.reload!(first).state == "discarded" + refute Custode.ProviderJobs.active_turn?(routine.id) assert :ok = Agents.stop_agent(routine.id, provider) + assert_receive {:DOWN, ^monitor, :process, ^old_pid, _reason}, 1_000 + refute Process.alive?(old_pid) + assert {:ok, :offline} = Agents.await(routine.id, provider, :offline, 1_000) assert ConversationArcs.read_model(routine.id).current.provider_session_id == native_id write_cli!(binary, provider, native_id, :complete) - assert {:ok, delivery} = Actions.message(routine.id, "continue that conversation") - assert delivery in [:started, :queued] - second = latest_job(routine.id, first.id) + assert {:ok, receipt, :created} = continue(routine, @ordering) + assert receipt.delivery in ~w(started queued admitting delivered) + second = correlated_job(receipt) + assert [second.id] == later_job_ids(routine.id, first.id) + [{new_pid, _state}] = Registry.lookup(registry(provider), routine.id) + refute new_pid == old_pid + assert is_binary(first.meta["arc_id"]) + assert is_binary(first.meta["agent_generation"]) + assert is_binary(second.meta["agent_generation"]) assert second.id != first.id assert second.meta["arc_id"] == first.meta["arc_id"] - assert (second.args["resume"] || second.args["session_id"]) == native_id + assert second.args[resume_key(provider)] == native_id assert second.meta["agent_generation"] != first.meta["agent_generation"] second = prepare_job!(second, binary, 2_000) assert :ok = worker.perform(second) mark_terminal!(second, "completed") assert {:ok, :idle} = Agents.await(routine.id, :idle, 1_000) + eventually(fn -> + completed = OperatorMessages.get(receipt.message_id) + assert completed.status == "completed" + assert completed.provider_correlation_id == receipt.provider_correlation_id + assert completed.agent_generation == second.meta["agent_generation"] + assert completed.arc_id == first.meta["arc_id"] + expected_delivery = if receipt.delivery == "started", do: "started", else: "delivered" + assert completed.delivery == expected_delivery + assert correlated_job(receipt).id == second.id + assert [second.id] == later_job_ids(routine.id, first.id) + end) + arguments = binary |> Kernel.<>(".args") |> File.read!() |> String.split("\n") assert native_id in arguments assert resume_switch(provider) in arguments @@ -93,24 +130,84 @@ defmodule Custode.EarlySessionRecoveryTest do refute Process.alive?(task.pid) end - defp latest_job(agent_id, after_id \\ 0) do + defp continue(routine, :ordinary), + do: Actions.message_with_receipt(routine.id, "continue that conversation") + + defp continue(routine, :replay_first) do + eventually(fn -> assert AgentHandoff.status(routine.id) == :ready end) + config = Routine.agent_config(routine, ConversationArcs.seed_map(routine)) + coordinator = Process.whereis(AgentHandoff) + + on_exit(fn -> + if Process.alive?(coordinator), do: :sys.resume(coordinator) + end) + + :ok = :sys.suspend(coordinator) + + # Suspend admission, not execution. The real replay function admits the + # durable row before its original caller enters the coordinator callback. + # Always release the shared coordinator, including assertion failures. + task = + Task.async(fn -> Actions.message_with_receipt(routine.id, "continue that conversation") end) + + try do + eventually(fn -> + {:messages, messages} = Process.info(coordinator, :messages) + + assert Enum.any?(messages, fn + {:"$gen_call", _, {:admit, id, _fun, opts}} -> + id == routine.id and opts[:current_queue_head] == true + + _other -> + false + end) + end) + + [queued] = OperatorMessages.queued_for(routine.id) + assert {:ok, _pid} = Agents.start_agent(routine.id, routine.provider, config) + assert {:ok, :idle} = Agents.await(routine.id, routine.provider, :idle, 1_000) + assert {:accepted, message_id} = Actions.replay_next(routine.id) + assert message_id == queued.message_id + after + :sys.resume(coordinator) + end + + assert {:ok, receipt, :created} = Task.await(task, 5_000) + assert receipt.delivery == "delivered" + assert AgentHandoff.status(routine.id) == :ready + {:ok, receipt, :created} + end + + defp correlated_job(receipt) do eventually(fn -> - job = - Repo.one( - from(j in Oban.Job, - where: j.worker in ["ObanClaude.Agent.Job", "ObanCodex.Agent.Job"], - where: j.id > ^after_id, - where: fragment("json_extract(?, '$.agent_id')", j.meta) == ^agent_id, - order_by: [desc: j.id], - limit: 1 - ) - ) - - assert %Oban.Job{} = job + assert [job] = + Repo.all( + from(j in Oban.Job, + where: j.worker in ["ObanClaude.Agent.Job", "ObanCodex.Agent.Job"], + where: + fragment("json_extract(?, '$.agent_id')", j.meta) == ^receipt.target_agent_id, + where: + fragment("json_extract(?, '$.correlation_id')", j.meta) == + ^receipt.provider_correlation_id + ) + ) + job end) end + defp later_job_ids(agent_id, after_id) do + Repo.all( + from(j in Oban.Job, + where: j.worker in ["ObanClaude.Agent.Job", "ObanCodex.Agent.Job"], + where: j.id > ^after_id, + where: fragment("json_extract(?, '$.agent_id')", j.meta) == ^agent_id, + order_by: j.id, + select: j.id + ) + ) + end + defp prepare_job!(job, binary, timeout) do args = Map.merge(job.args, %{"binary" => binary, "timeout" => timeout}) job |> Ecto.Changeset.change(args: args, state: "executing", attempt: 1) |> Repo.update!() @@ -157,6 +254,10 @@ defmodule Custode.EarlySessionRecoveryTest do defp worker(:claude), do: ObanClaude.Agent.Job defp worker(:codex), do: ObanCodex.Agent.Job + defp registry(:claude), do: ObanClaude.Agent.Registry + defp registry(:codex), do: ObanCodex.Agent.Registry + defp resume_key(:claude), do: "resume" + defp resume_key(:codex), do: "session_id" defp resume_switch(:claude), do: "--resume" defp resume_switch(:codex), do: "resume" end