fix: store session agent in chat.message for prometheus-md-only hook

The prometheus-md-only hook was not enforcing file restrictions because getSessionAgent() returned undefined - setSessionAgent was only called in message.updated event which doesn't always provide agent info. - Add setSessionAgent call in chat.message hook when input.agent exists - Add session state tests for setSessionAgent/getSessionAgent - Add clearSessionAgent cleanup to prometheus-md-only tests This ensures prometheus-md-only hook can reliably identify Prometheus sessions and enforce .sisyphus/*.md write restrictions.
2026-01-16 11:33:25 +09:00
parent 75925d5433
commit c282244439
3 changed files with 133 additions and 0 deletions
@@ -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()
+    })
+  })
+})