diff --git a/src/hooks/prometheus-md-only/agent-matcher.ts b/src/hooks/prometheus-md-only/agent-matcher.ts new file mode 100644 index 00000000..1932b5d4 --- /dev/null +++ b/src/hooks/prometheus-md-only/agent-matcher.ts @@ -0,0 +1,5 @@ +import { PROMETHEUS_AGENT } from "./constants" + +export function isPrometheusAgent(agentName: string | undefined): boolean { + return agentName?.toLowerCase().includes(PROMETHEUS_AGENT) ?? false +} diff --git a/src/hooks/prometheus-md-only/hook.ts b/src/hooks/prometheus-md-only/hook.ts index d693836c..b0b5a01a 100644 --- a/src/hooks/prometheus-md-only/hook.ts +++ b/src/hooks/prometheus-md-only/hook.ts @@ -1,9 +1,10 @@ import type { PluginInput } from "@opencode-ai/plugin" -import { HOOK_NAME, PROMETHEUS_AGENT, BLOCKED_TOOLS, PLANNING_CONSULT_WARNING, PROMETHEUS_WORKFLOW_REMINDER } from "./constants" +import { HOOK_NAME, BLOCKED_TOOLS, PLANNING_CONSULT_WARNING, PROMETHEUS_WORKFLOW_REMINDER } from "./constants" import { log } from "../../shared/logger" import { SYSTEM_DIRECTIVE_PREFIX } from "../../shared/system-directive" import { getAgentDisplayName } from "../../shared/agent-display-names" import { getAgentFromSession } from "./agent-resolution" +import { isPrometheusAgent } from "./agent-matcher" import { isAllowedFile } from "./path-policy" const TASK_TOOLS = ["task", "call_omo_agent"] @@ -16,7 +17,7 @@ export function createPrometheusMdOnlyHook(ctx: PluginInput) { ): Promise => { const agentName = getAgentFromSession(input.sessionID, ctx.directory) - if (agentName !== PROMETHEUS_AGENT) { + if (!isPrometheusAgent(agentName)) { return } diff --git a/src/hooks/prometheus-md-only/index.test.ts b/src/hooks/prometheus-md-only/index.test.ts index ca9cd32f..2c99b6e5 100644 --- a/src/hooks/prometheus-md-only/index.test.ts +++ b/src/hooks/prometheus-md-only/index.test.ts @@ -30,11 +30,11 @@ describe("prometheus-md-only", () => { } as never } - function setupMessageStorage(sessionID: string, agent: string): void { + function setupMessageStorage(sessionID: string, agent: string | undefined): void { testMessageDir = join(MESSAGE_STORAGE, sessionID) mkdirSync(testMessageDir, { recursive: true }) const messageContent = { - agent, + ...(agent ? { agent } : {}), model: { providerID: "test", modelID: "test-model" }, } writeFileSync( @@ -55,6 +55,122 @@ describe("prometheus-md-only", () => { rmSync(TEST_STORAGE_ROOT, { recursive: true, force: true }) }) + describe("agent name matching", () => { + test("should enforce md-only restriction for exact prometheus agent name", async () => { + //#given + setupMessageStorage(TEST_SESSION_ID, "prometheus") + const hook = createPrometheusMdOnlyHook(createMockPluginInput()) + const input = { + tool: "Write", + sessionID: TEST_SESSION_ID, + callID: "call-1", + } + const output = { + args: { filePath: "/path/to/file.ts" }, + } + + //#when //#then + await expect( + hook["tool.execute.before"](input, output) + ).rejects.toThrow("can only write/edit .md files") + }) + + test("should enforce md-only restriction for Prometheus display name Plan Builder", async () => { + //#given + setupMessageStorage(TEST_SESSION_ID, "Prometheus (Plan Builder)") + const hook = createPrometheusMdOnlyHook(createMockPluginInput()) + const input = { + tool: "Write", + sessionID: TEST_SESSION_ID, + callID: "call-1", + } + const output = { + args: { filePath: "/path/to/file.ts" }, + } + + //#when //#then + await expect( + hook["tool.execute.before"](input, output) + ).rejects.toThrow("can only write/edit .md files") + }) + + test("should enforce md-only restriction for Prometheus display name Planner", async () => { + //#given + setupMessageStorage(TEST_SESSION_ID, "Prometheus (Planner)") + const hook = createPrometheusMdOnlyHook(createMockPluginInput()) + const input = { + tool: "Write", + sessionID: TEST_SESSION_ID, + callID: "call-1", + } + const output = { + args: { filePath: "/path/to/file.ts" }, + } + + //#when //#then + await expect( + hook["tool.execute.before"](input, output) + ).rejects.toThrow("can only write/edit .md files") + }) + + test("should enforce md-only restriction for uppercase PROMETHEUS", async () => { + //#given + setupMessageStorage(TEST_SESSION_ID, "PROMETHEUS") + const hook = createPrometheusMdOnlyHook(createMockPluginInput()) + const input = { + tool: "Write", + sessionID: TEST_SESSION_ID, + callID: "call-1", + } + const output = { + args: { filePath: "/path/to/file.ts" }, + } + + //#when //#then + await expect( + hook["tool.execute.before"](input, output) + ).rejects.toThrow("can only write/edit .md files") + }) + + test("should not enforce restriction for non-Prometheus agent", async () => { + //#given + setupMessageStorage(TEST_SESSION_ID, "sisyphus") + const hook = createPrometheusMdOnlyHook(createMockPluginInput()) + const input = { + tool: "Write", + sessionID: TEST_SESSION_ID, + callID: "call-1", + } + const output = { + args: { filePath: "/path/to/file.ts" }, + } + + //#when //#then + await expect( + hook["tool.execute.before"](input, output) + ).resolves.toBeUndefined() + }) + + test("should not enforce restriction when agent name is undefined", async () => { + //#given + setupMessageStorage(TEST_SESSION_ID, undefined) + const hook = createPrometheusMdOnlyHook(createMockPluginInput()) + const input = { + tool: "Write", + sessionID: TEST_SESSION_ID, + callID: "call-1", + } + const output = { + args: { filePath: "/path/to/file.ts" }, + } + + //#when //#then + await expect( + hook["tool.execute.before"](input, output) + ).resolves.toBeUndefined() + }) + }) + describe("with Prometheus agent in message storage", () => { beforeEach(() => { setupMessageStorage(TEST_SESSION_ID, "prometheus")