fix: update grok session after new (#2683)

* fix: update grok session after new

refs #2681

* fix: avoid powershell args shadowing

refs #2681

---------

Co-authored-by: akbash-bot <300245827+akbash-bot@users.noreply.github.com>

5 files changed, 129 insertions, 13 deletions

diff --git a/src/integration/assets/grok/herdr-agent-state.ps1 b/src/integration/assets/grok/herdr-agent-state.ps1
index 69c4f515..1b183a3b 100644
--- a/src/integration/assets/grok/herdr-agent-state.ps1
+++ b/src/integration/assets/grok/herdr-agent-state.ps1
@@ -2,7 +2,7 @@
 # managed by herdr; reinstalling or updating the integration overwrites this file.
 # add custom hooks beside this file instead of editing it.
 # HERDR_INTEGRATION_ID=grok
-# HERDR_INTEGRATION_VERSION=1
+# HERDR_INTEGRATION_VERSION=2
 
 param([string]$Action = "")
 
@@ -26,6 +26,12 @@ $event = if ($null -ne $payload -and $payload.hook_event_name -is [string]) {
 }
 if ($null -ne $event -and $event -notin @("session_start", "SessionStart", "sessionStart")) { exit 0 }
 
+$sessionStartSource = if ($null -ne $payload -and $payload.source -is [string]) {
+    $payload.source
+} else {
+    $null
+}
+
 $sessionId = $env:GROK_SESSION_ID
 if ([string]::IsNullOrWhiteSpace($sessionId) -and $null -ne $payload) {
     if ($payload.session_id -is [string]) { $sessionId = $payload.session_id }
@@ -35,7 +41,17 @@ if ([string]::IsNullOrWhiteSpace($sessionId)) { exit 0 }
 
 $seq = [DateTime]::UtcNow.Ticks
 $herdr = if ([string]::IsNullOrWhiteSpace($env:HERDR_BIN_PATH)) { "herdr" } else { $env:HERDR_BIN_PATH }
+$herdrArgs = @(
+    "pane", "report-agent-session", $env:HERDR_PANE_ID,
+    "--source", "herdr:grok",
+    "--agent", "grok",
+    "--seq", "$seq",
+    "--agent-session-id", "$sessionId"
+)
+if (-not [string]::IsNullOrWhiteSpace($sessionStartSource)) {
+    $herdrArgs += @("--session-start-source", "$sessionStartSource")
+}
 try {
-    & $herdr pane report-agent-session $env:HERDR_PANE_ID --source herdr:grok --agent grok --seq $seq --agent-session-id $sessionId 2>$null | Out-Null
+    & $herdr @herdrArgs 2>$null | Out-Null
 } catch {
 }
diff --git a/src/integration/assets/grok/herdr-agent-state.sh b/src/integration/assets/grok/herdr-agent-state.sh
index bdf0f671..3eaa3431 100644
--- a/src/integration/assets/grok/herdr-agent-state.sh
+++ b/src/integration/assets/grok/herdr-agent-state.sh
@@ -3,7 +3,7 @@
 # managed by herdr; reinstalling or updating the integration overwrites this file.
 # add custom hooks beside this file instead of editing it.
 # HERDR_INTEGRATION_ID=grok
-# HERDR_INTEGRATION_VERSION=1
+# HERDR_INTEGRATION_VERSION=2
 
 set -eu
 
@@ -62,6 +62,7 @@ def first_text(*keys):
 hook_event_name = first_text("hook_event_name", "hookEventName")
 if hook_event_name not in (None, "session_start", "SessionStart", "sessionStart"):
     raise SystemExit(0)
+session_start_source = first_text("source")
 
 # Grok injects GROK_SESSION_ID into every hook process; prefer it and fall
 # back to the event payload's session id fields.
@@ -72,16 +73,19 @@ if not agent_session_id:
 
 request_id = f"{source}:{int(time.time() * 1000)}:{random.randrange(1_000_000):06d}"
 report_seq = time.time_ns()
+params = {
+    "pane_id": pane_id,
+    "source": source,
+    "agent": "grok",
+    "seq": report_seq,
+    "agent_session_id": agent_session_id,
+}
+if session_start_source:
+    params["session_start_source"] = session_start_source
 request = {
     "id": request_id,
     "method": "pane.report_agent_session",
-    "params": {
-        "pane_id": pane_id,
-        "source": source,
-        "agent": "grok",
-        "seq": report_seq,
-        "agent_session_id": agent_session_id,
-    },
+    "params": params,
 }
 
 try:
diff --git a/src/integration/mod.rs b/src/integration/mod.rs
index 49808ca2..55dcb2be 100644
--- a/src/integration/mod.rs
+++ b/src/integration/mod.rs
@@ -296,7 +296,7 @@ const GROK_HOOK_ASSET: &str = if cfg!(windows) {
 } else {
     include_str!("assets/grok/herdr-agent-state.sh")
 };
-const GROK_INTEGRATION_VERSION: u32 = 1;
+const GROK_INTEGRATION_VERSION: u32 = 2;
 
 pub(crate) const INSTALL_WARNING_PREFIX: &str = "warning:";
 
diff --git a/src/integration/tests.rs b/src/integration/tests.rs
index dc597d1d..e8aad72b 100644
--- a/src/integration/tests.rs
+++ b/src/integration/tests.rs
@@ -2897,6 +2897,7 @@ fn bundled_integration_asset_versions_match_expected_versions() {
             MASTRACODE_HOOK_ASSET,
             MASTRACODE_INTEGRATION_VERSION,
         ),
+        ("grok", GROK_HOOK_ASSET, GROK_INTEGRATION_VERSION),
     ] {
         assert_eq!(
             parse_integration_version(asset),
@@ -4056,7 +4057,39 @@ fn install_antigravity_cli_errors_when_config_dir_missing() {
 }
 
 #[test]
-fn grok_v1_integration_status_is_current() {
+fn grok_v1_integration_status_is_outdated() {
+    let _lock = integration_env_lock();
+    let base = unique_base();
+    let grok_dir = base.join(".grok");
+    let hooks_dir = grok_dir.join("hooks");
+    fs::create_dir_all(&hooks_dir).unwrap();
+    std::env::set_var(GROK_CONFIG_DIR_ENV_VAR, &grok_dir);
+    let hook_path = hooks_dir.join(GROK_HOOK_INSTALL_NAME);
+    fs::write(
+        &hook_path,
+        "#!/bin/sh\n# HERDR_INTEGRATION_ID=grok\n# HERDR_INTEGRATION_VERSION=1\n",
+    )
+    .unwrap();
+    fs::write(
+        hooks_dir.join(GROK_HOOK_CONFIG_INSTALL_NAME),
+        serde_json::to_string(&grok_hook_config(&hook_path)).unwrap(),
+    )
+    .unwrap();
+
+    let grok = installed_integration_statuses()
+        .into_iter()
+        .find(|status| status.target == crate::api::schema::IntegrationTarget::Grok)
+        .expect("grok integration status");
+    assert_eq!(grok.installed_version, Some(1));
+    assert_eq!(grok.expected_version, GROK_INTEGRATION_VERSION);
+    assert_eq!(grok.state, IntegrationStatusKind::Outdated);
+
+    clear_integration_path_env();
+    let _ = fs::remove_dir_all(base);
+}
+
+#[test]
+fn grok_v2_integration_status_is_current() {
     let _lock = integration_env_lock();
     let base = unique_base();
     let grok_dir = base.join(".grok");
diff --git a/src/terminal/state.rs b/src/terminal/state.rs
index dc28af10..a7c84574 100644
--- a/src/terminal/state.rs
+++ b/src/terminal/state.rs
@@ -1331,6 +1331,7 @@ impl TerminalState {
                 | ("herdr:hermes", "hermes", Some("startup" | "new" | "resume"))
                 | ("herdr:opencode", "opencode", Some("select"))
                 | ("herdr:pi", "pi", Some("new" | "resume" | "fork"))
+                | ("herdr:grok", "grok", Some("new"))
                 | (
                     "herdr:omp",
                     "omp",
@@ -1635,7 +1636,8 @@ impl TerminalState {
         session_ref: &crate::agent_resume::AgentSessionRef,
         session_start_source: Option<&str>,
     ) -> bool {
-        Self::session_start_source_is_recognized(session_start_source)
+        (source, agent_label) != ("herdr:grok", "grok")
+            && Self::session_start_source_is_recognized(session_start_source)
             && self.foreground_agent_confirms_session_owner(source, agent_label, session_ref)
     }
 
@@ -4653,6 +4655,38 @@ mod tests {
     }
 
     #[test]
+    fn grok_new_session_ref_replaces_existing_session_ref() {
+        let mut terminal = test_terminal();
+        terminal
+            .set_agent_session_ref(
+                "herdr:grok".into(),
+                "grok".into(),
+                crate::agent_resume::AgentSessionRef::id("grok-old"),
+                Some(20),
+            )
+            .expect("initial session should be accepted");
+
+        let mutation = terminal
+            .set_agent_session_ref_for_session_start(
+                "herdr:grok".into(),
+                "grok".into(),
+                crate::agent_resume::AgentSessionRef::id("grok-new"),
+                Some(21),
+                Some("new".into()),
+            )
+            .expect("new should replace the grok session");
+
+        assert!(mutation.session_ref_changed);
+        assert_eq!(
+            terminal
+                .persisted_agent_session
+                .as_ref()
+                .map(|session| session.session_ref.value.as_str()),
+            Some("grok-new")
+        );
+    }
+
+    #[test]
     fn opencode_server_new_does_not_replace_existing_session_ref() {
         let mut terminal = test_terminal();
         terminal.set_detected_state(Some(Agent::OpenCode), AgentState::Idle);
@@ -5106,6 +5140,35 @@ mod tests {
     }
 
     #[test]
+    fn grok_new_session_does_not_replace_a_different_owner() {
+        let mut terminal = test_terminal();
+        terminal.set_persisted_agent_session(crate::agent_resume::PersistedAgentSession {
+            source: "herdr:claude".into(),
+            agent: "claude".into(),
+            session_ref: crate::agent_resume::AgentSessionRef::id("claude-session").unwrap(),
+        });
+        terminal.set_detected_state(Some(Agent::Grok), AgentState::Idle);
+
+        let mutation = terminal.set_agent_session_ref_for_session_start(
+            "herdr:grok".into(),
+            "grok".into(),
+            crate::agent_resume::AgentSessionRef::id("grok-session"),
+            Some(21),
+            Some("new".into()),
+        );
+
+        assert!(mutation.is_none());
+        assert_eq!(
+            terminal.persisted_agent_session.as_ref().map(|session| (
+                session.source.as_str(),
+                session.agent.as_str(),
+                session.session_ref.value.as_str()
+            )),
+            Some(("herdr:claude", "claude", "claude-session"))
+        );
+    }
+
+    #[test]
     fn foreground_agent_session_replaces_stale_different_owner_session_ref() {
         for session_start_source in ["resume", "startup"] {
             let mut terminal = test_terminal();