fix(upstreams): avoid token refresh loops on fresh usage probes

Treat usage-probe 401/403 responses as quota unavailable when the stored access token is still fresh, so reconciliation does not mark routable accounts refresh_failed or call OAuth unnecessarily.

Keep retryable account-reconciliation refresh failures recoverable through the existing token refresh worker and tighten job read-model redaction for recovery metadata.

Author: masterkain

lib/codex_pooler/jobs/read_model/failure_presentation.ex

@@ -2,6 +2,15 @@ defmodule CodexPooler.Jobs.ReadModel.FailurePresentation do
   @moduledoc false
 
   @sensitive_projection_keys [:args, :meta, :errors, "args", "meta", "errors"]
+  @sensitive_secret_key_pattern "(?:secret(?:[_-][a-z0-9]+)*|[a-z0-9]+(?:[_-][a-z0-9]+)*[_-]secret(?:[_-][a-z0-9]+)*)"
+  @sensitive_body_key_pattern "(?:auth[_-]?json|provider[_-]?body|request[_-]?body|response[_-]?body|body)"
+  @sensitive_failure_key_pattern "(?:authorization|cookie|set-cookie|api[_-]?key|(?:access|refresh|id|session|api)[_-]?token|[a-z0-9]+(?:[_-][a-z0-9]+)*[_-]token|#{@sensitive_body_key_pattern}|password|prompt|#{@sensitive_secret_key_pattern}|token)"
+  @sensitive_spaced_secret_key_pattern "(?:password|#{@sensitive_secret_key_pattern})"
+  @sensitive_quoted_string_body_failure_fragment ~r/(?i)(?:"?\b#{@sensitive_body_key_pattern}\b"?\s*[:=]\s*)"(?:\\.|[^"\\])*"/
+  @sensitive_jsonish_failure_fragment ~r/(?i)(?:"?\b#{@sensitive_failure_key_pattern}\b"?\s*[:=]\s*)(?:\{(?:[^{}]|\{[^{}]*\})*\}|\[(?:[^\[\]]|\[[^\[\]]*\])*\]|"[^"]*"|'[^']*')/
+  @sensitive_body_failure_fragment ~r/(?i)(?:"?\b#{@sensitive_body_key_pattern}\b"?\s*[:=]\s*).*?(?=\s+"?\b#{@sensitive_failure_key_pattern}\b"?\s*[:=]|\z)/
+  @sensitive_spaced_secret_failure_fragment ~r/(?i)(?:"?\b#{@sensitive_spaced_secret_key_pattern}\b"?\s*[:=]\s*).*?(?=[,;]|\.\s+|\s+"?\b#{@sensitive_failure_key_pattern}\b"?\s*[:=]|\z)/
+  @sensitive_text_failure_fragment ~r/(?i)(?:"?\b#{@sensitive_failure_key_pattern}\b"?\s*[:=]\s*)[^,;\s]+/
 
   @type failure_summary :: %{
           required(:title) => String.t(),
@@ -21,6 +30,12 @@ defmodule CodexPooler.Jobs.ReadModel.FailurePresentation do
   def sanitize_projection(%NaiveDateTime{} = value), do: value
   def sanitize_projection(%Date{} = value), do: value
 
+  def sanitize_projection(%{trigger_kind: trigger_kind} = value) when trigger_kind in [nil, ""] do
+    value
+    |> Map.delete(:trigger_kind)
+    |> sanitize_projection()
+  end
+
   def sanitize_projection(%{errors: errors} = value) when is_list(errors) do
     value
     |> Map.drop(@sensitive_projection_keys)
@@ -109,11 +124,12 @@ defmodule CodexPooler.Jobs.ReadModel.FailurePresentation do
   defp redact_failure_secrets(message) do
     message
     |> String.replace(~r/(?i)bearer\s+[a-z0-9._~+\/=:-]+/, "Bearer [redacted]")
+    |> String.replace(@sensitive_quoted_string_body_failure_fragment, "[redacted]")
+    |> String.replace(@sensitive_jsonish_failure_fragment, "[redacted]")
+    |> String.replace(@sensitive_body_failure_fragment, "[redacted]")
+    |> String.replace(@sensitive_spaced_secret_failure_fragment, "[redacted]")
+    |> String.replace(@sensitive_text_failure_fragment, "[redacted]")
     |> String.replace(~r/(?i)\bsecret[-_a-z0-9]*\b/, "[redacted]")
-    |> String.replace(
-      ~r/(?i)\b(authorization|cookie|set-cookie|api[_-]?key|access[_-]?token|refresh[_-]?token|password|prompt|secret|token)\b\s*[:=]\s*[^,;\s]+/,
-      "[redacted]"
-    )
   end
 
   defp unwrap_oban_failure_message(message) do

lib/codex_pooler/jobs/read_model/query.ex

@@ -29,17 +29,29 @@ defmodule CodexPooler.Jobs.ReadModel.Query do
 
   def with_attention(jobs, opts) when is_list(jobs) do
     now = attention_now(opts)
-    Enum.map(jobs, &HealthPolicy.put_attention(&1, now: now))
+    Enum.map(jobs, &put_attention(&1, now))
   end
 
   def with_attention(nil, _opts), do: nil
 
   def with_attention(job, opts) when is_map(job) do
-    HealthPolicy.put_attention(job, now: attention_now(opts))
+    put_attention(job, attention_now(opts))
   end
 
   def attention_now(opts), do: Keyword.get_lazy(opts, :now, &DateTime.utc_now/0)
 
+  defp put_attention(job, now) do
+    job
+    |> drop_blank_trigger_kind()
+    |> HealthPolicy.put_attention(now: now)
+  end
+
+  defp drop_blank_trigger_kind(%{trigger_kind: trigger_kind} = job)
+       when trigger_kind in [nil, ""],
+       do: Map.delete(job, :trigger_kind)
+
+  defp drop_blank_trigger_kind(job), do: job
+
   def group_worker_rows_by_index(worker_groups) do
     worker_groups = Enum.with_index(worker_groups)
 
@@ -71,6 +83,7 @@ defmodule CodexPooler.Jobs.ReadModel.Query do
         worker: unquote(job).worker,
         queue: unquote(job).queue,
         state: unquote(job).state,
+        trigger_kind: fragment("?->>?", unquote(job).args, "trigger_kind"),
         errors: unquote(job).errors,
         attempt: unquote(job).attempt,
         max_attempts: unquote(job).max_attempts,

lib/codex_pooler/upstreams/reconciliation/pool_reconciliation.ex

@@ -3,6 +3,7 @@ defmodule CodexPooler.Upstreams.Reconciliation.PoolReconciliation do
 
   import Ecto.Query
 
+  alias CodexPooler.Jobs
   alias CodexPooler.Pools.Pool
   alias CodexPooler.Repo
   alias CodexPooler.Upstreams.Auth.TokenRefresh
@@ -17,6 +18,7 @@ defmodule CodexPooler.Upstreams.Reconciliation.PoolReconciliation do
   @eligible PoolUpstreamAssignment.eligible_status()
   @health_active PoolUpstreamAssignment.active_health_status()
   @account_quota_key "account"
+  @usage_auth_refresh_skew_seconds 5 * 60
   @codex_usage_paths [
     "/api/codex/usage",
     "/backend-api/codex/usage",
@@ -204,7 +206,7 @@ defmodule CodexPooler.Upstreams.Reconciliation.PoolReconciliation do
           {:usage, identity, payload, windows}
 
         {:error, {:upstream_status, status}} when status in [401, 403] ->
-          retry_codex_usage_after_token_refresh(identity, assignment, opts)
+          maybe_retry_codex_usage_after_token_refresh(identity, assignment, observed_at, opts)
 
         _error ->
           :usage_unavailable
@@ -214,13 +216,62 @@ defmodule CodexPooler.Upstreams.Reconciliation.PoolReconciliation do
     end
   end
 
+  defp maybe_retry_codex_usage_after_token_refresh(identity, assignment, observed_at, opts) do
+    if access_token_refresh_due_after_usage_auth_failure?(identity, observed_at) do
+      retry_codex_usage_after_token_refresh(identity, assignment, opts)
+    else
+      :usage_unavailable
+    end
+  end
+
+  defp access_token_refresh_due_after_usage_auth_failure?(
+         %UpstreamIdentity{} = identity,
+         %DateTime{} = observed_at
+       ) do
+    case access_token_expires_at(identity.metadata) do
+      {:ok, expires_at} ->
+        refresh_at = DateTime.add(observed_at, @usage_auth_refresh_skew_seconds, :second)
+        DateTime.compare(expires_at, refresh_at) in [:lt, :eq]
+
+      :unknown ->
+        true
+    end
+  end
+
+  defp access_token_expires_at(%{} = metadata) do
+    case metadata["access_token_expires_at"] do
+      expires_at when is_binary(expires_at) ->
+        case DateTime.from_iso8601(expires_at) do
+          {:ok, parsed, _offset} -> {:ok, DateTime.truncate(parsed, :microsecond)}
+          _invalid -> :unknown
+        end
+
+      _value ->
+        :unknown
+    end
+  end
+
+  defp access_token_expires_at(_metadata), do: :unknown
+
   defp retry_codex_usage_after_token_refresh(identity, assignment, opts) do
-    with {:ok, %{status: :active, identity: refreshed_identity}} <-
-           TokenRefresh.refresh_access_token(identity,
-             trigger_kind: "account_reconciliation",
-             receive_timeout: Keyword.get(opts, :receive_timeout, 30_000)
-           ),
-         {:ok, access_token} <- Secrets.decrypt_active_secret(refreshed_identity, "access_token"),
+    case TokenRefresh.refresh_access_token(identity,
+           trigger_kind: "account_reconciliation",
+           receive_timeout: Keyword.get(opts, :receive_timeout, 30_000)
+         ) do
+      {:ok, %{status: :active, identity: refreshed_identity}} ->
+        fetch_codex_usage_after_successful_token_refresh(refreshed_identity, assignment, opts)
+
+      {:ok, %{status: :refresh_failed, retryable?: true, identity: failed_identity}} ->
+        maybe_enqueue_account_reconciliation_token_refresh_recovery(failed_identity)
+        :auth_unavailable
+
+      _unavailable ->
+        :auth_unavailable
+    end
+  end
+
+  defp fetch_codex_usage_after_successful_token_refresh(refreshed_identity, assignment, opts) do
+    with {:ok, access_token} <- Secrets.decrypt_active_secret(refreshed_identity, "access_token"),
          observed_at <- now(),
          {:ok, payload, _url, windows} <-
            fetch_codex_usage_payload(
@@ -236,6 +287,26 @@ defmodule CodexPooler.Upstreams.Reconciliation.PoolReconciliation do
     end
   end
 
+  defp maybe_enqueue_account_reconciliation_token_refresh_recovery(
+         %UpstreamIdentity{} = failed_identity
+       ) do
+    if account_reconciliation_refresh_failure?(failed_identity) do
+      case Jobs.enqueue_token_refresh(failed_identity,
+             trigger_kind: "account_reconciliation_recovery"
+           ) do
+        {:ok, %Oban.Job{}} -> :ok
+        {:error, _reason} -> :ok
+      end
+    end
+  end
+
+  defp account_reconciliation_refresh_failure?(%UpstreamIdentity{} = identity) do
+    case identity.metadata["token_refresh"] do
+      %{"status" => "failed", "trigger_kind" => "account_reconciliation"} -> true
+      _metadata -> false
+    end
+  end
+
   defp fetch_codex_usage_payload(identity, assignment, access_token, observed_at, opts) do
     base = upstream_usage_base_url(identity, assignment)
     timeout = Keyword.get(opts, :receive_timeout, 30_000)