fix: harden provider recovery and session lifecycle

lib/kodo/agent/loop.ex

@@ -159,28 +159,17 @@ defmodule Kodo.Agent.Loop do
     review = ModelMapping.role!(context.mapping, :review)
     review_model = ModelMapping.request_model(review)
     contract = Roles.fetch!(:review, review["role_contract"])
+    review_invocation_id = Ecto.UUID.generate()
 
     with {:ok, capability_validation} <-
            context.adapter.validate_model(review_model, review, contract),
-         {:ok, request} <- resolve_request(context.session_id, review),
          :ok <-
            Phoenix.PubSub.subscribe(
              Kodo.PubSub,
              "runner_responses:#{context.projection.runner_id}"
            ),
          :ok <- Phoenix.PubSub.subscribe(Kodo.PubSub, "session:#{context.session_id}"),
          :ok <- Runners.subscribe_lifecycle(),
-         {:ok, {review_invocation_id, admission}} <-
-           start_review_invocation(
-             context.session_id,
-             primary_invocation_id,
-             context.mapping,
-             review,
-             contract,
-             capability_validation,
-             request,
-             context.ownership
-           ),
          {:ok, diff} <-
            execute_tool(
              context.session_id,
@@ -192,6 +181,17 @@ defmodule Kodo.Agent.Loop do
              context.ownership
            ),
          :ok <- rehoming_boundary(),
+         {:ok, request} <- resolve_request(context.session_id, review),
+         {:ok, admission} <-
+           start_review_invocation(
+             context,
+             primary_invocation_id,
+             review_invocation_id,
+             review,
+             contract,
+             capability_validation,
+             request
+           ),
          {:ok, generated} <-
            Sessions.dispatch_if_owner(context.ownership, fn ->
              task = original_task(context.events)
@@ -237,19 +237,16 @@ defmodule Kodo.Agent.Loop do
   end
 
   defp start_review_invocation(
-         session_id,
+         context,
          primary_invocation_id,
-         mapping,
+         invocation_id,
          review,
          contract,
          capability_validation,
-         request,
-         ownership
+         request
        ) do
-    invocation_id = Ecto.UUID.generate()
-
     case Sessions.append_admitted_model_event(
-           session_id,
+           context.session_id,
            "review_invocation_started",
            %{
              "invocation_id" => invocation_id,
@@ -264,15 +261,15 @@ defmodule Kodo.Agent.Loop do
              "role_contract" => contract.id,
              "toolset_version" => contract.toolset_version,
              "capability_validation" => capability_validation,
-             "model_mapping" => mapping
+             "model_mapping" => context.mapping
            },
            request.model,
            request.credential,
            version: 2,
            parent_id: primary_invocation_id,
-           ownership: ownership
+           ownership: context.ownership
          ) do
-      {:ok, {_event, admission}} -> {:ok, {invocation_id, admission}}
+      {:ok, {_event, admission}} -> {:ok, admission}
       error -> normalize_admission_error(error, request)
     end
   end

lib/kodo/agent/model_mapping.ex

@@ -67,11 +67,14 @@ defmodule Kodo.Agent.ModelMapping do
     roles =
       Map.new(mapping["roles"], fn {role, role_mapping} ->
         contract = Roles.fetch!(role_atom(role), role_mapping["role_contract"])
+        {route, selector} = model_route_and_selector(role_mapping)
 
         {role,
          role_mapping
-         |> Map.put("execution_route", role_mapping["provider"])
-         |> Map.put("model_selector", model_selector(role_mapping["model"]))
+         |> Map.put("provider", route)
+         |> Map.put("execution_route", route)
+         |> Map.put("model", canonical_model(route, selector, role_mapping["model"]))
+         |> Map.put("model_selector", selector)
          |> Map.put_new("capability_contract", capability_contract(contract))}
       end)
 
@@ -249,10 +252,41 @@ defmodule Kodo.Agent.ModelMapping do
     end
   end
 
-  defp model_selector(model) do
-    case String.split(model, ":", parts: 2) do
-      [_provider, selector] -> selector
-      [selector] -> selector
+  defp model_route_and_selector(%{
+         "execution_route" => route,
+         "model_selector" => selector
+       })
+       when is_binary(route) and route != "" and is_binary(selector) and selector != "",
+       do: {route, selector}
+
+  defp model_route_and_selector(%{"model" => model, "provider" => provider})
+       when is_binary(model) and is_binary(provider) and provider != "" do
+    accepted_prefix =
+      [provider, String.replace(provider, "_", "-")]
+      |> Enum.uniq()
+      |> Enum.find(&String.starts_with?(model, &1 <> ":"))
+
+    if accepted_prefix do
+      {provider, String.replace_prefix(model, accepted_prefix <> ":", "")}
+    else
+      parse_model_route_and_selector(model, provider)
+    end
+  end
+
+  defp model_route_and_selector(%{"model" => model, "provider" => provider}) do
+    parse_model_route_and_selector(model, provider)
+  end
+
+  defp parse_model_route_and_selector(model, provider) do
+    case LLMDB.Spec.parse_spec(model) do
+      {:ok, {parsed_provider, selector}} -> {Atom.to_string(parsed_provider), selector}
+      {:error, _reason} -> {provider, model}
     end
   end
+
+  defp canonical_model(route, selector, _original)
+       when is_binary(route) and route != "" and is_binary(selector) and selector != "",
+       do: "#{route}:#{selector}"
+
+  defp canonical_model(_route, _selector, original), do: original
 end

lib/kodo/integrations.ex

@@ -251,6 +251,30 @@ defmodule Kodo.Integrations do
     end
   end
 
+  @doc "Returns bounded terminal authorization feedback without loading encrypted payloads."
+  def get_device_authorization_outcome(%Scope{user: user}, integration_id) do
+    with {:ok, integration_id} <- cast_uuid(integration_id) do
+      DeviceAuthorizationAttempt
+      |> join(:inner, [attempt], integration in assoc(attempt, :integration))
+      |> where(
+        [attempt, integration],
+        attempt.integration_id == ^integration_id and integration.user_id == ^user.id
+      )
+      |> order_by([attempt], desc: attempt.attempt_generation)
+      |> limit(1)
+      |> select([attempt], %{
+        attempt_generation: attempt.attempt_generation,
+        state: attempt.state,
+        terminal_error_code: attempt.terminal_error_code
+      })
+      |> Repo.one()
+      |> case do
+        %{state: state} = outcome when state in ["failed", "expired"] -> {:ok, outcome}
+        _missing_or_non_failure -> {:error, :device_authorization_outcome_not_found}
+      end
+    end
+  end
+
   @doc false
   def expire_device_authorization(%Scope{user: user}, integration_id) do
     with {:ok, integration_id} <- cast_uuid(integration_id) do
@@ -441,6 +465,7 @@ defmodule Kodo.Integrations do
     do: {:error, :authentication_type_mismatch}
 
   defp insert_oauth_integration(user_id, changeset) do
+    lock_user!(user_id)
     lock_provider_identity(user_id, "openai_codex")
 
     case Repo.insert(changeset) do
@@ -782,6 +807,8 @@ defmodule Kodo.Integrations do
 
   defp activate_transaction(user_id, id, generation, audit?) do
     Repo.transaction(fn ->
+      lock_user!(user_id)
+
       case lock_provider_integrations(user_id, id) do
         {:ok, integration, provider_integrations} ->
           activate_locked(user_id, integration, provider_integrations, generation, audit?)
@@ -1561,6 +1588,7 @@ defmodule Kodo.Integrations do
     case Ecto.UUID.cast(id) do
       {:ok, id} ->
         Repo.transaction(fn ->
+          lock_user!(user.id)
           maybe_lock_transition(user.id, id, audit_event_type)
 
           {restore_active, changes} = Map.pop(changes, :restore_active, false)
@@ -1617,6 +1645,7 @@ defmodule Kodo.Integrations do
 
   defp insert_connected(changeset, actor_user_id, provider, event_type) do
     Repo.transaction(fn ->
+      lock_user!(actor_user_id)
       lock_provider_identity(actor_user_id, provider)
       first_account? = !provider_account_exists?(actor_user_id, provider)
 
@@ -1655,7 +1684,14 @@ defmodule Kodo.Integrations 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)
+
+        case lock_provider_rows(user_id, provider, integration_id) do
+          {:ok, _integration, _provider_integrations} ->
+            invalidate_provider_selection_recovery(user_id, provider)
+
+          {:error, reason} ->
+            Repo.rollback(reason)
+        end
 
       nil ->
         :ok

lib/kodo/sessions.ex

@@ -588,13 +588,14 @@ defmodule Kodo.Sessions do
   end
 
   defp preflight_turn(scope, session_id) do
-    mapping =
-      session_id
-      |> get_session!()
-      |> turn_route_snapshot()
-      |> Map.fetch!("model_mapping")
+    case get_session(scope, session_id) do
+      %Session{} = session ->
+        mapping = session |> turn_route_snapshot() |> Map.fetch!("model_mapping")
+        preflight_mapping(scope, mapping)
 
-    preflight_mapping(scope, mapping)
+      nil ->
+        {:error, :session_not_found}
+    end
   end
 
   defp preflight_mapping(scope, mapping) do
@@ -643,7 +644,7 @@ defmodule Kodo.Sessions do
   end
 
   defp append_admitted_model_event_locked(session_id, type, payload, model, credential, opts) do
-    session = lock_session!(session_id, opts)
+    session = lock_session_after_owner!(session_id, opts)
     integration = admit_model_credential!(session, credential)
     billing_path = IntegrationRef.from_integration(integration).billing_path
 
@@ -732,8 +733,17 @@ defmodule Kodo.Sessions do
   end
 
   defp reconcile_cancel({:error, reason}, scope, session_id)
-       when reason in [:not_running, :already_finished, :coordinator_unavailable] do
-    if get_session(scope, session_id).status == "cancelled", do: :ok, else: {:error, reason}
+       when reason in [
+              :not_running,
+              :already_finished,
+              :coordinator_unavailable,
+              :session_not_found
+            ] do
+    case get_session(scope, session_id) do
+      %Session{status: "cancelled"} -> :ok
+      %Session{} -> {:error, reason}
+      nil -> {:error, :not_found}
+    end
   end
 
   defp reconcile_cancel(result, _scope, _session_id), do: result
@@ -1090,7 +1100,7 @@ defmodule Kodo.Sessions do
   end
 
   defp begin_turn_locked(session_id, content, client_request_id, opts) do
-    session = lock_session!(session_id, opts)
+    session = lock_session_after_owner!(session_id, opts)
 
     if turn_request_recorded?(session_id, client_request_id) do
       []
@@ -1149,6 +1159,7 @@ defmodule Kodo.Sessions do
   end
 
   defp change_execution_route_locked(scope, session_id, destination) do
+    lock_scope_user!(scope)
     session = lock_user_session!(scope, session_id)
     projection = session.id |> events_after() |> Projection.from_events()
 
@@ -1498,17 +1509,55 @@ defmodule Kodo.Sessions do
       Session
       |> where([session], session.id == ^session_id)
       |> lock("FOR UPDATE")
-      |> Repo.one!()
+      |> Repo.one()
 
     ownership = Keyword.get(opts, :ownership)
 
     cond do
+      is_nil(session) -> Repo.rollback(:session_not_found)
       Keyword.get(opts, :allow_unowned, false) -> session
       ownership_matches?(session, ownership) and owner_alive?(ownership) -> session
       true -> Repo.rollback(:stale_ownership)
     end
   end
 
+  defp lock_session_after_owner!(session_id, opts) do
+    owner =
+      Session
+      |> where([session], session.id == ^session_id)
+      |> select([session], {session.id, session.user_id})
+      |> Repo.one()
+
+    case owner do
+      {^session_id, user_id} when is_integer(user_id) ->
+        lock_user!(user_id)
+        session = lock_session!(session_id, opts)
+
+        if session.user_id == user_id,
+          do: session,
+          else: Repo.rollback(:session_not_found)
+
+      {^session_id, nil} ->
+        lock_session!(session_id, opts)
+
+      nil ->
+        Repo.rollback(:session_not_found)
+    end
+  end
+
+  defp lock_scope_user!(%Scope{user: user}), do: lock_user!(user.id)
+
+  defp lock_user!(user_id) do
+    case User
+         |> where([user], user.id == ^user_id)
+         |> select([user], user.id)
+         |> lock("FOR KEY SHARE")
+         |> Repo.one() do
+      ^user_id -> :ok
+      nil -> Repo.rollback(:session_not_found)
+    end
+  end
+
   defp lock_user_session!(%Scope{user: user}, session_id) do
     case Ecto.UUID.cast(session_id) do
       {:ok, session_id} ->
@@ -1583,11 +1632,13 @@ defmodule Kodo.Sessions do
           from session in Session, where: session.id == ^event.session_id, select: session.user_id
         )
 
-      Phoenix.PubSub.broadcast(
-        Kodo.PubSub,
-        session_index_topic(user_id),
-        {:session_index_changed, event.session_id}
-      )
+      if user_id do
+        Phoenix.PubSub.broadcast(
+          Kodo.PubSub,
+          session_index_topic(user_id),
+          {:session_index_changed, event.session_id}
+        )
+      end
     end
   end