fix(gateway): release neutral stream circuit probes

Terminal-missing HTTP SSE closes are health-neutral, but they still need to complete any half-open routing circuit probe started before dispatch.

Release the probe slot without counting success or failure so OMP retry storms do not collapse into no_eligible_backend while an interrupted stream is in flight.

Author: masterkain

lib/codex_pooler/gateway/routing/circuit_state.ex

@@ -233,6 +233,45 @@ defmodule CodexPooler.Gateway.Routing.CircuitState do
       ),
       do: {:error, :invalid_route_class}
 
+  @spec record_neutral_completion(auth(), Model.t(), PoolUpstreamAssignment.t(), String.t()) ::
+          {:ok, :ok | RoutingCircuitState.t()} | {:error, term()}
+  def record_neutral_completion(
+        %{pool: %Pool{}, api_key: %APIKey{}} = auth,
+        %Model{} = model,
+        %PoolUpstreamAssignment{} = assignment,
+        route_class
+      )
+      when is_binary(route_class) and route_class != "" do
+    now = now()
+
+    Repo.transaction(fn ->
+      case latest_for_update(auth, model, assignment, route_class) do
+        %RoutingCircuitState{status: @half_open_status} = state ->
+          state
+          |> RoutingCircuitState.changeset(%{
+            metadata: probe_metadata(state, max(probe_in_flight_count(state) - 1, 0)),
+            updated_at: now
+          })
+          |> persist_or_rollback(:update)
+
+        %RoutingCircuitState{} = state ->
+          state
+
+        nil ->
+          :ok
+      end
+    end)
+    |> unwrap_transaction()
+  end
+
+  def record_neutral_completion(
+        %{pool: %Pool{}, api_key: %APIKey{}},
+        %Model{},
+        %PoolUpstreamAssignment{},
+        _route_class
+      ),
+      do: {:error, :invalid_route_class}
+
   defp begin_attempt_with_snapshot(
          _auth,
          _model,

lib/codex_pooler/gateway/routing/route_lifecycle.ex

@@ -49,6 +49,28 @@ defmodule CodexPooler.Gateway.Routing.RouteLifecycle do
     end
   end
 
+  @spec selection_neutral_completion(map(), CodexPooler.Catalog.Model.t(), RoutingSelection.t()) ::
+          success_result()
+  def selection_neutral_completion(
+        auth,
+        model,
+        %RoutingSelection{
+          circuit_state: %RoutingCircuitState{}
+        } = selection
+      ) do
+    case CircuitState.record_neutral_completion(
+           auth,
+           model,
+           selection.assignment,
+           selection.route_class
+         ) do
+      {:ok, _state} -> :ok
+      {:error, reason} -> lifecycle_failure(:record_route_circuit_neutral_completion, reason)
+    end
+  end
+
+  def selection_neutral_completion(_auth, _model, %RoutingSelection{}), do: :ok
+
   @spec log_optional_result(String.t(), keyword(), success_result() | failure_result()) :: :ok
   def log_optional_result(_operation, _metadata, :ok), do: :ok
   def log_optional_result(_operation, _metadata, {:ok, _value}), do: :ok

lib/codex_pooler/gateway/runtime/finalization/streaming.ex

@@ -204,9 +204,9 @@ defmodule CodexPooler.Gateway.Runtime.Finalization.Streaming do
          "upstream_stream_error",
          nil,
          _headers,
-         _context
+         context
        ),
-       do: :ok
+       do: RouteLifecycle.neutral_completion(context)
 
   defp record_stream_failure_health(reason, code, nil, _headers, context) do
     record_health_failure(reason, code, context)

lib/codex_pooler/gateway/runtime/routing/route_lifecycle.ex

@@ -38,6 +38,15 @@ defmodule CodexPooler.Gateway.Runtime.Routing.RouteLifecycle do
     end
   end
 
+  @spec neutral_completion(DispatchContext.t()) :: success_result()
+  def neutral_completion(%DispatchContext{} = context) do
+    RoutingRouteLifecycle.selection_neutral_completion(
+      context.auth,
+      context.model,
+      routing_selection(context)
+    )
+  end
+
   defp merge_failure_metadata(%DispatchContext{} = context, demotion_reason) do
     case Accounting.merge_request_metadata(
            context.reserved.request,

test/codex_pooler/gateway/persistence/routing_circuit_state_test.exs

@@ -175,6 +175,22 @@ defmodule CodexPooler.Gateway.Persistence.RoutingCircuitStateTest do
     assert DateTime.compare(resumed.updated_at, state.updated_at) == :gt
   end
 
+  test "neutral completions release half-open probes without counting success or failure" do
+    {auth, model, assignment} = routing_fixture()
+    state = half_open_circuit!(auth, model, assignment, updated_at: now(), probe_count: 1)
+
+    assert {:ok, %RoutingCircuitState{} = updated} =
+             CircuitState.record_neutral_completion(auth, model, assignment, "proxy_websocket")
+
+    assert updated.id == state.id
+    assert updated.status == "half_open"
+    assert updated.reason_code == "test_probe"
+    assert updated.failure_count == state.failure_count
+    assert updated.success_count == state.success_count
+    assert updated.metadata["probe_in_flight_count"] == 0
+    assert CircuitState.eligible?(auth, model, assignment, "proxy_websocket")
+  end
+
   defp open_circuit!(auth, model, assignment, attrs) do
     now = now()
     next_probe_at = Keyword.fetch!(attrs, :next_probe_at)

test/codex_pooler/gateway/runtime/streaming/stream_lifecycle_test.exs

@@ -300,6 +300,83 @@ defmodule CodexPooler.Gateway.Runtime.Streaming.StreamLifecycleTest do
     assert Repo.all(from(c in RoutingCircuitState)) == []
   end
 
+  test "terminal-missing upstream SSE close releases half-open route probe" do
+    {setup, _first_upstream, _second_upstream} =
+      stream_retry_setup(
+        FakeUpstream.sse_stream([]),
+        FakeUpstream.sse_stream([])
+      )
+
+    {:ok, auth} = Access.authenticate_authorization_header(setup.authorization)
+    payload = payload(setup)
+    request_options = request_options(auth, payload, setup)
+
+    now = DateTime.utc_now() |> DateTime.truncate(:microsecond)
+
+    circuit =
+      %RoutingCircuitState{
+        pool_id: auth.pool.id,
+        pool_upstream_assignment_id: setup.assignment.id,
+        upstream_identity_id: setup.identity.id,
+        model_identifier: setup.model.exposed_model_id,
+        route_class: request_options.transport.route_class,
+        status: "half_open",
+        reason_code: "upstream_5xx",
+        failure_count: 3,
+        success_count: 0,
+        opened_at: DateTime.add(now, -120, :second),
+        half_opened_at: now,
+        metadata: %{"probe_in_flight_count" => 1},
+        created_at: DateTime.add(now, -120, :second),
+        updated_at: now
+      }
+      |> Repo.insert!()
+
+    assert {:ok, reserved} =
+             Accounting.reserve(auth, setup.model, payload, %{
+               endpoint: @endpoint_path,
+               transport: "http_sse",
+               correlation_id:
+                 "upstream-stream-neutral-probe-#{System.unique_integer([:positive])}",
+               request_metadata: %{}
+             })
+
+    assert {:ok, attempt} = Accounting.create_attempt(reserved.request, setup.assignment)
+
+    context =
+      retry_context(setup, auth, request_options, reserved.request,
+        candidates: [{setup.assignment, setup.identity}],
+        attempt: attempt,
+        routing_circuit_state: circuit
+      )
+
+    response_context = %ResponseContext{context: context, response: %Req.Response{status: 200}}
+
+    assert {:ok, _finalized} =
+             Streaming.finalize_failure(
+               "event: response.created\n\n",
+               :upstream_stream_interrupted,
+               response_context
+             )
+
+    assert [request] = Repo.all(from(r in Request, where: r.id == ^reserved.request.id))
+    assert request.status == "failed"
+    assert request.last_error_code == "upstream_stream_error"
+
+    assert [attempt] = Repo.all(from(a in Attempt, where: a.request_id == ^request.id))
+    assert attempt.status == "failed"
+    assert attempt.network_error_code == "upstream_stream_error"
+
+    assert Repo.all(from(d in BridgeDemotion)) == []
+
+    assert %RoutingCircuitState{} = updated = Repo.get!(RoutingCircuitState, circuit.id)
+    assert updated.status == "half_open"
+    assert updated.reason_code == "upstream_5xx"
+    assert updated.failure_count == 3
+    assert updated.success_count == 0
+    assert updated.metadata["probe_in_flight_count"] == 0
+  end
+
   defp retry_context(setup, auth, request_options, request, opts \\ []) do
     candidates =
       Keyword.get(opts, :candidates, [
@@ -330,6 +407,7 @@ defmodule CodexPooler.Gateway.Runtime.Streaming.StreamLifecycleTest do
       allow_retry?: true,
       routing_attempt_metadata: %{},
       route_class: request_options.transport.route_class,
+      routing_circuit_state: Keyword.get(opts, :routing_circuit_state),
       attempt: Keyword.get(opts, :attempt),
       started: System.monotonic_time(:millisecond)
     }