Connect Code Suggestions LSP directly to AI Gateway for Code Completion

Description

To improve Code Suggestions performance, let's bypass the Rails monolith and connect directly to the AI Model Gateway for code completion requests. We'll need to solve some authentication pro

  • See this thread about the selection of this approach and discussion of authentication methods.

Proposal

Since the Rails app already generates a JWT, we don't need to touch the AI Gateway for this iteration and we can keep the Rails app as the generator for this authentication token.

For the time being, let's also treat the Rails app as a service registry for the AI Gateway. Currently the client knows nothing about the AI Gateway. Let's return this information along with the jwt token when the client requests it from the Rails app.

sequenceDiagram
    autonumber

    participant C as Client
    participant G as GitLab
    participant AI as AI Gateway

    C->>G: POST GitLab<br>with "/ai_gateway_auth"
    G->>C: return `jwt`, `ai_url`, `expires`
    loop Each code suggestions request
      alt token is about to expire?
        C->>G: POST GitLab<br>with "/ai_gateway_auth"
        G->>C: return `jwt`, `ai_url`, `expires`
      end

      C->>AI: POST `ai_url`/completions<br>with `jwt`
      AI->>C: Response
    end
Original proposal

From this comment:

what if GitLab was an OAuth App for the AI Gateway? This way the AI Gateway doesn't have to ping back to GitLab to find out the who?

Here's a diagram of what I'm thinking:

sequenceDiagram
    autonumber

    participant C as Client
    participant G as GitLab
    participant AI as AI Gateway

    C->>G: POST GitLab<br>with OAuth/PAT for AI Gateway OAuth token
    G->>AI: POST /oauth/token as OAuth app for User X
    AI->>G: return AI Gateway OAuth Token
    G->>C: AI Gateway OAuth Token
    loop Each subsequent request
      C->>AI: POST AI Gateway/completions<br>with AI Gateway OAuth Token
      AI->>C: Response
    end

See also this comment:

This could work as long as we use OIDC extension here, and the token returned is JWT, otherwise the standard OAuth bearer token might not contain authorization claims. The JWT would still expire after a while, so the VScode extension would need to refresh the token 🤔

Tasks

  • Add Rails endpoint to return { ai_model_gateway_url: string, jwt: string, expires_ms: number }. Possibly put this behind feature flag 🤷? (DRI @jessieay)
  • Add aiDirectToModelGateway feature flag in Extension Client and LSP. (DRI @pslaughter)
  • Under aiDirectToModelGateway FF:
    • Hit Rails endpoint to get AI model gateway URL and token info (DRI @pslaughter)
    • Use AI model gateway URL to make code completion requests (DRI @pslaughter)
    • Handle token refresh seamlessly (DRI @pslaughter)

Related MR's

These MR's were being worked on until this effort was put on pause. If we pick this up, please reference these implementations:

Edited by Paul Slaughter