Persist Duo Workflow model metadata

What does this MR do and why?

Implements the Rails side of provider stickiness for Duo Workflow token caching. This is the Rails counterpart to the DWS change in gitlab-org/modelops/applied-ml/code-suggestions/ai-assist!5407 (merged).

Problem: With the LiteLLM load balancer distributing requests across multiple providers (e.g., Vertex, Bedrock, Anthropic), subsequent workflow sessions from the same user may be routed to different providers. This breaks token caching, which relies on the same provider receiving consecutive turns.

Solution: DWS returns the selected model's metadata in NewCheckpoint.model_metadata_json (a JSON string matching the x-gitlab-agent-platform-model-metadata header format). Rails stores this on the workflow record and re-injects it as the model metadata on subsequent sessions, ensuring the same provider is used.

There are two distinct paths that need to be covered:

  1. CI pipeline path (ResumeWorkflowService): AGENT_PLATFORM_MODEL_METADATA is set as a CI env var. Covered by the agent_platform_model_metadata_json override.
  2. WebSocket/browser path (GET /api/v4/ai/duo_workflows/ws): The Duo Chat web UI connects via WebSocket. This endpoint calls DuoAgentPlatformModelMetadataService to build the x-gitlab-agent-platform-model-metadata gRPC header. This path now also checks for stored metadata when a workflow_id is provided.

Changes

File Change
vendor/gems/gitlab-duo-workflow-service-client/lib/gitlab/duo_workflow_service/version.rb Bump gem version to 0.8
vendor/gems/gitlab-duo-workflow-service-client/lib/proto/contract_pb.rb Regenerated Ruby proto bindings with model_metadata_json field in NewCheckpoint
Gemfile / Gemfile.lock Update gem constraint to ~> 0.8
db/migrate/20260501000000_add_selected_model_metadata_json_to_duo_workflows_workflows.rb Add selected_model_metadata_json text column to duo_workflows_workflows
db/structure.sql Updated schema
ee/lib/api/ai/duo_workflows/workflows_internal.rb Accept optional model_metadata_json param in the checkpoint REST API
ee/app/services/ai/duo_workflows/create_checkpoint_service.rb Persist model_metadata_json to workflow.selected_model_metadata_json when received
ee/app/services/ai/duo_workflows/resume_workflow_service.rb Override agent_platform_model_metadata_json to use stored metadata on resume (gated by FF); fix FF check to handle namespace-level workflows
ee/lib/api/ai/duo_workflows/workflows.rb Add optional workflow_id param to /ws endpoint; when present and stored metadata exists, use it as the model metadata header (gated by FF)
ee/app/assets/javascripts/ai/constants.js Add DUO_WORKFLOW_WEBSOCKET_PARAM_WORKFLOW_ID constant
ee/app/assets/javascripts/ai/duo_agentic_chat/websocket/workflow_utils.js Accept and forward workflowId in buildWebsocketUrl
ee/app/assets/javascripts/ai/duo_agentic_chat/components/duo_agentic_chat_state_manager.vue Pass workflowIid to buildWebsocketUrl so the browser sends the current workflow ID
ee/config/feature_flags/development/duo_workflow_provider_stickiness.yml Feature flag definition

How it works

CI pipeline path:

  1. DWS selects a model via random.choice(default_models) and returns its metadata in NewCheckpoint.model_metadata_json
  2. The Node executor (duo-cli) reads model_metadata_json from the checkpoint action and passes it to the Rails checkpoint REST API
  3. CreateCheckpointService persists it to workflow.selected_model_metadata_json
  4. On workflow resume, ResumeWorkflowService#agent_platform_model_metadata_json returns the stored metadata instead of calling DuoAgentPlatformModelMetadataService, so AGENT_PLATFORM_MODEL_METADATA is set to the previously selected provider
  5. DWS ModelMetadataInterceptor sees the explicit provider and bypasses random.choice(), routing to the same provider → token cache hit

WebSocket/browser path:

  1. Same steps 1-3 as above (DWS returns metadata, Node executor or DWS HTTP POST persists it)
  2. When the browser reconnects, duo_agentic_chat_state_manager.vue passes workflowIid to buildWebsocketUrl
  3. The /ws endpoint receives workflow_id, looks up the workflow, and if selected_model_metadata_json is present, uses it directly as the x-gitlab-agent-platform-model-metadata gRPC header
  4. DWS routes to the same provider → token cache hit

The feature is gated behind duo_workflow_provider_stickiness FF (default: disabled).

References

Screenshots or screen recordings

N/A — backend-only change.

Update Query

UPDATE "duo_workflows_workflows" SET "updated_at" = '2026-05-09 18:12:49.183682', "model_metadata_json" = '{"model":"claude-3"}' WHERE "duo_workflows_workflows"."id" = 21 /*application:test,correlation_id:abbc1cfbffe6ed8c1232efcbe088cb4a,db_config_database:gitlabhq_test,db_config_name:main,line:/ee/app/services/ai/duo_workflows/create_checkpoint_service.rb:44:in `persist_model_metadata_if_present'*/

How to set up and validate locally

  1. Apply the DWS MR (!5407) and run DWS locally.
  2. Check out this branch and run the Rails migration:
    bundle exec rails db:migrate
  3. Enable the feature flag:
    Feature.enable(:duo_workflow_provider_stickiness)
  4. Start a Duo Chat workflow session in the browser. After the first checkpoint is created, verify that duo_workflows_workflows.selected_model_metadata_json is populated.
  5. Reload the page (which triggers a new WebSocket connection). Verify that the /ws request includes workflow_id as a query param and that the gRPC header x-gitlab-agent-platform-model-metadata matches the stored metadata.
  6. For the CI pipeline path: resume the workflow via the API and verify that AGENT_PLATFORM_MODEL_METADATA in the new workload's CI variables matches the stored metadata.

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 Alejandro Rodríguez

Merge request reports

Loading