Cache code suggestions direct access tokens per user

Tian Gaorequested to merge
635137-cache-direct-access-token into master

What does this MR do and why?

POST /api/v4/code_suggestions/direct_access asks the AI Gateway to mint a user access token on every call, as a synchronous HTTP request from the Puma thread that serves the API call. A token lives for one hour, so a client needs about one mint per hour, but the Language Server asks for connection details on every completion while it has none, and nothing on the server side reuses the token it minted a moment ago.

Behind the cache_code_suggestions_direct_access_token feature flag (default off), Gitlab::Llm::AiGateway::CodeSuggestionsClient#direct_access_token now:

  • returns the cached token when one exists for the same user, project, governing namespace, code completions feature setting and AI Gateway token URL (the inputs that select which gateway signs the token and whose billing context it carries; a cache key built from a hash, so a nil part cannot be read as another one);
  • otherwise mints as before and caches a successful result until five minutes before expires_at (TOKEN_CACHE_EXPIRY_MARGIN), so a client that receives a cached token still has time to use it;
  • never caches a failed mint;
  • stores the cached token encrypted with the application's db_key_base through Authn::TokenField::EncryptionHelper, which derives a nonce from each token and stores it with the ciphertext (the scheme the encrypted token columns use), so a copy of the shared cache alone is not usable and no two entries share a keystream; an entry that no longer decrypts is treated as a miss.

The route itself is unchanged: all access checks and the rate limit still run before the client is asked for a token, and the response body is the same (a repeated call now returns the same token and expires_at).

The log_info call that Danger flags as AI logging is the existing mint log line (Creating user access token, event name and component only), moved unchanged into mint_direct_access_token; nothing new is logged.

How to reproduce (on master)

Apply only the spec change and run the new request example against master's application code:

bin/rspec ee/spec/requests/api/code_suggestions_spec.rb -e 'mints the token once for repeated requests'
# expected: The request POST ... /v1/code/user_access_token was expected to execute 1 time but it executed 2 times

Or in bin/rails console against a GDK with the AI Gateway running: call the route twice for the same user and watch the AI Gateway log two user_access_token requests.

How to verify (on this branch)

bin/rspec ee/spec/lib/gitlab/llm/ai_gateway/code_suggestions_client_spec.rb
# 44 examples, 0 failures

bin/rspec ee/spec/requests/api/code_suggestions_spec.rb -e 'POST /code_suggestions/direct_access'
# 60 examples, 0 failures, 4 pending

bundle exec rubocop on the changed files: no offenses.

How to set up and validate locally

  1. Run GDK with the AI Gateway (gdk start gitlab-ai-gateway) and a user with a Duo Pro or Duo Enterprise seat.

  2. Enable the flag in bin/rails console:

    Feature.enable(:cache_code_suggestions_direct_access_token)

  3. Call the route twice with the same personal access token:

    curl -s -X POST -H "PRIVATE-TOKEN: $TOKEN" http://gdk.test:3000/api/v4/code_suggestions/direct_access | jq '{token, expires_at}'

    The second call returns the same token and expires_at, and the AI Gateway log shows one POST /v1/code/user_access_token.

  4. Disable the flag and call the route twice again: each call mints a new token.

Screenshots or screen recordings

No UI change.

References

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.

Edited by Tian Gao

Merge request reports