Draft: Let a foundational flow react to a session that never finished

Wanderson Policarporequested to merge
ai-session-lifecycle-hooks into master

What does this MR do and why?

A foundational flow that tracks something outside its own session has no way to learn that the session died. The only tool available today is a timeout worker: schedule a job, wait, and assume the worst if nothing reported back. That is a guess, and it gets the edges wrong - a stale timer can fail a run that already succeeded, and a run that dies in the first second still waits out the whole clock before anything notices.

A Duo session now publishes an event when it ends without finishing, and a flow definition can declare an on_session_dropped hook to react to it.

The event carries the terminal status and the status it came from, because those cases are not the same. A user cancelling a session is not a failure, and a session dropped straight from created never ran at all, so a hook has no earlier work to undo. Getting either wrong shows the user an error for something that did not fail.

Nothing declares the hook yet, which is deliberate - this is the foundation, and the risk assessment work will be the first consumer.

References

  • Related: #609292 (closed)
  • Grew out of the discussion on !255592 (closed), which drove the risk assessment lifecycle from session events for that one flow. This generalizes the idea so any foundational flow can use it.

Screenshots or screen recordings

No UI change.

How to set up and validate locally

Enable risk_classification/v1 flow:

  1. Enable the existing flag in the Rails console: 
    Feature.enable(:duo_mr_risk_classification)
Feature.enable(:show_duo_mr_risk_classification_widget)
  2. Enable the risk_classification/v1 foundational flow for the root namespace. Go to the group's Duo settings at http://gdk.test:3000/groups/gitlab-duo/-/edit#js-gitlab-duo-settings. Turn the toggle on.

Simulate session hooks

Session lifecycle events are only fired if the flow has hooks enabled. This saves us from enqueuing a Sidekiq job that doesn't have a hook session to execute.

Apply the following patch to gitlab project to simulate a session hook for the risk_classification/v1 flow:

diff --git a/ee/app/models/ai/catalog/foundational_flow/risk_classification/definition.rb b/ee/app/models/ai/catalog/foundational_flow/risk_classification/definition.rb
index 17a29d1c8e1e..9901d87d7cba 100644
--- a/ee/app/models/ai/catalog/foundational_flow/risk_classification/definition.rb
+++ b/ee/app/models/ai/catalog/foundational_flow/risk_classification/definition.rb
@@ -42,6 +42,20 @@ def configuration
                 # enqueue is refused while a run is already queued; touching keeps the
                 # timeout counting from this attempt instead of the earlier one's start.
                 assessment.enqueue || assessment.touch
+              end,
+              on_session_failed: ->(workflow:) do
+                Gitlab::AppLogger.info(
+                  message: "session_hook_test",
+                  hook: "on_session_failed",
+                  workflow: workflow
+                )
+              end,
+              on_session_stopped: ->(workflow:) do
+                Gitlab::AppLogger.info(
+                  message: "session_hook_test",
+                  hook: "on_session_stopped",
+                  workflow: workflow
+                )
               end
             }
           end

Testing on_session_stopped: stop a running flow

Restart impacted services:

gdk restart rails-web rails-background-jobs ai-gateway

  1. Open a merge request in gitlab-duo/test with an actual code change in it

  2. Watch the session at AI > Sessions in the project and cancel the session before it moves to running

  3. Confirm the event was published and dispatched:

    tail -f log/* | grep session_hook_test
Example of log 
{
    "severity": "INFO",
    "time": "2026-09-21T17:29:27.962Z",
    "correlation_id": "01M32G7Z211M8E5CNMC7YM1PMA",
    "meta.caller_id": "Ai::DuoWorkflows::SessionLifecycleWorker",
    "meta.remote_ip": "172.16.123.1",
    "meta.feature_category": "duo_agent_platform",
    "meta.user": "root",
    "meta.gl_user_id": 1,
    "meta.client_id": "user/1",
    "meta.organization_id": 1,
    "meta.root_caller_id": "PATCH /api/:version/ai/duo_workflows/workflows/:id",
    "message": "session_hook_test",
    "hook": "on_session_stopped",
    "workflow": {
        "id": 47,
        "user_id": 1,
        "project_id": 1000000,
        "created_at": "2026-09-21T17:29:22.691Z",
        "updated_at": "2026-09-21T17:29:27.935Z",
        "status": 5,
        "goal": "http://gdk.test:3000/gitlab-duo/test/-/merge_requests/25",
        "agent_privileges": [
            1,
            2,
            3,
            4,
            5,
            6,
            7
        ],
        "workflow_definition": "risk_classification/v1",
        "allow_agent_to_request_user": true,
        "pre_approved_agent_privileges": [
            1,
            2,
            3,
            4,
            5,
            6,
            7
        ],
        "image": null,
        "environment": "web",
        "namespace_id": null,
        "ai_catalog_item_version_id": 2,
        "issue_id": null,
        "merge_request_id": 161,
        "service_account_id": 71,
        "tool_call_approvals": {},
        "ai_catalog_item_id": null,
        "summary": null,
        "messaging_callback_context": null,
        "title": "Review merge request 25",
        "model_metadata_json": null,
        "incremental_checkpoints_enabled": true,
        "agent_type": null,
        "jsonl_sha256": null,
        "idempotency_key": null,
        "sync_type": null,
        "agent_identity_id": null,
        "flow_metadata_json": null,
        "web_search_enabled": false,
        "trigger_source": "human",
        "trigger_flow_trigger_id": null,
        "execution_mode": null,
        "source_type": null,
        "source_link": null,
        "trigger_flow_schedule_id": null,
        "trigger_event_type": null
    }
}

Testing on_session_failed: break a Code Review run on purpose

Apply the following patch to gitlab-ai-gateway to simulate a random error during flow execution:

diff --git a/duo_workflow_service/agent_platform/v1/flows/configs/code_review/1.0.0.yml b/duo_workflow_service/agent_platform/v1/flows/configs/code_review/1.0.0.yml
index e97b702db..825dffa76 100644
--- a/duo_workflow_service/agent_platform/v1/flows/configs/code_review/1.0.0.yml
+++ b/duo_workflow_service/agent_platform/v1/flows/configs/code_review/1.0.0.yml
@@ -20,7 +20,7 @@ components:
     tool_name: "build_review_merge_request_context"
     inputs:
       - from: "context:project_id"
-        as: "project_id"
+        as: "projekt_id"
       - from: "context:goal"
         as: "merge_request_iid"
       - from: "true"

Then restart impacted services:

gdk restart rails-web rails-background-jobs ai-gateway

  1. Open a merge request in gitlab-duo/test with an actual code change in it

  2. Watch the session at AI > Sessions in the project. It moves to running, the gitlab--duo CI job starts, then the session ends as failed once the executor cannot load its flow config. gdk tail duo-workflow-service shows the resolution error if you want to see the cause.

  3. Confirm the event was published and dispatched. This needs no hook and no code change - the worker only runs because the event reached it:

    tail -f log/* | grep session_hook_test
Example of log 
{
    "severity": "INFO",
    "time": "2026-09-21T17:15:25.247Z",
    "correlation_id": "01M32FE7BCKGEZTZM0G8DNX856",
    "meta.duo_workflow_id": "46",
    "meta.scoped_user": "root",
    "meta.scoped_user_id": 1,
    "meta.user": "root",
    "meta.gl_user_id": 1,
    "meta.client_id": "user/71",
    "meta.caller_id": "Ai::DuoWorkflows::SessionLifecycleWorker",
    "meta.remote_ip": "172.16.123.1",
    "meta.feature_category": "duo_agent_platform",
    "meta.organization_id": 1,
    "meta.root_caller_id": "PATCH /api/:version/ai/duo_workflows/workflows/:id",
    "message": "session_hook_test",
    "hook": "on_session_failed",
    "workflow": {
        "id": 46,
        "user_id": 1,
        "project_id": 1000000,
        "created_at": "2026-09-21T17:15:16.813Z",
        "updated_at": "2026-09-21T17:15:25.228Z",
        "status": 4,
        "goal": "http://gdk.test:3000/gitlab-duo/test/-/merge_requests/24",
        "agent_privileges": [
            1,
            2,
            3,
            4,
            5,
            6,
            7
        ],
        "workflow_definition": "risk_classification/v1",
        "allow_agent_to_request_user": true,
        "pre_approved_agent_privileges": [
            1,
            2,
            3,
            4,
            5,
            6,
            7
        ],
        "image": null,
        "environment": "web",
        "namespace_id": null,
        "ai_catalog_item_version_id": 2,
        "issue_id": null,
        "merge_request_id": 160,
        "service_account_id": 71,
        "tool_call_approvals": {},
        "ai_catalog_item_id": null,
        "summary": null,
        "messaging_callback_context": null,
        "title": "Review MR !24 risk classification",
        "model_metadata_json": null,
        "incremental_checkpoints_enabled": true,
        "agent_type": null,
        "jsonl_sha256": null,
        "idempotency_key": null,
        "sync_type": null,
        "agent_identity_id": null,
        "flow_metadata_json": null,
        "web_search_enabled": false,
        "trigger_source": "human",
        "trigger_flow_trigger_id": null,
        "execution_mode": null,
        "source_type": null,
        "source_link": null,
        "trigger_flow_schedule_id": null,
        "trigger_event_type": null
    }
}

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 Wanderson Policarpo

Merge request reports