From c350fc85873b533ec16e40e8c8a229cd5c6c0294 Mon Sep 17 00:00:00 2001 From: Darwin Wu Date: Fri, 4 Sep 2026 02:32:18 +0000 Subject: [PATCH 1/2] fix(agent): persist canonical invocation providers --- lib/kodo/agent/loop.ex | 29 +++++++--- lib/kodo/llm/credential_resolver.ex | 62 +++++++++++++++++----- lib/kodo/llm/integration_ref.ex | 28 ++++++++-- test/kodo/agent/loop_test.exs | 7 ++- test/kodo/llm/credential_resolver_test.exs | 12 +++++ 5 files changed, 112 insertions(+), 26 deletions(-) diff --git a/lib/kodo/agent/loop.ex b/lib/kodo/agent/loop.ex index a446468..1220e48 100644 --- a/lib/kodo/agent/loop.ex +++ b/lib/kodo/agent/loop.ex @@ -185,6 +185,7 @@ defmodule Kodo.Agent.Loop do review, contract, capability_validation, + request, context.ownership ), {:ok, diff} <- @@ -249,6 +250,7 @@ defmodule Kodo.Agent.Loop do review, contract, capability_validation, + request, ownership ) do invocation_id = Ecto.UUID.generate() @@ -260,15 +262,18 @@ defmodule Kodo.Agent.Loop do "invocation_id" => invocation_id, "primary_invocation_id" => primary_invocation_id, "role" => "review", - "provider" => review["provider"], + "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), "reasoning" => review["reasoning"], "role_contract" => contract.id, "toolset_version" => contract.toolset_version, "capability_validation" => capability_validation, "model_mapping" => mapping }, - version: 1, + version: 2, parent_id: primary_invocation_id, ownership: ownership ) do @@ -476,6 +481,7 @@ defmodule Kodo.Agent.Loop do mapping, primary, capability_validation, + request, ownership ), :ok <- rehoming_boundary(), @@ -559,6 +565,7 @@ defmodule Kodo.Agent.Loop do mapping, primary, capability_validation, + request, ownership ) do invocation_id = Ecto.UUID.generate() @@ -570,15 +577,18 @@ defmodule Kodo.Agent.Loop do "invocation_id" => invocation_id, "continuation" => continuation, "role" => "primary", - "provider" => primary["provider"], + "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), "reasoning" => primary["reasoning"], "role_contract" => primary["role_contract"], "toolset_version" => primary["toolset_version"], "capability_validation" => capability_validation, "model_mapping" => mapping }, - version: 3, + version: 4, ownership: ownership ) do {:ok, _event} -> {:ok, invocation_id} @@ -927,6 +937,7 @@ defmodule Kodo.Agent.Loop do state.contract, state.capability_validation, continuation, + request, state.context ), :ok <- rehoming_boundary(), @@ -1045,6 +1056,7 @@ defmodule Kodo.Agent.Loop do contract, capability_validation, continuation, + request, context ) do invocation_id = Ecto.UUID.generate() @@ -1057,15 +1069,18 @@ defmodule Kodo.Agent.Loop do "delegation_tool_call_id" => parent_call["id"], "continuation" => continuation, "role" => "search", - "provider" => search["provider"], + "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), "reasoning" => search["reasoning"], "role_contract" => contract.id, "toolset_version" => contract.toolset_version, "capability_validation" => capability_validation, "model_mapping" => context.mapping }, - version: 1, + version: 2, parent_id: context.invocation_id, ownership: context.ownership ) do @@ -1633,6 +1648,8 @@ defmodule Kodo.Agent.Loop do ) end + defp canonical_provider(request), do: Atom.to_string(request.model.provider) + defp within_budget(invocations, tokens, budgets) do cond do invocations > budgets[:max_continuations] -> {:error, :continuation_budget_exceeded} diff --git a/lib/kodo/llm/credential_resolver.ex b/lib/kodo/llm/credential_resolver.ex index 901ccd4..034dba9 100644 --- a/lib/kodo/llm/credential_resolver.ex +++ b/lib/kodo/llm/credential_resolver.ex @@ -44,9 +44,11 @@ defmodule Kodo.LLM.CredentialResolver do :ok <- require_generation(integration, reference.credential_generation), :ok <- require_provider(reference.provider, provider), :ok <- require_provider(integration.provider, provider), + :ok <- require_authentication_type(integration, reference), + :ok <- require_billing_path(integration, reference), :ok <- require_usable(integration), {:ok, payload} <- CredentialEncryption.decrypt(integration) do - build_credential(integration, payload) + build_credential(integration, reference, payload) end end @@ -58,10 +60,13 @@ defmodule Kodo.LLM.CredentialResolver do defp require_reference(%IntegrationRef{ integration_id: id, provider: provider, - credential_generation: generation + authentication_type: authentication_type, + credential_generation: generation, + billing_path: billing_path }) - when is_binary(provider) and provider in @providers and is_integer(generation) and - generation >= 0 do + when is_binary(provider) and provider in @providers and + authentication_type in ["api_key", "oauth"] and is_integer(generation) and + generation >= 0 and billing_path in [:platform, :subscription, :aggregator] do case Ecto.UUID.cast(id) do {:ok, _id} -> :ok :error -> {:error, :invalid_integration_reference} @@ -78,6 +83,34 @@ defmodule Kodo.LLM.CredentialResolver do defp require_generation(%Integration{}, _generation), do: {:error, :stale_credential_generation} + defp require_authentication_type( + %Integration{authentication_type: authentication_type}, + %IntegrationRef{authentication_type: authentication_type} + ), + do: :ok + + defp require_authentication_type(%Integration{}, %IntegrationRef{}), + do: {:error, :integration_authentication_mismatch} + + defp require_billing_path(%Integration{provider: "openai_codex"}, %IntegrationRef{ + billing_path: :subscription + }), + do: :ok + + defp require_billing_path(%Integration{provider: "openrouter"}, %IntegrationRef{ + billing_path: :aggregator + }), + do: :ok + + defp require_billing_path(%Integration{provider: provider}, %IntegrationRef{ + billing_path: :platform + }) + when provider in ~w(openai anthropic), + do: :ok + + defp require_billing_path(%Integration{}, %IntegrationRef{}), + do: {:error, :integration_billing_mismatch} + defp require_usable(%Integration{connection_status: "disconnected"}), do: {:error, :integration_disconnected} @@ -89,32 +122,37 @@ defmodule Kodo.LLM.CredentialResolver do defp require_usable(%Integration{connection_status: "connected"}), do: :ok - defp build_credential(%Integration{authentication_type: "api_key"} = integration, payload) do + defp build_credential( + %Integration{authentication_type: "api_key"} = integration, + reference, + payload + ) do with {:ok, api_key} <- fetch_secret(payload, "api_key") do - {:ok, credential(integration, api_key, nil)} + {:ok, credential(integration, reference, api_key, nil)} end end defp build_credential( %Integration{provider: "openai_codex", authentication_type: "oauth"} = integration, + reference, payload ) do with {:ok, access_token} <- fetch_secret(payload, "access_token"), {:ok, account_id} <- fetch_secret(payload, "account_id") do - {:ok, credential(integration, access_token, account_id)} + {:ok, credential(integration, reference, access_token, account_id)} end end - defp build_credential(%Integration{}, _payload), + defp build_credential(%Integration{}, %IntegrationRef{}, _payload), do: {:error, :credential_payload_invalid} - defp credential(integration, token, account_id) do + defp credential(integration, reference, token, account_id) do %Credential{ integration_id: integration.id, provider: integration.provider, authentication_type: integration.authentication_type, credential_generation: integration.credential_generation, - billing_path: billing_path(integration.provider), + billing_path: reference.billing_path, token: token, account_id: account_id } @@ -126,8 +164,4 @@ defmodule Kodo.LLM.CredentialResolver do _missing -> {:error, :credential_payload_invalid} end end - - defp billing_path("openai_codex"), do: :subscription - defp billing_path("openrouter"), do: :aggregator - defp billing_path(_provider), do: :platform end diff --git a/lib/kodo/llm/integration_ref.ex b/lib/kodo/llm/integration_ref.ex index 91438b3..09cb97f 100644 --- a/lib/kodo/llm/integration_ref.ex +++ b/lib/kodo/llm/integration_ref.ex @@ -8,13 +8,27 @@ defmodule Kodo.LLM.IntegrationRef do alias Kodo.Integrations.Integration - @enforce_keys [:integration_id, :provider, :credential_generation] - defstruct [:integration_id, :provider, :credential_generation] + @enforce_keys [ + :integration_id, + :provider, + :authentication_type, + :credential_generation, + :billing_path + ] + defstruct [ + :integration_id, + :provider, + :authentication_type, + :credential_generation, + :billing_path + ] @opaque t :: %__MODULE__{ integration_id: Ecto.UUID.t(), provider: String.t(), - credential_generation: non_neg_integer() + authentication_type: String.t(), + credential_generation: non_neg_integer(), + billing_path: :platform | :subscription | :aggregator } @doc "Builds a reference from scoped integration metadata." @@ -22,7 +36,13 @@ defmodule Kodo.LLM.IntegrationRef do %__MODULE__{ integration_id: integration.id, provider: integration.provider, - credential_generation: integration.credential_generation + authentication_type: integration.authentication_type, + credential_generation: integration.credential_generation, + billing_path: billing_path(integration.provider) } end + + defp billing_path("openai_codex"), do: :subscription + defp billing_path("openrouter"), do: :aggregator + defp billing_path(_provider), do: :platform end diff --git a/test/kodo/agent/loop_test.exs b/test/kodo/agent/loop_test.exs index 0687889..0b2103c 100644 --- a/test/kodo/agent/loop_test.exs +++ b/test/kodo/agent/loop_test.exs @@ -114,7 +114,10 @@ defmodule Kodo.Agent.LoopTest do assert invocation.payload["provider"] == "openai" assert invocation.payload["model"] == "openai:gpt-4o-mini" assert invocation.payload["reasoning"] == "none" - assert invocation.version == 3 + assert invocation.version == 4 + assert invocation.payload["model_identity"] == "gpt-4o-mini" + assert invocation.payload["authentication_type"] == "api_key" + assert invocation.payload["billing_path"] == "platform" assert invocation.payload["role_contract"] == "alpha-v1" refute Map.has_key?(invocation.payload, "role_prompt_version") assert invocation.payload["toolset_version"] == "workspace-v5" @@ -177,7 +180,7 @@ defmodule Kodo.Agent.LoopTest do invocation = Enum.find(Sessions.events_after(session.id), &(&1.type == "model_invocation_started")) - assert invocation.version == 3 + assert invocation.version == 4 assert invocation.payload["model"] == "openai:gpt-4o-mini" assert invocation.payload["role_contract"] == "alpha-v1" assert invocation.payload["toolset_version"] == "workspace-v5" diff --git a/test/kodo/llm/credential_resolver_test.exs b/test/kodo/llm/credential_resolver_test.exs index 667c088..d00eed7 100644 --- a/test/kodo/llm/credential_resolver_test.exs +++ b/test/kodo/llm/credential_resolver_test.exs @@ -183,6 +183,18 @@ defmodule Kodo.LLM.CredentialResolverTest do CredentialResolver.resolve(context.scope, unsupported_model, context.reference) end + test "rejects forged authentication and billing provenance", context do + forged_authentication = %{context.reference | authentication_type: "oauth"} + + assert {:error, :integration_authentication_mismatch} = + resolve(%{context | reference: forged_authentication}) + + forged_billing = %{context.reference | billing_path: :subscription} + + assert {:error, :integration_billing_mismatch} = + resolve(%{context | reference: forged_billing}) + end + test "rejects malformed decrypted payloads with a bounded error" do other_scope = AccountsFixtures.user_scope_fixture() From e895fac65c0a29be53bb5bdf29ce6ef81a7ba5f8 Mon Sep 17 00:00:00 2001 From: Darwin Wu Date: Fri, 4 Sep 2026 02:37:00 +0000 Subject: [PATCH 2/2] refactor(agent): resolve canonical model providers --- lib/kodo/agent/model_mapping.ex | 7 ++++++- test/kodo/agent/model_mapping_test.exs | 7 +++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/lib/kodo/agent/model_mapping.ex b/lib/kodo/agent/model_mapping.ex index 79f2492..0413f10 100644 --- a/lib/kodo/agent/model_mapping.ex +++ b/lib/kodo/agent/model_mapping.ex @@ -118,5 +118,10 @@ defmodule Kodo.Agent.ModelMapping do defp maybe_update_provider(mapping, :reasoning, _reasoning), do: mapping - defp provider(model), do: model |> String.split(~r/[:\/]/, parts: 2) |> hd() + defp provider(model) do + case ReqLLM.model(model) do + {:ok, %LLMDB.Model{provider: provider}} -> Atom.to_string(provider) + {:error, _reason} -> nil + end + end end diff --git a/test/kodo/agent/model_mapping_test.exs b/test/kodo/agent/model_mapping_test.exs index 921cc78..8ca5a66 100644 --- a/test/kodo/agent/model_mapping_test.exs +++ b/test/kodo/agent/model_mapping_test.exs @@ -51,4 +51,11 @@ defmodule Kodo.Agent.ModelMappingTest do assert search["reasoning"] == "high" assert search["sources"] == %{"model" => "repository", "reasoning" => "session"} end + + test "derives display providers from resolved models rather than string splitting" do + mapping = ModelMapping.balanced([{"user", %{search: %{model: "not-a-model"}}}]) + + assert mapping["roles"]["search"]["model"] == "not-a-model" + assert mapping["roles"]["search"]["provider"] == nil + end end