Add session lifecycle hooks for foundational flows

What does this MR do and why?

This MR introduces execution hooks to foundational flows. A flow definition can now declare two optional hooks, on_session_failed and on_session_stopped - each taking workflow: that a consumer resolves from the session and calls the hook.

A messaging session still publishes the aggregated event: the payload travels as Sidekiq job arguments, and an older release cannot read a payload a newer one split. Nothing older reads a non-messaging session's events, so those are safe to split now. An older release also cannot resolve the stopped event class at all. Moving messaging sessions over is tracked in #630318

No flow declares a hook yet, so nothing changes for users in this MR.

References

How to set up and validate locally

Enable the risk_classification/v1 flow

  1. Enable the flags 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 and turn the toggle on.

Declare a hook to observe

A session lifecycle event only reaches the worker when the flow declares a hook for it. That guard is what keeps a Sidekiq job off the queue for a flow with nothing to run. No flow declares one yet, so apply this patch to the gitlab project to add a pair that only logs:

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

Then restart the affected services:

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

Test on_session_stopped: cancel a running session

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

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

  3. Confirm the hook ran:

    tail -f log/* | grep session_hook_test

    The line reports "hook": "on_session_stopped" and "meta.caller_id": "Ai::DuoWorkflows::SessionLifecycleWorker".

Test on_session_failed: break the run on purpose

  1. In gitlab-ai-gateway, break the flow config the session runs, at duo_workflow_service/agent_platform/v1/flows/configs/<flow>/<version>.yml. Renaming one input key is enough, for example changing as: "project_id" to as: "projekt_id". The executor then fails to resolve its config.

  2. Restart the affected services again:

    gdk restart rails-web rails-background-jobs ai-gateway
  3. Open another merge request in gitlab-duo/test with a real code change in it.

  4. Watch the session at AI > Sessions. It moves to running, the gitlab--duo CI job starts, then the session ends as failed. gdk tail duo-workflow-service shows the resolution error if you want the cause.

  5. Confirm the hook ran, as above. The line reports "hook": "on_session_failed".

Check a messaging session is untouched

Start a session from a @GitLabDuo note on an issue and cancel it. The reply on the note still reports a cancellation, and no session_hook_test line appears, because a messaging session still publishes the aggregated event.

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

Loading
Loading