my-pi/pi-hermes-memory/tests/subagent-mode.test.ts

/**
 * Subagent (pi-subagents child) sessions are read-only: memory/skill write
 * tools return a fixed error and all automatic writes are skipped.
 *
 * Child detection keys on the `<active_agent name="X"/>` tag pi-subagents
 * stamps into `systemPromptOptions.customPrompt`; membership is tracked per
 * sessionManager so a child in the same process never flips the parent.
 */
import { describe, it, beforeEach, afterEach } from "node:test";
import assert from "node:assert/strict";
import fs from "node:fs";
import os from "node:os";
import path from "node:path";
import {
  isSubagentCustomPrompt,
  isSubagentSession,
  markSubagentSession,
  markSubagentSessionFromPrompt,
  SUBAGENT_WRITE_DISABLED_MESSAGE,
} from "../src/subagent-mode.js";
import { registerMemoryTool } from "../src/tools/memory-tool.js";
import { registerSkillTool } from "../src/tools/skill-tool.js";
import { setupSessionFlush } from "../src/handlers/session-flush.js";
import { MemoryStore } from "../src/store/memory-store.js";
import { SkillStore } from "../src/store/skill-store.js";
import type { MemoryConfig } from "../src/types.js";

const CHILD_PROMPT = 'You are the Explore agent.\n<active_agent name="Explore"/>';

function beforeAgentStartEvent(customPrompt?: string) {
  return { systemPromptOptions: customPrompt === undefined ? {} : { customPrompt } } as any;
}

describe("subagent detection", () => {
  it("matches the pi-subagents <active_agent> contract on customPrompt only", () => {
    assert.ok(isSubagentCustomPrompt(CHILD_PROMPT));
    assert.ok(isSubagentCustomPrompt("intro <active_agent name=\"General Purpose\" extra=\"1\"> tail"));
    assert.ok(!isSubagentCustomPrompt("a normal custom prompt"));
    assert.ok(!isSubagentCustomPrompt(""));
    assert.ok(!isSubagentCustomPrompt(undefined));
    assert.ok(!isSubagentCustomPrompt(null));
  });

  it("marks only sessions whose before_agent_start customPrompt carries the tag", () => {
    const childManager = {};
    const mainManager = {};
    markSubagentSession(beforeAgentStartEvent(CHILD_PROMPT), { sessionManager: childManager } as any);
    markSubagentSession(beforeAgentStartEvent(), { sessionManager: mainManager } as any);
    markSubagentSession(beforeAgentStartEvent("no tag here"), { sessionManager: {} } as any);

    assert.ok(isSubagentSession({ sessionManager: childManager }));
    assert.ok(!isSubagentSession({ sessionManager: mainManager }));
    assert.ok(!isSubagentSession(undefined));
    assert.ok(!isSubagentSession({}));
  });

  it("detects tag-leading rendered prompts early (session_start), not mid-prompt tags", () => {
    const childManager = {};
    const mainManager = {};
    markSubagentSessionFromPrompt(
      `<active_agent name="Explore"/>\n\nYou are the Explore agent.\n\nrest`,
      { sessionManager: childManager } as any,
    );
    markSubagentSessionFromPrompt(
      `preamble\n\nliteral <active_agent name="X"/> inside project docs`,
      { sessionManager: mainManager } as any,
    );
    assert.ok(isSubagentSession({ sessionManager: childManager }));
    assert.ok(!isSubagentSession({ sessionManager: mainManager }));
  });
});

describe("subagent write gating", () => {
  let tmpDir: string;

  beforeEach(() => {
    tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "subagent-mode-test-"));
  });

  afterEach(() => {
    fs.rmSync(tmpDir, { recursive: true, force: true });
  });

  function makeMemoryStore(): MemoryStore {
    return new MemoryStore({
      memoryMode: "policy-only",
      memoryCharLimit: 5000,
      userCharLimit: 5000,
      projectCharLimit: 5000,
      nudgeInterval: 10,
      reviewEnabled: false,
      flushOnCompact: false,
      flushOnShutdown: false,
      flushMinTurns: 6,
      autoConsolidate: false,
      correctionDetection: false,
      failureInjectionEnabled: false,
      failureInjectionMaxAgeDays: 7,
      failureInjectionMaxEntries: 5,
      nudgeToolCalls: 15,
      consolidationTimeoutMs: 60000,
      memoryDir: tmpDir,
    });
  }

  it("memory_add/replace/remove return the disabled error and perform no writes", async () => {
    const tools = new Map<string, any>();
    const mockPi = { registerTool: (def: any) => tools.set(def.name, def) } as any;
    const store = makeMemoryStore();
    await store.loadFromDisk();
    let bindCalls = 0;
    registerMemoryTool(mockPi, store, null, null, null, () => { bindCalls++; });

    const sessionManager = {};
    markSubagentSession(beforeAgentStartEvent(CHILD_PROMPT), { sessionManager } as any);
    const ctx = { cwd: tmpDir, sessionManager };

    for (const name of ["memory_add", "memory_replace", "memory_remove"]) {
      const result = await tools.get(name).execute("tc", { target: "memory", content: "x", old_text: "x" }, undefined, undefined, ctx);
      assert.strictEqual(result.content[0].text, SUBAGENT_WRITE_DISABLED_MESSAGE, name);
      assert.strictEqual(result.details.success, false, name);
      assert.strictEqual(result.details.error, SUBAGENT_WRITE_DISABLED_MESSAGE, name);
    }
    assert.strictEqual(bindCalls, 0, "blocked writes must not bind projects");
    assert.deepStrictEqual(store.getMemoryEntries(), [], "blocked writes must not mutate the store");
    assert.ok(!fs.existsSync(path.join(tmpDir, "MEMORY.md")), "blocked writes must not write files");
  });

  it("memory tools still write in an unmarked (main) session", async () => {
    const tools = new Map<string, any>();
    const mockPi = { registerTool: (def: any) => tools.set(def.name, def) } as any;
    const store = makeMemoryStore();
    await store.loadFromDisk();
    registerMemoryTool(mockPi, store, null);

    const result = await tools.get("memory_add").execute(
      "tc",
      { target: "memory", content: "main session entry" },
      undefined,
      undefined,
      { cwd: tmpDir, sessionManager: {} },
    );
    assert.strictEqual(result.details.success, true);
    assert.deepStrictEqual(store.getMemoryEntries(), ["main session entry"]);
  });

  it("skill_manage returns the disabled error in a subagent session", async () => {
    let tool: any;
    const mockPi = { registerTool: (def: any) => { tool = def; } } as any;
    const store = new SkillStore({
      globalSkillsDir: path.join(tmpDir, "skills"),
      piGlobalSkillsDir: path.join(tmpDir, "pi-skills"),
      projectSkillsDir: null,
      projectName: null,
      legacySkillsDir: path.join(tmpDir, "legacy-skills"),
      migrationSentinelPath: path.join(tmpDir, ".migrated"),
    });
    registerSkillTool(mockPi, store);

    const sessionManager = {};
    markSubagentSession(beforeAgentStartEvent(CHILD_PROMPT), { sessionManager } as any);
    const result = await tool.execute(
      "tc",
      { action: "create", name: "demo", description: "d", scope: "global", content: "body" },
      undefined,
      undefined,
      { sessionManager },
    );
    assert.strictEqual(result.content[0].text, SUBAGENT_WRITE_DISABLED_MESSAGE);
    assert.strictEqual(result.details.success, false);
    assert.ok(!fs.existsSync(path.join(tmpDir, "skills")), "blocked skill writes must not create the skills dir");
  });

  it("session flush skips compact and shutdown flushes in a subagent session", async () => {
    const handlers: Record<string, Function[]> = {};
    const mockPi = {
      on(event: string, handler: Function) {
        (handlers[event] ??= []).push(handler);
      },
      async exec() {
        throw new Error("subagent flush must not spawn a child pi");
      },
      registerTool() {},
      registerCommand() {},
    } as any;

    let directCalls = 0;
    const config = {
      flushOnCompact: true,
      flushOnShutdown: true,
      flushMinTurns: 1,
      flushCompactTimeoutMs: 60000,
      reviewTransport: "direct",
    } as MemoryConfig;

    const sessionManager = { getBranch: () => [] };
    markSubagentSession(beforeAgentStartEvent(CHILD_PROMPT), { sessionManager } as any);

    setupSessionFlush(mockPi, { getMemoryEntries: () => [], getUserEntries: () => [] } as any, null, config, null, null, {
      runDirectMemoryCompletion: async () => {
        directCalls++;
        return { ok: true } as any;
      },
    });

    const ctx = { sessionManager, cwd: tmpDir };
    for (const handler of handlers["message_end"] ?? []) {
      await handler({ message: { role: "user", content: "hi" } }, ctx);
    }
    for (const handler of handlers["session_before_compact"] ?? []) {
      await handler({ signal: undefined }, ctx);
    }
    for (const handler of handlers["session_shutdown"] ?? []) {
      await handler({ reason: "exit" }, ctx);
    }
    assert.strictEqual(directCalls, 0, "subagent sessions must not flush");

    // Control: the same handlers must still flush an unmarked (main) session.
    const mainCtx = { sessionManager: { getBranch: () => [] }, cwd: tmpDir };
    for (const handler of handlers["message_end"] ?? []) {
      await handler({ message: { role: "user", content: "hi" } }, mainCtx);
    }
    for (const handler of handlers["session_before_compact"] ?? []) {
      await handler({ signal: undefined }, mainCtx);
    }
    assert.ok(directCalls > 0, "main sessions must still flush");
  });
});