feat(integrations): complete ChatGPT subscription support

lib/kodo/integrations/audit_event.ex

@@ -14,9 +14,13 @@ defmodule Kodo.Integrations.AuditEvent do
     integration_disconnected
     device_authorization_started
     device_authorization_cancelled
+    device_authorization_completed
+    device_authorization_failed
     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/device_authorization.ex

@@ -0,0 +1,217 @@
+defmodule Kodo.Integrations.DeviceAuthorization do
+  @moduledoc "Runs finite, durable ChatGPT device authorization tasks."
+
+  alias Kodo.Accounts.Scope
+  alias Kodo.Cluster.InstanceManager
+  alias Kodo.Integrations
+  alias Kodo.Integrations.DeviceAuthorizationAttempt
+  alias Kodo.Integrations.DeviceAuthorizationTokens
+
+  @lease_renewal_interval_ms 15_000
+  @verification_url "https://auth.openai.com/codex/device"
+
+  def begin(%Scope{} = scope, integration_id, expected_generation, opts \\ []) do
+    client = Keyword.get(opts, :client, configured_client())
+    client_options = Keyword.get(opts, :client_options, [])
+
+    with {:ok, created} <- safe_create(client, client_options),
+         {:ok, attempt} <-
+           Integrations.begin_device_authorization(
+             scope,
+             integration_id,
+             expected_generation,
+             created.payload,
+             created.polling_interval_ms
+           ),
+         {:ok, task} <- resume(scope, integration_id, opts) do
+      {:ok,
+       %{
+         attempt: attempt,
+         task: task,
+         verification_url: created.verification_url
+       }}
+    end
+  end
+
+  def resume(%Scope{} = scope, integration_id, opts \\ []) do
+    claim_owner_id = Keyword.get_lazy(opts, :claim_owner_id, &InstanceManager.current_boot_id/0)
+    supervisor = Keyword.get(opts, :supervisor, Kodo.ControlPlaneTaskSupervisor)
+
+    with owner when is_binary(owner) <- claim_owner_id,
+         {:ok, {claim, payload}} <-
+           Integrations.claim_device_authorization(scope, integration_id, owner) do
+      task =
+        Task.Supervisor.async_nolink(supervisor, fn ->
+          run(scope, claim, payload, opts)
+        end)
+
+      {:ok, task}
+    else
+      nil -> {:error, :device_authorization_owner_unavailable}
+      {:error, _reason} = error -> error
+    end
+  end
+
+  @doc false
+  def run(
+        %Scope{} = scope,
+        %DeviceAuthorizationAttempt{} = claim,
+        payload,
+        opts \\ []
+      ) do
+    client = Keyword.get(opts, :client, configured_client())
+    client_options = Keyword.get(opts, :client_options, [])
+    continue(scope, claim, payload, client, client_options)
+  end
+
+  defp continue(scope, claim, %{"device_auth_id" => _, "user_code" => _}, client, options) do
+    with {:ok, {admission, payload}} <- await_poll_admission(scope, claim),
+         result <- safe_poll(client, payload, options) do
+      case result do
+        :pending ->
+          continue(scope, admission, payload, client, options)
+
+        {:ok, exchange_payload} ->
+          persist_and_exchange(scope, admission, exchange_payload, client, options)
+
+        {:error, reason} ->
+          fail(scope, admission, reason)
+      end
+    else
+      {:error, reason} -> stop(reason)
+    end
+  end
+
+  defp continue(scope, claim, %{"authorization_code" => _} = _payload, client, options) do
+    exchange(scope, claim, client, options)
+  end
+
+  defp continue(scope, claim, _payload, _client, _options) do
+    fail(scope, claim, :device_authorization_response_invalid)
+  end
+
+  defp await_poll_admission(scope, claim) do
+    delay = max(DateTime.diff(claim.next_poll_at, DateTime.utc_now(), :millisecond), 0)
+
+    if delay == 0 do
+      case Integrations.admit_device_authorization_poll(scope, claim) do
+        {:error, :device_authorization_poll_not_due} -> wait_and_renew(scope, claim, 1)
+        result -> result
+      end
+    else
+      wait_and_renew(scope, claim, min(delay, @lease_renewal_interval_ms))
+    end
+  end
+
+  defp wait_and_renew(scope, claim, delay) do
+    receive do
+    after
+      delay ->
+        case Integrations.renew_device_authorization_claim(
+               scope,
+               claim.id,
+               claim.attempt_generation,
+               claim.claim_owner_id,
+               claim.claim_epoch
+             ) do
+          {:ok, {renewed, _payload}} -> await_poll_admission(scope, renewed)
+          {:error, reason} -> {:error, reason}
+        end
+    end
+  end
+
+  defp persist_and_exchange(scope, claim, payload, client, options) do
+    case Integrations.store_device_authorization_exchange(scope, claim, payload) do
+      {:ok, persisted} -> exchange(scope, persisted, client, options)
+      {:error, reason} -> stop(reason)
+    end
+  end
+
+  defp exchange(scope, claim, client, options) do
+    with {:ok, {admission, payload}} <-
+           Integrations.admit_device_authorization_exchange(scope, claim),
+         {:ok, tokens} <- safe_exchange(client, payload, options),
+         {:ok, normalized} <- DeviceAuthorizationTokens.normalize(tokens),
+         {:ok, _integration} <-
+           Integrations.complete_device_authorization(
+             scope,
+             admission,
+             normalized.credentials,
+             normalized.expires_at
+           ) do
+      :ok
+    else
+      {:error, :stale_device_authorization_claim} = error -> stop(error)
+      {:error, reason} -> fail(scope, claim, reason)
+    end
+  end
+
+  defp fail(scope, claim, reason) do
+    case Integrations.fail_device_authorization(scope, claim, terminal_error(reason)) do
+      {:ok, _attempt} -> {:error, reason}
+      {:error, stale_reason} -> stop(stale_reason)
+    end
+  end
+
+  defp stop({:error, reason}), do: {:error, reason}
+  defp stop(reason), do: {:error, reason}
+
+  defp terminal_error(:redirect), do: "redirect"
+  defp terminal_error(:device_authorization_rejected), do: "provider_rejected"
+  defp terminal_error(:device_authorization_response_invalid), do: "response_invalid"
+  defp terminal_error(_reason), do: "provider_unavailable"
+
+  defp safe_create(client, options) do
+    case safe_client_call(fn -> client.create(options) end) do
+      {:ok,
+       %{
+         payload: payload,
+         polling_interval_ms: interval,
+         verification_url: verification_url
+       } = created}
+      when is_map(payload) and is_integer(interval) and verification_url == @verification_url ->
+        {:ok, created}
+
+      {:error, reason} when is_atom(reason) ->
+        {:error, reason}
+
+      _invalid ->
+        {:error, :device_authorization_response_invalid}
+    end
+  end
+
+  defp safe_poll(client, payload, options) do
+    case safe_client_call(fn -> client.poll(payload, options) end) do
+      :pending -> :pending
+      {:ok, exchange_payload} when is_map(exchange_payload) -> {:ok, exchange_payload}
+      {:error, reason} when is_atom(reason) -> {:error, reason}
+      _invalid -> {:error, :device_authorization_response_invalid}
+    end
+  end
+
+  defp safe_exchange(client, payload, options) do
+    case safe_client_call(fn -> client.exchange(payload, options) end) do
+      {:ok, tokens} when is_map(tokens) -> {:ok, tokens}
+      {:error, reason} when is_atom(reason) -> {:error, reason}
+      _invalid -> {:error, :device_authorization_response_invalid}
+    end
+  end
+
+  defp safe_client_call(operation) do
+    try do
+      operation.()
+    rescue
+      _exception -> {:error, :provider_unavailable}
+    catch
+      _kind, _reason -> {:error, :provider_unavailable}
+    end
+  end
+
+  defp configured_client do
+    Application.get_env(
+      :kodo,
+      :device_authorization_client,
+      Kodo.Integrations.ReqDeviceAuthorizationClient
+    )
+  end
+end

lib/kodo/integrations/device_authorization_tokens.ex

@@ -0,0 +1,70 @@
+defmodule Kodo.Integrations.DeviceAuthorizationTokens do
+  @moduledoc false
+
+  @openai_auth_claim "https://api.openai.com/auth"
+
+  def normalize(tokens, now \\ DateTime.utc_now())
+
+  def normalize(
+        %{
+          "access_token" => access_token,
+          "refresh_token" => refresh_token,
+          "id_token" => id_token
+        },
+        now
+      ) do
+    with {:ok, access_claims} <- decode_claims(access_token),
+         {:ok, identity_claims} <- decode_claims(id_token),
+         {:ok, expires_at} <- expiration(access_claims, now),
+         {:ok, account_id} <- account_id(identity_claims) do
+      {:ok,
+       %{
+         credentials: %{
+           "access_token" => access_token,
+           "refresh_token" => refresh_token,
+           "id_token" => id_token,
+           "account_id" => account_id
+         },
+         expires_at: expires_at
+       }}
+    else
+      _invalid -> {:error, :device_authorization_response_invalid}
+    end
+  end
+
+  def normalize(_tokens, _now), do: {:error, :device_authorization_response_invalid}
+
+  defp decode_claims(token) when is_binary(token) do
+    case String.split(token, ".") do
+      [_header, payload, _signature] ->
+        with {:ok, decoded} <- Base.url_decode64(payload, padding: false),
+             {:ok, claims} when is_map(claims) <- Jason.decode(decoded) do
+          {:ok, claims}
+        else
+          _invalid -> :error
+        end
+
+      _invalid ->
+        :error
+    end
+  end
+
+  defp decode_claims(_token), do: :error
+
+  defp expiration(%{"exp" => unix}, now) when is_integer(unix) do
+    with {:ok, expires_at} <- DateTime.from_unix(unix),
+         true <- DateTime.after?(expires_at, now) do
+      {:ok, %{expires_at | microsecond: {0, 6}}}
+    else
+      _invalid -> :error
+    end
+  end
+
+  defp expiration(_claims, _now), do: :error
+
+  defp account_id(%{@openai_auth_claim => %{"chatgpt_account_id" => account_id}})
+       when is_binary(account_id) and account_id != "",
+       do: {:ok, account_id}
+
+  defp account_id(_claims), do: :error
+end

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)