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
- Builds on the merged !256961 (merged), which shipped the event class and the wider schema
- Tracking issue: #630318
- Related work item: #609292 (closed)
- Supersedes the combined change in !255876 (closed)
How to set up and validate locally
Enable the risk_classification/v1 flow
- Enable the flags in the Rails console:
Feature.enable(:duo_mr_risk_classification) Feature.enable(:show_duo_mr_risk_classification_widget) - Enable the
risk_classification/v1foundational 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
}
endThen restart the affected services:
gdk restart rails-web rails-background-jobs ai-gatewayTest on_session_stopped: cancel a running session
-
Open a merge request in
gitlab-duo/testwith a real code change in it. -
Watch the session at AI > Sessions in the project, and cancel it before it moves to running.
-
Confirm the hook ran:
tail -f log/* | grep session_hook_testThe line reports
"hook": "on_session_stopped"and"meta.caller_id": "Ai::DuoWorkflows::SessionLifecycleWorker".
Test on_session_failed: break the run on purpose
-
In
gitlab-ai-gateway, break the flow config the session runs, atduo_workflow_service/agent_platform/v1/flows/configs/<flow>/<version>.yml. Renaming one input key is enough, for example changingas: "project_id"toas: "projekt_id". The executor then fails to resolve its config. -
Restart the affected services again:
gdk restart rails-web rails-background-jobs ai-gateway -
Open another merge request in
gitlab-duo/testwith a real code change in it. -
Watch the session at AI > Sessions. It moves to running, the
gitlab--duoCI job starts, then the session ends as failed.gdk tail duo-workflow-serviceshows the resolution error if you want the cause. -
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.