Add model allowlist claim to third-party agent tokens

Justin Ho Tuan Duongrequested to merge
fix/third-party-agent-token-model-allowlist into master

What does this MR do and why?

A group admin can turn on the model allowlist for Agentic Chat. The setting is under Settings > GitLab Duo > Model selection. The Duo Chat model picker honors this setting. The direct_access endpoint does not. Third-party agents, such as OpenCode, get a token from this endpoint. The token can call any model that the AI gateway proxy supports. A customer removed Claude Fable 5 from the allowlist. The customer could still use that model from OpenCode.

Rails cannot enforce this rule alone. The client picks the model on each proxy request. Rails never sees these requests. This MR adds the allowlist to the token itself. Ai::ThirdPartyAgents::TokenService gets the Agentic Chat feature setting for the governing namespace. It then asks Ai::ModelSelection::UserModelSelection for the new enforced_model_refs. When the result is not nil, the service adds it to the signed Cloud Connector token as a gitlab_allowed_model_refs claim. The value is the stored model refs plus the current or default model. This is the same set the Duo Chat picker offers.

Design points:

  • Why UserModelSelection: this class already decides which models the Duo Chat picker offers. The proxy must enforce the same set, so the rule stays in one place. TokenService only builds the claim.
  • Why not selectable_models: that method returns [] when model definitions cannot be fetched. The gateway cannot tell an empty list from "block all models". enforced_model_refs returns nil for "no restriction".
  • Why a signed claim, not a header: the AI gateway user_access_token endpoint copies every non-basic claim from the incoming token into the token it mints. The claim reaches the proxy with no change to the mint step. A client cannot forge a signed claim. A header could be changed by the client.
  • GitLab.com only, for now: TokenService builds the claim only when gitlab_duo_saas_only is available. Self-managed instances do not send a self-signed token on this path, so a claim would be dropped. Self-managed enforcement is a follow-up.
  • If model definitions cannot be fetched: the default model is unknown, so the claim holds only the stored refs. If no refs are stored, the claim is an empty list and the gateway allows no model until the next token. This fails closed on purpose.
  • Out of scope: an admin-pinned model with the allowlist turned off still gives no claim, so all models stay allowed.
  • No behavior change yet: the claim does nothing until the AI gateway proxy checks it. A companion AI gateway MR must map each ref to the upstream model name, refuse other models with a clear 403 error, and allow all models when the claim is absent.
  • Feature flag: the claim is behind third_party_agent_token_model_allowlist (gitlab_com_derisk, off by default). The actor is the governing namespace. The rollout is tracked in #632588.
  • Known gap: Code Suggestions direct access tokens (POST /api/v4/code_suggestions/direct_access) also work on the AI gateway proxy. These tokens do not get the claim yet, so a user can still use one to call other models. A follow-up MR covers this path.
  • Related work: !255682 blocks the token when Duo is off for the whole namespace. This MR carries the per-model allowlist. The two MRs are independent.

References

Screenshots or screen recordings

This MR changes backend token logic only. There is no UI, so there are no screenshots.

How to set up and validate locally

  1. Start GDK in SaaS mode: set GITLAB_SIMULATE_SAAS=1.

  2. Pick a group that has a Duo add-on. Set it as your default Duo namespace.

  3. In rails console, turn on the feature flag for the group:

    Feature.enable(:third_party_agent_token_model_allowlist, group)

  4. In rails console, turn on the model allowlist for the group:

    Ai::ModelSelection::NamespaceFeatureSetting.find_or_initialize_by_feature(group, :duo_agent_platform_agentic_chat)
  .update!(model_allowlist_enabled: true, model_allowlist_gitlab_model_refs: %w[claude_sonnet_4_5_20250929])

  5. Call the endpoint with a personal access token:

    curl --request POST --header "PRIVATE-TOKEN: <pat>" "http://gdk.test:3000/api/v4/ai/third_party_agents/direct_access"

  6. Decode the token field from the response. It is a JWT:

    echo "<token>" | cut -d. -f2 | base64 -d

  7. Check that the decoded payload has a gitlab_allowed_model_refs claim.

  8. Turn the allowlist off for the group. Call the endpoint again. Check that the claim is absent.

MR acceptance checklist

Evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.

🤖 AI-generated

Edited by Justin Ho Tuan Duong

Merge request reports