From e9651e6203886db79175b5767e90a89e05724959 Mon Sep 17 00:00:00 2001 From: Darwin Wu Date: Sat, 12 Sep 2026 19:56:34 +0000 Subject: [PATCH 01/14] fix(integrations): classify permanent refresh failures --- lib/kodo/integrations/req_refresh_client.ex | 19 ++++++++++++++----- .../integrations/req_refresh_client_test.exs | 14 +++++++++++++- 2 files changed, 27 insertions(+), 6 deletions(-) diff --git a/lib/kodo/integrations/req_refresh_client.ex b/lib/kodo/integrations/req_refresh_client.ex index 8e78832..8515121 100644 --- a/lib/kodo/integrations/req_refresh_client.ex +++ b/lib/kodo/integrations/req_refresh_client.ex @@ -51,8 +51,8 @@ defmodule Kodo.Integrations.ReqRefreshClient do {:ok, %Req.Response{status: status, body: body}} when status in 200..299 -> decode_success(body) - {:ok, %Req.Response{status: 400, body: body}} -> - if invalid_grant?(body), + {:ok, %Req.Response{status: status, body: body}} when status in [400, 401] -> + if permanent_refresh_failure?(status, body), do: {:error, :invalid_grant}, else: {:error, :provider_unavailable} @@ -77,10 +77,19 @@ defmodule Kodo.Integrations.ReqRefreshClient do end end - defp invalid_grant?(body) do + defp permanent_refresh_failure?(401, _body), do: true + + defp permanent_refresh_failure?(400, body) do case decode_json(body) do - {:ok, %{"error" => "invalid_grant"}} -> true - _other -> false + {:ok, %{"error" => "invalid_grant"}} -> + true + + {:ok, %{"code" => code}} + when code in ~w(refresh_token_expired refresh_token_reused refresh_token_invalidated) -> + true + + _other -> + false end end diff --git a/test/kodo/integrations/req_refresh_client_test.exs b/test/kodo/integrations/req_refresh_client_test.exs index 53a7b58..cbca3bf 100644 --- a/test/kodo/integrations/req_refresh_client_test.exs +++ b/test/kodo/integrations/req_refresh_client_test.exs @@ -29,7 +29,7 @@ defmodule Kodo.Integrations.ReqRefreshClientTest do ReqRefreshClient.refresh("refresh-secret", plug: plug) end - test "classifies only an exact invalid_grant response as invalid credentials" do + test "classifies the pinned permanent refresh failures as invalid credentials" do invalid_grant = fn conn -> conn |> Plug.Conn.put_status(400) |> Req.Test.json(%{"error" => "invalid_grant"}) end @@ -37,6 +37,18 @@ defmodule Kodo.Integrations.ReqRefreshClientTest do assert {:error, :invalid_grant} = ReqRefreshClient.refresh("refresh-secret", plug: invalid_grant) + for {status, body} <- [ + {401, %{"code" => "token_expired"}}, + {400, %{"code" => "refresh_token_expired"}}, + {400, %{"code" => "refresh_token_reused"}}, + {400, %{"code" => "refresh_token_invalidated"}} + ] do + plug = fn conn -> conn |> Plug.Conn.put_status(status) |> Req.Test.json(body) end + + assert {:error, :invalid_grant} = + ReqRefreshClient.refresh("refresh-secret", plug: plug) + end + for body <- [%{"error" => "temporarily_unavailable"}, %{"message" => "invalid_grant"}] do plug = fn conn -> conn |> Plug.Conn.put_status(400) |> Req.Test.json(body) end From 2fe11bbd03f35f92fa69d86180513f59d609e7e7 Mon Sep 17 00:00:00 2001 From: Darwin Wu Date: Sat, 12 Sep 2026 19:57:44 +0000 Subject: [PATCH 02/14] fix(integrations): preserve refresh completion lineage --- lib/kodo/integrations.ex | 95 ++++++++++++++----- lib/kodo/integrations/integration.ex | 4 + lib/kodo/integrations/oauth_refresh.ex | 19 +++- ...ce_generation_to_provider_integrations.exs | 15 +++ priv/repo/structure.sql | 7 +- test/kodo/integrations/oauth_refresh_test.exs | 73 +++++++++++++- 6 files changed, 179 insertions(+), 34 deletions(-) create mode 100644 priv/repo/migrations/20260912195208_add_refresh_source_generation_to_provider_integrations.exs diff --git a/lib/kodo/integrations.ex b/lib/kodo/integrations.ex index bb3959c..5b5c6b0 100644 --- a/lib/kodo/integrations.ex +++ b/lib/kodo/integrations.ex @@ -428,7 +428,7 @@ defmodule Kodo.Integrations do id, generation, credentials, - opts, + Keyword.put(opts, :refresh_source_generation, generation), ~w(connected reauthorization_required), "oauth", "refresh_succeeded" @@ -442,11 +442,9 @@ defmodule Kodo.Integrations do 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, + as: :integration, where: integration.id == ^id and integration.user_id == ^user.id and integration.provider == "openai_codex" and @@ -454,21 +452,31 @@ defmodule Kodo.Integrations do 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) + integration.refresh_lease_expires_at <= fragment("clock_timestamp()")) and + not exists( + from attempt in DeviceAuthorizationAttempt, + where: + attempt.integration_id == parent_as(:integration).id and + attempt.state == "active" + ), + update: [ + set: [ + refresh_claim_owner_id: ^claim_owner_id, + refresh_claim_generation: ^generation, + refresh_lease_expires_at: + fragment( + "clock_timestamp() + (? * interval '1 millisecond')", + ^@refresh_claim_lease_ms + ), + updated_at: fragment("clock_timestamp()") + ], + inc: [refresh_claim_epoch: 1] + ], + select: integration + + case Repo.update_all(query, []) do + {1, [claim]} -> {:ok, claim} + {0, []} -> refresh_claim_conflict(user.id, id, generation) end end end @@ -476,6 +484,27 @@ defmodule Kodo.Integrations do def claim_refresh(%Scope{}, _id, _generation, _claim_owner_id), do: {:error, :stale_credential_generation} + @doc "Revalidates an exact refresh claim immediately before credential use." + def admit_refresh_claim(%Scope{user: user}, %Integration{} = claim) do + active_attempt = + from attempt in DeviceAuthorizationAttempt, + where: attempt.integration_id == ^claim.id and attempt.state == "active" + + query = + exact_refresh_claim_query(user.id, claim) + |> where( + [integration], + integration.connection_status == "connected" and + integration.refresh_lease_expires_at > fragment("clock_timestamp()") and + not exists(subquery(active_attempt)) + ) + + case Repo.one(query) do + %Integration{} = admitted -> {:ok, admitted} + nil -> {:error, :stale_refresh_claim} + end + end + @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) @@ -496,6 +525,8 @@ defmodule Kodo.Integrations do 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 -> + lock_user!(user.id) + changes = [ connection_status: "reauthorization_required", active: false, @@ -526,6 +557,8 @@ defmodule Kodo.Integrations do def fail_refresh(%Scope{user: user}, %Integration{} = claim) do Repo.transaction(fn -> + lock_user!(user.id) + case Repo.update_all(exact_refresh_claim_query(user.id, claim), set: [ refresh_claim_owner_id: nil, @@ -1249,6 +1282,7 @@ defmodule Kodo.Integrations do connection_status: "connected", validation_status: "unverified", credential_generation: generation + 1, + refresh_source_generation: opts[:refresh_source_generation], expires_at: opts[:expires_at], validated_at: nil, refreshed_at: opts[:refreshed_at], @@ -1276,15 +1310,26 @@ defmodule Kodo.Integrations do 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} + active_authorization? = + Repo.exists?( + from attempt in DeviceAuthorizationAttempt, + where: + attempt.integration_id == ^id and attempt.user_id == ^user_id and + attempt.state == "active" + ) + + case {Repo.get_by(Integration, id: id, user_id: user_id), active_authorization?} do + {%Integration{credential_generation: ^generation}, true} -> + {:error, :integration_reauthorization_required} + + {%Integration{credential_generation: ^generation, refresh_claim_owner_id: owner}, false} when not is_nil(owner) -> {:error, :refresh_in_progress} - %Integration{} -> + {%Integration{}, _active?} -> {:error, :stale_credential_generation} - nil -> + {nil, _active?} -> {:error, :integration_not_found} end end @@ -1430,6 +1475,10 @@ defmodule Kodo.Integrations do end end + defp maybe_lock_transition(user_id, _integration_id, "refresh_succeeded") do + lock_user!(user_id) + end + defp maybe_lock_transition(_user_id, _integration_id, _event_type), do: :ok defp maybe_terminalize_device_authorization(user_id, integration_id, "integration_disconnected") do diff --git a/lib/kodo/integrations/integration.ex b/lib/kodo/integrations/integration.ex index 2e80063..967c478 100644 --- a/lib/kodo/integrations/integration.ex +++ b/lib/kodo/integrations/integration.ex @@ -33,6 +33,7 @@ defmodule Kodo.Integrations.Integration do field :expires_at, :utc_datetime_usec field :validated_at, :utc_datetime_usec field :refreshed_at, :utc_datetime_usec + field :refresh_source_generation, :integer field :validation_error_code, :string field :refresh_claim_owner_id, Ecto.UUID field :refresh_claim_epoch, :integer, default: 0 @@ -79,6 +80,9 @@ defmodule Kodo.Integrations.Integration do |> check_constraint(:refresh_claim_epoch, name: :provider_integrations_refresh_claim_epoch_valid ) + |> check_constraint(:refresh_source_generation, + name: :provider_integrations_refresh_source_generation_valid + ) |> check_constraint(:refresh_claim_owner_id, name: :provider_integrations_refresh_claim_valid ) diff --git a/lib/kodo/integrations/oauth_refresh.ex b/lib/kodo/integrations/oauth_refresh.ex index 76705d4..9a5a541 100644 --- a/lib/kodo/integrations/oauth_refresh.ex +++ b/lib/kodo/integrations/oauth_refresh.ex @@ -2,7 +2,6 @@ defmodule Kodo.Integrations.OAuthRefresh do @moduledoc "Runs generation-fenced, just-in-time OAuth refreshes as finite tasks." alias Kodo.Accounts.Scope - alias Kodo.Cluster.InstanceManager alias Kodo.Integrations alias Kodo.Integrations.CredentialEncryption alias Kodo.Integrations.Integration @@ -24,7 +23,7 @@ defmodule Kodo.Integrations.OAuthRefresh do end defp refresh(scope, integration, opts) do - owner_id = Keyword.get_lazy(opts, :claim_owner_id, &InstanceManager.current_boot_id/0) + owner_id = Keyword.get_lazy(opts, :claim_owner_id, &Ecto.UUID.generate/0) deadline = System.monotonic_time(:millisecond) + Keyword.get(opts, :wait_timeout, @wait_timeout_ms) @@ -52,6 +51,9 @@ defmodule Kodo.Integrations.OAuthRefresh do {:error, :stale_credential_generation} -> current_result(scope, integration) + {:error, :integration_reauthorization_required} -> + {:error, :integration_reauthorization_required} + {:error, _reason} -> {:error, :provider_unavailable} end @@ -84,7 +86,10 @@ defmodule Kodo.Integrations.OAuthRefresh do {:error, :provider_unavailable} nil -> - _result = Task.shutdown(task, :brutal_kill) + # The caller is bounded, but an admitted worker must retain ownership of a rotated-token + # response long enough to persist it. Ignoring detaches the eventual task reply without + # cancelling the finite supervised operation. + _result = Task.ignore(task) {:error, :provider_unavailable} end end @@ -93,7 +98,8 @@ defmodule Kodo.Integrations.OAuthRefresh do client = Keyword.get(opts, :client, configured_client()) client_options = Keyword.get(opts, :client_options, []) - with {:ok, current} <- CredentialEncryption.decrypt(claim), + with {:ok, admitted} <- Integrations.admit_refresh_claim(scope, claim), + {:ok, current} <- CredentialEncryption.decrypt(admitted), {:ok, refresh_token} <- fetch_refresh_token(current), {:ok, response} <- safe_refresh(client, refresh_token, client_options), {:ok, normalized} <- RefreshTokens.normalize(response, current), @@ -117,6 +123,9 @@ defmodule Kodo.Integrations.OAuthRefresh do {:error, :stale_credential_generation} -> current_result(scope, claim) + {:error, :stale_refresh_claim} -> + current_result(scope, claim) + {:error, _reason} -> _result = Integrations.fail_refresh(scope, claim) {:error, :provider_unavailable} @@ -134,7 +143,7 @@ defmodule Kodo.Integrations.OAuthRefresh do defp current_result(scope, integration) do case Integrations.get_integration(scope, integration.id) do {:ok, %{connection_status: "connected"} = current} - when current.credential_generation > integration.credential_generation -> + when current.refresh_source_generation == integration.credential_generation -> {:ok, current} {:ok, %{connection_status: "reauthorization_required"}} -> diff --git a/priv/repo/migrations/20260912195208_add_refresh_source_generation_to_provider_integrations.exs b/priv/repo/migrations/20260912195208_add_refresh_source_generation_to_provider_integrations.exs new file mode 100644 index 0000000..09f600e --- /dev/null +++ b/priv/repo/migrations/20260912195208_add_refresh_source_generation_to_provider_integrations.exs @@ -0,0 +1,15 @@ +defmodule Kodo.Repo.Migrations.AddRefreshSourceGenerationToProviderIntegrations do + use Ecto.Migration + + def change do + alter table(:provider_integrations) do + add :refresh_source_generation, :bigint + end + + create constraint( + :provider_integrations, + :provider_integrations_refresh_source_generation_valid, + check: "refresh_source_generation IS NULL OR refresh_source_generation >= 0" + ) + end +end diff --git a/priv/repo/structure.sql b/priv/repo/structure.sql index bef1f4c..cd098cc 100644 --- a/priv/repo/structure.sql +++ b/priv/repo/structure.sql @@ -2,7 +2,7 @@ -- PostgreSQL database dump -- -\restrict VoWETYGZi4vEwyliUNWQG3BVneag5OunvKzcS2NOBFzVtuTzyBWpDfhYIneu2ek +\restrict cMSeT1HADB9gn2fd6oRSfRA6k31lS6KseACcNAs16bTriITaRrwptynxWKofvCj -- Dumped from database version 15.19 (Debian 15.19-0+deb12u1) -- Dumped by pg_dump version 15.19 (Debian 15.19-0+deb12u1) @@ -187,6 +187,7 @@ CREATE TABLE public.provider_integrations ( refresh_claim_epoch bigint DEFAULT 0 NOT NULL, refresh_claim_generation bigint, refresh_lease_expires_at timestamp without time zone, + refresh_source_generation bigint, CONSTRAINT provider_integrations_active_connected CHECK (((NOT active) OR ((connection_status)::text = 'connected'::text))), CONSTRAINT provider_integrations_authentication_type_valid CHECK (((authentication_type)::text = ANY ((ARRAY['api_key'::character varying, 'oauth'::character varying])::text[]))), CONSTRAINT provider_integrations_credential_generation_valid CHECK ((credential_generation >= 0)), @@ -194,6 +195,7 @@ CREATE TABLE public.provider_integrations ( CONSTRAINT provider_integrations_provider_valid CHECK (((provider)::text = ANY ((ARRAY['openai'::character varying, 'openai_codex'::character varying, 'anthropic'::character varying, 'openrouter'::character varying])::text[]))), CONSTRAINT provider_integrations_refresh_claim_epoch_valid CHECK ((refresh_claim_epoch >= 0)), CONSTRAINT provider_integrations_refresh_claim_valid CHECK ((((refresh_claim_owner_id IS NULL) AND (refresh_claim_generation IS NULL) AND (refresh_lease_expires_at IS NULL)) OR ((refresh_claim_owner_id IS NOT NULL) AND (refresh_claim_generation IS NOT NULL) AND (refresh_claim_generation >= 0) AND (refresh_lease_expires_at IS NOT NULL)))), + CONSTRAINT provider_integrations_refresh_source_generation_valid CHECK (((refresh_source_generation IS NULL) OR (refresh_source_generation >= 0))), CONSTRAINT provider_integrations_state_valid CHECK (((((connection_status)::text = 'disconnected'::text) AND ((validation_status)::text = 'unverified'::text) AND (encrypted_credentials IS NULL) AND (encryption_key_version IS NULL) AND (credential_format_version IS NULL) AND (expires_at IS NULL) AND (validated_at IS NULL) AND (refreshed_at IS NULL) AND (validation_error_code IS NULL)) OR (((connection_status)::text = 'connected'::text) AND ((validation_status)::text = ANY ((ARRAY['unverified'::character varying, 'valid'::character varying, 'invalid'::character varying, 'unavailable'::character varying])::text[])) AND (encrypted_credentials IS NOT NULL) AND (encryption_key_version IS NOT NULL) AND (credential_format_version IS NOT NULL)) OR (((connection_status)::text = 'reauthorization_required'::text) AND ((validation_status)::text = 'unverified'::text) AND (encrypted_credentials IS NOT NULL) AND (encryption_key_version IS NOT NULL) AND (credential_format_version IS NOT NULL)))) ); @@ -780,7 +782,7 @@ ALTER TABLE ONLY public.users_tokens -- PostgreSQL database dump complete -- -\unrestrict VoWETYGZi4vEwyliUNWQG3BVneag5OunvKzcS2NOBFzVtuTzyBWpDfhYIneu2ek +\unrestrict cMSeT1HADB9gn2fd6oRSfRA6k31lS6KseACcNAs16bTriITaRrwptynxWKofvCj INSERT INTO public."schema_migrations" (version) VALUES (20260807073017); INSERT INTO public."schema_migrations" (version) VALUES (20260808062115); @@ -800,3 +802,4 @@ INSERT INTO public."schema_migrations" (version) VALUES (20260904051620); INSERT INTO public."schema_migrations" (version) VALUES (20260907083117); INSERT INTO public."schema_migrations" (version) VALUES (20260907161929); INSERT INTO public."schema_migrations" (version) VALUES (20260910044734); +INSERT INTO public."schema_migrations" (version) VALUES (20260912195208); diff --git a/test/kodo/integrations/oauth_refresh_test.exs b/test/kodo/integrations/oauth_refresh_test.exs index ba2aff2..f57cf40 100644 --- a/test/kodo/integrations/oauth_refresh_test.exs +++ b/test/kodo/integrations/oauth_refresh_test.exs @@ -164,7 +164,44 @@ defmodule Kodo.Integrations.OAuthRefreshTest do assert persisted["access_token"] == second_tokens["access_token"] end - test "user deletion cascades OAuth state and fences an in-flight refresh", context do + test "persists a rotated response after its caller stops waiting", context do + owner = self() + + blocked = fn _refresh_token -> + send(owner, {:rotated_response_ready, self()}) + + receive do + :persist_rotated_response -> {:ok, rotated_tokens("account", "detached")} + end + end + + client = fake_client([blocked]) + + caller = + Task.async(fn -> + OAuthRefresh.ensure_fresh(context.scope, context.integration, + client: FakeRefreshClient, + client_options: [agent: client], + supervisor: context.supervisor, + claim_owner_id: Ecto.UUID.generate(), + wait_timeout: 1 + ) + end) + + assert_receive {:rotated_response_ready, worker} + worker_ref = Process.monitor(worker) + assert {:error, :provider_unavailable} = Task.await(caller) + send(worker, :persist_rotated_response) + assert_receive {:DOWN, ^worker_ref, :process, ^worker, :normal} + + refreshed = Repo.reload!(context.integration) + assert refreshed.refresh_source_generation == context.integration.credential_generation + assert refreshed.credential_generation == context.integration.credential_generation + 1 + assert {:ok, credentials} = CredentialEncryption.decrypt(refreshed) + assert credentials["refresh_token"] == "detached-refresh" + end + + test "an active reauthorization attempt blocks refresh without using its credential", context do assert {:ok, _attempt} = Integrations.begin_device_authorization( context.scope, @@ -174,7 +211,25 @@ defmodule Kodo.Integrations.OAuthRefreshTest do 0 ) - integration = Repo.reload!(context.integration) + reauthorizing = Repo.reload!(context.integration) + client = fake_client([]) + + assert {:error, :integration_reauthorization_required} = + OAuthRefresh.ensure_fresh(context.scope, reauthorizing, + force: true, + client: FakeRefreshClient, + client_options: [agent: client], + supervisor: context.supervisor, + claim_owner_id: Ecto.UUID.generate() + ) + + assert Agent.get(client, & &1.refresh_tokens) == [] + + assert Repo.reload!(reauthorizing).credential_generation == + reauthorizing.credential_generation + end + + test "user deletion cascades OAuth state and fences an in-flight refresh", context do owner = self() blocked = fn _refresh_token -> @@ -186,10 +241,20 @@ defmodule Kodo.Integrations.OAuthRefreshTest do end client = fake_client([blocked]) - refresh = Task.async(fn -> refresh(%{context | integration: integration}, client) end) + refresh = Task.async(fn -> refresh(context, client) end) assert_receive {:deletion_refresh_started, worker} - assert Repo.reload!(integration).refresh_claim_owner_id + assert Repo.reload!(context.integration).refresh_claim_owner_id + + assert {:ok, _attempt} = + Integrations.begin_device_authorization( + context.scope, + context.integration.id, + context.integration.credential_generation, + %{"device_auth_id" => "device", "user_code" => "CODE"}, + 0 + ) + Repo.delete!(context.scope.user) refute Repo.exists?(Integration) refute Repo.exists?(DeviceAuthorizationAttempt) From ea0974774f9d81379aefd48bfe5279e0d4be6dda Mon Sep 17 00:00:00 2001 From: Darwin Wu Date: Sat, 12 Sep 2026 20:01:36 +0000 Subject: [PATCH 03/14] fix(integrations): recover device authorization workers --- config/config.exs | 1 + lib/kodo/integrations.ex | 23 +++++- lib/kodo/integrations/device_authorization.ex | 26 +++--- lib/kodo_web/live/integrations_live.ex | 60 ++++++++++++-- .../device_authorization_test.exs | 24 ++++++ test/kodo_web/live/integrations_live_test.exs | 82 ++++++++++++++++++- 6 files changed, 190 insertions(+), 26 deletions(-) diff --git a/config/config.exs b/config/config.exs index 374931e..dc06f2c 100644 --- a/config/config.exs +++ b/config/config.exs @@ -48,6 +48,7 @@ config :phoenix, :filter_parameters, [ "refresh_token", "id_token", "device_code", + "device_authorization_attempt_id", "device_auth_id", "user_code", "code_verifier", diff --git a/lib/kodo/integrations.ex b/lib/kodo/integrations.ex index 5b5c6b0..e448531 100644 --- a/lib/kodo/integrations.ex +++ b/lib/kodo/integrations.ex @@ -157,6 +157,7 @@ defmodule Kodo.Integrations do polling_interval_ms ) end) + |> notify_device_authorization_change() end end @@ -279,6 +280,7 @@ defmodule Kodo.Integrations do audit!(user.id, integration, "device_authorization_failed") attempt end) + |> notify_device_authorization_change() end def fail_device_authorization(%Scope{}, %DeviceAuthorizationAttempt{}, _error_code), @@ -290,6 +292,7 @@ defmodule Kodo.Integrations do Repo.transaction(fn -> cancel_device_authorization_locked(user.id, attempt_id, attempt_generation) end) + |> notify_device_authorization_change() end end @@ -752,7 +755,7 @@ defmodule Kodo.Integrations do lease_expires_at = DateTime.add(timestamp, @device_authorization_claim_lease_ms, :millisecond) query = - claimable_device_authorization_query(user_id, integration_id, claim_owner_id, timestamp) + claimable_device_authorization_query(user_id, integration_id, timestamp) |> select([attempt, _integration], attempt) case Repo.update_all( @@ -772,7 +775,7 @@ defmodule Kodo.Integrations do end end - defp claimable_device_authorization_query(user_id, integration_id, claim_owner_id, timestamp) do + defp claimable_device_authorization_query(user_id, integration_id, timestamp) do from attempt in DeviceAuthorizationAttempt, join: integration in Integration, on: integration.id == attempt.integration_id and integration.user_id == attempt.user_id, @@ -780,8 +783,7 @@ defmodule Kodo.Integrations do attempt.user_id == ^user_id and attempt.integration_id == ^integration_id and attempt.state == "active" and attempt.provider_deadline > ^timestamp and integration.credential_generation == attempt.expected_integration_generation and - (is_nil(attempt.claim_owner_id) or attempt.claim_lease_expires_at <= ^timestamp or - attempt.claim_owner_id == ^claim_owner_id) + (is_nil(attempt.claim_owner_id) or attempt.claim_lease_expires_at <= ^timestamp) end # The boolean terms below are one atomic database fence, not control-flow branches. @@ -1567,6 +1569,19 @@ defmodule Kodo.Integrations do defp notify_integration_change(error), do: error + defp notify_device_authorization_change({:ok, attempt} = result) do + Phoenix.PubSub.broadcast( + Kodo.PubSub, + "integration:#{attempt.user_id}", + {:device_authorization_changed, attempt.integration_id, attempt.attempt_generation, + attempt.state, attempt.terminal_error_code} + ) + + result + end + + defp notify_device_authorization_change(error), do: error + defp constraint_error?(changeset, type) do Enum.any?(changeset.errors, fn {_field, {_message, metadata}} -> metadata[:constraint] == type diff --git a/lib/kodo/integrations/device_authorization.ex b/lib/kodo/integrations/device_authorization.ex index ab3722b..d56695a 100644 --- a/lib/kodo/integrations/device_authorization.ex +++ b/lib/kodo/integrations/device_authorization.ex @@ -2,7 +2,6 @@ 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 @@ -34,21 +33,22 @@ defmodule Kodo.Integrations.DeviceAuthorization do end def resume(%Scope{} = scope, integration_id, opts \\ []) do - claim_owner_id = Keyword.get_lazy(opts, :claim_owner_id, &InstanceManager.current_boot_id/0) + # A claim identifies one finite worker, not its VM. Reusing a node boot ID + # would let a second browser tab replace a healthy worker on the same node. + claim_owner_id = Keyword.get_lazy(opts, :claim_owner_id, &Ecto.UUID.generate/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) + case Integrations.claim_device_authorization(scope, integration_id, claim_owner_id) do + {:ok, {claim, payload}} -> + 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 + {:ok, task} + + {:error, _reason} = error -> + error end end diff --git a/lib/kodo_web/live/integrations_live.ex b/lib/kodo_web/live/integrations_live.ex index 06de0cb..ac3b6cc 100644 --- a/lib/kodo_web/live/integrations_live.ex +++ b/lib/kodo_web/live/integrations_live.ex @@ -44,6 +44,7 @@ defmodule KodoWeb.IntegrationsLive do @providers Enum.map(@provider_configs, & &1.id) @actions ~w(connect replace reauthorize disconnect) @max_api_key_bytes 4_096 + @device_authorization_retry_ms 30_100 @impl true def mount(_params, _session, socket) do @@ -65,7 +66,14 @@ defmodule KodoWeb.IntegrationsLive do |> assign(:device_authorization_tasks, %{}) |> load_integrations() - socket = if connected?(socket), do: resume_device_authorizations(socket), else: socket + socket = + if connected?(socket) do + {:ok, _deleted_count} = Integrations.cleanup_device_authorizations() + resume_device_authorizations(socket) + else + socket + end + {:ok, socket} end @@ -251,7 +259,7 @@ defmodule KodoWeb.IntegrationsLive do def handle_event( "cancel_device_authorization", - %{"attempt" => attempt_id, "generation" => generation}, + %{"device_authorization_attempt_id" => attempt_id, "generation" => generation}, socket ) do with {attempt_generation, ""} <- Integer.parse(generation), @@ -352,6 +360,23 @@ defmodule KodoWeb.IntegrationsLive do {:noreply, load_integrations(socket)} end + def handle_info( + {:device_authorization_changed, _integration_id, _attempt_generation, state, _error}, + socket + ) do + socket = load_integrations(socket) + + if state == "failed" do + {:noreply, put_flash(socket, :error, "Authorization failed. Try again.")} + else + {:noreply, socket} + end + end + + def handle_info({:retry_device_authorization, integration_id}, socket) do + {:noreply, resume_device_authorization(socket, integration_id, false)} + end + defp save_api_key( %{assigns: %{action: "connect", action_target: :new}} = socket, %{"api_key" => api_key} = params @@ -479,15 +504,34 @@ defmodule KodoWeb.IntegrationsLive do end defp resume_device_authorizations(socket) do - Enum.reduce(socket.assigns.device_authorizations, socket, fn {integration_id, _attempt}, - acc -> - case DeviceAuthorization.resume(acc.assigns.current_scope, integration_id) do - {:ok, task} -> track_device_authorization_task(acc, task, integration_id) - {:error, _reason} -> acc + Enum.reduce(socket.assigns.integrations, socket, fn integration, acc -> + if integration.provider == "openai_codex" do + resume_device_authorization(acc, integration.id, true) + else + acc end end) end + defp resume_device_authorization(socket, integration_id, schedule_retry?) do + case DeviceAuthorization.resume(socket.assigns.current_scope, integration_id) do + {:ok, task} -> + track_device_authorization_task(socket, task, integration_id) + + {:error, :device_authorization_not_claimable} when schedule_retry? -> + Process.send_after( + self(), + {:retry_device_authorization, integration_id}, + @device_authorization_retry_ms + ) + + socket + + {:error, _reason} -> + socket + end + end + defp load_device_authorizations(scope, integrations) do integrations |> Enum.filter(&(&1.provider == "openai_codex")) @@ -871,7 +915,7 @@ defmodule KodoWeb.IntegrationsLive do id={dom_id(integration, "cancel-device-authorization")} type="button" phx-click="cancel_device_authorization" - phx-value-attempt={authorization.attempt_id} + phx-value-device_authorization_attempt_id={authorization.attempt_id} phx-value-generation={authorization.attempt_generation} phx-disable-with="Cancelling…" class="rounded-lg px-3 py-1.5 text-xs font-semibold text-red-700 transition hover:bg-red-100 disabled:opacity-60 dark:text-red-400 dark:hover:bg-red-950/40" diff --git a/test/kodo/integrations/device_authorization_test.exs b/test/kodo/integrations/device_authorization_test.exs index 7003fd9..452f3f0 100644 --- a/test/kodo/integrations/device_authorization_test.exs +++ b/test/kodo/integrations/device_authorization_test.exs @@ -118,6 +118,30 @@ defmodule Kodo.Integrations.DeviceAuthorizationTest do assert calls(client) == [] end + test "a healthy claim cannot be replaced by another task using the same node owner", %{ + scope: scope, + integration: integration + } do + owner = Ecto.UUID.generate() + + assert {:ok, _attempt} = + Integrations.begin_device_authorization( + scope, + integration.id, + integration.credential_generation, + %{"device_auth_id" => "device", "user_code" => "CODE"}, + 0 + ) + + assert {:ok, {claim, _payload}} = + Integrations.claim_device_authorization(scope, integration.id, owner) + + assert {:error, :device_authorization_not_claimable} = + Integrations.claim_device_authorization(scope, integration.id, owner) + + assert Repo.reload!(claim).claim_epoch == claim.claim_epoch + end + test "a poll admitted before takeover may finish but cannot persist its result", %{ scope: scope, integration: integration diff --git a/test/kodo_web/live/integrations_live_test.exs b/test/kodo_web/live/integrations_live_test.exs index 7248721..872d665 100644 --- a/test/kodo_web/live/integrations_live_test.exs +++ b/test/kodo_web/live/integrations_live_test.exs @@ -163,7 +163,7 @@ defmodule KodoWeb.IntegrationsLiveTest do refute inspect(:sys.get_state(view.pid)) =~ "private-device" render_click(view, "cancel_device_authorization", %{ - "attempt" => Ecto.UUID.generate(), + "device_authorization_attempt_id" => Ecto.UUID.generate(), "generation" => "1" }) @@ -254,6 +254,86 @@ defmodule KodoWeb.IntegrationsLiveTest do assert Repo.reload!(attempt).state == "expired" end + test "resumes an exchange-stage authorization after its worker lease expires", %{ + conn: conn, + scope: scope + } do + _client = configure_device_client(exchange: [{:ok, device_tokens()}]) + + assert {:ok, integration} = + Integrations.create_oauth_integration(scope, "openai_codex", + display_name: "Recovering subscription" + ) + + assert {:ok, _attempt} = + Integrations.begin_device_authorization( + scope, + integration.id, + integration.credential_generation, + %{"device_auth_id" => "device", "user_code" => "CODE"}, + 0 + ) + + assert {:ok, {claim, _payload}} = + Integrations.claim_device_authorization(scope, integration.id, Ecto.UUID.generate()) + + assert {:ok, persisted} = + Integrations.store_device_authorization_exchange(scope, claim, %{ + "authorization_code" => "authorization", + "code_challenge" => "challenge", + "code_verifier" => "verifier" + }) + + Repo.update!( + Ecto.Changeset.change(persisted, + claim_lease_expires_at: DateTime.add(DateTime.utc_now(), -1, :second) + ) + ) + + {:ok, view, _html} = live(conn, ~p"/integrations") + monitor = Process.monitor(view.pid) + _ = :sys.get_state(view.pid) + + assert {:ok, connected} = Integrations.get_integration(scope, integration.id) + assert connected.connection_status == "connected" + refute has_element?(view, "#integration-#{integration.id}-device-authorization") + Process.demonitor(monitor, [:flush]) + end + + test "cancelling authorization removes its code from another open tab", %{ + conn: conn, + scope: scope, + user: user + } do + assert {:ok, integration} = + Integrations.create_oauth_integration(scope, "openai_codex", + display_name: "Shared subscription" + ) + + assert {:ok, attempt} = + Integrations.begin_device_authorization( + scope, + integration.id, + integration.credential_generation, + %{"device_auth_id" => "device", "user_code" => "SHARED"}, + 30_000 + ) + + {:ok, first_view, _html} = live(conn, ~p"/integrations") + {:ok, second_view, _html} = live(log_in_user(build_conn(), user), ~p"/integrations") + panel = "#integration-#{integration.id}-device-authorization" + assert has_element?(first_view, panel, "SHARED") + assert has_element?(second_view, panel, "SHARED") + + first_view + |> element("#integration-#{integration.id}-cancel-device-authorization") + |> render_click() + + _ = :sys.get_state(second_view.pid) + refute has_element?(second_view, panel) + assert Repo.reload!(attempt).state == "cancelled" + end + test "shows completion and reauthorization after the device flow succeeds", %{ conn: conn, scope: scope From 01ed0b3ad429cc3c15dda8c3742a8a1cbd0ea7dc Mon Sep 17 00:00:00 2001 From: Darwin Wu Date: Sat, 12 Sep 2026 20:04:42 +0000 Subject: [PATCH 04/14] fix(integrations): preserve OAuth account lifecycle intent --- lib/kodo/integrations.ex | 23 ++++---- .../device_authorization_attempt.ex | 3 ++ lib/kodo_web/live/integrations_live.ex | 25 ++++++--- ...etion_to_device_authorization_attempts.exs | 9 ++++ priv/repo/structure.sql | 6 ++- test/kodo/integrations_test.exs | 54 +++++++++++++++++++ test/kodo_web/live/integrations_live_test.exs | 31 +++++++++++ 7 files changed, 131 insertions(+), 20 deletions(-) create mode 100644 priv/repo/migrations/20260912200217_add_activate_on_completion_to_device_authorization_attempts.exs diff --git a/lib/kodo/integrations.ex b/lib/kodo/integrations.ex index e448531..44c8325 100644 --- a/lib/kodo/integrations.ex +++ b/lib/kodo/integrations.ex @@ -932,7 +932,13 @@ defmodule Kodo.Integrations do }) ) |> Repo.update!() - |> maybe_activate_completed_oauth(user_id) + + integration = + if attempt.activate_on_completion do + integration |> change(active: true, updated_at: timestamp) |> Repo.update!() + else + integration + end attempt |> change(terminal_attempt_changes("completed", nil, timestamp)) @@ -983,14 +989,6 @@ defmodule Kodo.Integrations do end end - defp maybe_activate_completed_oauth(integration, user_id) do - if integration.active or active_provider_account_exists?(user_id, integration.provider) do - integration - else - integration |> change(active: true, updated_at: now()) |> Repo.update!() - end - end - defp terminalize_claimed_device_authorization(user_id, claim, state, error_code) do timestamp = database_now() @@ -1060,6 +1058,9 @@ defmodule Kodo.Integrations do timestamp = database_now() expected_integration_generation = integration.credential_generation + 1 + activate_on_completion = + integration.active or !active_provider_account_exists?(user_id, integration.provider) + supersede_active_device_authorization(integration.id, timestamp) integration = @@ -1079,7 +1080,8 @@ defmodule Kodo.Integrations do integration_id: integration.id, provider: "openai_codex", attempt_generation: next_device_authorization_generation(integration.id), - expected_integration_generation: expected_integration_generation + expected_integration_generation: expected_integration_generation, + activate_on_completion: activate_on_completion } case DeviceAuthorizationEncryption.encrypt(attempt, payload) do @@ -1094,6 +1096,7 @@ defmodule Kodo.Integrations do provider: "openai_codex", attempt_generation: attempt.attempt_generation, expected_integration_generation: expected_integration_generation, + activate_on_completion: activate_on_completion, provider_deadline: deadline, polling_interval_ms: polling_interval_ms, next_poll_at: next_poll_at diff --git a/lib/kodo/integrations/device_authorization_attempt.ex b/lib/kodo/integrations/device_authorization_attempt.ex index 409b1ee..0808fd4 100644 --- a/lib/kodo/integrations/device_authorization_attempt.ex +++ b/lib/kodo/integrations/device_authorization_attempt.ex @@ -15,6 +15,7 @@ defmodule Kodo.Integrations.DeviceAuthorizationAttempt do field :state, :string, default: "active" field :attempt_generation, :integer field :expected_integration_generation, :integer + field :activate_on_completion, :boolean, default: false field :encrypted_payload, :binary, redact: true field :encryption_key_version, :string field :payload_format_version, :integer @@ -39,6 +40,7 @@ defmodule Kodo.Integrations.DeviceAuthorizationAttempt do :provider, :attempt_generation, :expected_integration_generation, + :activate_on_completion, :provider_deadline, :polling_interval_ms, :next_poll_at @@ -51,6 +53,7 @@ defmodule Kodo.Integrations.DeviceAuthorizationAttempt do :state, :attempt_generation, :expected_integration_generation, + :activate_on_completion, :encrypted_payload, :encryption_key_version, :payload_format_version, diff --git a/lib/kodo_web/live/integrations_live.ex b/lib/kodo_web/live/integrations_live.ex index ac3b6cc..47a7521 100644 --- a/lib/kodo_web/live/integrations_live.ex +++ b/lib/kodo_web/live/integrations_live.ex @@ -308,8 +308,9 @@ defmodule KodoWeb.IntegrationsLive do def handle_event("disconnect", %{"modal-token" => modal_token}, socket) do with true <- socket.assigns.action == "disconnect", true <- modal_token == socket.assigns.modal_token, - %{id: id, credential_generation: generation, connection_status: "connected"} <- + %{id: id, credential_generation: generation, connection_status: connection_status} <- socket.assigns.action_target, + true <- connection_status in ~w(connected reauthorization_required), {:ok, integration} <- Integrations.disconnect(socket.assigns.current_scope, id, generation) do {:noreply, @@ -601,8 +602,12 @@ defmodule KodoWeb.IntegrationsLive do defp action_target( "reauthorize", - %{provider: "openai_codex", connection_status: "connected"} = integration - ), + %{ + provider: "openai_codex", + connection_status: connection_status + } = integration + ) + when connection_status in ~w(connected reauthorization_required), do: {:ok, target_metadata(integration)} defp action_target( @@ -611,8 +616,9 @@ defmodule KodoWeb.IntegrationsLive do ), do: {:ok, target_metadata(integration)} - defp action_target("disconnect", %{connection_status: "connected"} = integration), - do: {:ok, target_metadata(integration)} + defp action_target("disconnect", %{connection_status: connection_status} = integration) + when connection_status in ~w(connected reauthorization_required), + do: {:ok, target_metadata(integration)} defp action_target(_action, _integration), do: :error @@ -929,7 +935,7 @@ defmodule KodoWeb.IntegrationsLive do
<.link - :if={!integration_connected?(integration)} + :if={integration.connection_status == "disconnected"} id={dom_id(integration, "reconnect")} patch={action_path(integration, "connect")} phx-click={JS.push_focus()} @@ -986,7 +992,10 @@ defmodule KodoWeb.IntegrationsLive do Replace key <.link - :if={integration_connected?(integration) and integration.provider == "openai_codex"} + :if={ + integration.provider == "openai_codex" and + integration.connection_status in ~w(connected reauthorization_required) + } id={dom_id(integration, "reauthorize")} patch={action_path(integration, "reauthorize")} phx-click={JS.push_focus()} @@ -995,7 +1004,7 @@ defmodule KodoWeb.IntegrationsLive do Reauthorize <.link - :if={integration_connected?(integration)} + :if={integration.connection_status in ~w(connected reauthorization_required)} id={dom_id(integration, "disconnect")} patch={action_path(integration, "disconnect")} phx-click={JS.push_focus()} diff --git a/priv/repo/migrations/20260912200217_add_activate_on_completion_to_device_authorization_attempts.exs b/priv/repo/migrations/20260912200217_add_activate_on_completion_to_device_authorization_attempts.exs new file mode 100644 index 0000000..28f85d9 --- /dev/null +++ b/priv/repo/migrations/20260912200217_add_activate_on_completion_to_device_authorization_attempts.exs @@ -0,0 +1,9 @@ +defmodule Kodo.Repo.Migrations.AddActivateOnCompletionToDeviceAuthorizationAttempts do + use Ecto.Migration + + def change do + alter table(:device_authorization_attempts) do + add :activate_on_completion, :boolean, null: false, default: false + end + end +end diff --git a/priv/repo/structure.sql b/priv/repo/structure.sql index cd098cc..140e388 100644 --- a/priv/repo/structure.sql +++ b/priv/repo/structure.sql @@ -2,7 +2,7 @@ -- PostgreSQL database dump -- -\restrict cMSeT1HADB9gn2fd6oRSfRA6k31lS6KseACcNAs16bTriITaRrwptynxWKofvCj +\restrict zlmayUQMxfTEc8sR5OS893sgwJPRmRbQbN8OxqjM51G1hsJJMUZbtQRN809vFYQ -- Dumped from database version 15.19 (Debian 15.19-0+deb12u1) -- Dumped by pg_dump version 15.19 (Debian 15.19-0+deb12u1) @@ -135,6 +135,7 @@ CREATE TABLE public.device_authorization_attempts ( terminal_error_code character varying(64), inserted_at timestamp without time zone NOT NULL, updated_at timestamp without time zone NOT NULL, + activate_on_completion boolean DEFAULT false NOT NULL, CONSTRAINT device_attempts_claim_epoch_valid CHECK ((claim_epoch >= 0)), CONSTRAINT device_attempts_claim_valid CHECK (((((state)::text = 'active'::text) AND (terminal_error_code IS NULL) AND (((claim_owner_id IS NULL) AND (claim_lease_expires_at IS NULL)) OR ((claim_owner_id IS NOT NULL) AND (claim_lease_expires_at IS NOT NULL)))) OR (((state)::text <> 'active'::text) AND (claim_owner_id IS NULL) AND (claim_lease_expires_at IS NULL)))), CONSTRAINT device_attempts_generations_valid CHECK (((attempt_generation > 0) AND (expected_integration_generation >= 0))), @@ -782,7 +783,7 @@ ALTER TABLE ONLY public.users_tokens -- PostgreSQL database dump complete -- -\unrestrict cMSeT1HADB9gn2fd6oRSfRA6k31lS6KseACcNAs16bTriITaRrwptynxWKofvCj +\unrestrict zlmayUQMxfTEc8sR5OS893sgwJPRmRbQbN8OxqjM51G1hsJJMUZbtQRN809vFYQ INSERT INTO public."schema_migrations" (version) VALUES (20260807073017); INSERT INTO public."schema_migrations" (version) VALUES (20260808062115); @@ -803,3 +804,4 @@ INSERT INTO public."schema_migrations" (version) VALUES (20260907083117); INSERT INTO public."schema_migrations" (version) VALUES (20260907161929); INSERT INTO public."schema_migrations" (version) VALUES (20260910044734); INSERT INTO public."schema_migrations" (version) VALUES (20260912195208); +INSERT INTO public."schema_migrations" (version) VALUES (20260912200217); diff --git a/test/kodo/integrations_test.exs b/test/kodo/integrations_test.exs index bb9022b..5063d8a 100644 --- a/test/kodo/integrations_test.exs +++ b/test/kodo/integrations_test.exs @@ -856,6 +856,51 @@ defmodule Kodo.IntegrationsTest do Repo.reload!(valid) end + test "a later authorization completion preserves its inactive account decision", %{ + scope: scope + } do + first = oauth_integration(scope) + + assert {:ok, active} = + Integrations.oauth_succeeded(scope, first.id, 0, %{"access_token" => "first"}) + + second = oauth_integration(scope) + + assert {:ok, _attempt} = + Integrations.begin_device_authorization( + scope, + second.id, + second.credential_generation, + %{"device_auth_id" => "second", "user_code" => "SECOND"}, + 0 + ) + + assert {:ok, {claim, _payload}} = + Integrations.claim_device_authorization( + scope, + second.id, + Ecto.UUID.generate() + ) + + assert {:ok, _disconnected} = + Integrations.disconnect(scope, active.id, active.credential_generation) + + assert {:ok, completed} = + Integrations.complete_device_authorization( + scope, + claim, + %{ + "access_token" => "second-access", + "refresh_token" => "second-refresh", + "id_token" => "second-id", + "account_id" => "second-account" + }, + DateTime.add(DateTime.utc_now(), 3_600, :second) + ) + + refute completed.active + end + test "attempt operations suppress identifiers and ciphertext from query observability", %{ scope: scope } do @@ -981,6 +1026,15 @@ defmodule Kodo.IntegrationsTest do assert {:error, :device_authorization_not_claimable} = Integrations.claim_device_authorization(scope, integration.id, second_owner) + assert {:error, :device_authorization_not_claimable} = + Integrations.claim_device_authorization(scope, integration.id, first_owner) + + Repo.update!( + change(first_claim, + claim_lease_expires_at: DateTime.add(DateTime.utc_now(), -1, :second) + ) + ) + assert {:ok, {replacement_claim, ^payload}} = Integrations.claim_device_authorization(scope, integration.id, first_owner) diff --git a/test/kodo_web/live/integrations_live_test.exs b/test/kodo_web/live/integrations_live_test.exs index 872d665..21aaf5d 100644 --- a/test/kodo_web/live/integrations_live_test.exs +++ b/test/kodo_web/live/integrations_live_test.exs @@ -392,6 +392,37 @@ defmodule KodoWeb.IntegrationsLiveTest do assert has_element?(view, "#{status(integration)} dd.text-green-700", "Valid") end + test "offers recovery and disconnect actions when ChatGPT authorization is required", %{ + conn: conn, + scope: scope + } do + assert {:ok, integration} = + Integrations.create_oauth_integration(scope, "openai_codex", + display_name: "Expired subscription" + ) + + assert {:ok, connected} = + Integrations.oauth_succeeded(scope, integration.id, 0, %{ + "access_token" => "access", + "refresh_token" => "refresh" + }) + + assert {:ok, required} = + Integrations.refresh_invalid_grant( + scope, + connected.id, + connected.credential_generation + ) + + {:ok, view, _html} = live(conn, ~p"/integrations") + assert has_element?(view, "#integration-#{required.id}-reauthorize", "Reauthorize") + assert has_element?(view, "#integration-#{required.id}-disconnect", "Disconnect") + refute has_element?(view, "#integration-#{required.id}-reconnect") + + view |> element("#integration-#{required.id}-reauthorize") |> render_click() + assert has_element?(view, "#device-authorization-form") + end + test "adds multiple accounts for one provider without exposing secrets", %{ conn: conn, scope: scope From 052bd997f4a9fa49760e80378f554f02cc9bcc79 Mon Sep 17 00:00:00 2001 From: Darwin Wu Date: Sat, 12 Sep 2026 20:10:35 +0000 Subject: [PATCH 05/14] fix(agent): enforce immutable turn capability snapshots --- lib/kodo/agent/execution_route_change.ex | 41 +++++++--- lib/kodo/agent/loop.ex | 52 +++++++------ lib/kodo/agent/model_capabilities.ex | 36 ++++++++- lib/kodo/agent/model_mapping.ex | 77 +++++++++++++++++++ lib/kodo/sessions/projection.ex | 6 +- .../agent/execution_route_change_test.exs | 41 +++++++++- test/kodo/agent/loop_test.exs | 46 ++++++++++- test/kodo/agent/model_mapping_test.exs | 51 ++++++++++++ test/kodo/sessions_test.exs | 10 ++- 9 files changed, 320 insertions(+), 40 deletions(-) diff --git a/lib/kodo/agent/execution_route_change.ex b/lib/kodo/agent/execution_route_change.ex index a24e671..e4002d6 100644 --- a/lib/kodo/agent/execution_route_change.ex +++ b/lib/kodo/agent/execution_route_change.ex @@ -12,15 +12,9 @@ defmodule Kodo.Agent.ExecutionRouteChange do def change(mapping, destination, records \\ @approved_records) def change(mapping, destination, records) when destination in @routes and is_list(records) do - mapping = ModelMapping.snapshot(mapping) - - affected = - Enum.filter(mapping["roles"], fn {_role, role_mapping} -> - role_mapping["execution_route"] in @routes and - role_mapping["execution_route"] != destination - end) - - with false <- affected == [], + with {:ok, mapping} <- change_mapping(mapping), + affected = affected_roles(mapping, destination), + false <- affected == [], [] <- incompatible_roles(affected, destination, records) do {:ok, apply_route(mapping, affected, destination)} else @@ -31,6 +25,33 @@ defmodule Kodo.Agent.ExecutionRouteChange do def change(_mapping, _destination, _records), do: {:error, :unsupported_execution_route} + defp change_mapping(mapping) do + case ModelMapping.validate_snapshot(mapping) do + {:ok, snapshot} -> + {:ok, snapshot} + + {:error, :invalid_model_mapping_snapshot} -> + if explicit_snapshot?(mapping), + do: {:error, :invalid_model_mapping_snapshot}, + else: {:ok, ModelMapping.snapshot(mapping)} + end + end + + defp explicit_snapshot?(%{"roles" => roles}) when is_map(roles) do + Enum.any?(roles, fn {_role, mapping} -> + is_map(mapping) and is_map_key(mapping, "capability_contract") + end) + end + + defp explicit_snapshot?(_mapping), do: false + + defp affected_roles(mapping, destination) do + Enum.filter(mapping["roles"], fn {_role, role_mapping} -> + role_mapping["execution_route"] in @routes and + role_mapping["execution_route"] != destination + end) + end + defp incompatible_roles(affected, destination, records) do for {role, role_mapping} <- affected, not compatible?(role, role_mapping, destination, records), @@ -44,6 +65,8 @@ defmodule Kodo.Agent.ExecutionRouteChange do record.destination == destination and record.selector == role_mapping["model_selector"] and record.role_contract == role_mapping["role_contract"] and + record.reasoning == role_mapping["reasoning"] and + record.capability_contract == role_mapping["capability_contract"] and record.status == :approved end) end diff --git a/lib/kodo/agent/loop.ex b/lib/kodo/agent/loop.ex index 7e89861..dfde343 100644 --- a/lib/kodo/agent/loop.ex +++ b/lib/kodo/agent/loop.ex @@ -33,29 +33,21 @@ defmodule Kodo.Agent.Loop do latest_response = List.last(responses) - case rehoming_boundary() do - :ok -> - case within_budget(invocations, tokens, budgets) do - :ok -> - action = next_action(latest_invocation, latest_response) - - resume_action( - action, - session_id, - projection, - events, - adapter, - budgets, - invocations, - ownership - ) - - error -> - error - end + with :ok <- validate_turn_mapping(projection), + :ok <- rehoming_boundary(), + :ok <- within_budget(invocations, tokens, budgets) do + action = next_action(latest_invocation, latest_response) - error -> - error + resume_action( + action, + session_id, + projection, + events, + adapter, + budgets, + invocations, + ownership + ) end end @@ -459,7 +451,21 @@ defmodule Kodo.Agent.Loop do end defp turn_mapping(projection) do - projection.turn_model_mapping || projection.model_mapping || legacy_mapping(projection.model) + case {projection.turn_route_snapshot_present, projection.turn_model_mapping} do + {false, _mapping} -> + mapping = projection.model_mapping || legacy_mapping(projection.model) + ModelMapping.snapshot(mapping) + + {true, snapshot} -> + snapshot + end + end + + defp validate_turn_mapping(projection) do + case ModelMapping.validate_snapshot(turn_mapping(projection)) do + {:ok, _mapping} -> :ok + {:error, _reason} = error -> error + end end defp infer(session_id, projection, events, adapter, budgets, invocation, ownership) do diff --git a/lib/kodo/agent/model_capabilities.ex b/lib/kodo/agent/model_capabilities.ex index adce20e..66c85f2 100644 --- a/lib/kodo/agent/model_capabilities.ex +++ b/lib/kodo/agent/model_capabilities.ex @@ -4,7 +4,7 @@ defmodule Kodo.Agent.ModelCapabilities do alias ReqLLM.ModelHelpers def validate(model_spec, role_mapping, contract) do - required = requirements(contract) + required = requirements(role_mapping, contract) with {:ok, model} <- resolve(model_spec), [] <- missing_capabilities(model, role_mapping, required) do @@ -47,9 +47,14 @@ defmodule Kodo.Agent.ModelCapabilities do defp require_capability(missing, false, _capability), do: missing defp require_modalities(missing, model, required) do - supported = model.modalities |> then(&(&1 && &1.input)) |> List.wrap() + supported = + model.modalities + |> then(&(&1 && &1.input)) + |> List.wrap() + |> Enum.map(&to_string/1) Enum.reduce(required, missing, fn modality, missing -> + modality = to_string(modality) require_capability(missing, modality not in supported, "input:#{modality}") end) end @@ -83,7 +88,28 @@ defmodule Kodo.Agent.ModelCapabilities do defp context_window(model), do: (model.limits && model.limits.context) || 0 - defp requirements(contract) do + defp requirements( + %{ + "capability_contract" => %{ + "requirements" => %{ + "tools" => tools, + "structured_output" => structured_output, + "min_context" => min_context, + "input_modalities" => input_modalities + } + } + }, + _contract + ) do + %{ + tools: tools, + structured_output: structured_output(structured_output), + min_context: min_context, + input_modalities: input_modalities + } + end + + defp requirements(_role_mapping, contract) do Map.get(contract, :capabilities, %{ tools: true, structured_output: false, @@ -91,4 +117,8 @@ defmodule Kodo.Agent.ModelCapabilities do input_modalities: [:text] }) end + + defp structured_output("json_schema"), do: :json_schema + defp structured_output("object"), do: :object + defp structured_output(value), do: value end diff --git a/lib/kodo/agent/model_mapping.ex b/lib/kodo/agent/model_mapping.ex index 9c2f9fa..a5f9d69 100644 --- a/lib/kodo/agent/model_mapping.ex +++ b/lib/kodo/agent/model_mapping.ex @@ -78,6 +78,22 @@ defmodule Kodo.Agent.ModelMapping do %{mapping | "roles" => roles} end + @doc "Validates a persisted turn snapshot without repairing it from current defaults." + def validate_snapshot( + %{"profile" => profile, "profile_version" => version, "roles" => roles} = snapshot + ) + when is_binary(profile) and profile != "" and is_binary(version) and version != "" and + is_map(roles) and map_size(roles) == 3 do + if Enum.sort(Map.keys(roles)) == ~w(primary review search) and + Enum.all?(roles, fn {role, mapping} -> valid_snapshot_role?(role, mapping) end) do + {:ok, snapshot} + else + {:error, :invalid_model_mapping_snapshot} + end + end + + def validate_snapshot(_snapshot), do: {:error, :invalid_model_mapping_snapshot} + @doc "Builds the provider-specific request model from an immutable route snapshot." def request_model(%{"execution_route" => route, "model_selector" => selector}) when is_binary(route) and route != "" and is_binary(selector) and selector != "" do @@ -99,6 +115,67 @@ defmodule Kodo.Agent.ModelMapping do } end + defp valid_snapshot_role?( + role, + %{ + "model" => model, + "reasoning" => reasoning, + "provider" => provider, + "execution_route" => route, + "model_selector" => selector, + "role_contract" => role_contract, + "toolset_version" => toolset_version, + "capability_contract" => %{ + "id" => capability_id, + "toolset_version" => capability_toolset, + "requirements" => %{ + "tools" => tools, + "structured_output" => structured_output, + "min_context" => min_context, + "input_modalities" => modalities + } + } + } + ) do + valid_snapshot_identity?(role, [ + model, + reasoning, + provider, + route, + selector, + role_contract, + toolset_version + ]) and + capability_id == role_contract and capability_toolset == toolset_version and + valid_snapshot_route?(model, provider, route, selector) and + valid_snapshot_requirements?(tools, structured_output, min_context, modalities) and + executable_contract?(role, role_contract, toolset_version) + end + + defp valid_snapshot_role?(_role, _mapping), do: false + + defp valid_snapshot_identity?(role, values) do + role in ~w(primary review search) and + Enum.all?(values, &(is_binary(&1) and &1 != "")) + end + + defp valid_snapshot_route?(model, provider, route, selector) do + model == "#{route}:#{selector}" and provider == route + end + + defp valid_snapshot_requirements?(tools, structured_output, min_context, modalities) do + is_boolean(tools) and structured_output in [false, "json_schema", "object"] and + is_integer(min_context) and min_context > 0 and is_list(modalities) and modalities != [] and + Enum.all?(modalities, &(is_binary(&1) and &1 != "")) + end + + defp executable_contract?(role, contract_id, toolset_version) do + contract = Roles.fetch!(role_atom(role), contract_id) + contract.toolset_version == toolset_version + rescue + KeyError -> false + end + defp role_atom("primary"), do: :primary defp role_atom("search"), do: :search defp role_atom("review"), do: :review diff --git a/lib/kodo/sessions/projection.ex b/lib/kodo/sessions/projection.ex index 8d25fb7..8c646bd 100644 --- a/lib/kodo/sessions/projection.ex +++ b/lib/kodo/sessions/projection.ex @@ -13,6 +13,7 @@ defmodule Kodo.Sessions.Projection do :model_mapping, route_revision: 1, turn_route_revision: nil, + turn_route_snapshot_present: false, turn_model_mapping: nil, approval_policy: "standard", status: "idle", @@ -90,7 +91,10 @@ defmodule Kodo.Sessions.Projection do defp reduce(projection, "turn_route_snapshot", payload) do %{ projection - | turn_model_mapping: normalize_model_mapping(payload["model_mapping"]), + | # Turn snapshots are immutable records. Repairing one with today's + # defaults would silently change a replayed request. + turn_route_snapshot_present: true, + turn_model_mapping: payload["model_mapping"], turn_route_revision: payload["route_revision"] } end diff --git a/test/kodo/agent/execution_route_change_test.exs b/test/kodo/agent/execution_route_change_test.exs index c35a952..184bf0b 100644 --- a/test/kodo/agent/execution_route_change_test.exs +++ b/test/kodo/agent/execution_route_change_test.exs @@ -40,6 +40,40 @@ defmodule Kodo.Agent.ExecutionRouteChangeTest do assert ModelMapping.role!(mapping, :primary)["provider"] == "openai" end + test "rejects evidence for a different reasoning or capability envelope" do + mapping = mapping_for_selector("gpt-4o-mini") + records = records_for(mapping, "openai_codex") + + reasoning_mismatch = + Enum.map(records, &Map.put(&1, :reasoning, "medium")) + + assert {:error, {:incompatible_execution_route, ["primary", "review", "search"]}} = + ExecutionRouteChange.change(mapping, "openai_codex", reasoning_mismatch) + + capability_mismatch = + Enum.map(records, fn record -> + put_in(record, [:capability_contract, "requirements", "tools"], false) + end) + + assert {:error, {:incompatible_execution_route, ["primary", "review", "search"]}} = + ExecutionRouteChange.change(mapping, "openai_codex", capability_mismatch) + end + + test "preserves and verifies an explicit mapping capability envelope" do + mapping = mapping_for_selector("gpt-4o-mini") |> ModelMapping.snapshot() + records = records_for(mapping, "openai_codex") + + changed_requirement = + put_in( + mapping, + ["roles", "primary", "capability_contract", "requirements", "min_context"], + 200_000 + ) + + assert {:error, {:incompatible_execution_route, ["primary"]}} = + ExecutionRouteChange.change(changed_requirement, "openai_codex", records) + end + defp mapping_for_selector(selector) do ModelMapping.balanced([ {"test", @@ -52,13 +86,18 @@ defmodule Kodo.Agent.ExecutionRouteChangeTest do end defp records_for(mapping, destination) do - Enum.map(mapping["roles"], fn {role, role_mapping} -> + mapping + |> ModelMapping.snapshot() + |> Map.fetch!("roles") + |> Enum.map(fn {role, role_mapping} -> %{ role: role, source: "openai", destination: destination, selector: role_mapping["model"] |> String.split(":", parts: 2) |> List.last(), role_contract: role_mapping["role_contract"], + reasoning: role_mapping["reasoning"], + capability_contract: role_mapping["capability_contract"], status: :approved } end) diff --git a/test/kodo/agent/loop_test.exs b/test/kodo/agent/loop_test.exs index dee42bf..90dc94a 100644 --- a/test/kodo/agent/loop_test.exs +++ b/test/kodo/agent/loop_test.exs @@ -99,10 +99,14 @@ defmodule Kodo.Agent.LoopTest do [created] = Sessions.events_after(session.id) mapping = - update_in(created.payload["model_mapping"], ["roles"], fn roles -> + created.payload["model_mapping"] + |> Kodo.Agent.ModelMapping.snapshot() + |> update_in(["roles"], fn roles -> Map.new(roles, fn {role, role_mapping} -> {role, role_mapping + |> Map.put("provider", "openai_codex") + |> Map.put("model", "openai_codex:gpt-5.4") |> Map.put("execution_route", "openai_codex") |> Map.put("model_selector", "gpt-5.4")} end) @@ -138,6 +142,46 @@ defmodule Kodo.Agent.LoopTest do refute Enum.any?(Sessions.events_after(session.id), &(&1.type == "model_invocation_started")) end + test "rejects a malformed persisted turn capability envelope before dispatch", %{ + session: session, + ownership: ownership + } do + [created] = Sessions.events_after(session.id) + + malformed = + update_in( + created.payload, + ["model_mapping", "roles", "primary"], + &Map.delete(&1, "capability_contract") + ) + + {:ok, _snapshot} = + Sessions.append_event( + session.id, + "turn_route_snapshot", + %{"route_revision" => 1, "model_mapping" => malformed["model_mapping"]}, + source: "system", + ownership: ownership + ) + + {:ok, _message} = + Sessions.append_event( + session.id, + "user_message", + %{"role" => "user", "content" => "do not repair this snapshot"}, + ownership: ownership + ) + + assert {:error, :invalid_model_mapping_snapshot} = + Loop.run(session.id, + adapter: Kodo.Test.FakeLLM, + budgets: budgets([]), + ownership: ownership + ) + + refute_received {:llm_request, _model, _system, _tools, _opts} + end + test "drops provider-private assistant state when replay crosses execution routes", %{ runner: runner, session: session, diff --git a/test/kodo/agent/model_mapping_test.exs b/test/kodo/agent/model_mapping_test.exs index 8ca5a66..7f28d43 100644 --- a/test/kodo/agent/model_mapping_test.exs +++ b/test/kodo/agent/model_mapping_test.exs @@ -58,4 +58,55 @@ defmodule Kodo.Agent.ModelMappingTest do assert mapping["roles"]["search"]["model"] == "not-a-model" assert mapping["roles"]["search"]["provider"] == nil end + + test "rejects malformed immutable snapshots instead of filling current defaults" do + snapshot = ModelMapping.snapshot(ModelMapping.balanced()) + + assert {:ok, ^snapshot} = ModelMapping.validate_snapshot(snapshot) + + malformed = + update_in(snapshot, ["roles", "primary"], &Map.delete(&1, "capability_contract")) + + assert {:error, :invalid_model_mapping_snapshot} = + ModelMapping.validate_snapshot(malformed) + end + + test "rejects malformed capability requirements instead of accepting a partial envelope" do + snapshot = ModelMapping.snapshot(ModelMapping.balanced()) + + malformed = + put_in( + snapshot, + ["roles", "primary", "capability_contract", "requirements", "min_context"], + "current default" + ) + + assert {:error, :invalid_model_mapping_snapshot} = + ModelMapping.validate_snapshot(malformed) + end + + test "rejects capability variants and toolsets with no executable implementation" do + snapshot = ModelMapping.snapshot(ModelMapping.balanced()) + + unsupported_output = + put_in( + snapshot, + ["roles", "primary", "capability_contract", "requirements", "structured_output"], + true + ) + + assert {:error, :invalid_model_mapping_snapshot} = + ModelMapping.validate_snapshot(unsupported_output) + + unsupported_toolset = + snapshot + |> put_in(["roles", "primary", "toolset_version"], "workspace-v999") + |> put_in( + ["roles", "primary", "capability_contract", "toolset_version"], + "workspace-v999" + ) + + assert {:error, :invalid_model_mapping_snapshot} = + ModelMapping.validate_snapshot(unsupported_toolset) + end end diff --git a/test/kodo/sessions_test.exs b/test/kodo/sessions_test.exs index ca98e7c..f1ce8db 100644 --- a/test/kodo/sessions_test.exs +++ b/test/kodo/sessions_test.exs @@ -2,6 +2,7 @@ defmodule Kodo.SessionsTest do use Kodo.DataCase alias Kodo.Agent.ExecutionRouteChange + alias Kodo.Agent.ModelMapping alias Kodo.Cluster.Instances alias Kodo.Integrations alias Kodo.Runners @@ -415,6 +416,11 @@ defmodule Kodo.SessionsTest do {:ok, [first_snapshot, _message, _status]} = Sessions.begin_turn(session.id, "First") + [created | _events] = Sessions.events_after(session.id) + + primary = + created.payload["model_mapping"] |> ModelMapping.snapshot() |> ModelMapping.role!(:primary) + records = [ %{ role: "primary", @@ -422,12 +428,12 @@ defmodule Kodo.SessionsTest do destination: "openai", selector: "gpt-5.4", role_contract: "alpha-v1", + reasoning: primary["reasoning"], + capability_contract: primary["capability_contract"], status: :approved } ] - [created | _events] = Sessions.events_after(session.id) - assert {:ok, changed_mapping} = ExecutionRouteChange.change( created.payload["model_mapping"], From f39d20afc9feff64be05e57b1b8d8e6b85677524 Mon Sep 17 00:00:00 2001 From: Darwin Wu Date: Sat, 12 Sep 2026 20:34:12 +0000 Subject: [PATCH 06/14] fix(integrations): harden authorization transitions --- lib/kodo/integrations.ex | 199 ++++++++++++------ lib/kodo/integrations/device_authorization.ex | 9 +- lib/kodo/integrations/oauth_refresh.ex | 3 +- lib/kodo_web/live/integrations_live.ex | 48 ++++- .../device_authorization_test.exs | 53 +++++ test/kodo/integrations_test.exs | 27 +++ test/kodo_web/live/integrations_live_test.exs | 28 ++- 7 files changed, 287 insertions(+), 80 deletions(-) diff --git a/lib/kodo/integrations.ex b/lib/kodo/integrations.ex index 44c8325..4d0f1a4 100644 --- a/lib/kodo/integrations.ex +++ b/lib/kodo/integrations.ex @@ -168,13 +168,21 @@ defmodule Kodo.Integrations do with {:ok, integration_id} <- cast_uuid(integration_id) do result = Repo.transaction(fn -> - expire_device_authorization(user.id, integration_id) - read_active_device_authorization(user.id, integration_id) + expired = expire_device_authorization(user.id, integration_id) + {expired, read_active_device_authorization(user.id, integration_id)} end) case result do - {:ok, :not_found} -> {:error, :device_authorization_not_found} - other -> other + {:ok, {expired, :not_found}} -> + Enum.each(expired, &broadcast_device_authorization_change/1) + {:error, :device_authorization_not_found} + + {:ok, {expired, authorization}} -> + Enum.each(expired, &broadcast_device_authorization_change/1) + {:ok, authorization} + + error -> + error end end end @@ -184,13 +192,21 @@ defmodule Kodo.Integrations do {:ok, claim_owner_id} <- cast_uuid(claim_owner_id) do result = Repo.transaction(fn -> - expire_device_authorization(user.id, integration_id) - claim_device_authorization_locked(user.id, integration_id, claim_owner_id) + expired = expire_device_authorization(user.id, integration_id) + {expired, claim_device_authorization_locked(user.id, integration_id, claim_owner_id)} end) case result do - {:ok, {:error, reason}} -> {:error, reason} - other -> other + {:ok, {expired, {:error, reason}}} -> + Enum.each(expired, &broadcast_device_authorization_change/1) + {:error, reason} + + {:ok, {expired, claim}} -> + Enum.each(expired, &broadcast_device_authorization_change/1) + {:ok, claim} + + error -> + error end end end @@ -225,6 +241,13 @@ defmodule Kodo.Integrations do Repo.transaction(fn -> admit_device_authorization_poll_locked(user.id, claim) end) end + def schedule_device_authorization_poll( + %Scope{user: user}, + %DeviceAuthorizationAttempt{} = claim + ) do + Repo.transaction(fn -> schedule_device_authorization_poll_locked(user.id, claim) end) + end + def store_device_authorization_exchange( %Scope{user: user}, %DeviceAuthorizationAttempt{} = claim, @@ -436,6 +459,7 @@ defmodule Kodo.Integrations do "oauth", "refresh_succeeded" ) + |> notify_integration_change() end @doc "Claims one generation of an OAuth integration for a bounded refresh operation." @@ -460,7 +484,8 @@ defmodule Kodo.Integrations do from attempt in DeviceAuthorizationAttempt, where: attempt.integration_id == parent_as(:integration).id and - attempt.state == "active" + attempt.state == "active" and + attempt.provider_deadline > fragment("clock_timestamp()") ), update: [ set: [ @@ -640,6 +665,7 @@ defmodule Kodo.Integrations do encryption_key_version: nil, credential_format_version: nil, credential_generation: generation + 1, + refresh_source_generation: nil, expires_at: nil, validated_at: nil, refreshed_at: nil, @@ -795,9 +821,6 @@ defmodule Kodo.Integrations do claim_owner_id, claim_epoch ) do - timestamp = database_now() - lease_expires_at = DateTime.add(timestamp, @device_authorization_claim_lease_ms, :millisecond) - query = from attempt in DeviceAuthorizationAttempt, join: integration in Integration, @@ -806,38 +829,45 @@ defmodule Kodo.Integrations do attempt.id == ^attempt_id and attempt.user_id == ^user_id and attempt.attempt_generation == ^attempt_generation and attempt.state == "active" and attempt.claim_owner_id == ^claim_owner_id and attempt.claim_epoch == ^claim_epoch and - attempt.claim_lease_expires_at > ^timestamp and attempt.provider_deadline > ^timestamp and + attempt.claim_lease_expires_at > fragment("clock_timestamp()") and + attempt.provider_deadline > fragment("clock_timestamp()") and integration.credential_generation == attempt.expected_integration_generation, + update: [ + set: [ + claim_lease_expires_at: + fragment( + "clock_timestamp() + (? * interval '1 millisecond')", + ^@device_authorization_claim_lease_ms + ), + updated_at: fragment("clock_timestamp()") + ] + ], select: attempt - case Repo.update_all( - query, - [set: [claim_lease_expires_at: lease_expires_at, updated_at: timestamp]], - @secret_repo_options - ) do + case Repo.update_all(query, [], @secret_repo_options) do {1, [attempt]} -> decrypt_device_authorization(attempt) {0, []} -> Repo.rollback(:stale_device_authorization_claim) end end defp admit_device_authorization_poll_locked(user_id, claim) do - timestamp = database_now() - lease_expires_at = DateTime.add(timestamp, @device_authorization_claim_lease_ms, :millisecond) - query = - claimed_device_authorization_query(user_id, claim, timestamp) - |> where([attempt, _integration], attempt.next_poll_at <= ^timestamp) + claimed_device_authorization_query(user_id, claim) + |> where([attempt, _integration], attempt.next_poll_at <= fragment("clock_timestamp()")) |> update( [attempt, _integration], set: [ next_poll_at: fragment( - "CAST(? AS timestamp) + (? * interval '1 millisecond')", - ^timestamp, + "clock_timestamp() + (? * interval '1 millisecond')", attempt.polling_interval_ms ), - claim_lease_expires_at: ^lease_expires_at, - updated_at: ^timestamp + claim_lease_expires_at: + fragment( + "clock_timestamp() + (? * interval '1 millisecond')", + ^@device_authorization_claim_lease_ms + ), + updated_at: fragment("clock_timestamp()") ] ) |> select([attempt, _integration], attempt) @@ -847,15 +877,37 @@ defmodule Kodo.Integrations do decrypt_device_authorization(attempt) {0, []} -> - if claimed_device_authorization?(user_id, claim, timestamp), + if claimed_device_authorization?(user_id, claim), do: Repo.rollback(:device_authorization_poll_not_due), else: Repo.rollback(:stale_device_authorization_claim) end end - defp claimed_device_authorization?(user_id, claim, timestamp) do + defp schedule_device_authorization_poll_locked(user_id, claim) do + query = + claimed_device_authorization_query(user_id, claim) + |> update( + [attempt, _integration], + set: [ + next_poll_at: + fragment( + "clock_timestamp() + (? * interval '1 millisecond')", + attempt.polling_interval_ms + ), + updated_at: fragment("clock_timestamp()") + ] + ) + |> select([attempt, _integration], attempt) + + case Repo.update_all(query, [], @secret_repo_options) do + {1, [attempt]} -> decrypt_device_authorization(attempt) + {0, []} -> Repo.rollback(:stale_device_authorization_claim) + end + end + + defp claimed_device_authorization?(user_id, claim) do claim - |> then(&claimed_device_authorization_query(user_id, &1, timestamp)) + |> then(&claimed_device_authorization_query(user_id, &1)) |> Repo.exists?(@secret_repo_options) end @@ -863,7 +915,7 @@ defmodule Kodo.Integrations do timestamp = database_now() query = - claimed_device_authorization_query(user_id, claim, timestamp) + claimed_device_authorization_query(user_id, claim) |> select([attempt, _integration], attempt) changes = @@ -878,18 +930,22 @@ defmodule Kodo.Integrations do end defp admit_device_authorization_exchange_locked(user_id, claim) do - timestamp = database_now() - lease_expires_at = DateTime.add(timestamp, @device_authorization_claim_lease_ms, :millisecond) - query = - claimed_device_authorization_query(user_id, claim, timestamp) + claimed_device_authorization_query(user_id, claim) + |> update( + [attempt, _integration], + set: [ + claim_lease_expires_at: + fragment( + "clock_timestamp() + (? * interval '1 millisecond')", + ^@device_authorization_claim_lease_ms + ), + updated_at: fragment("clock_timestamp()") + ] + ) |> select([attempt, _integration], attempt) - case Repo.update_all( - query, - [set: [claim_lease_expires_at: lease_expires_at, updated_at: timestamp]], - @secret_repo_options - ) do + case Repo.update_all(query, [], @secret_repo_options) do {1, [attempt]} -> {attempt, payload} = decrypt_device_authorization(attempt) @@ -904,16 +960,16 @@ defmodule Kodo.Integrations do end defp complete_device_authorization_locked(user_id, claim, credentials, expires_at) do - timestamp = database_now() lock_user!(user_id) provider = owned_integration_provider!(user_id, claim.integration_id) lock_provider_identity(user_id, provider) integration = lock_owned_integration!(user_id, claim.integration_id) + timestamp = database_now() if !DateTime.after?(expires_at, timestamp), do: Repo.rollback(:device_authorization_response_invalid) - attempt = lock_claimed_device_authorization!(user_id, integration, claim, timestamp) + attempt = lock_claimed_device_authorization!(user_id, integration, claim) case CredentialEncryption.encrypt(integration, credentials) do {:ok, encrypted} -> @@ -924,6 +980,7 @@ defmodule Kodo.Integrations do connection_status: "connected", validation_status: "unverified", credential_generation: integration.credential_generation + 1, + refresh_source_generation: nil, expires_at: expires_at, validated_at: nil, refreshed_at: timestamp, @@ -934,7 +991,8 @@ defmodule Kodo.Integrations do |> Repo.update!() integration = - if attempt.activate_on_completion do + if attempt.activate_on_completion and !integration.active and + !active_provider_account_exists?(user_id, integration.provider) do integration |> change(active: true, updated_at: timestamp) |> Repo.update!() else integration @@ -967,7 +1025,7 @@ defmodule Kodo.Integrations do # The boolean terms below are one atomic database fence, not control-flow branches. # credo:disable-for-next-line Credo.Check.Refactor.CyclomaticComplexity - defp lock_claimed_device_authorization!(user_id, integration, claim, timestamp) do + defp lock_claimed_device_authorization!(user_id, integration, claim) do query = DeviceAuthorizationAttempt |> where( @@ -977,8 +1035,8 @@ defmodule Kodo.Integrations do attempt.attempt_generation == ^claim.attempt_generation and attempt.state == "active" and attempt.claim_owner_id == ^claim.claim_owner_id and attempt.claim_epoch == ^claim.claim_epoch and - attempt.claim_lease_expires_at > ^timestamp and - attempt.provider_deadline > ^timestamp and + attempt.claim_lease_expires_at > fragment("clock_timestamp()") and + attempt.provider_deadline > fragment("clock_timestamp()") and attempt.expected_integration_generation == ^integration.credential_generation ) |> lock("FOR UPDATE") @@ -992,7 +1050,7 @@ defmodule Kodo.Integrations do defp terminalize_claimed_device_authorization(user_id, claim, state, error_code) do timestamp = database_now() - query = claimed_device_authorization_query(user_id, claim, timestamp) + query = claimed_device_authorization_query(user_id, claim) case Repo.update_all( query, @@ -1006,7 +1064,7 @@ defmodule Kodo.Integrations do # The boolean terms below are one atomic database fence, not control-flow branches. # credo:disable-for-next-line Credo.Check.Refactor.CyclomaticComplexity - defp claimed_device_authorization_query(user_id, claim, timestamp) do + defp claimed_device_authorization_query(user_id, claim) do from attempt in DeviceAuthorizationAttempt, join: integration in Integration, on: integration.id == attempt.integration_id and integration.user_id == attempt.user_id, @@ -1015,8 +1073,8 @@ defmodule Kodo.Integrations do attempt.attempt_generation == ^claim.attempt_generation and attempt.state == "active" and attempt.claim_owner_id == ^claim.claim_owner_id and attempt.claim_epoch == ^claim.claim_epoch and - attempt.claim_lease_expires_at > ^timestamp and - attempt.provider_deadline > ^timestamp and + attempt.claim_lease_expires_at > fragment("clock_timestamp()") and + attempt.provider_deadline > fragment("clock_timestamp()") and integration.credential_generation == attempt.expected_integration_generation end @@ -1059,7 +1117,7 @@ defmodule Kodo.Integrations do expected_integration_generation = integration.credential_generation + 1 activate_on_completion = - integration.active or !active_provider_account_exists?(user_id, integration.provider) + integration.active or provider_account_count(user_id, integration.provider) == 1 supersede_active_device_authorization(integration.id, timestamp) @@ -1114,6 +1172,7 @@ defmodule Kodo.Integrations do defp device_authorization_integration_changes(%Integration{}, generation, timestamp) do [ credential_generation: generation, + refresh_source_generation: nil, refresh_claim_owner_id: nil, refresh_claim_generation: nil, refresh_lease_expires_at: nil, @@ -1139,14 +1198,24 @@ defmodule Kodo.Integrations do end defp expire_device_authorization(user_id, integration_id) do - DeviceAuthorizationAttempt - |> where( - [attempt], - attempt.user_id == ^user_id and attempt.integration_id == ^integration_id and - attempt.state == "active" and - attempt.provider_deadline <= fragment("timezone('UTC', clock_timestamp())") - ) - |> Repo.update_all(set: terminal_attempt_changes("expired", "deadline_exceeded")) + query = + DeviceAuthorizationAttempt + |> where( + [attempt], + attempt.user_id == ^user_id and attempt.integration_id == ^integration_id and + attempt.state == "active" and + attempt.provider_deadline <= fragment("timezone('UTC', clock_timestamp())") + ) + |> select([attempt], attempt) + + case Repo.update_all( + query, + [set: terminal_attempt_changes("expired", "deadline_exceeded")], + @secret_repo_options + ) do + {0, []} -> [] + {_count, attempts} -> attempts + end end defp terminal_attempt_changes(state, error_code, timestamp \\ now()) do @@ -1573,18 +1642,22 @@ defmodule Kodo.Integrations do defp notify_integration_change(error), do: error defp notify_device_authorization_change({:ok, attempt} = result) do + broadcast_device_authorization_change(attempt) + + result + end + + defp notify_device_authorization_change(error), do: error + + defp broadcast_device_authorization_change(attempt) do Phoenix.PubSub.broadcast( Kodo.PubSub, "integration:#{attempt.user_id}", {:device_authorization_changed, attempt.integration_id, attempt.attempt_generation, attempt.state, attempt.terminal_error_code} ) - - result end - defp notify_device_authorization_change(error), do: error - defp constraint_error?(changeset, type) do Enum.any?(changeset.errors, fn {_field, {_message, metadata}} -> metadata[:constraint] == type diff --git a/lib/kodo/integrations/device_authorization.ex b/lib/kodo/integrations/device_authorization.ex index d56695a..1f6763e 100644 --- a/lib/kodo/integrations/device_authorization.ex +++ b/lib/kodo/integrations/device_authorization.ex @@ -69,7 +69,7 @@ defmodule Kodo.Integrations.DeviceAuthorization do result <- safe_poll(client, payload, options) do case result do :pending -> - continue(scope, admission, payload, client, options) + continue_after_pending(scope, admission, client, options) {:ok, exchange_payload} -> persist_and_exchange(scope, admission, exchange_payload, client, options) @@ -90,6 +90,13 @@ defmodule Kodo.Integrations.DeviceAuthorization do fail(scope, claim, :device_authorization_response_invalid) end + defp continue_after_pending(scope, admission, client, options) do + case Integrations.schedule_device_authorization_poll(scope, admission) do + {:ok, {scheduled, payload}} -> continue(scope, scheduled, payload, client, options) + {:error, reason} -> stop(reason) + end + end + defp await_poll_admission(scope, claim) do delay = max(DateTime.diff(claim.next_poll_at, DateTime.utc_now(), :millisecond), 0) diff --git a/lib/kodo/integrations/oauth_refresh.ex b/lib/kodo/integrations/oauth_refresh.ex index 9a5a541..3badca2 100644 --- a/lib/kodo/integrations/oauth_refresh.ex +++ b/lib/kodo/integrations/oauth_refresh.ex @@ -143,7 +143,8 @@ defmodule Kodo.Integrations.OAuthRefresh do defp current_result(scope, integration) do case Integrations.get_integration(scope, integration.id) do {:ok, %{connection_status: "connected"} = current} - when current.refresh_source_generation == integration.credential_generation -> + when current.refresh_source_generation == integration.credential_generation and + current.credential_generation == integration.credential_generation + 1 -> {:ok, current} {:ok, %{connection_status: "reauthorization_required"}} -> diff --git a/lib/kodo_web/live/integrations_live.ex b/lib/kodo_web/live/integrations_live.ex index 47a7521..7d74e7f 100644 --- a/lib/kodo_web/live/integrations_live.ex +++ b/lib/kodo_web/live/integrations_live.ex @@ -333,7 +333,12 @@ defmodule KodoWeb.IntegrationsLive do Map.has_key?(socket.assigns.device_authorization_tasks, reference) -> Process.demonitor(reference, [:flush]) - {:noreply, socket |> drop_device_authorization_task(reference) |> load_integrations()} + + {:noreply, + socket + |> drop_device_authorization_task(reference) + |> load_integrations() + |> resume_device_authorizations()} true -> {:noreply, socket} @@ -346,7 +351,11 @@ defmodule KodoWeb.IntegrationsLive do {:noreply, socket |> drop_validation_task(reference) |> load_integrations()} Map.has_key?(socket.assigns.device_authorization_tasks, reference) -> - {:noreply, socket |> drop_device_authorization_task(reference) |> load_integrations()} + {:noreply, + socket + |> drop_device_authorization_task(reference) + |> load_integrations() + |> resume_device_authorizations()} true -> {:noreply, socket} @@ -375,7 +384,7 @@ defmodule KodoWeb.IntegrationsLive do end def handle_info({:retry_device_authorization, integration_id}, socket) do - {:noreply, resume_device_authorization(socket, integration_id, false)} + {:noreply, resume_device_authorization(socket, integration_id, true)} end defp save_api_key( @@ -505,12 +514,9 @@ defmodule KodoWeb.IntegrationsLive do end defp resume_device_authorizations(socket) do - Enum.reduce(socket.assigns.integrations, socket, fn integration, acc -> - if integration.provider == "openai_codex" do - resume_device_authorization(acc, integration.id, true) - else - acc - end + Enum.reduce(socket.assigns.device_authorizations, socket, fn {integration_id, _attempt}, + acc -> + resume_device_authorization(acc, integration_id, true) end) end @@ -542,10 +548,20 @@ defmodule KodoWeb.IntegrationsLive do Map.put(attempts, integration.id, %{ attempt_id: attempt.id, attempt_generation: attempt.attempt_generation, + stage: :polling, user_code: user_code, provider_deadline: attempt.provider_deadline }) + {:ok, {attempt, %{"authorization_code" => _code}}} -> + Map.put(attempts, integration.id, %{ + attempt_id: attempt.id, + attempt_generation: attempt.attempt_generation, + stage: :exchanging, + user_code: nil, + provider_deadline: attempt.provider_deadline + }) + _missing_or_exchanging -> attempts end @@ -869,9 +885,12 @@ defmodule KodoWeb.IntegrationsLive do class="mt-4 max-w-xl rounded-2xl border border-sky-200 bg-sky-50/80 p-4 dark:border-sky-900 dark:bg-sky-950/30" >

- Waiting for OpenAI authorization + {if(authorization.stage == :polling, + do: "Waiting for OpenAI authorization", + else: "Completing OpenAI authorization" + )}

-
+

Open this verification page

<.link @@ -905,6 +924,13 @@ defmodule KodoWeb.IntegrationsLive do
+

+ Authorization was approved. Securely exchanging credentials… +

+ send(owner, {:pending_poll_started, self()}) + + receive do + :finish_pending -> :pending + end + end + + second_poll = fn :poll, _payload -> + send(owner, :second_poll_started) + {:error, :device_authorization_rejected} + end + + client = fake_client(poll: [blocking_poll, second_poll]) + + assert {:ok, _attempt} = + Integrations.begin_device_authorization( + scope, + integration.id, + integration.credential_generation, + %{"device_auth_id" => "device", "user_code" => "CODE"}, + 100 + ) + + assert {:ok, {claim, payload}} = + Integrations.claim_device_authorization( + scope, + integration.id, + Ecto.UUID.generate() + ) + + Repo.update!(change(claim, next_poll_at: DateTime.add(DateTime.utc_now(), -1, :second))) + claim = Repo.reload!(claim) + supervisor = start_supervised!(Task.Supervisor) + + task = + Task.Supervisor.async_nolink(supervisor, fn -> + DeviceAuthorization.run(scope, claim, payload, client_opts(client)) + end) + + assert_receive {:pending_poll_started, poller} + send(poller, :finish_pending) + refute_receive :second_poll_started, 50 + assert_receive :second_poll_started, 200 + assert {:error, :device_authorization_rejected} = Task.await(task) + end + test "begin persists the provider response and launches a supervised finite task", %{ scope: scope, integration: integration diff --git a/test/kodo/integrations_test.exs b/test/kodo/integrations_test.exs index 5063d8a..4827a2a 100644 --- a/test/kodo/integrations_test.exs +++ b/test/kodo/integrations_test.exs @@ -823,6 +823,7 @@ defmodule Kodo.IntegrationsTest do assert after_start.connection_status == "connected" assert after_start.validation_status == "valid" assert after_start.active + assert is_nil(after_start.refresh_source_generation) assert {:ok, _cancelled} = Integrations.cancel_device_authorization( @@ -901,6 +902,32 @@ defmodule Kodo.IntegrationsTest do refute completed.active end + test "an inactive account does not become active merely because its sibling was disconnected", + %{ + scope: scope + } do + first = oauth_integration(scope) + + assert {:ok, active} = + Integrations.oauth_succeeded(scope, first.id, 0, %{"access_token" => "first"}) + + second = oauth_integration(scope) + + assert {:ok, _disconnected} = + Integrations.disconnect(scope, active.id, active.credential_generation) + + assert {:ok, attempt} = + Integrations.begin_device_authorization( + scope, + second.id, + second.credential_generation, + %{"device_auth_id" => "second", "user_code" => "SECOND"}, + 0 + ) + + refute attempt.activate_on_completion + end + test "attempt operations suppress identifiers and ciphertext from query observability", %{ scope: scope } do diff --git a/test/kodo_web/live/integrations_live_test.exs b/test/kodo_web/live/integrations_live_test.exs index 21aaf5d..ef9597a 100644 --- a/test/kodo_web/live/integrations_live_test.exs +++ b/test/kodo_web/live/integrations_live_test.exs @@ -258,7 +258,18 @@ defmodule KodoWeb.IntegrationsLiveTest do conn: conn, scope: scope } do - _client = configure_device_client(exchange: [{:ok, device_tokens()}]) + owner = self() + + blocking_exchange = fn :exchange, _payload -> + send(owner, {:exchange_resumed, self()}) + + receive do + :finish_exchange -> {:ok, device_tokens()} + end + end + + _client = configure_device_client(exchange: [blocking_exchange]) + Phoenix.PubSub.subscribe(Kodo.PubSub, "integration:#{scope.user.id}") assert {:ok, integration} = Integrations.create_oauth_integration(scope, "openai_codex", @@ -291,13 +302,22 @@ defmodule KodoWeb.IntegrationsLiveTest do ) {:ok, view, _html} = live(conn, ~p"/integrations") - monitor = Process.monitor(view.pid) + assert_receive {:exchange_resumed, worker} + + panel = "#integration-#{integration.id}-device-authorization" + assert has_element?(view, panel, "Completing OpenAI authorization") + assert has_element?(view, panel, "Securely exchanging credentials") + refute has_element?(view, "#integration-#{integration.id}-copy-device-code") + + send(worker, :finish_exchange) + integration_id = integration.id + assert_receive message = {:integration_changed, ^integration_id, _generation} + send(view.pid, message) _ = :sys.get_state(view.pid) assert {:ok, connected} = Integrations.get_integration(scope, integration.id) assert connected.connection_status == "connected" - refute has_element?(view, "#integration-#{integration.id}-device-authorization") - Process.demonitor(monitor, [:flush]) + refute has_element?(view, panel) end test "cancelling authorization removes its code from another open tab", %{ From 449a728ce0d4857c004f50330450e1f3d7c6377d Mon Sep 17 00:00:00 2001 From: Darwin Wu Date: Sat, 12 Sep 2026 20:36:49 +0000 Subject: [PATCH 07/14] fix(sessions): preflight every turn provider --- lib/kodo/sessions.ex | 23 ++++++++++++++++++++++- rfd/0005/IMPLEMENTATION.org | 2 +- test/kodo/sessions_test.exs | 27 +++++++++++++++++++++++++++ 3 files changed, 50 insertions(+), 2 deletions(-) diff --git a/lib/kodo/sessions.ex b/lib/kodo/sessions.ex index 4e0b9f2..af51d3b 100644 --- a/lib/kodo/sessions.ex +++ b/lib/kodo/sessions.ex @@ -14,6 +14,7 @@ defmodule Kodo.Sessions do alias Kodo.Cluster.Placement alias Kodo.ControlPlaneTelemetry alias Kodo.Integrations + alias Kodo.LLM alias Kodo.LLM.ProviderError alias Kodo.Repo alias Kodo.Runners.Runner @@ -546,7 +547,9 @@ defmodule Kodo.Sessions do if turn_request_recorded?(session_id, client_request_id) do :ok else - start_turn(session_id, content, client_request_id) + with :ok <- preflight_turn(scope, session_id) do + start_turn(session_id, content, client_request_id) + end end nil -> @@ -554,6 +557,24 @@ defmodule Kodo.Sessions do end end + defp preflight_turn(scope, session_id) do + mapping = + session_id + |> get_session!() + |> turn_route_snapshot() + |> Map.fetch!("model_mapping") + + mapping["roles"] + |> Map.values() + |> Enum.uniq_by(& &1["execution_route"]) + |> Enum.reduce_while(:ok, fn role, :ok -> + case LLM.resolve_integration(scope, ModelMapping.request_model(role)) do + {:ok, _model, _reference} -> {:cont, :ok} + {:error, reason} -> {:halt, {:error, reason}} + end + end) + end + def cancel(session_id) do case get_session(session_id) do %Session{status: status} when status in ["running", "awaiting_approval"] -> diff --git a/rfd/0005/IMPLEMENTATION.org b/rfd/0005/IMPLEMENTATION.org index 10f72c5..3780fcd 100644 --- a/rfd/0005/IMPLEMENTATION.org +++ b/rfd/0005/IMPLEMENTATION.org @@ -430,7 +430,7 @@ manufactured after that no-go decision. cross-route replay, no automatic route fallback, same-origin and cross-origin redirect rejection, malicious verification URLs, absence of a credential-bearing follow-up request, and secret redaction. -- [X] Run the Phase 0 live compatibility suite against the production path and +- [ ] Run the Phase 0 live compatibility suite against the production path and do not declare ChatGPT/Codex implementation complete while a required capability fails; classify later contract failures as regressions and never fall back to Platform billing. diff --git a/test/kodo/sessions_test.exs b/test/kodo/sessions_test.exs index f1ce8db..da4eb90 100644 --- a/test/kodo/sessions_test.exs +++ b/test/kodo/sessions_test.exs @@ -3,8 +3,10 @@ defmodule Kodo.SessionsTest do alias Kodo.Agent.ExecutionRouteChange alias Kodo.Agent.ModelMapping + alias Kodo.Agent.ModelSettings alias Kodo.Cluster.Instances alias Kodo.Integrations + alias Kodo.LLM.ProviderError alias Kodo.Runners alias Kodo.Sessions alias Kodo.Sessions.Session @@ -359,6 +361,31 @@ defmodule Kodo.SessionsTest do assert Enum.count(Sessions.events_after(session.id), &(&1.type == "user_message")) == 1 end + test "rejects a turn before acceptance when any configured role lacks access", %{ + runner: runner, + scope: scope + } do + assert {:ok, _override} = + ModelSettings.put_user_override(scope, :review, %{ + model: "anthropic:claude-3-5-haiku-latest" + }) + + assert {:ok, _integration} = + Integrations.connect(scope, "openai", "api_key", %{"api_key" => "primary-key"}) + + assert {:ok, session} = + Sessions.create_session(scope, %{ + runner_id: runner.id, + title: "Missing review provider", + model: "openai:gpt-4o-mini" + }) + + assert {:error, %ProviderError{provider: "anthropic", kind: :integration_required}} = + Sessions.start_turn(scope, session.id, "Do not accept", Ecto.UUID.generate()) + + assert Enum.map(Sessions.events_after(session.id), & &1.type) == ["session_created"] + end + test "rejects route changes without approved exact-identity evidence", %{ runner: runner, scope: scope From 08e8ecd3a92bf573c3f7a2fddb755da8d41c7c6a Mon Sep 17 00:00:00 2001 From: Darwin Wu Date: Sat, 12 Sep 2026 20:38:53 +0000 Subject: [PATCH 08/14] fix(agent): atomically admit model invocations --- lib/kodo/agent/loop.ex | 48 +++++++------ lib/kodo/llm.ex | 30 ++++++++ lib/kodo/sessions.ex | 52 ++++++++++++++ test/kodo/llm/credential_resolver_test.exs | 14 ++++ test/kodo/sessions_test.exs | 82 ++++++++++++++++++++++ 5 files changed, 206 insertions(+), 20 deletions(-) diff --git a/lib/kodo/agent/loop.ex b/lib/kodo/agent/loop.ex index dfde343..61a7419 100644 --- a/lib/kodo/agent/loop.ex +++ b/lib/kodo/agent/loop.ex @@ -6,6 +6,7 @@ defmodule Kodo.Agent.Loop do alias Kodo.Agent.Roles alias Kodo.Agent.Tools alias Kodo.LLM + alias Kodo.LLM.ProviderError alias Kodo.RunnerProtocol alias Kodo.Runners alias Kodo.Sessions @@ -247,7 +248,7 @@ defmodule Kodo.Agent.Loop do ) do invocation_id = Ecto.UUID.generate() - case Sessions.append_event( + case Sessions.append_admitted_model_event( session_id, "review_invocation_started", %{ @@ -257,20 +258,21 @@ defmodule Kodo.Agent.Loop do "provider" => canonical_provider(request), "model" => review["model"], "model_identity" => request.model.id, - "authentication_type" => request.reference.authentication_type, - "billing_path" => Atom.to_string(request.reference.billing_path), + "authentication_type" => request.credential.authentication_type, + "billing_path" => Atom.to_string(request.credential.billing_path), "reasoning" => review["reasoning"], "role_contract" => contract.id, "toolset_version" => contract.toolset_version, "capability_validation" => capability_validation, "model_mapping" => mapping }, + request.credential, version: 2, parent_id: primary_invocation_id, ownership: ownership ) do {:ok, _event} -> {:ok, invocation_id} - error -> error + error -> normalize_admission_error(error, request) end end @@ -620,7 +622,7 @@ defmodule Kodo.Agent.Loop do ) do invocation_id = Ecto.UUID.generate() - case Sessions.append_event( + case Sessions.append_admitted_model_event( session_id, "model_invocation_started", %{ @@ -630,19 +632,20 @@ defmodule Kodo.Agent.Loop do "provider" => canonical_provider(request), "model" => primary["model"], "model_identity" => request.model.id, - "authentication_type" => request.reference.authentication_type, - "billing_path" => Atom.to_string(request.reference.billing_path), + "authentication_type" => request.credential.authentication_type, + "billing_path" => Atom.to_string(request.credential.billing_path), "reasoning" => primary["reasoning"], "role_contract" => primary["role_contract"], "toolset_version" => primary["toolset_version"], "capability_validation" => capability_validation, "model_mapping" => mapping }, + request.credential, version: 4, ownership: ownership ) do {:ok, _event} -> {:ok, invocation_id} - error -> error + error -> normalize_admission_error(error, request) end end @@ -1125,7 +1128,7 @@ defmodule Kodo.Agent.Loop do ) do invocation_id = Ecto.UUID.generate() - case Sessions.append_event( + case Sessions.append_admitted_model_event( context.session_id, "subagent_invocation_started", %{ @@ -1136,20 +1139,21 @@ defmodule Kodo.Agent.Loop do "provider" => canonical_provider(request), "model" => search["model"], "model_identity" => request.model.id, - "authentication_type" => request.reference.authentication_type, - "billing_path" => Atom.to_string(request.reference.billing_path), + "authentication_type" => request.credential.authentication_type, + "billing_path" => Atom.to_string(request.credential.billing_path), "reasoning" => search["reasoning"], "role_contract" => contract.id, "toolset_version" => contract.toolset_version, "capability_validation" => capability_validation, "model_mapping" => context.mapping }, + request.credential, version: 2, parent_id: context.invocation_id, ownership: context.ownership ) do {:ok, _event} -> {:ok, invocation_id} - error -> error + error -> normalize_admission_error(error, request) end end @@ -1686,16 +1690,21 @@ defmodule Kodo.Agent.Loop do defp resolve_request(session_id, role) do with {:ok, scope} <- Sessions.owner_scope(session_id), {:ok, model, reference} <- - LLM.resolve_integration(scope, ModelMapping.request_model(role)) do - {:ok, %{scope: scope, model: model, reference: reference}} + LLM.resolve_integration(scope, ModelMapping.request_model(role)), + {:ok, credential} <- LLM.admit_credential(scope, model, reference) do + {:ok, %{model: model, credential: credential}} end end + defp normalize_admission_error({:error, :stale_credential_generation}, request), + do: {:error, ProviderError.from_integration(:stale_credential_generation, request.model)} + + defp normalize_admission_error(error, _request), do: error + defp generate(adapter, request, messages, tools, opts) do - LLM.generate( - request.scope, + LLM.generate_admitted( request.model, - request.reference, + request.credential, messages, tools, Keyword.put(opts, :adapter, adapter) @@ -1703,10 +1712,9 @@ defmodule Kodo.Agent.Loop do end defp generate_object(adapter, request, messages, schema, opts) do - LLM.generate_object( - request.scope, + LLM.generate_object_admitted( request.model, - request.reference, + request.credential, messages, schema, Keyword.put(opts, :adapter, adapter) diff --git a/lib/kodo/llm.ex b/lib/kodo/llm.ex index 2029a02..484640f 100644 --- a/lib/kodo/llm.ex +++ b/lib/kodo/llm.ex @@ -90,6 +90,20 @@ defmodule Kodo.LLM do end end + @doc false + def admit_credential(%Scope{} = scope, %LLMDB.Model{} = model, %IntegrationRef{} = reference), + do: resolve_credential(scope, model, reference) + + @doc false + def generate_admitted(%LLMDB.Model{} = model, %Credential{} = credential, messages, tools, opts) do + adapter = Keyword.get(opts, :adapter, adapter()) + adapter_opts = Keyword.delete(opts, :adapter) + + with :ok <- reject_credential_options(adapter_opts) do + adapter.generate(model, messages, tools, credential, adapter_opts) + end + end + @doc "Generates a structured object after scoped credential admission." def generate_object( %Scope{} = scope, @@ -108,6 +122,22 @@ defmodule Kodo.LLM do end end + @doc false + def generate_object_admitted( + %LLMDB.Model{} = model, + %Credential{} = credential, + messages, + schema, + opts + ) do + adapter = Keyword.get(opts, :adapter, adapter()) + adapter_opts = Keyword.delete(opts, :adapter) + + with :ok <- reject_credential_options(adapter_opts) do + adapter.generate_object(model, messages, schema, credential, adapter_opts) + end + end + def adapter, do: Application.get_env(:kodo, :llm_adapter, Kodo.LLM.ReqLLM) defp resolve_model(model) do diff --git a/lib/kodo/sessions.ex b/lib/kodo/sessions.ex index af51d3b..8513fe9 100644 --- a/lib/kodo/sessions.ex +++ b/lib/kodo/sessions.ex @@ -15,6 +15,7 @@ defmodule Kodo.Sessions do alias Kodo.ControlPlaneTelemetry alias Kodo.Integrations alias Kodo.LLM + alias Kodo.LLM.Credential alias Kodo.LLM.ProviderError alias Kodo.Repo alias Kodo.Runners.Runner @@ -541,6 +542,39 @@ defmodule Kodo.Sessions do end end + @doc false + def append_admitted_model_event( + session_id, + type, + payload, + %Credential{} = credential, + opts \\ [] + ) do + case Repo.transaction(fn -> + session = lock_session!(session_id, opts) + admit_model_credential!(session, credential) + + payload = + Map.merge(payload, %{ + "provider" => credential.provider, + "authentication_type" => credential.authentication_type, + "billing_path" => Atom.to_string(credential.billing_path) + }) + + case append_locked(session, type, payload, opts) do + {:ok, event} -> event + {:error, changeset} -> Repo.rollback(changeset) + end + end) do + {:ok, event} = result -> + broadcast(event) + result + + error -> + error + end + end + def start_turn(%Scope{} = scope, session_id, content, client_request_id) do case get_session(scope, session_id) do %Session{} -> @@ -575,6 +609,24 @@ defmodule Kodo.Sessions do end) end + defp admit_model_credential!(session, credential) do + admitted? = + Integrations.Integration + |> where( + [integration], + integration.id == ^credential.integration_id and + integration.user_id == ^session.user_id and integration.provider == ^credential.provider and + integration.authentication_type == ^credential.authentication_type and + integration.credential_generation == ^credential.credential_generation and + integration.connection_status == "connected" and integration.active and + integration.validation_status != "invalid" + ) + |> lock("FOR SHARE") + |> Repo.exists?() + + if !admitted?, do: Repo.rollback(:stale_credential_generation) + end + def cancel(session_id) do case get_session(session_id) do %Session{status: status} when status in ["running", "awaiting_approval"] -> diff --git a/test/kodo/llm/credential_resolver_test.exs b/test/kodo/llm/credential_resolver_test.exs index ef0500c..da386ed 100644 --- a/test/kodo/llm/credential_resolver_test.exs +++ b/test/kodo/llm/credential_resolver_test.exs @@ -36,6 +36,8 @@ defmodule Kodo.LLM.CredentialResolverTest do assert {:ok, resolved_model, reference} = LLM.resolve_integration(context.scope, "openai:gpt-4o-mini") + assert {:ok, credential} = LLM.admit_credential(context.scope, resolved_model, reference) + for unsafe_option <- [ {:api_key, "caller-secret"}, {:base_url, "https://attacker.example"}, @@ -55,6 +57,18 @@ defmodule Kodo.LLM.CredentialResolverTest do [unsafe_option] ) ) + + assert {:error, :credential_options_not_allowed} = + LLM.generate_admitted( + resolved_model, + credential, + [%{"role" => "user", "content" => "final answer"}], + [], + Keyword.merge( + [adapter: Kodo.Test.FakeLLM, timeout: 1_000], + [unsafe_option] + ) + ) end assert {:ok, _replaced} = diff --git a/test/kodo/sessions_test.exs b/test/kodo/sessions_test.exs index da4eb90..32c2ec4 100644 --- a/test/kodo/sessions_test.exs +++ b/test/kodo/sessions_test.exs @@ -6,6 +6,7 @@ defmodule Kodo.SessionsTest do alias Kodo.Agent.ModelSettings alias Kodo.Cluster.Instances alias Kodo.Integrations + alias Kodo.LLM alias Kodo.LLM.ProviderError alias Kodo.Runners alias Kodo.Sessions @@ -386,6 +387,87 @@ defmodule Kodo.SessionsTest do assert Enum.map(Sessions.events_after(session.id), & &1.type) == ["session_created"] end + test "atomically records invocation admission for the exact active credential generation", %{ + runner: runner, + scope: scope + } do + assert {:ok, integration} = + Integrations.connect(scope, "openai", "api_key", %{"api_key" => "admitted-key"}) + + assert {:ok, model, reference} = LLM.resolve_integration(scope, "openai:gpt-4o-mini") + assert {:ok, credential} = LLM.admit_credential(scope, model, reference) + + assert {:ok, session} = + Sessions.create_session(scope, %{ + runner_id: runner.id, + title: "Admitted generation", + model: "openai:gpt-4o-mini" + }) + + payload = %{ + "invocation_id" => Ecto.UUID.generate(), + "provider" => "forged-provider", + "authentication_type" => "forged-authentication", + "billing_path" => "forged-billing" + } + + assert {:ok, event} = + Sessions.append_admitted_model_event( + session.id, + "model_invocation_started", + payload, + credential + ) + + assert event.payload == %{ + payload + | "provider" => credential.provider, + "authentication_type" => credential.authentication_type, + "billing_path" => Atom.to_string(credential.billing_path) + } + + refute inspect(event) =~ "admitted-key" + assert integration.credential_generation == credential.credential_generation + end + + test "rejects a decrypted credential replaced before invocation admission", %{ + runner: runner, + scope: scope + } do + assert {:ok, integration} = + Integrations.connect(scope, "openai", "api_key", %{"api_key" => "stale-key"}) + + assert {:ok, model, reference} = LLM.resolve_integration(scope, "openai:gpt-4o-mini") + assert {:ok, credential} = LLM.admit_credential(scope, model, reference) + + assert {:ok, _replacement} = + Integrations.replace_credentials( + scope, + integration.id, + integration.credential_generation, + %{"api_key" => "replacement-key"} + ) + + assert {:ok, session} = + Sessions.create_session(scope, %{ + runner_id: runner.id, + title: "Stale generation", + model: "openai:gpt-4o-mini" + }) + + events_before = Sessions.events_after(session.id) + + assert {:error, :stale_credential_generation} = + Sessions.append_admitted_model_event( + session.id, + "model_invocation_started", + %{"invocation_id" => Ecto.UUID.generate()}, + credential + ) + + assert Sessions.events_after(session.id) == events_before + end + test "rejects route changes without approved exact-identity evidence", %{ runner: runner, scope: scope From b52031b92d0383fa736a8fb6f1a247154cd26ebf Mon Sep 17 00:00:00 2001 From: Darwin Wu Date: Sat, 12 Sep 2026 20:42:45 +0000 Subject: [PATCH 09/14] refactor(sessions): simplify turn admission control flow --- lib/kodo/sessions.ex | 56 ++++++++++++++++++++++++++------------------ 1 file changed, 33 insertions(+), 23 deletions(-) diff --git a/lib/kodo/sessions.ex b/lib/kodo/sessions.ex index 8513fe9..9ba88c9 100644 --- a/lib/kodo/sessions.ex +++ b/lib/kodo/sessions.ex @@ -550,22 +550,11 @@ defmodule Kodo.Sessions do %Credential{} = credential, opts \\ [] ) do - case Repo.transaction(fn -> - session = lock_session!(session_id, opts) - admit_model_credential!(session, credential) - - payload = - Map.merge(payload, %{ - "provider" => credential.provider, - "authentication_type" => credential.authentication_type, - "billing_path" => Atom.to_string(credential.billing_path) - }) - - case append_locked(session, type, payload, opts) do - {:ok, event} -> event - {:error, changeset} -> Repo.rollback(changeset) - end - end) do + transaction = fn -> + append_admitted_model_event_locked(session_id, type, payload, credential, opts) + end + + case Repo.transaction(transaction) do {:ok, event} = result -> broadcast(event) result @@ -578,19 +567,23 @@ defmodule Kodo.Sessions do def start_turn(%Scope{} = scope, session_id, content, client_request_id) do case get_session(scope, session_id) do %Session{} -> - if turn_request_recorded?(session_id, client_request_id) do - :ok - else - with :ok <- preflight_turn(scope, session_id) do - start_turn(session_id, content, client_request_id) - end - end + start_scoped_turn(scope, session_id, content, client_request_id) nil -> {:error, :not_found} end end + defp start_scoped_turn(scope, session_id, content, client_request_id) do + if turn_request_recorded?(session_id, client_request_id) do + :ok + else + with :ok <- preflight_turn(scope, session_id) do + start_turn(session_id, content, client_request_id) + end + end + end + defp preflight_turn(scope, session_id) do mapping = session_id @@ -627,6 +620,23 @@ defmodule Kodo.Sessions do if !admitted?, do: Repo.rollback(:stale_credential_generation) end + defp append_admitted_model_event_locked(session_id, type, payload, credential, opts) do + session = lock_session!(session_id, opts) + admit_model_credential!(session, credential) + + payload = + Map.merge(payload, %{ + "provider" => credential.provider, + "authentication_type" => credential.authentication_type, + "billing_path" => Atom.to_string(credential.billing_path) + }) + + case append_locked(session, type, payload, opts) do + {:ok, event} -> event + {:error, changeset} -> Repo.rollback(changeset) + end + end + def cancel(session_id) do case get_session(session_id) do %Session{status: status} when status in ["running", "awaiting_approval"] -> From 3b0913c7d9c194001b1dbb05f13fa847cd2eaf63 Mon Sep 17 00:00:00 2001 From: Darwin Wu Date: Sat, 12 Sep 2026 21:13:35 +0000 Subject: [PATCH 10/14] fix(llm): harden provider request metadata --- lib/kodo/llm/req_llm.ex | 1 + lib/kodo/llm/safe_req_adapter.ex | 9 ++++++++- .../req_device_authorization_client_test.exs | 5 ++++- test/kodo/integrations/req_refresh_client_test.exs | 12 +++++++++++- test/kodo/llm/req_llm_test.exs | 1 + test/kodo/llm/safe_req_adapter_test.exs | 7 +++++++ 6 files changed, 32 insertions(+), 3 deletions(-) diff --git a/lib/kodo/llm/req_llm.ex b/lib/kodo/llm/req_llm.ex index b6fbf59..6554c9d 100644 --- a/lib/kodo/llm/req_llm.ex +++ b/lib/kodo/llm/req_llm.ex @@ -248,6 +248,7 @@ defmodule Kodo.LLM.ReqLLM do |> Keyword.put(:auth_mode, :oauth) |> Keyword.put(:access_token, token) |> Keyword.put(:chatgpt_account_id, account_id) + |> Keyword.put(:provider_options, codex_originator: "kodo") end # ReqLLM emits terminal errors to telemetry before returning control to Kodo. diff --git a/lib/kodo/llm/safe_req_adapter.ex b/lib/kodo/llm/safe_req_adapter.ex index 6ebfbe6..05db977 100644 --- a/lib/kodo/llm/safe_req_adapter.ex +++ b/lib/kodo/llm/safe_req_adapter.ex @@ -111,7 +111,7 @@ defmodule Kodo.LLM.SafeReqAdapter do defp send_request(request, conn, timeout) do method = request.method |> to_string() |> String.upcase() - path = URI.to_string(%{request.url | scheme: nil, host: nil, port: nil, userinfo: nil}) + path = request_target(request.url) headers = Req.Fields.get_list(request.headers) case Mint.HTTP.request(conn, method, path, headers, request.body) do @@ -120,6 +120,13 @@ defmodule Kodo.LLM.SafeReqAdapter do end end + @doc false + def request_target(url) do + # Build a fresh URI because URI.parse/1 retains a legacy authority field. + # Mutating the absolute URI could therefore send `//host/path` to Mint. + URI.to_string(%URI{path: url.path || "/", query: url.query}) + end + defp receive_response(request, conn, ref, timeout, response) do case Mint.HTTP.recv(conn, 0, timeout) do {:ok, conn, entries} -> diff --git a/test/kodo/integrations/req_device_authorization_client_test.exs b/test/kodo/integrations/req_device_authorization_client_test.exs index c8fca92..53f5938 100644 --- a/test/kodo/integrations/req_device_authorization_client_test.exs +++ b/test/kodo/integrations/req_device_authorization_client_test.exs @@ -175,7 +175,10 @@ defmodule Kodo.Integrations.ReqDeviceAuthorizationClientTest do plug = fn conn -> Agent.update(counter, &(&1 + 1)) - Plug.Conn.send_resp(conn, 302, "redirect without a location") + + conn + |> Plug.Conn.put_resp_header("location", "https://attacker.example/collect") + |> Plug.Conn.send_resp(302, "redirect") end payloads = [ diff --git a/test/kodo/integrations/req_refresh_client_test.exs b/test/kodo/integrations/req_refresh_client_test.exs index cbca3bf..71ce119 100644 --- a/test/kodo/integrations/req_refresh_client_test.exs +++ b/test/kodo/integrations/req_refresh_client_test.exs @@ -58,10 +58,20 @@ defmodule Kodo.Integrations.ReqRefreshClientTest do end test "rejects redirects, malformed success, and unsafe caller input with bounded errors" do - redirect = fn conn -> Plug.Conn.send_resp(conn, 302, "private") end + counter = start_supervised!({Agent, fn -> 0 end}) + + redirect = fn conn -> + Agent.update(counter, &(&1 + 1)) + + conn + |> Plug.Conn.put_resp_header("location", "https://auth.openai.com/oauth/redirected") + |> Plug.Conn.send_resp(302, "private") + end + malformed = fn conn -> Req.Test.json(conn, %{"refresh_token" => "private"}) end assert {:error, :redirect} = ReqRefreshClient.refresh("refresh-secret", plug: redirect) + assert Agent.get(counter, & &1) == 1 assert {:error, :oauth_refresh_response_invalid} = ReqRefreshClient.refresh("refresh-secret", plug: malformed) diff --git a/test/kodo/llm/req_llm_test.exs b/test/kodo/llm/req_llm_test.exs index b960079..1901267 100644 --- a/test/kodo/llm/req_llm_test.exs +++ b/test/kodo/llm/req_llm_test.exs @@ -93,6 +93,7 @@ defmodule Kodo.LLM.ReqLLMTest do assert options[:auth_mode] == :oauth assert options[:access_token] == "request-local-access" assert options[:chatgpt_account_id] == "request-local-account" + assert options[:provider_options] == [codex_originator: "kodo"] refute Keyword.has_key?(options, :refresh_token) end diff --git a/test/kodo/llm/safe_req_adapter_test.exs b/test/kodo/llm/safe_req_adapter_test.exs index c95827e..26d5d4e 100644 --- a/test/kodo/llm/safe_req_adapter_test.exs +++ b/test/kodo/llm/safe_req_adapter_test.exs @@ -5,6 +5,13 @@ defmodule Kodo.LLM.SafeReqAdapterTest do alias Kodo.LLM.SafeReqAdapter + test "uses an origin-form request target for absolute provider URLs" do + url = URI.parse("https://chatgpt.com/backend-api/codex/responses?stream=true") + + assert SafeReqAdapter.request_target(url) == + "/backend-api/codex/responses?stream=true" + end + test "removes operation secrets and untrusted details before error steps" do request = Req.new( From bd62d664526d05c0d0dad299d9d86be3fd12519f Mon Sep 17 00:00:00 2001 From: Darwin Wu Date: Sat, 12 Sep 2026 21:13:35 +0000 Subject: [PATCH 11/14] fix(integrations): preserve admitted OAuth results --- lib/kodo/integrations.ex | 96 ++++++++++-- lib/kodo/integrations/device_authorization.ex | 21 ++- lib/kodo/integrations/oauth_refresh.ex | 39 +++-- .../device_authorization_test.exs | 65 ++++++++ test/kodo/integrations/oauth_refresh_test.exs | 146 ++++++++++++++++++ test/kodo/integrations_test.exs | 52 +++++++ 6 files changed, 389 insertions(+), 30 deletions(-) diff --git a/lib/kodo/integrations.ex b/lib/kodo/integrations.ex index 4d0f1a4..f2ad691 100644 --- a/lib/kodo/integrations.ex +++ b/lib/kodo/integrations.ex @@ -164,11 +164,57 @@ defmodule Kodo.Integrations do def begin_device_authorization(%Scope{}, _id, _generation, _payload, _interval), do: {:error, :device_authorization_invalid} + @doc false + def begin_claimed_device_authorization( + %Scope{user: user}, + integration_id, + expected_generation, + payload, + polling_interval_ms, + claim_owner_id + ) + when is_integer(expected_generation) and expected_generation >= 0 and + is_integer(polling_interval_ms) and polling_interval_ms >= 0 and + polling_interval_ms <= 900_000 do + with {:ok, integration_id} <- cast_uuid(integration_id), + {:ok, claim_owner_id} <- cast_uuid(claim_owner_id), + :ok <- validate_device_authorization_payload(payload) do + result = + Repo.transaction(fn -> + attempt = + begin_device_authorization_locked( + user.id, + integration_id, + expected_generation, + payload, + polling_interval_ms + ) + + {claim, decrypted} = + claim_device_authorization_locked(user.id, integration_id, claim_owner_id) + + {attempt, claim, decrypted} + end) + + case result do + {:ok, {attempt, _claim, _decrypted}} -> + broadcast_device_authorization_change(attempt) + result + + error -> + error + end + end + end + + def begin_claimed_device_authorization(%Scope{}, _id, _generation, _payload, _interval, _owner), + do: {:error, :device_authorization_invalid} + def get_active_device_authorization(%Scope{user: user}, integration_id) do with {:ok, integration_id} <- cast_uuid(integration_id) do result = Repo.transaction(fn -> - expired = expire_device_authorization(user.id, integration_id) + expired = expire_device_authorization_rows(user.id, integration_id) {expired, read_active_device_authorization(user.id, integration_id)} end) @@ -187,12 +233,21 @@ defmodule Kodo.Integrations do end end + @doc false + def expire_device_authorization(%Scope{user: user}, integration_id) do + with {:ok, integration_id} <- cast_uuid(integration_id) do + expired = expire_device_authorization_rows(user.id, integration_id) + Enum.each(expired, &broadcast_device_authorization_change/1) + {:ok, expired} + end + end + def claim_device_authorization(%Scope{user: user}, integration_id, claim_owner_id) do with {:ok, integration_id} <- cast_uuid(integration_id), {:ok, claim_owner_id} <- cast_uuid(claim_owner_id) do result = Repo.transaction(fn -> - expired = expire_device_authorization(user.id, integration_id) + expired = expire_device_authorization_rows(user.id, integration_id) {expired, claim_device_authorization_locked(user.id, integration_id, claim_owner_id)} end) @@ -454,7 +509,9 @@ defmodule Kodo.Integrations do id, generation, credentials, - Keyword.put(opts, :refresh_source_generation, generation), + opts + |> Keyword.put(:refresh_source_generation, generation) + |> Keyword.put_new(:restore_active, false), ~w(connected reauthorization_required), "oauth", "refresh_succeeded" @@ -469,6 +526,10 @@ defmodule Kodo.Integrations do when is_integer(generation) and generation >= 0 do with {:ok, id} <- cast_uuid(id), {:ok, claim_owner_id} <- cast_uuid(claim_owner_id) do + user.id + |> expire_device_authorization_rows(id) + |> Enum.each(&broadcast_device_authorization_change/1) + query = from integration in Integration, as: :integration, @@ -969,7 +1030,7 @@ defmodule Kodo.Integrations do if !DateTime.after?(expires_at, timestamp), do: Repo.rollback(:device_authorization_response_invalid) - attempt = lock_claimed_device_authorization!(user_id, integration, claim) + attempt = lock_device_authorization_completion!(user_id, integration, claim) case CredentialEncryption.encrypt(integration, credentials) do {:ok, encrypted} -> @@ -1025,7 +1086,7 @@ defmodule Kodo.Integrations do # The boolean terms below are one atomic database fence, not control-flow branches. # credo:disable-for-next-line Credo.Check.Refactor.CyclomaticComplexity - defp lock_claimed_device_authorization!(user_id, integration, claim) do + defp lock_device_authorization_completion!(user_id, integration, claim) do query = DeviceAuthorizationAttempt |> where( @@ -1035,8 +1096,6 @@ defmodule Kodo.Integrations do attempt.attempt_generation == ^claim.attempt_generation and attempt.state == "active" and attempt.claim_owner_id == ^claim.claim_owner_id and attempt.claim_epoch == ^claim.claim_epoch and - attempt.claim_lease_expires_at > fragment("clock_timestamp()") and - attempt.provider_deadline > fragment("clock_timestamp()") and attempt.expected_integration_generation == ^integration.credential_generation ) |> lock("FOR UPDATE") @@ -1197,7 +1256,7 @@ defmodule Kodo.Integrations do |> Kernel.+(1) end - defp expire_device_authorization(user_id, integration_id) do + defp expire_device_authorization_rows(user_id, integration_id) do query = DeviceAuthorizationAttempt |> where( @@ -1365,6 +1424,7 @@ defmodule Kodo.Integrations do refresh_claim_generation: nil, refresh_lease_expires_at: nil }) + |> Map.put(:restore_active, opts[:restore_active]) update_fenced(scope, id, generation, allowed_connections, audit_event_type, changes) else @@ -1474,11 +1534,14 @@ defmodule Kodo.Integrations do Repo.transaction(fn -> maybe_lock_transition(user.id, id, audit_event_type) + {restore_active, changes} = Map.pop(changes, :restore_active, false) + integration = execute_fenced_update(user.id, id, generation, allowed_connections, changes) maybe_terminalize_device_authorization(user.id, id, audit_event_type) integration = maybe_activate_initial_oauth(user.id, integration, audit_event_type) + integration = maybe_restore_refresh_selection(user.id, integration, restore_active) audit!(user.id, integration, audit_event_type) integration end) @@ -1549,8 +1612,11 @@ defmodule Kodo.Integrations do end end - defp maybe_lock_transition(user_id, _integration_id, "refresh_succeeded") do - lock_user!(user_id) + defp maybe_lock_transition(user_id, integration_id, "refresh_succeeded") do + case Repo.get_by(Integration, id: integration_id, user_id: user_id) do + %Integration{provider: provider} -> lock_provider_identity(user_id, provider) + nil -> lock_user!(user_id) + end end defp maybe_lock_transition(_user_id, _integration_id, _event_type), do: :ok @@ -1582,6 +1648,16 @@ defmodule Kodo.Integrations do defp maybe_activate_initial_oauth(_user_id, integration, _event_type), do: integration + defp maybe_restore_refresh_selection(user_id, integration, true) do + if !integration.active and !active_provider_account_exists?(user_id, integration.provider) do + integration |> change(active: true, updated_at: now()) |> Repo.update!() + else + integration + end + end + + defp maybe_restore_refresh_selection(_user_id, integration, _restore_active), do: integration + defp provider_account_exists?(user_id, provider) do Repo.exists?( from integration in Integration, diff --git a/lib/kodo/integrations/device_authorization.ex b/lib/kodo/integrations/device_authorization.ex index 1f6763e..41bd255 100644 --- a/lib/kodo/integrations/device_authorization.ex +++ b/lib/kodo/integrations/device_authorization.ex @@ -12,17 +12,24 @@ defmodule Kodo.Integrations.DeviceAuthorization do def begin(%Scope{} = scope, integration_id, expected_generation, opts \\ []) do client = Keyword.get(opts, :client, configured_client()) client_options = Keyword.get(opts, :client_options, []) + claim_owner_id = Keyword.get_lazy(opts, :claim_owner_id, &Ecto.UUID.generate/0) + supervisor = Keyword.get(opts, :supervisor, Kodo.ControlPlaneTaskSupervisor) with {:ok, created} <- safe_create(client, client_options), - {:ok, attempt} <- - Integrations.begin_device_authorization( + {:ok, {attempt, claim, payload}} <- + Integrations.begin_claimed_device_authorization( scope, integration_id, expected_generation, created.payload, - created.polling_interval_ms - ), - {:ok, task} <- resume(scope, integration_id, opts) do + created.polling_interval_ms, + claim_owner_id + ) do + task = + Task.Supervisor.async_nolink(supervisor, fn -> + run(scope, claim, payload, opts) + end) + {:ok, %{ attempt: attempt, @@ -78,7 +85,9 @@ defmodule Kodo.Integrations.DeviceAuthorization do fail(scope, admission, reason) end else - {:error, reason} -> stop(reason) + {:error, reason} -> + _result = Integrations.expire_device_authorization(scope, claim.integration_id) + stop(reason) end end diff --git a/lib/kodo/integrations/oauth_refresh.ex b/lib/kodo/integrations/oauth_refresh.ex index 3badca2..9676297 100644 --- a/lib/kodo/integrations/oauth_refresh.ex +++ b/lib/kodo/integrations/oauth_refresh.ex @@ -49,7 +49,7 @@ defmodule Kodo.Integrations.OAuthRefresh do wait_for_claim(scope, integration, owner_id, deadline, opts) {:error, :stale_credential_generation} -> - current_result(scope, integration) + current_result(scope, integration, opts) {:error, :integration_reauthorization_required} -> {:error, :integration_reauthorization_required} @@ -66,7 +66,7 @@ defmodule Kodo.Integrations.OAuthRefresh do @wait_interval_ms -> claim_or_wait(scope, integration, owner_id, deadline, opts) end else - current_result(scope, integration) + current_result(scope, integration, opts) end end @@ -97,12 +97,13 @@ defmodule Kodo.Integrations.OAuthRefresh do defp perform(scope, claim, opts) do client = Keyword.get(opts, :client, configured_client()) client_options = Keyword.get(opts, :client_options, []) + refresh_started_at = Keyword.get(opts, :now, DateTime.utc_now()) with {:ok, admitted} <- Integrations.admit_refresh_claim(scope, claim), {:ok, current} <- CredentialEncryption.decrypt(admitted), {:ok, refresh_token} <- fetch_refresh_token(current), {:ok, response} <- safe_refresh(client, refresh_token, client_options), - {:ok, normalized} <- RefreshTokens.normalize(response, current), + {:ok, normalized} <- RefreshTokens.normalize(response, current, refresh_started_at), {:ok, integration} <- Integrations.refresh_succeeded( scope, @@ -110,21 +111,23 @@ defmodule Kodo.Integrations.OAuthRefresh do claim.refresh_claim_generation, normalized.credentials, expires_at: normalized.expires_at, - refreshed_at: DateTime.utc_now() - ) do - {:ok, integration} + refreshed_at: DateTime.utc_now(), + restore_active: claim.active + ), + {:ok, usable} <- ensure_persisted_access_is_fresh(scope, integration, opts) do + {:ok, usable} else {:error, :invalid_grant} -> - require_reauthorization(scope, claim, "refresh_invalid_grant") + require_reauthorization(scope, claim, "refresh_invalid_grant", opts) {:error, :refresh_account_identity_mismatch} -> - require_reauthorization(scope, claim, "refresh_account_identity_mismatch") + require_reauthorization(scope, claim, "refresh_account_identity_mismatch", opts) {:error, :stale_credential_generation} -> - current_result(scope, claim) + current_result(scope, claim, opts) {:error, :stale_refresh_claim} -> - current_result(scope, claim) + current_result(scope, claim, opts) {:error, _reason} -> _result = Integrations.fail_refresh(scope, claim) @@ -132,20 +135,28 @@ defmodule Kodo.Integrations.OAuthRefresh do end end - defp require_reauthorization(scope, claim, error_code) do + defp ensure_persisted_access_is_fresh(scope, integration, opts) do + if DateTime.after?(integration.expires_at, DateTime.utc_now()) do + {:ok, integration} + else + refresh(scope, integration, Keyword.delete(opts, :now)) + end + end + + defp require_reauthorization(scope, claim, error_code, opts) do case Integrations.require_refresh_reauthorization(scope, claim, error_code) do {:ok, _integration} -> {:error, :integration_reauthorization_required} - {:error, :stale_refresh_claim} -> current_result(scope, claim) + {:error, :stale_refresh_claim} -> current_result(scope, claim, opts) {:error, _reason} -> {:error, :provider_unavailable} end end - defp current_result(scope, integration) do + defp current_result(scope, integration, opts) do case Integrations.get_integration(scope, integration.id) do {:ok, %{connection_status: "connected"} = current} when current.refresh_source_generation == integration.credential_generation and current.credential_generation == integration.credential_generation + 1 -> - {:ok, current} + ensure_persisted_access_is_fresh(scope, current, opts) {:ok, %{connection_status: "reauthorization_required"}} -> {:error, :integration_reauthorization_required} diff --git a/test/kodo/integrations/device_authorization_test.exs b/test/kodo/integrations/device_authorization_test.exs index 4d74d08..32e3fd7 100644 --- a/test/kodo/integrations/device_authorization_test.exs +++ b/test/kodo/integrations/device_authorization_test.exs @@ -118,6 +118,22 @@ defmodule Kodo.Integrations.DeviceAuthorizationTest do assert calls(client) == [] end + test "terminalizes an attempt whose deadline expires while its worker is paused", %{ + scope: scope, + integration: integration + } do + client = fake_client(poll: [{:ok, exchange_payload()}]) + {claim, payload} = claimed_attempt(scope, integration) + + Repo.update!(change(claim, provider_deadline: DateTime.add(DateTime.utc_now(), -1, :second))) + + assert {:error, :stale_device_authorization_claim} = + DeviceAuthorization.run(scope, claim, payload, client_opts(client)) + + assert %{state: "expired", encrypted_payload: nil} = Repo.reload!(claim) + assert calls(client) == [] + end + test "a healthy claim cannot be replaced by another task using the same node owner", %{ scope: scope, integration: integration @@ -272,6 +288,55 @@ defmodule Kodo.Integrations.DeviceAuthorizationTest do assert Repo.reload!(started.attempt).state == "failed" end + test "begin claims the new attempt before another worker can take it", %{ + scope: scope, + integration: integration + } do + owner = self() + supervisor = start_supervised!(Task.Supervisor) + + blocking_poll = fn :poll, _payload -> + send(owner, {:begin_worker_claimed, self()}) + + receive do + :finish_begin_worker -> {:error, :device_authorization_rejected} + end + end + + client = + fake_client( + create: [ + {:ok, + %{ + payload: %{"device_auth_id" => "device", "user_code" => "CODE"}, + polling_interval_ms: 0, + verification_url: "https://auth.openai.com/codex/device" + }} + ], + poll: [blocking_poll] + ) + + assert {:ok, started} = + DeviceAuthorization.begin( + scope, + integration.id, + integration.credential_generation, + client_opts(client, supervisor: supervisor, claim_owner_id: Ecto.UUID.generate()) + ) + + assert_receive {:begin_worker_claimed, worker} + + assert {:error, :device_authorization_not_claimable} = + Integrations.claim_device_authorization( + scope, + integration.id, + Ecto.UUID.generate() + ) + + send(worker, :finish_begin_worker) + assert {:error, :device_authorization_rejected} = Task.await(started.task) + end + test "begin rejects a client-selected verification destination before persistence", %{ scope: scope, integration: integration diff --git a/test/kodo/integrations/oauth_refresh_test.exs b/test/kodo/integrations/oauth_refresh_test.exs index f57cf40..7a2431a 100644 --- a/test/kodo/integrations/oauth_refresh_test.exs +++ b/test/kodo/integrations/oauth_refresh_test.exs @@ -164,6 +164,112 @@ defmodule Kodo.Integrations.OAuthRefreshTest do assert persisted["access_token"] == second_tokens["access_token"] end + test "a delayed success restores its active selection after provisional invalid grant", + context do + owner = self() + + first_response = fn _refresh_token -> + send(owner, {:successful_refresh_started, self()}) + + receive do + :finish_successful_refresh -> {:ok, rotated_tokens("account", "recovered")} + end + end + + client = fake_client([first_response, {:error, :invalid_grant}]) + first = Task.async(fn -> refresh(context, client) end) + assert_receive {:successful_refresh_started, first_worker} + + context.integration + |> Repo.reload!() + |> change(refresh_lease_expires_at: DateTime.add(DateTime.utc_now(), -1, :second)) + |> Repo.update!() + + assert {:error, :integration_reauthorization_required} = refresh(context, client) + refute Repo.reload!(context.integration).active + + send(first_worker, :finish_successful_refresh) + assert {:ok, recovered} = Task.await(first) + assert recovered.active + + assert {:ok, active} = + Integrations.get_active_integration_by_provider(context.scope, "openai_codex") + + assert active.id == context.integration.id + end + + test "a delayed success does not replace a newer explicit active selection", context do + owner = self() + + first_response = fn _refresh_token -> + send(owner, {:selectable_refresh_started, self()}) + + receive do + :finish_selectable_refresh -> {:ok, rotated_tokens("account", "late")} + end + end + + client = fake_client([first_response, {:error, :invalid_grant}]) + first = Task.async(fn -> refresh(context, client) end) + assert_receive {:selectable_refresh_started, first_worker} + + context.integration + |> Repo.reload!() + |> change(refresh_lease_expires_at: DateTime.add(DateTime.utc_now(), -1, :second)) + |> Repo.update!() + + assert {:error, :integration_reauthorization_required} = refresh(context, client) + + sibling = oauth_integration(context.scope, "Second subscription") + + assert {:ok, sibling} = + Integrations.oauth_succeeded(context.scope, sibling.id, 0, current_tokens(), + expires_at: DateTime.add(DateTime.utc_now(), 3_600, :second) + ) + + assert {:ok, sibling} = + Integrations.activate( + context.scope, + sibling.id, + sibling.credential_generation + ) + + send(first_worker, :finish_selectable_refresh) + assert {:ok, recovered} = Task.await(first) + refute recovered.active + assert Repo.reload!(sibling).active + end + + test "persists a rotated response that expires during a worker pause and refreshes it again", + context do + admitted_at = DateTime.add(DateTime.utc_now(), -2, :second) + + expired_rotation = + rotated_tokens("account", "short") + |> put_in( + ["access_token"], + jwt(%{"exp" => DateTime.utc_now() |> DateTime.to_unix() |> Kernel.-(1)}) + ) + + client = fake_client([{:ok, expired_rotation}, {:ok, rotated_tokens("account", "usable")}]) + + assert {:ok, refreshed} = + OAuthRefresh.ensure_fresh(context.scope, context.integration, + force: true, + now: admitted_at, + client: FakeRefreshClient, + client_options: [agent: client], + supervisor: context.supervisor, + claim_owner_id: Ecto.UUID.generate() + ) + + assert refreshed.credential_generation == context.integration.credential_generation + 2 + assert Agent.get(client, &Enum.reverse(&1.refresh_tokens)) == ["old-refresh", "short-refresh"] + + assert {:ok, credentials} = CredentialEncryption.decrypt(refreshed) + assert credentials["refresh_token"] == "usable-refresh" + end + test "persists a rotated response after its caller stops waiting", context do owner = self() @@ -229,6 +335,36 @@ defmodule Kodo.Integrations.OAuthRefreshTest do reauthorizing.credential_generation end + test "an expired abandoned authorization attempt is terminalized before refresh", context do + assert {:ok, attempt} = + Integrations.begin_device_authorization( + context.scope, + context.integration.id, + context.integration.credential_generation, + %{"device_auth_id" => "abandoned", "user_code" => "OLD"}, + 0 + ) + + Repo.update!( + change(attempt, provider_deadline: DateTime.add(DateTime.utc_now(), -1, :second)) + ) + + reauthorizing = Repo.reload!(context.integration) + client = fake_client([{:ok, rotated_tokens("account")}]) + + assert {:ok, _refreshed} = + OAuthRefresh.ensure_fresh(context.scope, reauthorizing, + force: true, + client: FakeRefreshClient, + client_options: [agent: client], + supervisor: context.supervisor, + claim_owner_id: Ecto.UUID.generate() + ) + + assert %{state: "expired", encrypted_payload: nil} = Repo.reload!(attempt) + assert Agent.get(client, &Enum.reverse(&1.refresh_tokens)) == ["old-refresh"] + end + test "user deletion cascades OAuth state and fences an in-flight refresh", context do owner = self() @@ -281,6 +417,16 @@ defmodule Kodo.Integrations.OAuthRefreshTest do ) end + defp oauth_integration(scope, display_name) do + %Integration{user_id: scope.user.id} + |> Integration.create_changeset(%{ + provider: "openai_codex", + authentication_type: "oauth", + display_name: display_name + }) + |> Repo.insert!() + end + defp current_tokens do %{ "access_token" => "old-access", diff --git a/test/kodo/integrations_test.exs b/test/kodo/integrations_test.exs index 4827a2a..67a822f 100644 --- a/test/kodo/integrations_test.exs +++ b/test/kodo/integrations_test.exs @@ -1389,6 +1389,58 @@ defmodule Kodo.IntegrationsTest do assert Repo.reload!(integration).connection_status == "disconnected" end + test "completion accepts an exchange admitted before its lease expires", %{scope: scope} do + integration = oauth_integration(scope) + + assert {:ok, _attempt} = + Integrations.begin_device_authorization( + scope, + integration.id, + integration.credential_generation, + %{"device_auth_id" => "device", "user_code" => "CODE"}, + 0 + ) + + assert {:ok, {claim, _payload}} = + Integrations.claim_device_authorization( + scope, + integration.id, + Ecto.UUID.generate() + ) + + assert {:ok, claim} = + Integrations.store_device_authorization_exchange(scope, claim, %{ + "authorization_code" => "code", + "code_challenge" => "challenge", + "code_verifier" => "verifier" + }) + + assert {:ok, {admission, _payload}} = + Integrations.admit_device_authorization_exchange(scope, claim) + + Repo.update!( + change(admission, claim_lease_expires_at: DateTime.add(DateTime.utc_now(), -1, :second)) + ) + + credentials = %{ + "access_token" => "access", + "refresh_token" => "refresh", + "id_token" => "identity", + "account_id" => "account" + } + + assert {:ok, completed} = + Integrations.complete_device_authorization( + scope, + admission, + credentials, + DateTime.add(DateTime.utc_now(), 3_600, :second) + ) + + assert completed.connection_status == "connected" + assert Repo.reload!(admission).state == "completed" + end + test "cancellation is owned and generation fenced and removes the one-time code", %{ scope: scope } do From 52122bc1e1eda2d5a0b78413a9dfe068a307019d Mon Sep 17 00:00:00 2001 From: Darwin Wu Date: Sat, 12 Sep 2026 21:13:35 +0000 Subject: [PATCH 12/14] fix(integrations): stop terminal authorization retries --- lib/kodo_web/live/integrations_live.ex | 72 +++++++++++++++---- test/kodo_web/live/integrations_live_test.exs | 44 ++++++++++++ 2 files changed, 103 insertions(+), 13 deletions(-) diff --git a/lib/kodo_web/live/integrations_live.ex b/lib/kodo_web/live/integrations_live.ex index 7d74e7f..009e322 100644 --- a/lib/kodo_web/live/integrations_live.ex +++ b/lib/kodo_web/live/integrations_live.ex @@ -64,6 +64,7 @@ defmodule KodoWeb.IntegrationsLive do |> assign(:api_key_form, empty_form()) |> assign(:validation_tasks, %{}) |> assign(:device_authorization_tasks, %{}) + |> assign(:device_authorization_retries, %{}) |> load_integrations() socket = @@ -371,10 +372,13 @@ defmodule KodoWeb.IntegrationsLive do end def handle_info( - {:device_authorization_changed, _integration_id, _attempt_generation, state, _error}, + {:device_authorization_changed, integration_id, attempt_generation, state, _error}, socket ) do - socket = load_integrations(socket) + socket = + socket + |> maybe_clear_device_authorization_retry(integration_id, attempt_generation, state) + |> load_integrations() if state == "failed" do {:noreply, put_flash(socket, :error, "Authorization failed. Try again.")} @@ -383,8 +387,19 @@ defmodule KodoWeb.IntegrationsLive do end end - def handle_info({:retry_device_authorization, integration_id}, socket) do - {:noreply, resume_device_authorization(socket, integration_id, true)} + def handle_info({:retry_device_authorization, integration_id, attempt_generation}, socket) do + socket = + socket + |> clear_device_authorization_retry(integration_id, attempt_generation) + |> load_integrations() + + case socket.assigns.device_authorizations[integration_id] do + %{attempt_generation: ^attempt_generation} -> + {:noreply, resume_device_authorization(socket, integration_id, attempt_generation, true)} + + _terminal_or_replaced -> + {:noreply, socket} + end end defp save_api_key( @@ -514,28 +529,59 @@ defmodule KodoWeb.IntegrationsLive do end defp resume_device_authorizations(socket) do - Enum.reduce(socket.assigns.device_authorizations, socket, fn {integration_id, _attempt}, - acc -> - resume_device_authorization(acc, integration_id, true) + Enum.reduce(socket.assigns.device_authorizations, socket, fn {integration_id, attempt}, acc -> + resume_device_authorization(acc, integration_id, attempt.attempt_generation, true) end) end - defp resume_device_authorization(socket, integration_id, schedule_retry?) do + defp resume_device_authorization(socket, integration_id, attempt_generation, schedule_retry?) do case DeviceAuthorization.resume(socket.assigns.current_scope, integration_id) do {:ok, task} -> - track_device_authorization_task(socket, task, integration_id) + socket + |> clear_device_authorization_retry(integration_id, attempt_generation) + |> track_device_authorization_task(task, integration_id) {:error, :device_authorization_not_claimable} when schedule_retry? -> + schedule_device_authorization_retry(socket, integration_id, attempt_generation) + + {:error, _reason} -> + socket + end + end + + defp schedule_device_authorization_retry(socket, integration_id, attempt_generation) do + key = {integration_id, attempt_generation} + + if Map.has_key?(socket.assigns.device_authorization_retries, key) do + socket + else + timer = Process.send_after( self(), - {:retry_device_authorization, integration_id}, + {:retry_device_authorization, integration_id, attempt_generation}, @device_authorization_retry_ms ) - socket + update(socket, :device_authorization_retries, &Map.put(&1, key, timer)) + end + end - {:error, _reason} -> - socket + defp maybe_clear_device_authorization_retry(socket, _integration_id, _generation, "active"), + do: socket + + defp maybe_clear_device_authorization_retry(socket, integration_id, generation, _state), + do: clear_device_authorization_retry(socket, integration_id, generation) + + defp clear_device_authorization_retry(socket, integration_id, attempt_generation) do + key = {integration_id, attempt_generation} + + case Map.pop(socket.assigns.device_authorization_retries, key) do + {nil, retries} -> + assign(socket, :device_authorization_retries, retries) + + {timer, retries} -> + Process.cancel_timer(timer) + assign(socket, :device_authorization_retries, retries) end end diff --git a/test/kodo_web/live/integrations_live_test.exs b/test/kodo_web/live/integrations_live_test.exs index ef9597a..a56112d 100644 --- a/test/kodo_web/live/integrations_live_test.exs +++ b/test/kodo_web/live/integrations_live_test.exs @@ -354,6 +354,50 @@ defmodule KodoWeb.IntegrationsLiveTest do assert Repo.reload!(attempt).state == "cancelled" end + test "stops remount retries when the exact authorization attempt becomes terminal", %{ + conn: conn, + scope: scope + } do + assert {:ok, integration} = + Integrations.create_oauth_integration(scope, "openai_codex", + display_name: "Contended subscription" + ) + + assert {:ok, attempt} = + Integrations.begin_device_authorization( + scope, + integration.id, + integration.credential_generation, + %{"device_auth_id" => "device", "user_code" => "CONTENDED"}, + 30_000 + ) + + assert {:ok, {_claim, _payload}} = + Integrations.claim_device_authorization( + scope, + integration.id, + Ecto.UUID.generate() + ) + + {:ok, view, _html} = live(conn, ~p"/integrations") + retry_key = {integration.id, attempt.attempt_generation} + assert Map.has_key?(live_assign(view, :device_authorization_retries), retry_key) + + assert {:ok, _cancelled} = + Integrations.cancel_device_authorization( + scope, + attempt.id, + attempt.attempt_generation + ) + + _ = :sys.get_state(view.pid) + assert live_assign(view, :device_authorization_retries) == %{} + + send(view.pid, {:retry_device_authorization, integration.id, attempt.attempt_generation}) + _ = :sys.get_state(view.pid) + assert live_assign(view, :device_authorization_retries) == %{} + end + test "shows completion and reauthorization after the device flow succeeds", %{ conn: conn, scope: scope From 9b9bb8d06099a689149a524b8d95e23a6dae2170 Mon Sep 17 00:00:00 2001 From: Darwin Wu Date: Sat, 12 Sep 2026 21:13:35 +0000 Subject: [PATCH 13/14] fix(integrations): fence recovery selection intent --- lib/kodo/integrations.ex | 75 ++++++++- lib/kodo/integrations/integration.ex | 1 + lib/kodo/integrations/oauth_refresh.ex | 17 +- lib/kodo_web/live/integrations_live.ex | 20 +++ ...estore_active_to_provider_integrations.exs | 9 ++ priv/repo/structure.sql | 6 +- test/kodo/integrations/oauth_refresh_test.exs | 146 +++++++++++++++++- test/kodo/integrations_test.exs | 62 ++++++++ test/kodo_web/live/integrations_live_test.exs | 85 ++++++++++ 9 files changed, 415 insertions(+), 6 deletions(-) create mode 100644 priv/repo/migrations/20260912213842_add_refresh_restore_active_to_provider_integrations.exs diff --git a/lib/kodo/integrations.ex b/lib/kodo/integrations.ex index f2ad691..346118a 100644 --- a/lib/kodo/integrations.ex +++ b/lib/kodo/integrations.ex @@ -90,6 +90,24 @@ defmodule Kodo.Integrations do end end + @doc false + def lock_turn_integrations(%Scope{user: user}, providers) when is_list(providers) do + providers = providers |> Enum.uniq() |> Enum.sort() + Enum.each(providers, &lock_provider_identity(user.id, &1)) + + Integration + |> where( + [integration], + integration.user_id == ^user.id and integration.provider in ^providers and + integration.active + ) + |> order_by([integration], asc: integration.provider, asc: integration.id) + |> lock("FOR SHARE") + |> Repo.all() + + :ok + end + @doc false def admit_execution_route_change(%Scope{user: user}, provider) do integration = @@ -313,6 +331,7 @@ defmodule Kodo.Integrations do Repo.transaction(fn -> update_claimed_device_authorization_payload(user.id, claim, encrypted) end) + |> notify_device_authorization_change() end end @@ -552,6 +571,7 @@ defmodule Kodo.Integrations do set: [ refresh_claim_owner_id: ^claim_owner_id, refresh_claim_generation: ^generation, + refresh_restore_active: integration.active, refresh_lease_expires_at: fragment( "clock_timestamp() + (? * interval '1 millisecond')", @@ -577,7 +597,9 @@ defmodule Kodo.Integrations do def admit_refresh_claim(%Scope{user: user}, %Integration{} = claim) do active_attempt = from attempt in DeviceAuthorizationAttempt, - where: attempt.integration_id == ^claim.id and attempt.state == "active" + where: + attempt.integration_id == ^claim.id and attempt.state == "active" and + attempt.provider_deadline > fragment("clock_timestamp()") query = exact_refresh_claim_query(user.id, claim) @@ -602,6 +624,7 @@ defmodule Kodo.Integrations do set: [ refresh_claim_owner_id: nil, refresh_claim_generation: nil, + refresh_restore_active: false, refresh_lease_expires_at: nil, updated_at: now() ] @@ -652,6 +675,7 @@ defmodule Kodo.Integrations do set: [ refresh_claim_owner_id: nil, refresh_claim_generation: nil, + refresh_restore_active: false, refresh_lease_expires_at: nil, updated_at: now() ] @@ -733,6 +757,7 @@ defmodule Kodo.Integrations do validation_error_code: nil, refresh_claim_owner_id: nil, refresh_claim_generation: nil, + refresh_restore_active: false, refresh_lease_expires_at: nil } ) @@ -770,6 +795,8 @@ defmodule Kodo.Integrations do defp activate_locked(user_id, integration, provider_integrations, generation, audit?) do with :ok <- require_generation(integration, generation), :ok <- require_connected(integration) do + invalidate_provider_selection_recovery(user_id, integration.provider) + cond do integration.active -> integration !audit? and Enum.any?(provider_integrations, & &1.active) -> integration @@ -1234,6 +1261,7 @@ defmodule Kodo.Integrations do refresh_source_generation: nil, refresh_claim_owner_id: nil, refresh_claim_generation: nil, + refresh_restore_active: false, refresh_lease_expires_at: nil, updated_at: timestamp ] @@ -1422,6 +1450,7 @@ defmodule Kodo.Integrations do validation_error_code: nil, refresh_claim_owner_id: nil, refresh_claim_generation: nil, + refresh_restore_active: false, refresh_lease_expires_at: nil }) |> Map.put(:restore_active, opts[:restore_active]) @@ -1536,6 +1565,9 @@ defmodule Kodo.Integrations do {restore_active, changes} = Map.pop(changes, :restore_active, false) + restore_active = + restore_active == true and refresh_selection_restorable?(user.id, id, generation) + integration = execute_fenced_update(user.id, id, generation, allowed_connections, changes) @@ -1619,6 +1651,17 @@ defmodule Kodo.Integrations do end end + defp maybe_lock_transition(user_id, integration_id, "integration_disconnected") do + case Repo.get_by(Integration, id: integration_id, user_id: user_id) do + %Integration{provider: provider} -> + lock_provider_identity(user_id, provider) + invalidate_provider_selection_recovery(user_id, provider) + + nil -> + :ok + end + end + defp maybe_lock_transition(_user_id, _integration_id, _event_type), do: :ok defp maybe_terminalize_device_authorization(user_id, integration_id, "integration_disconnected") do @@ -1658,6 +1701,36 @@ defmodule Kodo.Integrations do defp maybe_restore_refresh_selection(_user_id, integration, _restore_active), do: integration + defp refresh_selection_restorable?(user_id, integration_id, generation) do + Repo.exists?( + from integration in Integration, + where: + integration.id == ^integration_id and integration.user_id == ^user_id and + integration.credential_generation == ^generation and + integration.refresh_restore_active + ) + end + + defp invalidate_provider_selection_recovery(user_id, provider) do + Integration + |> where( + [integration], + integration.user_id == ^user_id and integration.provider == ^provider and + integration.refresh_restore_active + ) + |> Repo.update_all(set: [refresh_restore_active: false, updated_at: now()]) + + DeviceAuthorizationAttempt + |> where( + [attempt], + attempt.user_id == ^user_id and attempt.provider == ^provider and attempt.state == "active" and + attempt.activate_on_completion + ) + |> Repo.update_all(set: [activate_on_completion: false, updated_at: now()]) + + :ok + end + defp provider_account_exists?(user_id, provider) do Repo.exists?( from integration in Integration, diff --git a/lib/kodo/integrations/integration.ex b/lib/kodo/integrations/integration.ex index 967c478..82831ef 100644 --- a/lib/kodo/integrations/integration.ex +++ b/lib/kodo/integrations/integration.ex @@ -39,6 +39,7 @@ defmodule Kodo.Integrations.Integration do field :refresh_claim_epoch, :integer, default: 0 field :refresh_claim_generation, :integer field :refresh_lease_expires_at, :utc_datetime_usec + field :refresh_restore_active, :boolean, default: false belongs_to :user, Kodo.Accounts.User, type: :id diff --git a/lib/kodo/integrations/oauth_refresh.ex b/lib/kodo/integrations/oauth_refresh.ex index 9676297..e902611 100644 --- a/lib/kodo/integrations/oauth_refresh.ex +++ b/lib/kodo/integrations/oauth_refresh.ex @@ -112,7 +112,7 @@ defmodule Kodo.Integrations.OAuthRefresh do normalized.credentials, expires_at: normalized.expires_at, refreshed_at: DateTime.utc_now(), - restore_active: claim.active + restore_active: claim.refresh_restore_active ), {:ok, usable} <- ensure_persisted_access_is_fresh(scope, integration, opts) do {:ok, usable} @@ -127,6 +127,7 @@ defmodule Kodo.Integrations.OAuthRefresh do current_result(scope, claim, opts) {:error, :stale_refresh_claim} -> + _result = Integrations.release_refresh_claim(scope, claim) current_result(scope, claim, opts) {:error, _reason} -> @@ -139,7 +140,19 @@ defmodule Kodo.Integrations.OAuthRefresh do if DateTime.after?(integration.expires_at, DateTime.utc_now()) do {:ok, integration} else - refresh(scope, integration, Keyword.delete(opts, :now)) + remaining = Keyword.get(opts, :expired_success_retries, 1) + + if remaining > 0 do + refresh( + scope, + integration, + opts + |> Keyword.delete(:now) + |> Keyword.put(:expired_success_retries, remaining - 1) + ) + else + {:error, :provider_unavailable} + end end end diff --git a/lib/kodo_web/live/integrations_live.ex b/lib/kodo_web/live/integrations_live.ex index 009e322..6185d72 100644 --- a/lib/kodo_web/live/integrations_live.ex +++ b/lib/kodo_web/live/integrations_live.ex @@ -482,6 +482,26 @@ defmodule KodoWeb.IntegrationsLive do socket |> assign(:integrations, integrations) |> assign(:device_authorizations, device_authorizations) + |> reconcile_device_authorization_retries() + end + + defp reconcile_device_authorization_retries(socket) do + active_keys = + MapSet.new(socket.assigns.device_authorizations, fn {integration_id, attempt} -> + {integration_id, attempt.attempt_generation} + end) + + retries = + Enum.reduce(socket.assigns.device_authorization_retries, %{}, fn {key, timer}, kept -> + if MapSet.member?(active_keys, key) do + Map.put(kept, key, timer) + else + Process.cancel_timer(timer) + kept + end + end) + + assign(socket, :device_authorization_retries, retries) end # Browser-facing state needs lifecycle metadata only. In particular, keeping diff --git a/priv/repo/migrations/20260912213842_add_refresh_restore_active_to_provider_integrations.exs b/priv/repo/migrations/20260912213842_add_refresh_restore_active_to_provider_integrations.exs new file mode 100644 index 0000000..15ce076 --- /dev/null +++ b/priv/repo/migrations/20260912213842_add_refresh_restore_active_to_provider_integrations.exs @@ -0,0 +1,9 @@ +defmodule Kodo.Repo.Migrations.AddRefreshRestoreActiveToProviderIntegrations do + use Ecto.Migration + + def change do + alter table(:provider_integrations) do + add :refresh_restore_active, :boolean, null: false, default: false + end + end +end diff --git a/priv/repo/structure.sql b/priv/repo/structure.sql index 140e388..ab75fb7 100644 --- a/priv/repo/structure.sql +++ b/priv/repo/structure.sql @@ -2,7 +2,7 @@ -- PostgreSQL database dump -- -\restrict zlmayUQMxfTEc8sR5OS893sgwJPRmRbQbN8OxqjM51G1hsJJMUZbtQRN809vFYQ +\restrict bhNtvgo2sLbQAwP0xzeErVvadl8Hb8oQ9uHPQt7qDfIjayod6OKCW93lzqbuYlP -- Dumped from database version 15.19 (Debian 15.19-0+deb12u1) -- Dumped by pg_dump version 15.19 (Debian 15.19-0+deb12u1) @@ -189,6 +189,7 @@ CREATE TABLE public.provider_integrations ( refresh_claim_generation bigint, refresh_lease_expires_at timestamp without time zone, refresh_source_generation bigint, + refresh_restore_active boolean DEFAULT false NOT NULL, CONSTRAINT provider_integrations_active_connected CHECK (((NOT active) OR ((connection_status)::text = 'connected'::text))), CONSTRAINT provider_integrations_authentication_type_valid CHECK (((authentication_type)::text = ANY ((ARRAY['api_key'::character varying, 'oauth'::character varying])::text[]))), CONSTRAINT provider_integrations_credential_generation_valid CHECK ((credential_generation >= 0)), @@ -783,7 +784,7 @@ ALTER TABLE ONLY public.users_tokens -- PostgreSQL database dump complete -- -\unrestrict zlmayUQMxfTEc8sR5OS893sgwJPRmRbQbN8OxqjM51G1hsJJMUZbtQRN809vFYQ +\unrestrict bhNtvgo2sLbQAwP0xzeErVvadl8Hb8oQ9uHPQt7qDfIjayod6OKCW93lzqbuYlP INSERT INTO public."schema_migrations" (version) VALUES (20260807073017); INSERT INTO public."schema_migrations" (version) VALUES (20260808062115); @@ -805,3 +806,4 @@ INSERT INTO public."schema_migrations" (version) VALUES (20260907161929); INSERT INTO public."schema_migrations" (version) VALUES (20260910044734); INSERT INTO public."schema_migrations" (version) VALUES (20260912195208); INSERT INTO public."schema_migrations" (version) VALUES (20260912200217); +INSERT INTO public."schema_migrations" (version) VALUES (20260912213842); diff --git a/test/kodo/integrations/oauth_refresh_test.exs b/test/kodo/integrations/oauth_refresh_test.exs index 7a2431a..2b97e81 100644 --- a/test/kodo/integrations/oauth_refresh_test.exs +++ b/test/kodo/integrations/oauth_refresh_test.exs @@ -5,6 +5,7 @@ defmodule Kodo.Integrations.OAuthRefreshTest do alias Kodo.Integrations alias Kodo.Integrations.CredentialEncryption alias Kodo.Integrations.DeviceAuthorizationAttempt + alias Kodo.Integrations.DeviceAuthorizationEncryption alias Kodo.Integrations.Integration alias Kodo.Integrations.OAuthAccessValidation alias Kodo.Integrations.OAuthRefresh @@ -234,10 +235,56 @@ defmodule Kodo.Integrations.OAuthRefreshTest do sibling.credential_generation ) + assert {:ok, _disconnected} = + Integrations.disconnect( + context.scope, + sibling.id, + sibling.credential_generation + ) + send(first_worker, :finish_selectable_refresh) assert {:ok, recovered} = Task.await(first) refute recovered.active - assert Repo.reload!(sibling).active + + assert {:error, :integration_not_found} = + Integrations.get_active_integration_by_provider(context.scope, "openai_codex") + end + + test "an ordinary delayed refresh does not undo a later decision to leave no active account", + context do + owner = self() + + response = fn _refresh_token -> + send(owner, {:ordinary_refresh_started, self()}) + + receive do + :finish_ordinary_refresh -> {:ok, rotated_tokens("account", "late")} + end + end + + client = fake_client([response]) + refresh = Task.async(fn -> refresh(context, client) end) + assert_receive {:ordinary_refresh_started, worker} + + sibling = oauth_integration(context.scope, "Temporary selection") + + assert {:ok, sibling} = + Integrations.oauth_succeeded(context.scope, sibling.id, 0, current_tokens(), + expires_at: DateTime.add(DateTime.utc_now(), 3_600, :second) + ) + + assert {:ok, sibling} = + Integrations.activate(context.scope, sibling.id, sibling.credential_generation) + + assert {:ok, _disconnected} = + Integrations.disconnect(context.scope, sibling.id, sibling.credential_generation) + + send(worker, :finish_ordinary_refresh) + assert {:ok, recovered} = Task.await(refresh) + refute recovered.active + + assert {:error, :integration_not_found} = + Integrations.get_active_integration_by_provider(context.scope, "openai_codex") end test "persists a rotated response that expires during a worker pause and refreshes it again", @@ -270,6 +317,32 @@ defmodule Kodo.Integrations.OAuthRefreshTest do assert credentials["refresh_token"] == "usable-refresh" end + test "bounds automatic recovery when refreshed access is already expired", context do + admitted_at = DateTime.add(DateTime.utc_now(), -10, :second) + + expired = + rotated_tokens("account", "expired") + |> put_in( + ["access_token"], + jwt(%{"exp" => DateTime.utc_now() |> DateTime.to_unix() |> Kernel.-(1)}) + ) + + client = fake_client([{:ok, expired}, {:ok, expired}]) + + assert {:error, :provider_unavailable} = + OAuthRefresh.ensure_fresh(context.scope, context.integration, + force: true, + now: admitted_at, + client: FakeRefreshClient, + client_options: [agent: client], + supervisor: context.supervisor, + claim_owner_id: Ecto.UUID.generate() + ) + + _ = :sys.get_state(client) + assert length(Agent.get(client, & &1.refresh_tokens)) == 2 + end + test "persists a rotated response after its caller stops waiting", context do owner = self() @@ -365,6 +438,44 @@ defmodule Kodo.Integrations.OAuthRefreshTest do assert Agent.get(client, &Enum.reverse(&1.refresh_tokens)) == ["old-refresh"] end + test "an attempt expiring after refresh claim does not veto admission", context do + owner_id = Ecto.UUID.generate() + + assert {:ok, claim} = + Integrations.claim_refresh( + context.scope, + context.integration.id, + context.integration.credential_generation, + owner_id + ) + + insert_device_attempt(context, DateTime.add(DateTime.utc_now(), -1, :second)) + + assert {:ok, admitted} = Integrations.admit_refresh_claim(context.scope, claim) + assert admitted.refresh_claim_owner_id == owner_id + assert :ok = Integrations.release_refresh_claim(context.scope, claim) + end + + test "a pre-network admission failure can release its exact refresh claim", context do + owner_id = Ecto.UUID.generate() + + assert {:ok, claim} = + Integrations.claim_refresh( + context.scope, + context.integration.id, + context.integration.credential_generation, + owner_id + ) + + insert_device_attempt(context, DateTime.add(DateTime.utc_now(), 60, :second)) + + assert {:error, :stale_refresh_claim} = + Integrations.admit_refresh_claim(context.scope, claim) + + assert :ok = Integrations.release_refresh_claim(context.scope, claim) + assert is_nil(Repo.reload!(context.integration).refresh_claim_owner_id) + end + test "user deletion cascades OAuth state and fences an in-flight refresh", context do owner = self() @@ -427,6 +538,39 @@ defmodule Kodo.Integrations.OAuthRefreshTest do |> Repo.insert!() end + defp insert_device_attempt(context, deadline) do + timestamp = DateTime.utc_now() + + attempt = %DeviceAuthorizationAttempt{ + id: Ecto.UUID.generate(), + user_id: context.scope.user.id, + integration_id: context.integration.id, + provider: "openai_codex", + attempt_generation: 1, + expected_integration_generation: context.integration.credential_generation, + activate_on_completion: false + } + + assert {:ok, encrypted} = + DeviceAuthorizationEncryption.encrypt(attempt, %{ + "device_auth_id" => "device", + "user_code" => "CODE" + }) + + attempt + |> Map.merge(encrypted) + |> DeviceAuthorizationAttempt.create_changeset(%{ + provider: "openai_codex", + attempt_generation: 1, + expected_integration_generation: context.integration.credential_generation, + activate_on_completion: false, + provider_deadline: deadline, + polling_interval_ms: 0, + next_poll_at: timestamp + }) + |> Repo.insert!() + end + defp current_tokens do %{ "access_token" => "old-access", diff --git a/test/kodo/integrations_test.exs b/test/kodo/integrations_test.exs index 67a822f..148fc67 100644 --- a/test/kodo/integrations_test.exs +++ b/test/kodo/integrations_test.exs @@ -1441,6 +1441,68 @@ defmodule Kodo.IntegrationsTest do assert Repo.reload!(admission).state == "completed" end + test "device completion does not undo a later decision to leave no active account", %{ + scope: scope + } do + integration = oauth_integration(scope) + + assert {:ok, integration} = + Integrations.oauth_succeeded(scope, integration.id, 0, %{ + "access_token" => "old-access", + "refresh_token" => "old-refresh", + "account_id" => "old-account" + }) + + assert {:ok, _attempt} = + Integrations.begin_device_authorization( + scope, + integration.id, + integration.credential_generation, + %{"device_auth_id" => "device", "user_code" => "CODE"}, + 0 + ) + + assert {:ok, {claim, _payload}} = + Integrations.claim_device_authorization( + scope, + integration.id, + Ecto.UUID.generate() + ) + + sibling = oauth_integration(scope) + + assert {:ok, sibling} = + Integrations.oauth_succeeded(scope, sibling.id, 0, %{ + "access_token" => "sibling-access", + "refresh_token" => "sibling-refresh", + "account_id" => "sibling-account" + }) + + assert {:ok, sibling} = + Integrations.activate(scope, sibling.id, sibling.credential_generation) + + assert {:ok, _disconnected} = + Integrations.disconnect(scope, sibling.id, sibling.credential_generation) + + assert {:ok, completed} = + Integrations.complete_device_authorization( + scope, + Repo.reload!(claim), + %{ + "access_token" => "new-access", + "refresh_token" => "new-refresh", + "id_token" => "new-identity", + "account_id" => "old-account" + }, + DateTime.add(DateTime.utc_now(), 3_600, :second) + ) + + refute completed.active + + assert {:error, :integration_not_found} = + Integrations.get_active_integration_by_provider(scope, "openai_codex") + end + test "cancellation is owned and generation fenced and removes the one-time code", %{ scope: scope } do diff --git a/test/kodo_web/live/integrations_live_test.exs b/test/kodo_web/live/integrations_live_test.exs index a56112d..2f5bcb4 100644 --- a/test/kodo_web/live/integrations_live_test.exs +++ b/test/kodo_web/live/integrations_live_test.exs @@ -398,6 +398,91 @@ defmodule KodoWeb.IntegrationsLiveTest do assert live_assign(view, :device_authorization_retries) == %{} end + test "drops an old retry immediately when a new authorization supersedes it", %{ + conn: conn, + scope: scope + } do + assert {:ok, integration} = + Integrations.create_oauth_integration(scope, "openai_codex", + display_name: "Superseded subscription" + ) + + assert {:ok, first} = + Integrations.begin_device_authorization( + scope, + integration.id, + integration.credential_generation, + %{"device_auth_id" => "first", "user_code" => "FIRST"}, + 30_000 + ) + + assert {:ok, {_claim, _payload}} = + Integrations.claim_device_authorization( + scope, + integration.id, + Ecto.UUID.generate() + ) + + {:ok, view, _html} = live(conn, ~p"/integrations") + retry_key = {integration.id, first.attempt_generation} + assert Map.has_key?(live_assign(view, :device_authorization_retries), retry_key) + + assert {:ok, second} = + Integrations.begin_device_authorization( + scope, + integration.id, + first.expected_integration_generation, + %{"device_auth_id" => "second", "user_code" => "SECOND"}, + 30_000 + ) + + _ = :sys.get_state(view.pid) + assert second.attempt_generation > first.attempt_generation + assert live_assign(view, :device_authorization_retries) == %{} + end + + test "updates an open authorization panel when polling advances to exchange", %{ + conn: conn, + scope: scope + } do + assert {:ok, integration} = + Integrations.create_oauth_integration(scope, "openai_codex", + display_name: "Advancing subscription" + ) + + assert {:ok, _attempt} = + Integrations.begin_device_authorization( + scope, + integration.id, + integration.credential_generation, + %{"device_auth_id" => "device", "user_code" => "ADVANCE"}, + 30_000 + ) + + assert {:ok, {claim, _payload}} = + Integrations.claim_device_authorization( + scope, + integration.id, + Ecto.UUID.generate() + ) + + {:ok, view, _html} = live(conn, ~p"/integrations") + panel = "#integration-#{integration.id}-device-authorization" + assert has_element?(view, panel, "Waiting for OpenAI authorization") + assert has_element?(view, "#integration-#{integration.id}-copy-device-code") + + assert {:ok, _exchange} = + Integrations.store_device_authorization_exchange(scope, claim, %{ + "authorization_code" => "code", + "code_challenge" => "challenge", + "code_verifier" => "verifier" + }) + + _ = :sys.get_state(view.pid) + assert has_element?(view, panel, "Completing OpenAI authorization") + refute has_element?(view, "#integration-#{integration.id}-copy-device-code") + end + test "shows completion and reauthorization after the device flow succeeds", %{ conn: conn, scope: scope From 5b9466456562e9f761edede5d948f4ff6375076c Mon Sep 17 00:00:00 2001 From: Darwin Wu Date: Sat, 12 Sep 2026 21:13:35 +0000 Subject: [PATCH 14/14] security(agent): bind inference to durable admission --- lib/kodo/agent/evaluation_runner.ex | 4 +- lib/kodo/agent/execution_route_change.ex | 1 + lib/kodo/agent/loop.ex | 31 ++-- lib/kodo/agent/model_mapping.ex | 2 +- lib/kodo/llm.ex | 25 +++- lib/kodo/llm/credential_resolver.ex | 4 + lib/kodo/llm/invocation_admission.ex | 20 +++ lib/kodo/sessions.ex | 77 +++++++--- .../agent/execution_route_change_test.exs | 10 ++ test/kodo/agent/loop_test.exs | 13 +- test/kodo/llm/credential_resolver_test.exs | 35 ++--- test/kodo/sessions_test.exs | 133 +++++++++++++++++- 12 files changed, 283 insertions(+), 72 deletions(-) create mode 100644 lib/kodo/llm/invocation_admission.ex diff --git a/lib/kodo/agent/evaluation_runner.ex b/lib/kodo/agent/evaluation_runner.ex index 0a518ca..479684b 100644 --- a/lib/kodo/agent/evaluation_runner.ex +++ b/lib/kodo/agent/evaluation_runner.ex @@ -488,7 +488,7 @@ defmodule Kodo.Agent.EvaluationRunner do defp generate(scope, adapter, role, messages, tools) do with {:ok, model, reference} <- LLM.resolve_integration(scope, role["model"]) do - LLM.generate(scope, model, reference, messages, tools, + LLM.generate_evaluation(scope, model, reference, messages, tools, adapter: adapter, timeout: @timeout, reasoning: role["reasoning"] @@ -498,7 +498,7 @@ defmodule Kodo.Agent.EvaluationRunner do defp generate_object(scope, adapter, role, messages, schema) do with {:ok, model, reference} <- LLM.resolve_integration(scope, role["model"]) do - LLM.generate_object(scope, model, reference, messages, schema, + LLM.generate_object_evaluation(scope, model, reference, messages, schema, adapter: adapter, timeout: @timeout, reasoning: role["reasoning"] diff --git a/lib/kodo/agent/execution_route_change.ex b/lib/kodo/agent/execution_route_change.ex index e4002d6..cc9b36b 100644 --- a/lib/kodo/agent/execution_route_change.ex +++ b/lib/kodo/agent/execution_route_change.ex @@ -18,6 +18,7 @@ defmodule Kodo.Agent.ExecutionRouteChange do [] <- incompatible_roles(affected, destination, records) do {:ok, apply_route(mapping, affected, destination)} else + {:error, _reason} = error -> error true -> {:error, :execution_route_unchanged} roles when is_list(roles) -> {:error, {:incompatible_execution_route, roles}} end diff --git a/lib/kodo/agent/loop.ex b/lib/kodo/agent/loop.ex index 61a7419..8569011 100644 --- a/lib/kodo/agent/loop.ex +++ b/lib/kodo/agent/loop.ex @@ -170,7 +170,7 @@ defmodule Kodo.Agent.Loop do ), :ok <- Phoenix.PubSub.subscribe(Kodo.PubSub, "session:#{context.session_id}"), :ok <- Runners.subscribe_lifecycle(), - {:ok, review_invocation_id} <- + {:ok, {review_invocation_id, admission}} <- start_review_invocation( context.session_id, primary_invocation_id, @@ -198,7 +198,7 @@ defmodule Kodo.Agent.Loop do generate_object( context.adapter, - request, + admission, [ %{"role" => "system", "content" => contract.prompt}, %{ @@ -266,12 +266,13 @@ defmodule Kodo.Agent.Loop do "capability_validation" => capability_validation, "model_mapping" => mapping }, + request.model, request.credential, version: 2, parent_id: primary_invocation_id, ownership: ownership ) do - {:ok, _event} -> {:ok, invocation_id} + {:ok, {_event, admission}} -> {:ok, {invocation_id, admission}} error -> normalize_admission_error(error, request) end end @@ -487,7 +488,7 @@ defmodule Kodo.Agent.Loop do {:ok, capability_validation} <- adapter.validate_model(primary_model, primary, contract), {:ok, request} <- resolve_request(session_id, primary), - {:ok, invocation_id} <- + {:ok, {invocation_id, admission}} <- start_invocation( session_id, invocation, @@ -502,7 +503,7 @@ defmodule Kodo.Agent.Loop do Sessions.dispatch_if_owner(ownership, fn -> generate( adapter, - request, + admission, transcript(events, contract, primary["execution_route"] || primary["provider"]), tools, timeout: budgets[:model_timeout], @@ -640,11 +641,12 @@ defmodule Kodo.Agent.Loop do "capability_validation" => capability_validation, "model_mapping" => mapping }, + request.model, request.credential, version: 4, ownership: ownership ) do - {:ok, _event} -> {:ok, invocation_id} + {:ok, {_event, admission}} -> {:ok, {invocation_id, admission}} error -> normalize_admission_error(error, request) end end @@ -997,7 +999,7 @@ defmodule Kodo.Agent.Loop do with :ok <- within_budget(continuation, tokens, state.contract.budget), :ok <- rehoming_boundary(), {:ok, request} <- resolve_request(state.context.session_id, state.search), - {:ok, invocation_id} <- + {:ok, {invocation_id, admission}} <- start_subagent_invocation( state.parent_call, state.search, @@ -1012,7 +1014,7 @@ defmodule Kodo.Agent.Loop do Sessions.dispatch_if_owner(state.context.ownership, fn -> generate( state.context.adapter, - request, + admission, messages, Tools.definitions_for_turn( state.contract.toolset_version, @@ -1147,12 +1149,13 @@ defmodule Kodo.Agent.Loop do "capability_validation" => capability_validation, "model_mapping" => context.mapping }, + request.model, request.credential, version: 2, parent_id: context.invocation_id, ownership: context.ownership ) do - {:ok, _event} -> {:ok, invocation_id} + {:ok, {_event, admission}} -> {:ok, {invocation_id, admission}} error -> normalize_admission_error(error, request) end end @@ -1701,20 +1704,18 @@ defmodule Kodo.Agent.Loop do defp normalize_admission_error(error, _request), do: error - defp generate(adapter, request, messages, tools, opts) do + defp generate(adapter, admission, messages, tools, opts) do LLM.generate_admitted( - request.model, - request.credential, + admission, messages, tools, Keyword.put(opts, :adapter, adapter) ) end - defp generate_object(adapter, request, messages, schema, opts) do + defp generate_object(adapter, admission, messages, schema, opts) do LLM.generate_object_admitted( - request.model, - request.credential, + admission, messages, schema, Keyword.put(opts, :adapter, adapter) diff --git a/lib/kodo/agent/model_mapping.ex b/lib/kodo/agent/model_mapping.ex index a5f9d69..c7fe6b3 100644 --- a/lib/kodo/agent/model_mapping.ex +++ b/lib/kodo/agent/model_mapping.ex @@ -72,7 +72,7 @@ defmodule Kodo.Agent.ModelMapping do role_mapping |> Map.put("execution_route", role_mapping["provider"]) |> Map.put("model_selector", model_selector(role_mapping["model"])) - |> Map.put("capability_contract", capability_contract(contract))} + |> Map.put_new("capability_contract", capability_contract(contract))} end) %{mapping | "roles" => roles} diff --git a/lib/kodo/llm.ex b/lib/kodo/llm.ex index 484640f..adb6167 100644 --- a/lib/kodo/llm.ex +++ b/lib/kodo/llm.ex @@ -11,6 +11,7 @@ defmodule Kodo.LLM do alias Kodo.LLM.Credential alias Kodo.LLM.CredentialResolver alias Kodo.LLM.IntegrationRef + alias Kodo.LLM.InvocationAdmission alias Kodo.LLM.ProviderError @credential_option_keys ~w( @@ -72,8 +73,8 @@ defmodule Kodo.LLM do end end - @doc "Generates text after rechecking and decrypting the referenced credential." - def generate( + @doc "Runs the out-of-band evaluation harness with a scoped credential." + def generate_evaluation( %Scope{} = scope, %LLMDB.Model{} = model, %IntegrationRef{} = reference, @@ -95,7 +96,12 @@ defmodule Kodo.LLM do do: resolve_credential(scope, model, reference) @doc false - def generate_admitted(%LLMDB.Model{} = model, %Credential{} = credential, messages, tools, opts) do + def generate_admitted( + %InvocationAdmission{model: model, credential: credential}, + messages, + tools, + opts + ) do adapter = Keyword.get(opts, :adapter, adapter()) adapter_opts = Keyword.delete(opts, :adapter) @@ -104,8 +110,11 @@ defmodule Kodo.LLM do end end - @doc "Generates a structured object after scoped credential admission." - def generate_object( + def generate_admitted(_admission, _messages, _tools, _opts), + do: {:error, :invocation_not_admitted} + + @doc "Runs structured generation for the out-of-band evaluation harness." + def generate_object_evaluation( %Scope{} = scope, %LLMDB.Model{} = model, %IntegrationRef{} = reference, @@ -124,8 +133,7 @@ defmodule Kodo.LLM do @doc false def generate_object_admitted( - %LLMDB.Model{} = model, - %Credential{} = credential, + %InvocationAdmission{model: model, credential: credential}, messages, schema, opts @@ -138,6 +146,9 @@ defmodule Kodo.LLM do end end + def generate_object_admitted(_admission, _messages, _schema, _opts), + do: {:error, :invocation_not_admitted} + def adapter, do: Application.get_env(:kodo, :llm_adapter, Kodo.LLM.ReqLLM) defp resolve_model(model) do diff --git a/lib/kodo/llm/credential_resolver.ex b/lib/kodo/llm/credential_resolver.ex index c5e400b..6cc59d1 100644 --- a/lib/kodo/llm/credential_resolver.ex +++ b/lib/kodo/llm/credential_resolver.ex @@ -48,6 +48,7 @@ defmodule Kodo.LLM.CredentialResolver do :ok <- require_authentication_type(integration, reference), :ok <- require_billing_path(integration, reference), :ok <- require_usable(integration), + :ok <- require_active(integration), {:ok, integration} <- OAuthRefresh.ensure_fresh(scope, integration), {:ok, payload} <- CredentialEncryption.decrypt(integration) do build_credential(integration, reference, payload) @@ -113,6 +114,9 @@ defmodule Kodo.LLM.CredentialResolver do defp require_billing_path(%Integration{}, %IntegrationRef{}), do: {:error, :integration_billing_mismatch} + defp require_active(%Integration{active: true}), do: :ok + defp require_active(%Integration{}), do: {:error, :stale_credential_generation} + defp require_usable(%Integration{connection_status: "disconnected"}), do: {:error, :integration_disconnected} diff --git a/lib/kodo/llm/invocation_admission.ex b/lib/kodo/llm/invocation_admission.ex new file mode 100644 index 0000000..ea883cd --- /dev/null +++ b/lib/kodo/llm/invocation_admission.ex @@ -0,0 +1,20 @@ +defmodule Kodo.LLM.InvocationAdmission do + @moduledoc false + + alias Kodo.LLM.Credential + + @enforce_keys [:model, :credential] + defstruct [:model, :credential] + + @opaque t :: %__MODULE__{model: LLMDB.Model.t(), credential: Credential.t()} + + # This wrapper is an internal call-graph boundary, not a bearer capability: Elixir structs are + # constructible inside the VM. Production agent paths receive it only from the session operation + # that atomically records invocation start; the out-of-band evaluation harness is deliberately + # named and kept separate. Avoid adding a database recheck at dispatch, because an admitted + # provider operation is allowed to finish after replacement or disconnection. + @doc false + def new(%LLMDB.Model{} = model, %Credential{} = credential) do + %__MODULE__{model: model, credential: credential} + end +end diff --git a/lib/kodo/sessions.ex b/lib/kodo/sessions.ex index 9ba88c9..2db0f59 100644 --- a/lib/kodo/sessions.ex +++ b/lib/kodo/sessions.ex @@ -16,6 +16,8 @@ defmodule Kodo.Sessions do alias Kodo.Integrations alias Kodo.LLM alias Kodo.LLM.Credential + alias Kodo.LLM.IntegrationRef + alias Kodo.LLM.InvocationAdmission alias Kodo.LLM.ProviderError alias Kodo.Repo alias Kodo.Runners.Runner @@ -547,15 +549,16 @@ defmodule Kodo.Sessions do session_id, type, payload, + %LLMDB.Model{} = model, %Credential{} = credential, opts \\ [] ) do transaction = fn -> - append_admitted_model_event_locked(session_id, type, payload, credential, opts) + append_admitted_model_event_locked(session_id, type, payload, model, credential, opts) end case Repo.transaction(transaction) do - {:ok, event} = result -> + {:ok, {event, _admission}} = result -> broadcast(event) result @@ -591,6 +594,10 @@ defmodule Kodo.Sessions do |> turn_route_snapshot() |> Map.fetch!("model_mapping") + preflight_mapping(scope, mapping) + end + + defp preflight_mapping(scope, mapping) do mapping["roles"] |> Map.values() |> Enum.uniq_by(& &1["execution_route"]) @@ -603,36 +610,55 @@ defmodule Kodo.Sessions do end defp admit_model_credential!(session, credential) do - admitted? = + integration = Integrations.Integration + |> where([integration], integration.id == ^credential.integration_id) + |> where([integration], integration.user_id == ^session.user_id) + |> where([integration], integration.provider == ^credential.provider) + |> where( + [integration], + integration.authentication_type == ^credential.authentication_type + ) |> where( [integration], - integration.id == ^credential.integration_id and - integration.user_id == ^session.user_id and integration.provider == ^credential.provider and - integration.authentication_type == ^credential.authentication_type and - integration.credential_generation == ^credential.credential_generation and - integration.connection_status == "connected" and integration.active and - integration.validation_status != "invalid" + integration.credential_generation == ^credential.credential_generation ) + |> where([integration], integration.connection_status == "connected") + |> where([integration], integration.active) + |> where([integration], integration.validation_status != "invalid") |> lock("FOR SHARE") - |> Repo.exists?() + |> Repo.one() + + case integration do + %Integrations.Integration{} = admitted -> + reference = IntegrationRef.from_integration(admitted) + + if reference.billing_path == credential.billing_path, + do: admitted, + else: Repo.rollback(:stale_credential_generation) - if !admitted?, do: Repo.rollback(:stale_credential_generation) + nil -> + Repo.rollback(:stale_credential_generation) + end end - defp append_admitted_model_event_locked(session_id, type, payload, credential, opts) do + defp append_admitted_model_event_locked(session_id, type, payload, model, credential, opts) do session = lock_session!(session_id, opts) - admit_model_credential!(session, credential) + integration = admit_model_credential!(session, credential) + billing_path = IntegrationRef.from_integration(integration).billing_path + + if Atom.to_string(model.provider) != integration.provider, + do: Repo.rollback(:stale_credential_generation) payload = Map.merge(payload, %{ - "provider" => credential.provider, - "authentication_type" => credential.authentication_type, - "billing_path" => Atom.to_string(credential.billing_path) + "provider" => integration.provider, + "authentication_type" => integration.authentication_type, + "billing_path" => Atom.to_string(billing_path) }) case append_locked(session, type, payload, opts) do - {:ok, event} -> event + {:ok, event} -> {event, InvocationAdmission.new(model, credential)} {:error, changeset} -> Repo.rollback(changeset) end end @@ -1082,8 +1108,12 @@ defmodule Kodo.Sessions do |> maybe_put_request_id(client_request_id) route_snapshot = turn_route_snapshot(session) + scope = Scope.for_user(Repo.get!(User, session.user_id)) + mapping = route_snapshot["model_mapping"] - with {:ok, snapshot} <- + with :ok <- lock_turn_integrations(scope, mapping), + :ok <- preflight_mapping(scope, mapping), + {:ok, snapshot} <- append_locked( session, "turn_route_snapshot", @@ -1104,11 +1134,20 @@ defmodule Kodo.Sessions do append_locked(session, "session_status_changed", %{"status" => "running"}, []) do [snapshot, message, status_event] else - {:error, changeset} -> Repo.rollback(changeset) + {:error, reason} -> Repo.rollback(reason) end end end + defp lock_turn_integrations(scope, mapping) do + providers = + mapping["roles"] + |> Map.values() + |> Enum.map(& &1["execution_route"]) + + Integrations.lock_turn_integrations(scope, providers) + end + defp change_execution_route_locked(scope, session_id, destination) do session = lock_user_session!(scope, session_id) projection = session.id |> events_after() |> Projection.from_events() diff --git a/test/kodo/agent/execution_route_change_test.exs b/test/kodo/agent/execution_route_change_test.exs index 184bf0b..2f4a5d8 100644 --- a/test/kodo/agent/execution_route_change_test.exs +++ b/test/kodo/agent/execution_route_change_test.exs @@ -74,6 +74,16 @@ defmodule Kodo.Agent.ExecutionRouteChangeTest do ExecutionRouteChange.change(changed_requirement, "openai_codex", records) end + test "returns a bounded error for a partially explicit malformed snapshot" do + malformed = + mapping_for_selector("gpt-4o-mini") + |> ModelMapping.snapshot() + |> update_in(["roles", "primary"], &Map.delete(&1, "capability_contract")) + + assert {:error, :invalid_model_mapping_snapshot} = + ExecutionRouteChange.change(malformed, "openai_codex", []) + end + defp mapping_for_selector(selector) do ModelMapping.balanced([ {"test", diff --git a/test/kodo/agent/loop_test.exs b/test/kodo/agent/loop_test.exs index 90dc94a..183a91a 100644 --- a/test/kodo/agent/loop_test.exs +++ b/test/kodo/agent/loop_test.exs @@ -2,6 +2,7 @@ defmodule Kodo.Agent.LoopTest do use Kodo.DataCase alias Kodo.Agent.Loop + alias Kodo.Agent.ModelMapping alias Kodo.Agent.ModelSettings alias Kodo.Agent.Tools alias Kodo.Integrations @@ -148,18 +149,16 @@ defmodule Kodo.Agent.LoopTest do } do [created] = Sessions.events_after(session.id) - malformed = - update_in( - created.payload, - ["model_mapping", "roles", "primary"], - &Map.delete(&1, "capability_contract") - ) + baseline = ModelMapping.snapshot(created.payload["model_mapping"]) + assert {:ok, ^baseline} = ModelMapping.validate_snapshot(baseline) + + malformed = update_in(baseline, ["roles", "primary"], &Map.delete(&1, "capability_contract")) {:ok, _snapshot} = Sessions.append_event( session.id, "turn_route_snapshot", - %{"route_revision" => 1, "model_mapping" => malformed["model_mapping"]}, + %{"route_revision" => 1, "model_mapping" => malformed}, source: "system", ownership: ownership ) diff --git a/test/kodo/llm/credential_resolver_test.exs b/test/kodo/llm/credential_resolver_test.exs index da386ed..8fa7e9b 100644 --- a/test/kodo/llm/credential_resolver_test.exs +++ b/test/kodo/llm/credential_resolver_test.exs @@ -36,8 +36,6 @@ defmodule Kodo.LLM.CredentialResolverTest do assert {:ok, resolved_model, reference} = LLM.resolve_integration(context.scope, "openai:gpt-4o-mini") - assert {:ok, credential} = LLM.admit_credential(context.scope, resolved_model, reference) - for unsafe_option <- [ {:api_key, "caller-secret"}, {:base_url, "https://attacker.example"}, @@ -46,7 +44,7 @@ defmodule Kodo.LLM.CredentialResolverTest do {:finch_request, fn _request -> :forged end} ] do assert {:error, :credential_options_not_allowed} = - LLM.generate( + LLM.generate_evaluation( context.scope, resolved_model, reference, @@ -57,20 +55,19 @@ defmodule Kodo.LLM.CredentialResolverTest do [unsafe_option] ) ) - - assert {:error, :credential_options_not_allowed} = - LLM.generate_admitted( - resolved_model, - credential, - [%{"role" => "user", "content" => "final answer"}], - [], - Keyword.merge( - [adapter: Kodo.Test.FakeLLM, timeout: 1_000], - [unsafe_option] - ) - ) end + assert {:ok, credential} = LLM.admit_credential(context.scope, resolved_model, reference) + + assert {:error, :invocation_not_admitted} = + LLM.generate_admitted( + credential, + [%{"role" => "user", "content" => "final answer"}], + [], + adapter: Kodo.Test.FakeLLM, + timeout: 1_000 + ) + assert {:ok, _replaced} = Integrations.replace_credentials( context.scope, @@ -80,7 +77,7 @@ defmodule Kodo.LLM.CredentialResolverTest do ) assert {:error, %Kodo.LLM.ProviderError{} = error} = - LLM.generate( + LLM.generate_evaluation( context.scope, resolved_model, reference, @@ -223,7 +220,7 @@ defmodule Kodo.LLM.CredentialResolverTest do assert invalid.credential_generation == reference.credential_generation end - test "keeps admitted references pinned while new preflight selects the activated account", + test "rejects a selected reference when another account becomes active", context do assert {:ok, second} = Integrations.connect(context.scope, "openai", "api_key", %{ @@ -233,9 +230,7 @@ defmodule Kodo.LLM.CredentialResolverTest do assert {:ok, _activated} = Integrations.activate(context.scope, second.id, second.credential_generation) - assert {:ok, admitted} = resolve(context) - assert admitted.integration_id == context.integration.id - assert admitted.token == "scoped-secret" + assert {:error, :stale_credential_generation} = resolve(context) assert {:ok, next_reference} = CredentialResolver.reference(context.scope, model()) assert next_reference.integration_id == second.id diff --git a/test/kodo/sessions_test.exs b/test/kodo/sessions_test.exs index 32c2ec4..50b840e 100644 --- a/test/kodo/sessions_test.exs +++ b/test/kodo/sessions_test.exs @@ -315,6 +315,9 @@ defmodule Kodo.SessionsTest do runner: runner, scope: scope } do + assert {:ok, _integration} = + Integrations.connect(scope, "openai", "api_key", %{"api_key" => "turn-key"}) + {:ok, session} = Sessions.create_session(scope, %{ runner_id: runner.id, @@ -387,6 +390,73 @@ defmodule Kodo.SessionsTest do assert Enum.map(Sessions.events_after(session.id), & &1.type) == ["session_created"] end + test "rechecks every role inside the transaction that accepts the turn", %{ + runner: runner, + scope: scope + } do + assert {:ok, _override} = + ModelSettings.put_user_override(scope, :review, %{ + model: "anthropic:claude-3-5-haiku-latest" + }) + + assert {:ok, _primary} = + Integrations.connect(scope, "openai", "api_key", %{"api_key" => "primary-key"}) + + assert {:ok, review} = + Integrations.connect(scope, "anthropic", "api_key", %{"api_key" => "review-key"}) + + assert {:ok, session} = + Sessions.create_session(scope, %{ + runner_id: runner.id, + title: "Atomic role preflight", + model: "openai:gpt-4o-mini" + }) + + assert {:ok, _disconnected} = + Integrations.disconnect(scope, review.id, review.credential_generation) + + assert {:error, %ProviderError{provider: "anthropic", kind: :integration_required}} = + Sessions.begin_turn(session.id, "Do not accept") + + assert Enum.map(Sessions.events_after(session.id), & &1.type) == ["session_created"] + end + + test "holds provider selection locks through turn acceptance", %{ + runner: runner, + scope: scope + } do + assert {:ok, _integration} = + Integrations.connect(scope, "openai", "api_key", %{"api_key" => "locked-key"}) + + assert {:ok, session} = + Sessions.create_session(scope, %{ + runner_id: runner.id, + title: "Atomic provider preflight", + model: "openai:gpt-4o-mini" + }) + + assert {:ok, _events} = Sessions.begin_turn(session.id, "Accept atomically") + + key = "provider-integration:#{scope.user.id}:openai" + + lock_available? = + Task.async(fn -> + Ecto.Adapters.SQL.Sandbox.unboxed_run(Repo, fn -> + %{rows: [[available?]]} = + Ecto.Adapters.SQL.query!( + Repo, + "SELECT pg_try_advisory_xact_lock(hashtextextended($1, 0))", + [key] + ) + + available? + end) + end) + |> Task.await() + + refute lock_available? + end + test "atomically records invocation admission for the exact active credential generation", %{ runner: runner, scope: scope @@ -411,11 +481,21 @@ defmodule Kodo.SessionsTest do "billing_path" => "forged-billing" } - assert {:ok, event} = + assert {:error, :stale_credential_generation} = + Sessions.append_admitted_model_event( + session.id, + "model_invocation_started", + payload, + model, + %{credential | billing_path: :aggregator} + ) + + assert {:ok, {event, admission}} = Sessions.append_admitted_model_event( session.id, "model_invocation_started", payload, + model, credential ) @@ -428,6 +508,25 @@ defmodule Kodo.SessionsTest do refute inspect(event) =~ "admitted-key" assert integration.credential_generation == credential.credential_generation + + assert {:error, :credential_options_not_allowed} = + LLM.generate_admitted(admission, [], [], + adapter: Kodo.Test.FakeLLM, + timeout: 1_000, + api_key: "forged" + ) + + assert {:ok, _disconnected} = + Integrations.disconnect(scope, integration.id, integration.credential_generation) + + assert {:ok, %{type: :final_answer}} = + LLM.generate_admitted( + admission, + [%{"role" => "user", "content" => "final answer"}], + [], + adapter: Kodo.Test.FakeLLM, + timeout: 1_000 + ) end test "rejects a decrypted credential replaced before invocation admission", %{ @@ -462,6 +561,7 @@ defmodule Kodo.SessionsTest do session.id, "model_invocation_started", %{"invocation_id" => Ecto.UUID.generate()}, + model, credential ) @@ -516,6 +616,21 @@ defmodule Kodo.SessionsTest do runner: runner, scope: scope } do + assert {:ok, codex} = + Integrations.create_oauth_integration(scope, "openai_codex", + display_name: "Route subscription" + ) + + assert {:ok, _codex} = + Integrations.oauth_succeeded(scope, codex.id, codex.credential_generation, %{ + "access_token" => "route-access", + "refresh_token" => "route-refresh", + "account_id" => "route-account" + }) + + assert {:ok, _platform} = + Integrations.connect(scope, "openai", "api_key", %{"api_key" => "route-key"}) + {:ok, session} = Sessions.create_session(scope, %{ runner_id: runner.id, @@ -550,6 +665,13 @@ defmodule Kodo.SessionsTest do records ) + changed_mapping = + put_in( + changed_mapping, + ["roles", "primary", "capability_contract", "requirements", "min_context"], + 200_000 + ) + assert {:ok, route_event} = Sessions.append_event( session.id, @@ -580,6 +702,15 @@ defmodule Kodo.SessionsTest do "primary", "execution_route" ]) == "openai" + + assert get_in(second_snapshot.payload, [ + "model_mapping", + "roles", + "primary", + "capability_contract", + "requirements", + "min_context" + ]) == 200_000 end test "advances ownership epochs and fences every stale state mutation", %{