fix: keep fresh stream sessions routable

Soft-pin fresh proxy stream Codex sessions so circuit-filtered pinned assignments can fall back until a strong continuity anchor exists. Also recover stale in-progress Codex turns with no active owner lease.

Author: masterkain

lib/codex_pooler/accounting/request_lifecycle/recovery.ex

@@ -5,7 +5,7 @@ defmodule CodexPooler.Accounting.RequestLifecycle.Recovery do
 
   alias CodexPooler.Accounting.{Attempt, LedgerEntry, Request}
   alias CodexPooler.Accounting.RequestLifecycle
-  alias CodexPooler.Gateway.Persistence.SessionReadModel
+  alias CodexPooler.Gateway.Persistence.{CodexTurn, SessionReadModel}
   alias CodexPooler.Repo
 
   @stale_after_seconds 6 * 60 * 60
@@ -77,25 +77,54 @@ defmodule CodexPooler.Accounting.RequestLifecycle.Recovery do
   end
 
   defp release_undispatched_request(%Request{} = request, now) do
-    RequestLifecycle.finalize_reserved_request_failure(request, %{
-      request_status: "failed",
-      response_status_code: 499,
-      last_error_code: @recovery_code,
-      usage_status: "not_applicable",
-      now: now
-    })
+    with {:ok, result} <-
+           RequestLifecycle.finalize_reserved_request_failure(request, %{
+             request_status: "failed",
+             response_status_code: 499,
+             last_error_code: @recovery_code,
+             usage_status: "not_applicable",
+             now: now
+           }) do
+      recover_stale_turn(request, nil, now)
+      {:ok, result}
+    end
   end
 
   defp settle_dispatched_request(%Request{} = request, %Attempt{} = attempt, now) do
-    RequestLifecycle.finalize_request(request, attempt, %{
-      request_status: "failed",
-      attempt_status: "failed",
-      response_status_code: 499,
-      last_error_code: @recovery_code,
-      error_message: "stale reservation recovered after request lifecycle was abandoned",
-      usage: %{status: "usage_unknown", source: @recovery_source},
-      now: now
-    })
+    with {:ok, result} <-
+           RequestLifecycle.finalize_request(request, attempt, %{
+             request_status: "failed",
+             attempt_status: "failed",
+             response_status_code: 499,
+             last_error_code: @recovery_code,
+             error_message: "stale reservation recovered after request lifecycle was abandoned",
+             usage: %{status: "usage_unknown", source: @recovery_source},
+             now: now
+           }) do
+      recover_stale_turn(request, attempt, now)
+      {:ok, result}
+    end
+  end
+
+  defp recover_stale_turn(%Request{id: request_id}, attempt, now) do
+    final_attempt_id = attempt && attempt.id
+
+    CodexTurn
+    |> where(
+      [turn],
+      turn.request_id == ^request_id and turn.status == ^CodexTurn.in_progress_status()
+    )
+    |> Repo.update_all(
+      set: [
+        status: CodexTurn.interrupted_status(),
+        error_code: @recovery_code,
+        final_attempt_id: final_attempt_id,
+        completed_at: now,
+        updated_at: now
+      ]
+    )
+
+    :ok
   end
 
   defp initial_summary do

lib/codex_pooler/gateway/persistence/session_read_model.ex

@@ -111,9 +111,13 @@ defmodule CodexPooler.Gateway.Persistence.SessionReadModel do
       from turn in CodexTurn,
         join: session in CodexSession,
         on: session.id == turn.codex_session_id,
+        left_join: lease in BridgeOwnerLease,
+        on:
+          lease.codex_session_id == session.id and
+            lease.status == ^@owner_lease_active and lease.expires_at > ^now,
         where:
           turn.request_id == ^request_id and turn.status == ^CodexTurn.in_progress_status() and
-            (is_nil(session.owner_lease_expires_at) or session.owner_lease_expires_at > ^now)
+            (session.owner_lease_expires_at > ^now or not is_nil(lease.id))
     )
   end
 

lib/codex_pooler/gateway/routing/session_continuity.ex

@@ -4,14 +4,16 @@ defmodule CodexPooler.Gateway.Routing.SessionContinuity do
   import Ecto.Query
 
   alias CodexPooler.Accounting.Request
+  alias CodexPooler.Catalog.Model
   alias CodexPooler.Files
   alias CodexPooler.Gateway.Contracts
   alias CodexPooler.Gateway.Payloads.ContinuityPayload
   alias CodexPooler.Gateway.Payloads.RequestOptions
-  alias CodexPooler.Gateway.Persistence.{BridgeSessionAlias, CodexSession}
+  alias CodexPooler.Gateway.Persistence.{BridgeSessionAlias, CodexSession, CodexTurn}
   alias CodexPooler.Gateway.Persistence.SessionContinuity, as: ContinuityStore
   alias CodexPooler.Gateway.Routing.BridgeRing
   alias CodexPooler.Repo
+  alias CodexPooler.RouteClass
   alias CodexPooler.Upstreams.Schemas.{PoolUpstreamAssignment, UpstreamIdentity}
 
   @type auth :: CodexPooler.Access.auth_context()
@@ -154,6 +156,26 @@ defmodule CodexPooler.Gateway.Routing.SessionContinuity do
 
   def filter_file_affinity(candidates, %RequestOptions{}), do: {:ok, candidates}
 
+  @spec filter_codex_session_assignment([BridgeRing.candidate()], RequestOptions.t(), Model.t()) ::
+          {:ok, [BridgeRing.candidate()]} | {:error, gateway_error()}
+  def filter_codex_session_assignment(
+        candidates,
+        %RequestOptions{
+          continuity: %{codex_session: %CodexSession{pool_upstream_assignment_id: assignment_id}}
+        } = request_options,
+        %Model{} = model
+      )
+      when is_binary(assignment_id) do
+    if hard_pin_codex_session_assignment?(request_options, model) do
+      filter_codex_session_assignment(candidates, request_options)
+    else
+      {:ok, prefer_codex_session_assignment(candidates, assignment_id)}
+    end
+  end
+
+  def filter_codex_session_assignment(candidates, %RequestOptions{} = request_options, %Model{}),
+    do: filter_codex_session_assignment(candidates, request_options)
+
   @spec filter_codex_session_assignment([BridgeRing.candidate()], RequestOptions.t()) ::
           {:ok, [BridgeRing.candidate()]} | {:error, gateway_error()}
   def filter_codex_session_assignment(
@@ -175,6 +197,72 @@ defmodule CodexPooler.Gateway.Routing.SessionContinuity do
 
   def filter_codex_session_assignment(candidates, %RequestOptions{}), do: {:ok, candidates}
 
+  @spec hard_pin_codex_session_assignment?(RequestOptions.t(), Model.t()) :: boolean()
+  defp hard_pin_codex_session_assignment?(%RequestOptions{} = request_options, %Model{} = model) do
+    RequestOptions.route_class(request_options) != RouteClass.proxy_stream() or
+      continuity_anchor?(request_options) or
+      file_affinity?(request_options) or
+      same_model_successful_turn?(request_options, model)
+  end
+
+  @spec continuity_anchor?(RequestOptions.t()) :: boolean()
+  defp continuity_anchor?(%RequestOptions{continuity: continuity}) do
+    [
+      continuity.accepted_turn_state,
+      continuity.previous_response_id,
+      continuity.session_header
+    ]
+    |> Enum.any?(&is_binary(clean_string(&1)))
+  end
+
+  @spec file_affinity?(RequestOptions.t()) :: boolean()
+  defp file_affinity?(%RequestOptions{routing: %{file_affinity_assignment_id: assignment_id}}),
+    do: is_binary(clean_string(assignment_id))
+
+  @spec same_model_successful_turn?(RequestOptions.t(), Model.t()) :: boolean()
+  defp same_model_successful_turn?(
+         %RequestOptions{continuity: %{codex_session: %CodexSession{id: session_id}}} =
+           request_options,
+         %Model{} = model
+       )
+       when is_binary(session_id) do
+    case requested_model_identifier(request_options, model) do
+      requested_model when is_binary(requested_model) ->
+        Repo.exists?(
+          from turn in CodexTurn,
+            join: request in Request,
+            on: request.id == turn.request_id,
+            where:
+              turn.codex_session_id == ^session_id and
+                turn.status == ^CodexTurn.succeeded_status() and
+                request.status == "succeeded" and
+                request.requested_model == ^requested_model
+        )
+
+      nil ->
+        false
+    end
+  end
+
+  defp same_model_successful_turn?(%RequestOptions{}, %Model{}), do: false
+
+  @spec requested_model_identifier(RequestOptions.t(), Model.t()) :: String.t() | nil
+  defp requested_model_identifier(
+         %RequestOptions{routing: %{requested_model: requested_model}},
+         %Model{exposed_model_id: exposed_model_id}
+       ) do
+    clean_string(requested_model) || clean_string(exposed_model_id)
+  end
+
+  @spec prefer_codex_session_assignment([BridgeRing.candidate()], Ecto.UUID.t() | String.t()) ::
+          [BridgeRing.candidate()]
+  defp prefer_codex_session_assignment(candidates, assignment_id) do
+    {pinned, fallback} =
+      Enum.split_with(candidates, fn {assignment, _identity} -> assignment.id == assignment_id end)
+
+    pinned ++ fallback
+  end
+
   @spec pinned_session_assignment_unavailable_error(Ecto.UUID.t() | String.t()) :: gateway_error()
   defp pinned_session_assignment_unavailable_error(assignment_id) do
     case persisted_pinned_reauth_assignment(assignment_id) do

lib/codex_pooler/gateway/runtime/dispatch/pre_dispatch.ex

@@ -101,7 +101,7 @@ defmodule CodexPooler.Gateway.Runtime.Dispatch.PreDispatch do
          {:ok, request_options} <-
            SessionContinuity.attach_codex_session(auth, payload, request_options),
          {:ok, candidates} <-
-           SessionContinuity.filter_codex_session_assignment(candidates, request_options) do
+           SessionContinuity.filter_codex_session_assignment(candidates, request_options, model) do
       route_state =
         route_state
         |> RouteState.put_candidates(candidates)

test/codex_pooler/accounting_test.exs

@@ -430,6 +430,71 @@ defmodule CodexPooler.AccountingTest do
              |> Enum.map(& &1.entry_kind) == ["reservation"]
     end
 
+    test "recovers stale in-progress codex turns when no owner lease is active" do
+      setup = accounting_setup()
+      now = DateTime.utc_now() |> DateTime.truncate(:microsecond)
+      stale_admitted_at = DateTime.add(now, -7, :hour)
+
+      assert {:ok, reserved} =
+               Accounting.reserve(
+                 setup.auth,
+                 setup.model,
+                 %{
+                   "model" => setup.model.exposed_model_id,
+                   "stream" => true,
+                   "max_output_tokens" => 10
+                 },
+                 %{correlation_id: "corr-stale-turn-no-owner", now: stale_admitted_at}
+               )
+
+      assert {:ok, attempt} =
+               Accounting.create_attempt(reserved.request, setup.assignment, %{
+                 now: stale_admitted_at
+               })
+
+      session =
+        %CodexSession{
+          pool_id: setup.pool.id,
+          api_key_id: setup.api_key.id,
+          session_key: "session-#{System.unique_integer([:positive])}",
+          pool_upstream_assignment_id: setup.assignment.id,
+          status: "active",
+          created_at: stale_admitted_at,
+          updated_at: stale_admitted_at
+        }
+        |> Repo.insert!()
+
+      turn =
+        %CodexTurn{
+          codex_session_id: session.id,
+          request_id: reserved.request.id,
+          turn_sequence: 1,
+          transport_kind: "http_sse",
+          status: "in_progress",
+          final_attempt_id: attempt.id,
+          started_at: stale_admitted_at,
+          created_at: stale_admitted_at,
+          updated_at: stale_admitted_at
+        }
+        |> Repo.insert!()
+
+      assert {:ok, %{stale_reservations_released: 0, stale_reservations_settled: 1}} =
+               Accounting.recover_stale_reservations(now)
+
+      assert Repo.get!(CodexPooler.Accounting.Request, reserved.request.id).last_error_code ==
+               "stale_reservation_recovered"
+
+      assert Repo.reload!(attempt).status == "failed"
+
+      assert %CodexTurn{
+               status: "interrupted",
+               error_code: "stale_reservation_recovered",
+               final_attempt_id: attempt_id
+             } = Repo.reload!(turn)
+
+      assert attempt_id == attempt.id
+    end
+
     test "skips already finalized reservations during stale recovery" do
       setup = accounting_setup()
       now = DateTime.utc_now() |> DateTime.truncate(:microsecond)