diff --git a/src/features/claude-code-session-state/state.test.ts b/src/features/claude-code-session-state/state.test.ts
new file mode 100644
index 000000000..7b72ebf0a
--- /dev/null
+++ b/src/features/claude-code-session-state/state.test.ts
@@ -0,0 +1,127 @@
+import { describe, test, expect, beforeEach } from "bun:test"
+import {
+  setSessionAgent,
+  getSessionAgent,
+  clearSessionAgent,
+  updateSessionAgent,
+  setMainSession,
+  getMainSessionID,
+  subagentSessions,
+} from "./state"
+
+describe("claude-code-session-state", () => {
+  beforeEach(() => {
+    // #given - clean state before each test
+    clearSessionAgent("test-session-1")
+    clearSessionAgent("test-session-2")
+    clearSessionAgent("test-prometheus-session")
+    setMainSession(undefined)
+    subagentSessions.clear()
+  })
+
+  describe("setSessionAgent", () => {
+    test("should store agent for session", () => {
+      // #given
+      const sessionID = "test-session-1"
+      const agent = "Prometheus (Planner)"
+
+      // #when
+      setSessionAgent(sessionID, agent)
+
+      // #then
+      expect(getSessionAgent(sessionID)).toBe(agent)
+    })
+
+    test("should NOT overwrite existing agent (first-write wins)", () => {
+      // #given
+      const sessionID = "test-session-1"
+      setSessionAgent(sessionID, "Prometheus (Planner)")
+
+      // #when - try to overwrite
+      setSessionAgent(sessionID, "Sisyphus")
+
+      // #then - first agent preserved
+      expect(getSessionAgent(sessionID)).toBe("Prometheus (Planner)")
+    })
+
+    test("should return undefined for unknown session", () => {
+      // #given - no session set
+
+      // #when / #then
+      expect(getSessionAgent("unknown-session")).toBeUndefined()
+    })
+  })
+
+  describe("updateSessionAgent", () => {
+    test("should overwrite existing agent", () => {
+      // #given
+      const sessionID = "test-session-1"
+      setSessionAgent(sessionID, "Prometheus (Planner)")
+
+      // #when - force update
+      updateSessionAgent(sessionID, "Sisyphus")
+
+      // #then
+      expect(getSessionAgent(sessionID)).toBe("Sisyphus")
+    })
+  })
+
+  describe("clearSessionAgent", () => {
+    test("should remove agent from session", () => {
+      // #given
+      const sessionID = "test-session-1"
+      setSessionAgent(sessionID, "Prometheus (Planner)")
+      expect(getSessionAgent(sessionID)).toBe("Prometheus (Planner)")
+
+      // #when
+      clearSessionAgent(sessionID)
+
+      // #then
+      expect(getSessionAgent(sessionID)).toBeUndefined()
+    })
+  })
+
+  describe("mainSessionID", () => {
+    test("should store and retrieve main session ID", () => {
+      // #given
+      const mainID = "main-session-123"
+
+      // #when
+      setMainSession(mainID)
+
+      // #then
+      expect(getMainSessionID()).toBe(mainID)
+    })
+
+    test("should return undefined when not set", () => {
+      // #given - not set
+
+      // #then
+      expect(getMainSessionID()).toBeUndefined()
+    })
+  })
+
+  describe("prometheus-md-only integration scenario", () => {
+    test("should correctly identify Prometheus agent for permission checks", () => {
+      // #given - Prometheus session
+      const sessionID = "test-prometheus-session"
+      const prometheusAgent = "Prometheus (Planner)"
+
+      // #when - agent is set (simulating chat.message hook)
+      setSessionAgent(sessionID, prometheusAgent)
+
+      // #then - getSessionAgent returns correct agent for prometheus-md-only hook
+      const agent = getSessionAgent(sessionID)
+      expect(agent).toBe("Prometheus (Planner)")
+      expect(["Prometheus (Planner)"].includes(agent!)).toBe(true)
+    })
+
+    test("should return undefined when agent not set (bug scenario)", () => {
+      // #given - session exists but no agent set (the bug)
+      const sessionID = "test-prometheus-session"
+
+      // #when / #then - this is the bug: agent is undefined
+      expect(getSessionAgent(sessionID)).toBeUndefined()
+    })
+  })
+})
diff --git a/src/hooks/prometheus-md-only/index.test.ts b/src/hooks/prometheus-md-only/index.test.ts
index c703c1dd6..c3a186f44 100644
--- a/src/hooks/prometheus-md-only/index.test.ts
+++ b/src/hooks/prometheus-md-only/index.test.ts
@@ -4,6 +4,7 @@ import { join } from "node:path"
 import { createPrometheusMdOnlyHook } from "./index"
 import { MESSAGE_STORAGE } from "../../features/hook-message-injector"
 import { SYSTEM_DIRECTIVE_PREFIX, createSystemDirective, SystemDirectiveTypes } from "../../shared/system-directive"
+import { clearSessionAgent } from "../../features/claude-code-session-state"
 
 describe("prometheus-md-only", () => {
   const TEST_SESSION_ID = "test-session-prometheus"
@@ -30,6 +31,7 @@ describe("prometheus-md-only", () => {
   }
 
   afterEach(() => {
+    clearSessionAgent(TEST_SESSION_ID)
     if (testMessageDir) {
       try {
         rmSync(testMessageDir, { recursive: true, force: true })
diff --git a/src/index.ts b/src/index.ts
index 78ca94758..9c23dee04 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -310,6 +310,10 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => {
     },
 
     "chat.message": async (input, output) => {
+      if (input.agent) {
+        setSessionAgent(input.sessionID, input.agent);
+      }
+
       const message = (output as { message: { variant?: string } }).message
       if (firstMessageVariantGate.shouldOverride(input.sessionID)) {
         const variant = resolveAgentVariant(pluginConfig, input.agent)