Skip to content

Add support for new universal tokens

Nikola Milojevic requested to merge 426572-support-new-audience into main

What does this merge request do and why?

This MR is first step of the Phase 1: Replacing CS JWT with Universal JWT (gitlab-org&11590 - closed)

It adds support for new universal tokens:

Since we are changing the token audience to gitlab-ai-gateway, the model gateway needs to maintain compatibility with legacy gitlab-code-suggestions audience

This Mr introduces also scopes in UserClaims, as we are going to use this to verify if customer has permissions to acess specific Model Gateway apis in the future.

How to set up and validate locally

  1. Checkout this branch

  2. Run https://gitlab.com/gitlab-org/customers-gitlab-com on local environment (localhost:5000)

  3. Get a JWT token by running the following command in rails console for local customers-gitlab-com main.

    jwt_token = ::Auth::JsonWebToken.new(
      audience: ::Auth::JsonWebToken::JWT_AUDIENCE_CODE_SUGGESTIONS
    ).encode
  4. Update the .env file in the root folder with the following variables.

    GITLAB_URL=https://gitlab.com
    CUSTOMER_PORTAL_BASE_URL=http://host.docker.internal:5000/
    AUTH_BYPASS_EXTERNAL=false
    # Testing & development
    USE_FAKE_MODELS=true
    FASTAPI_API_PORT=5052
  5. Run model gateway

      poetry run ai_gateway
  6. Run a cURL command to check the authentication, using the jwt_token generated in Step 3.

      curl -vvv --request POST \
      --url 'http://localhost:5052/v2/completions' \
      --header 'X-Gitlab-Authentication-Type: oidc' \
      --header 'Authorization: Bearer <jwt_token>' \
      --header 'Content-Type: application/json' \
      --data '{
        "prompt_version": 1,
        "project_path": "awesome_project",
        "project_id": 23,
        "current_file": {
          "file_name": "main.py",
          "content_above_cursor": "\"\"\"\nImplement fastapi middleware to log all incoming requests\"\"\"\n",
          "content_below_cursor": "scoopy doo"
        }
      }'
  7. You should get the response:

    {"id":"id","model":{"engine":"","name":""},"object":"text_completion","created":1689862537,"choices":[{"text":"fake code suggestion from GitLab Codegen","index":0,"finish_reason":"length"}]} 

To test with the new universal token, checkout https://gitlab.com/gitlab-org/customers-gitlab-com/-/merge_requests/8336+, and repeat steps above:

Step 3 will now look like this:

  1. Get a JWT token by running the following command in rails console for local customers-gitlab-com main.
jwt_token = ::Auth::JsonWebToken.new(
  audience: ::Auth::JsonWebToken::JWT_AUDIENCE
).encode

Merge request checklist

  • Tests added for new functionality. If not, please raise an issue to follow up.
  • Documentation added/updated, if needed.

Related to gitlab-org/gitlab#426572 (closed)

Edited by Nikola Milojevic

Merge request reports