Add support for new universal tokens
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
-
Checkout this branch
-
Run https://gitlab.com/gitlab-org/customers-gitlab-com on local environment (localhost:5000)
-
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
-
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
-
Run model gateway
poetry run ai_gateway
-
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" } }'
-
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:
- 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)