fix(gateway): keep refreshing identities route-visible

Author: masterkain

lib/codex_pooler/gateway/routing/candidate_eligibility.ex

@@ -96,6 +96,10 @@ defmodule CodexPooler.Gateway.Routing.CandidateEligibility do
     PoolUpstreamAssignment.disabled_health_status(),
     PoolUpstreamAssignment.errored_health_status()
   ]
+  @visible_identity_statuses [
+    UpstreamIdentity.active_status(),
+    UpstreamIdentity.refreshing_status()
+  ]
 
   @type candidate :: {PoolUpstreamAssignment.t(), UpstreamIdentity.t()}
   @type gateway_error :: Contracts.gateway_error()
@@ -392,7 +396,6 @@ defmodule CodexPooler.Gateway.Routing.CandidateEligibility do
     else
       assignment_active_status = PoolUpstreamAssignment.active_status()
       assignment_eligible_status = PoolUpstreamAssignment.eligible_status()
-      identity_active_status = UpstreamIdentity.active_status()
 
       Repo.all(
         from assignment in PoolUpstreamAssignment,
@@ -402,7 +405,7 @@ defmodule CodexPooler.Gateway.Routing.CandidateEligibility do
             assignment.id in ^assignment_ids and assignment.status == ^assignment_active_status and
               assignment.eligibility_status == ^assignment_eligible_status and
               assignment.health_status not in ^@health_excluded and
-              identity.status == ^identity_active_status and
+              identity.status in ^@visible_identity_statuses and
               (is_nil(assignment.cooldown_until) or assignment.cooldown_until <= ^timestamp),
           order_by: [asc: assignment.created_at, asc: assignment.id],
           select: {assignment, identity}

test/codex_pooler/gateway/routing/candidate_eligibility_test.exs

@@ -57,6 +57,7 @@ defmodule CodexPooler.Gateway.Routing.CandidateEligibilityTest do
       active = upstream_assignment_fixture(pool, %{})
       disabled = upstream_assignment_fixture(pool, %{assignment_status: "disabled"})
       stale_identity = upstream_assignment_fixture(pool, %{identity_status: "refresh_due"})
+      refreshing_identity = upstream_assignment_fixture(pool, %{identity_status: "refreshing"})
 
       visible_model =
         model_fixture(pool, %{
@@ -65,7 +66,8 @@ defmodule CodexPooler.Gateway.Routing.CandidateEligibilityTest do
             "source_assignment_ids" => [
               active.assignment.id,
               disabled.assignment.id,
-              stale_identity.assignment.id
+              stale_identity.assignment.id,
+              refreshing_identity.assignment.id
             ]
           }
         })
@@ -83,7 +85,12 @@ defmodule CodexPooler.Gateway.Routing.CandidateEligibilityTest do
       refute Enum.any?(context.visible_models, &(&1.id == hidden_model.id))
 
       assert {:ok, candidates} = CandidateEligibility.routable_candidates(context, visible_model)
-      assert candidate_ids(candidates) == [active.assignment.id]
+
+      assert candidate_ids(candidates) == [
+               active.assignment.id,
+               refreshing_identity.assignment.id
+             ]
+
       refute disabled.assignment.id in candidate_ids(candidates)
       refute stale_identity.assignment.id in candidate_ids(candidates)