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.
This commit is contained in:
127
src/features/claude-code-session-state/state.test.ts
Normal file
127
src/features/claude-code-session-state/state.test.ts
Normal file
@@ -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()
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -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 })
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user