Add MLflow Stop and SessionStart hooks to setup and docs (#17)

Replace old session-start.sh hook references with inline MLflow hooks
(Stop for trace flushing, SessionStart for session ID capture).
Update interactive setup in SKILL.md to auto-add hooks when MLflow is
configured. Rename check_session_hook to check_mlflow_hooks in setup.py
to validate both hooks are present.

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: Shreyanand

.claude/settings.example.json

@@ -18,22 +18,22 @@
     "MLFLOW_TAG_USER": "<your-username>"
   },
   "hooks": {
-    "SessionStart": [
+    "Stop": [
       {
         "hooks": [
           {
             "type": "command",
-            "command": "bash .claude/hooks/session-start.sh"
+            "command": "python -c \"from mlflow.claude_code.hooks import stop_hook_handler; stop_hook_handler()\""
           }
         ]
       }
     ],
-    "Stop": [
+    "SessionStart": [
       {
         "hooks": [
           {
             "type": "command",
-            "command": "python -c \"from mlflow.claude_code.hooks import stop_hook_handler; stop_hook_handler()\""
+            "command": "INPUT=$(cat); SESSION_ID=$(echo \"$INPUT\" | jq -r '.session_id'); echo \"export CLAUDE_SESSION_ID='$SESSION_ID'\" >> \"$CLAUDE_ENV_FILE\""
           }
         ]
       }

README.md

@@ -176,22 +176,30 @@ Step 5   [Claude]  Analyze and summarize root cause
 ```json
 {
   "hooks": {
+    "Stop": [
+      {
+        "hooks": [
+          {
+            "type": "command",
+            "command": "python -c \"from mlflow.claude_code.hooks import stop_hook_handler; stop_hook_handler()\""
+          }
+        ]
+      }
+    ],
     "SessionStart": [
       {
-        "matcher": "*",
-        "hooks": [   
+        "hooks": [
           {
             "type": "command",
-            "command": "./.claude/hooks/session-start.sh"
+            "command": "INPUT=$(cat); SESSION_ID=$(echo \"$INPUT\" | jq -r '.session_id'); echo \"export CLAUDE_SESSION_ID='$SESSION_ID'\" >> \"$CLAUDE_ENV_FILE\""
           }
         ]
       }
     ]
   }
 }
 ```
----
-**Note:** Run the following to make the script executable: `chmod +x ./.claude/hooks/session-start.sh`
+**Note:** These hooks are required for MLflow tracing. The `Stop` hook flushes traces and the `SessionStart` hook captures the session ID.
 
 ### context-fetcher
 
@@ -366,17 +374,29 @@ Add the following environment variables to your Claude settings file at `~/.clau
 **Note**: Replace `<your-username>@<your-jumpbox> -p <port>` with your actual jumpbox connection details.
 Replace `<experiment-name>` to a experiment name to the desired experiment name. This will automatically create the experiment for you if it does not exist.
 
-### Step 3: Add hooks SessionStart to hooks in .claude/settings.json
+### Step 3: Add MLflow hooks to `.claude/settings.json`
+
+Add both `Stop` and `SessionStart` hooks to your settings:
+
 ```json
 {
   "hooks": {
+    "Stop": [
+      {
+        "hooks": [
+          {
+            "type": "command",
+            "command": "python -c \"from mlflow.claude_code.hooks import stop_hook_handler; stop_hook_handler()\""
+          }
+        ]
+      }
+    ],
     "SessionStart": [
       {
-        "matcher": "*",
         "hooks": [
           {
             "type": "command",
-            "command": "~/.claude/hooks/session-start.sh"
+            "command": "INPUT=$(cat); SESSION_ID=$(echo \"$INPUT\" | jq -r '.session_id'); echo \"export CLAUDE_SESSION_ID='$SESSION_ID'\" >> \"$CLAUDE_ENV_FILE\""
           }
         ]
       }
@@ -385,10 +405,8 @@ Replace `<experiment-name>` to a experiment name to the desired experiment name.
 }
 ```
 
-Make the script executable:
-```bash
-chmod +x ./.claude/hooks/session-start.sh
-```
+- **Stop hook**: Flushes MLflow traces when Claude stops
+- **SessionStart hook**: Captures the session ID for trace correlation
 
 ### Step 4: cd into AIOPS-SKILLS/ dir (where claude will be running)
 

skills/root-cause-analysis/README.md

@@ -177,7 +177,14 @@ MLflow tracing is supported for debugging and performance analysis but is **not
 To enable tracing:
 1. Install MLflow: `pip install "mlflow[genai]>=3.4"` (or `pip install -e ".[mlflow]"` from the repo root)
 2. Configure `MLFLOW_PORT`, `MLFLOW_EXPERIMENT_NAME`, and optionally `JUMPBOX_URI` in `.claude/settings.json`
-3. The `session-start.sh` hook will auto-start the SSH tunnel and configure tracing when `MLFLOW_PORT` is set
+3. Add the required MLflow hooks to `.claude/settings.json`:
+   ```json
+   "hooks": {
+     "Stop": [{ "hooks": [{ "type": "command", "command": "python -c \"from mlflow.claude_code.hooks import stop_hook_handler; stop_hook_handler()\"" }] }],
+     "SessionStart": [{ "hooks": [{ "type": "command", "command": "INPUT=$(cat); SESSION_ID=$(echo \"$INPUT\" | jq -r '.session_id'); echo \"export CLAUDE_SESSION_ID='$SESSION_ID'\" >> \"$CLAUDE_ENV_FILE\"" }] }]
+   }
+   ```
+   The `Stop` hook flushes traces and the `SessionStart` hook captures the session ID for correlation.
 
 When enabled, pipeline-level traces (analysis steps 1-4) appear in the MLflow UI. When not installed, tracing is silently skipped.
 

skills/root-cause-analysis/SKILL.md

@@ -61,9 +61,16 @@ If any checks have `"status": "missing"` and `"configurable": true`, offer to he
      - Then set `REMOTE_HOST` to the alias name
 4. After collecting all values, read the project's `.claude/settings.json` file
 5. Merge the new values into the `"env"` block (create it if it doesn't exist)
-6. Write the updated settings file
-7. Tell the user to **restart the Claude Code session** for env vars to take effect
-8. **Important**: Write secrets (tokens, passwords) to `.claude/settings.json` -- ensure this file is in `.gitignore`
+6. If MLflow env vars were configured (MLFLOW_PORT, MLFLOW_EXPERIMENT_NAME), also add the required MLflow hooks to the `"hooks"` block (create it if it doesn't exist):
+   ```json
+   "hooks": {
+     "Stop": [{ "hooks": [{ "type": "command", "command": "python -c \"from mlflow.claude_code.hooks import stop_hook_handler; stop_hook_handler()\"" }] }],
+     "SessionStart": [{ "hooks": [{ "type": "command", "command": "INPUT=$(cat); SESSION_ID=$(echo \"$INPUT\" | jq -r '.session_id'); echo \"export CLAUDE_SESSION_ID='$SESSION_ID'\" >> \"$CLAUDE_ENV_FILE\"" }] }]
+   }
+   ```
+7. Write the updated settings file
+8. Tell the user to **restart the Claude Code session** for env vars and hooks to take effect
+9. **Important**: Write secrets (tokens, passwords) to `.claude/settings.json` -- ensure this file is in `.gitignore`
 
 If checks show non-configurable errors (e.g., venv issues, rsync not found), provide the fix command instead.
 

skills/root-cause-analysis/scripts/setup.py

@@ -507,37 +507,40 @@ def check_mlflow_server() -> dict:
     }
 
 
-def check_session_hook(repo_root: Path) -> dict:
-    """Check if session-start hook is configured."""
-    hook_path = repo_root / ".claude" / "hooks" / "session-start.sh"
-    if not hook_path.exists():
-        return {
-            "name": "Session hook",
-            "status": "missing",
-            "message": ".claude/hooks/session-start.sh not found",
-        }
-
-    # Check if settings.json has the SessionStart hook registered
+def check_mlflow_hooks(repo_root: Path) -> dict:
+    """Check if MLflow Stop and SessionStart hooks are configured in settings.json."""
     for settings_name in ["settings.json", "settings.local.json"]:
         settings_path = repo_root / ".claude" / settings_name
         if settings_path.exists():
             try:
                 with open(settings_path) as f:
                     settings = json.load(f)
                 hooks = settings.get("hooks", {})
-                if "SessionStart" in hooks:
+                has_stop = "Stop" in hooks
+                has_session = "SessionStart" in hooks
+                if has_stop and has_session:
                     return {
-                        "name": "Session hook",
+                        "name": "MLflow hooks",
                         "status": "ok",
-                        "message": f"registered in {settings_name}",
+                        "message": f"Stop + SessionStart registered in {settings_name}",
                     }
+                missing = []
+                if not has_stop:
+                    missing.append("Stop")
+                if not has_session:
+                    missing.append("SessionStart")
+                return {
+                    "name": "MLflow hooks",
+                    "status": "missing",
+                    "message": f"Missing hooks: {', '.join(missing)}. See settings.example.json",
+                }
             except (json.JSONDecodeError, OSError):
                 pass
 
     return {
-        "name": "Session hook",
+        "name": "MLflow hooks",
         "status": "missing",
-        "message": "SessionStart hook not registered in settings. See settings.example.json",
+        "message": "Stop and SessionStart hooks not found in settings. See settings.example.json",
     }
 
 
@@ -556,7 +559,7 @@ def run_checks(base_dir: Path, repo_root: Path | None = None) -> list[dict]:
         check_github_mcp(repo_root),
         check_mlflow(),
         check_mlflow_server(),
-        check_session_hook(repo_root),
+        check_mlflow_hooks(repo_root),
     ]