Ability to slowly rollout of Code Suggestions via Monolith Instance

Problem

For the architecture agreed on in our AI gateway blueprint we need to have Code Suggestions go through the monolith. While we have the API endpoint for this already running there are many concerns around the impact of having the code suggestions requests going through the monolith (see e.g. #418955 (closed) or https://gitlab.com/gitlab-org/modelops/applied-ml/code-suggestions/ai-assist/-/issues/212#note_1481970629) when it comes to latency or the load that this puts the system under.

We are trying to assess some of these impacts within controlled experiments (see #419530 (closed)), but this does not give us completely realistic data. We have the ability to shut down the code suggestions endpoint completely via Feature flag but this would mean disabling code suggestions for many users if something does not work as expected on the monolith.

Desired Outcome

We have a way to slowly roll out the switch to have code suggestion requests go via gitlab.com for the VSCode Extension, as we do with feature flags in the main codebase already. (e.g. 1% of users going via monolith, 5% etc)

Chosen Solution

  1. Create a feature flag called :use_instance_for_code_suggestions so that /code-suggestions/completions works for whom the feature flag is enabled but returns a 403 if the feature flag is not enabled
  2. Change the VSCode extension to:
    1. Go to /code-suggestions/completions by default for any code completion request
    2. If /code-suggestions/completions returns 403 call the Gateway directly with the old method instead
    3. Cache the fact that /code-suggestions/completions returned a 403 for the next hour and use the old method for the next request.

This will lead to us:

  1. Control the rollout of the switch to /code-suggestions/completions via a flag on Gitlab.com
  2. Have up to 1 completion request per user/hour that is still supposed to go to the Gateway directly but will incur additional latency because it receives a 403 from /code-suggestions/completions first.

Diagram

Current Setup

sequenceDiagram
    Extension->>Gitlab.com: /code_suggestions/tokens with PAT
    Gitlab.com->>Extension: JWT token
    Extension->>Gateway: v2/completions with JWT
    Gateway->>Extension: completion

Intended Setup

sequenceDiagram
    Extension->>Gitlab.com: /code_suggestions/tokens with PAT
    Gitlab.com-->>Extension: JWT token
    Extension->>Gitlab.com: /code_suggestions/completions with JWT and PAT
    alt feature flag is enabled
        Gitlab.com->>Gateway: v2/completions with JWT
        Gateway-->>Gitlab.com: completion
        Gitlab.com-->>Extension: completion
    else feature flag is disabled
        Gitlab.com->>Extension: 403 forbidden
        Extension->>Gateway: v2/completions with JWT
        Gateway-->>Extension: completion
    end
Edited by Sebastian Rehm