feat(integrations): refresh ChatGPT OAuth credentials

lib/kodo/integrations.ex

@@ -27,6 +27,7 @@
   @device_authorization_token_max_bytes 8_192
   @device_authorization_cleanup_age_seconds 24 * 60 * 60
   @device_authorization_claim_lease_ms 30_000
+  @refresh_claim_lease_ms 30_000
   @device_authorization_terminal_errors ~w(
     provider_unavailable
     provider_rejected
@@ -294,7 +295,7 @@
       Repo.transaction(fn ->
         lock_provider_identity(user.id, "openai_codex")
 
         case Repo.insert(changeset) do
           {:ok, integration} -> integration
           {:error, reason} -> Repo.rollback(reason)
         end
@@ -393,6 +394,114 @@
     )
   end
 
+  @doc "Claims one generation of an OAuth integration for a bounded refresh operation."
+  def claim_refresh(%Scope{user: user}, id, generation, claim_owner_id)
+      when is_integer(generation) and generation >= 0 do
+    with {:ok, id} <- cast_uuid(id),
+         {:ok, claim_owner_id} <- cast_uuid(claim_owner_id) do
+      lease_expires_at =
+        DateTime.add(database_now(), @refresh_claim_lease_ms, :millisecond)
+
+      query =
+        from integration in Integration,
+          where:
+            integration.id == ^id and integration.user_id == ^user.id and
+              integration.provider == "openai_codex" and
+              integration.authentication_type == "oauth" and
+              integration.connection_status == "connected" and
+              integration.credential_generation == ^generation and
+              (is_nil(integration.refresh_claim_owner_id) or
+                 integration.refresh_lease_expires_at <= fragment("clock_timestamp()"))
+
+      updates = [
+        set: [
+          refresh_claim_owner_id: claim_owner_id,
+          refresh_claim_generation: generation,
+          refresh_lease_expires_at: lease_expires_at,
+          updated_at: now()
+        ],
+        inc: [refresh_claim_epoch: 1]
+      ]
+
+      case Repo.update_all(query, updates) do
+        {1, nil} -> {:ok, Repo.get_by!(Integration, id: id, user_id: user.id)}
+        {0, nil} -> refresh_claim_conflict(user.id, id, generation)
+      end
+    end
+  end
+
+  def claim_refresh(%Scope{}, _id, _generation, _claim_owner_id),
+    do: {:error, :stale_credential_generation}
+
+  @doc "Releases only the exact refresh claim represented by the supplied snapshot."
+  def release_refresh_claim(%Scope{user: user}, %Integration{} = claim) do
+    query = exact_refresh_claim_query(user.id, claim)
+
+    case Repo.update_all(query,
+           set: [
+             refresh_claim_owner_id: nil,
+             refresh_claim_generation: nil,
+             refresh_lease_expires_at: nil,
+             updated_at: now()
+           ]
+         ) do
+      {1, nil} -> :ok
+      {0, nil} -> {:error, :stale_refresh_claim}
+    end
+  end
+
+  def require_refresh_reauthorization(%Scope{user: user}, %Integration{} = claim, error_code)
+      when error_code in ~w(refresh_invalid_grant refresh_account_identity_mismatch) do
+    Repo.transaction(fn ->
+      changes = [
+        connection_status: "reauthorization_required",
+        active: false,
+        validation_status: "unverified",
+        validated_at: nil,
+        validation_error_code: error_code,
+        refresh_claim_owner_id: nil,
+        refresh_claim_generation: nil,
+        refresh_lease_expires_at: nil,
+        updated_at: now()
+      ]
+
+      case Repo.update_all(exact_refresh_claim_query(user.id, claim), set: changes) do
+        {1, nil} ->
+          integration = Repo.get_by!(Integration, id: claim.id, user_id: user.id)
+          audit!(user.id, integration, "refresh_reauthorization_required")
+          integration
+
+        {0, nil} ->
+          Repo.rollback(:stale_refresh_claim)
+      end
+    end)
+    |> notify_integration_change()
+  end
+
+  def require_refresh_reauthorization(%Scope{}, %Integration{}, _error_code),
+    do: {:error, :unsafe_refresh_error}
+
+  def fail_refresh(%Scope{user: user}, %Integration{} = claim) do
+    Repo.transaction(fn ->
+      case Repo.update_all(exact_refresh_claim_query(user.id, claim),
+             set: [
+               refresh_claim_owner_id: nil,
+               refresh_claim_generation: nil,
+               refresh_lease_expires_at: nil,
+               updated_at: now()
+             ]
+           ) do
+        {1, nil} ->
+          integration = Repo.get_by!(Integration, id: claim.id, user_id: user.id)
+          audit!(user.id, integration, "refresh_failed")
+          integration
+
+        {0, nil} ->
+          Repo.rollback(:stale_refresh_claim)
+      end
+    end)
+  end
+
   def validation_succeeded(%Scope{} = scope, id, generation) do
     update_fenced(scope, id, generation, ["connected"], "validation_succeeded", %{
       validation_status: "valid",
@@ -455,7 +564,10 @@
         expires_at: nil,
         validated_at: nil,
         refreshed_at: nil,
-        validation_error_code: nil
+        validation_error_code: nil,
+        refresh_claim_owner_id: nil,
+        refresh_claim_generation: nil,
+        refresh_lease_expires_at: nil
       }
     )
     |> notify_integration_change()
@@ -596,7 +708,7 @@
              attempt.claim_owner_id == ^claim_owner_id)
   end
 
   defp renew_device_authorization_claim_locked(
          user_id,
          attempt_id,
          attempt_generation,
@@ -768,7 +880,7 @@
     end
   end
 
   defp lock_claimed_device_authorization!(user_id, integration, claim, timestamp) do
     query =
       DeviceAuthorizationAttempt
       |> where(
@@ -914,7 +1026,13 @@
   end
 
   defp device_authorization_integration_changes(%Integration{}, generation, timestamp) do
-    [credential_generation: generation, updated_at: timestamp]
+    [
+      credential_generation: generation,
+      refresh_claim_owner_id: nil,
+      refresh_claim_generation: nil,
+      refresh_lease_expires_at: nil,
+      updated_at: timestamp
+    ]
   end
 
   defp supersede_active_device_authorization(integration_id, timestamp) do
@@ -1086,7 +1204,10 @@
           expires_at: opts[:expires_at],
           validated_at: nil,
           refreshed_at: opts[:refreshed_at],
-          validation_error_code: nil
+          validation_error_code: nil,
+          refresh_claim_owner_id: nil,
+          refresh_claim_generation: nil,
+          refresh_lease_expires_at: nil
         })
 
       update_fenced(scope, id, generation, allowed_connections, audit_event_type, changes)
@@ -1106,6 +1227,29 @@
   defp require_connected(%Integration{connection_status: "connected"}), do: :ok
   defp require_connected(%Integration{}), do: {:error, :integration_not_connected}
 
+  defp refresh_claim_conflict(user_id, id, generation) do
+    case Repo.get_by(Integration, id: id, user_id: user_id) do
+      %Integration{credential_generation: ^generation, refresh_claim_owner_id: owner}
+      when not is_nil(owner) ->
+        {:error, :refresh_in_progress}
+
+      %Integration{} ->
+        {:error, :stale_credential_generation}
+
+      nil ->
+        {:error, :integration_not_found}
+    end
+  end
+
+  defp exact_refresh_claim_query(user_id, %Integration{} = claim) do
+    from integration in Integration,
+      where:
+        integration.id == ^claim.id and integration.user_id == ^user_id and
+          integration.credential_generation == ^claim.refresh_claim_generation and
+          integration.refresh_claim_owner_id == ^claim.refresh_claim_owner_id and
+          integration.refresh_claim_epoch == ^claim.refresh_claim_epoch
+  end
+
   defp lock_provider_integrations(user_id, integration_id) do
     provider =
       Integration

lib/kodo/integrations/audit_event.ex

@@ -19,6 +19,8 @@ defmodule Kodo.Integrations.AuditEvent do
     oauth_succeeded
     refresh_succeeded
     refresh_invalid_grant
+    refresh_reauthorization_required
+    refresh_failed
   )
 
   @primary_key {:id, :binary_id, autogenerate: true}

lib/kodo/integrations/device_authorization_req_adapter.ex → lib/kodo/integrations/credential_req_adapter.ex

@@ -1,6 +1,6 @@
-defmodule Kodo.Integrations.DeviceAuthorizationReqAdapter do
+defmodule Kodo.Integrations.CredentialReqAdapter do
   @moduledoc """
-  Req adapter for credential-bearing device authorization requests.
+  Req adapter for credential-bearing OAuth requests.
 
   It deliberately uses Mint directly: Finch telemetry includes request and
   response values, which are secrets throughout this protocol.
@@ -10,7 +10,7 @@ defmodule Kodo.Integrations.DeviceAuthorizationReqAdapter do
   @max_response_bytes 64 * 1024
 
   def run(%Req.Request{} = request) do
-    timeout = request.options[:device_authorization_timeout] || @default_timeout
+    timeout = request.options[:credential_request_timeout] || @default_timeout
     started_at = System.monotonic_time(:millisecond)
 
     task = Task.async(fn -> safe_perform(request, started_at, timeout) end)
@@ -30,7 +30,7 @@ defmodule Kodo.Integrations.DeviceAuthorizationReqAdapter do
   end
 
   defp perform(
-         %Req.Request{options: %{device_authorization_transport: transport}} = request,
+         %Req.Request{options: %{credential_request_transport: transport}} = request,
          started_at,
          timeout
        )

lib/kodo/integrations/integration.ex

@@ -34,6 +34,10 @@ defmodule Kodo.Integrations.Integration do
     field :validated_at, :utc_datetime_usec
     field :refreshed_at, :utc_datetime_usec
     field :validation_error_code, :string
+    field :refresh_claim_owner_id, Ecto.UUID
+    field :refresh_claim_epoch, :integer, default: 0
+    field :refresh_claim_generation, :integer
+    field :refresh_lease_expires_at, :utc_datetime_usec
 
     belongs_to :user, Kodo.Accounts.User, type: :id
 
@@ -72,6 +76,12 @@ defmodule Kodo.Integrations.Integration do
     |> check_constraint(:credential_generation,
       name: :provider_integrations_credential_generation_valid
     )
+    |> check_constraint(:refresh_claim_epoch,
+      name: :provider_integrations_refresh_claim_epoch_valid
+    )
+    |> check_constraint(:refresh_claim_owner_id,
+      name: :provider_integrations_refresh_claim_valid
+    )
     |> check_constraint(:connection_status, name: :provider_integrations_state_valid)
     |> check_constraint(:active, name: :provider_integrations_active_connected)
     |> unique_constraint(:active, name: :provider_integrations_one_active_index)

lib/kodo/integrations/oauth_access_validation.ex

@@ -0,0 +1,67 @@
+defmodule Kodo.Integrations.OAuthAccessValidation do
+  @moduledoc "Validates ChatGPT subscription access through the bounded refresh contract."
+
+  alias Kodo.Accounts.Scope
+  alias Kodo.Integrations
+  alias Kodo.Integrations.Integration
+  alias Kodo.Integrations.OAuthRefresh
+
+  def start(%Scope{} = scope, %{id: id, credential_generation: generation}) do
+    Task.Supervisor.async_nolink(Kodo.ControlPlaneTaskSupervisor, fn ->
+      validate(scope, id, generation)
+    end)
+  end
+
+  @doc false
+  def validate(%Scope{} = scope, id, generation, opts \\ []) do
+    with {:ok, integration} <- Integrations.get_integration(scope, id),
+         :ok <- admit(integration, generation) do
+      case OAuthRefresh.ensure_fresh(scope, integration, Keyword.put(opts, :force, true)) do
+        {:ok, refreshed} ->
+          scope
+          |> Integrations.validation_succeeded(refreshed.id, refreshed.credential_generation)
+          |> broadcast_result(refreshed)
+
+        {:error, :provider_unavailable} = error ->
+          result =
+            Integrations.validation_unavailable(
+              scope,
+              integration.id,
+              integration.credential_generation,
+              "provider_unavailable"
+            )
+
+          _result = broadcast_result(result, integration)
+          error
+
+        {:error, _reason} = error ->
+          error
+      end
+    end
+  end
+
+  defp admit(
+         %Integration{
+           provider: "openai_codex",
+           authentication_type: "oauth",
+           connection_status: "connected",
+           credential_generation: generation
+         },
+         generation
+       ),
+       do: :ok
+
+  defp admit(%Integration{}, _generation), do: {:error, :stale_credential_generation}
+
+  defp broadcast_result({:ok, validated} = result, integration) do
+    Phoenix.PubSub.broadcast(
+      Kodo.PubSub,
+      "integration:#{integration.user_id}",
+      {:integration_validation_finished, validated.id, validated.credential_generation}
+    )
+
+    result
+  end
+
+  defp broadcast_result(error, _integration), do: error
+end