Populate duoBranches

Kaveh Nejadrequested to merge
feature/duo-workflow-branches-read into master

What does this MR do and why?

Allows the front end to request the alternative messages to a specific user message. This will be used by the retry feature, and allows the front end to display all branches to the user.

Each group of alternatives also returns forkThreadTs which is the ID the front end should pass as resume_checkpoint_ts to select this branch as the branch to continue from.

References

Full PoC: gitlab-org/modelops/applied-ml/code-suggestions/ai-assist!6124 (closed)

Epic: gitlab-org#21289

Python implementation for retries: gitlab-org/modelops/applied-ml/code-suggestions/ai-assist!6356 (merged)

Depends on: !249131 (merged)

Video demo: gitlab-org#21289 (comment 3666671057)

Skeleton was added in !249895 (merged), and both MRs were extracted from !249143 (closed)

How to set up and validate locally

Query to get thread_ts and counts
{
  duoWorkflowWorkflows(workflowId: "gid://gitlab/Ai::DuoWorkflows::Workflow/11") {
    nodes { latestCheckpoint { duoMessages { content messageType threadTs parentTs alternativeCount } } }
  }
}
Query to get alternatives
{
  duoWorkflowBranches(
    workflowId: "gid://gitlab/Ai::DuoWorkflows::Workflow/52",
    threadTs: "1f196488-56dc-6fb7-8002-12836bae9b61"
  )
  {
    forkThreadTs
    messages {
      content
      messageType
      threadTs
      parentTs
      additionalContext { id category content metadata }
    }
  }
}
  1. Enable flags for incremental checkpoints:
    1. Feature.enable(:duo_workflow_incremental_checkpoints)
    2. Feature.enable(:duo_workflow_read_incremental_checkpoints)
    3. Feature.enable(:dw_read_blobs_graphql)
  2. Checkout this branch
    1. From gdk/gitlab run git checkout feature/duo-workflow-branches-read
  3. Restart
    1. gdk restart
  4. Start a conversation "Say 1" => "1" => "Say 2" => "2"
  5. From http://gdk.test:8080/-/graphql-explorer (replace port) run the query "Query to get thread_ts and counts` above with the workflow ID (You can get the workflow ID from dev tools => Application => Session storage => http://gdk.test:8080 => key).
    1. You should get something like
GraphQL response 
{
  "data": {
    "duoWorkflowWorkflows": {
      "nodes": [
        {
          "latestCheckpoint": {
            "duoMessages": [
              {
                "content": "Say 1",
                "messageType": "user",
                "threadTs": "1f194af9-1e18-67ed-8000-179ed51f8f07",
                "parentTs": "1f194af9-1e15-6170-bfff-0db48fd15212",
                "alternativeCount": 0
              },
              {
                "content": "1",
                "messageType": "agent",
                "threadTs": "1f194af9-3081-6a93-8001-7df329fe8fe2",
                "parentTs": "1f194af9-1e18-67ed-8000-179ed51f8f07",
                "alternativeCount": null
              },
              {
                "content": "Say 2",
                "messageType": "user",
                "threadTs": "1f194af9-6ffa-6fcf-8002-69189ce6196a",
                "parentTs": "1f194af9-3081-6a93-8001-7df329fe8fe2",
                "alternativeCount": 0
              },
              {
                "content": "2",
                "messageType": "agent",
                "threadTs": "1f194af9-6ffa-6fcf-8002-69189ce6196a",
                "parentTs": "1f194af9-3081-6a93-8001-7df329fe8fe2",
                "alternativeCount": null
              }
            ]
          }
        }
      ]
    }
  },
  "correlationId": "01KZNQ9YCPF5K0BS2KQQQADJEC"
}
  1. Copy the parentTs of "Say 2" into gitlab-ai-gateway/duo_workflow_service/server.py replacing resume_checkpoint_ts=start_workflow_request.startRequest.resume_checkpoint_ts with resume_checkpoint_ts="1f194af9-3081-6a93-8001-7df329fe8fe2"
  2. Run gdk restart duo-workflow-service
  3. Send new messages "Say 3", "Say 4"
    1. 3 will replace 2, and then 4 will replace 3
  4. Run the query again, and you will see that 4 has 2 alternatives
  5. Run "Query to get thread_ts and counts" with the thread_ts of "say 4"
    1. You should see messages 2 and 3
GraphQL response 
{
  "data": {
    "duoWorkflowBranches": [
      {
        "forkThreadTs": "1f194af9-6ffa-6fcf-8002-69189ce6196a",
        "messages": [
          {
            "content": "Say 2",
            "messageType": "user",
            "threadTs": "1f194af9-6ffa-6fcf-8002-69189ce6196a",
            "parentTs": "1f194af9-3081-6a93-8001-7df329fe8fe2",
            "alternativeCount": 2
          },
          {
            "content": "2",
            "messageType": "agent",
            "threadTs": "1f194af9-6ffa-6fcf-8002-69189ce6196a",
            "parentTs": "1f194af9-3081-6a93-8001-7df329fe8fe2",
            "alternativeCount": null
          }
        ]
      },
      {
        "forkThreadTs": "1f194b00-309b-69f8-8003-8f17e5216602",
        "messages": [
          {
            "content": "Say 3",
            "messageType": "user",
            "threadTs": "1f194b00-216f-6995-8002-85db5d02e91a",
            "parentTs": "1f194af9-3081-6a93-8001-7df329fe8fe2",
            "alternativeCount": 2
          },
          {
            "content": "3",
            "messageType": "agent",
            "threadTs": "1f194b00-309b-69f8-8003-8f17e5216602",
            "parentTs": "1f194b00-216f-6995-8002-85db5d02e91a",
            "alternativeCount": null
          }
        ]
      }
    ]
  },
  "correlationId": "01KZNQPFTKXC02MNRQXQZR918Z"
}

Query

SELECT
    "p_duo_workflows_checkpoint_blobs\".\ "thread_ts ",
    "p_duo_workflows_checkpoint_blobs\".\ "data "
FROM
    "p_duo_workflows_checkpoint_blobs "
WHERE
    "p_duo_workflows_checkpoint_blobs\".\ "workflow_id " = 60
    AND "p_duo_workflows_checkpoint_blobs\".\ "workflow_created_at " = '2026-08-19 16:25:07.711908'
    AND "p_duo_workflows_checkpoint_blobs\".\ "project_id " = 1000000
    AND "p_duo_workflows_checkpoint_blobs\".\ "channel " = 'status'
    AND "p_duo_workflows_checkpoint_blobs\".\ "thread_ts " IN ('1f19beed-11e0-620d-8003-39a76db1468d', '1f19beed-2773-6b47-8004-f135be3a1817', '1f19bef0-b834-60de-8005-b266d256a794', '1f19bef0-c6e6-60d2-8006-26face071f75', '1f19beed-5a25-63c2-8003-3eb2bfd34517', '1f19beed-6988-6bed-8004-761701d06512')

I have 2 issues with creating the query plan, that I am not sure how to resolve. If a plan is needed for this MR, I would appreciate some advice on how to set it up.

  1. The query uses checkpoint_blobs which is scoped to the created_at needing sub second accuracy. But graphql only returns second accuracy, so I cannot replicate the query. The only possible way I could think of is binary search on https://console.postgres.ai/ using a timerange, but that will take too long.
  2. It returns data of retried turns, but there are no examples of this on production. The code to allow this will be added in !250054 (merged) but still blocked by a feature flag. Should I send a custom websocket request somehow to create the data I need, or copy it from my gdk somehow?

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 Kaveh Nejad

Merge request reports