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:
- CI pipeline path (
ResumeWorkflowService):AGENT_PLATFORM_MODEL_METADATAis set as a CI env var. Covered by theagent_platform_model_metadata_jsonoverride. - WebSocket/browser path (
GET /api/v4/ai/duo_workflows/ws): The Duo Chat web UI connects via WebSocket. This endpoint callsDuoAgentPlatformModelMetadataServiceto build thex-gitlab-agent-platform-model-metadatagRPC header. This path now also checks for stored metadata when aworkflow_idis 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:
- DWS selects a model via
random.choice(default_models)and returns its metadata inNewCheckpoint.model_metadata_json - The Node executor (duo-cli) reads
model_metadata_jsonfrom the checkpoint action and passes it to the Rails checkpoint REST API CreateCheckpointServicepersists it toworkflow.selected_model_metadata_json- On workflow resume,
ResumeWorkflowService#agent_platform_model_metadata_jsonreturns the stored metadata instead of callingDuoAgentPlatformModelMetadataService, soAGENT_PLATFORM_MODEL_METADATAis set to the previously selected provider - DWS
ModelMetadataInterceptorsees the explicit provider and bypassesrandom.choice(), routing to the same provider → token cache hit
WebSocket/browser path:
- Same steps 1-3 as above (DWS returns metadata, Node executor or DWS HTTP POST persists it)
- When the browser reconnects,
duo_agentic_chat_state_manager.vuepassesworkflowIidtobuildWebsocketUrl - The
/wsendpoint receivesworkflow_id, looks up the workflow, and ifselected_model_metadata_jsonis present, uses it directly as thex-gitlab-agent-platform-model-metadatagRPC header - DWS routes to the same provider → token cache hit
The feature is gated behind duo_workflow_provider_stickiness FF (default: disabled).
References
- DWS MR: gitlab-org/modelops/applied-ml/code-suggestions/ai-assist!5407 (merged)
- Issue: gitlab-org/modelops/applied-ml/code-suggestions/ai-assist#2174
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
- Apply the DWS MR (!5407) and run DWS locally.
- Check out this branch and run the Rails migration:
bundle exec rails db:migrate - Enable the feature flag:
Feature.enable(:duo_workflow_provider_stickiness) - Start a Duo Chat workflow session in the browser. After the first checkpoint is created, verify that
duo_workflows_workflows.selected_model_metadata_jsonis populated. - Reload the page (which triggers a new WebSocket connection). Verify that the
/wsrequest includesworkflow_idas a query param and that the gRPC headerx-gitlab-agent-platform-model-metadatamatches the stored metadata. - For the CI pipeline path: resume the workflow via the API and verify that
AGENT_PLATFORM_MODEL_METADATAin 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.