@illuma-ai/agents 1.4.0-alpha.6 → 1.5.0

package/src/hooks/__tests__/HookRegistry.test.ts

@@ -0,0 +1,190 @@
+// src/hooks/__tests__/HookRegistry.test.ts
+import { HookRegistry } from '../HookRegistry';
+import type {
+  HookMatcher,
+  HookCallback,
+  PreToolUseHookOutput,
+  PostToolUseHookOutput,
+} from '../types';
+
+const noop: HookCallback<
+  'PreToolUse'
+> = async (): Promise<PreToolUseHookOutput> => ({});
+const noopPost: HookCallback<
+  'PostToolUse'
+> = async (): Promise<PostToolUseHookOutput> => ({});
+
+function makePreToolUseMatcher(pattern?: string): HookMatcher<'PreToolUse'> {
+  return {
+    pattern,
+    hooks: [noop],
+  };
+}
+
+function makePostToolUseMatcher(): HookMatcher<'PostToolUse'> {
+  return {
+    hooks: [noopPost],
+  };
+}
+
+describe('HookRegistry', () => {
+  describe('global registration', () => {
+    it('stores a matcher and returns it via getMatchers', () => {
+      const registry = new HookRegistry();
+      const matcher = makePreToolUseMatcher('Bash');
+      registry.register('PreToolUse', matcher);
+      const matchers = registry.getMatchers('PreToolUse');
+      expect(matchers).toHaveLength(1);
+      expect(matchers[0]).toBe(matcher);
+    });
+
+    it('keeps registrations isolated across events', () => {
+      const registry = new HookRegistry();
+      const pre = makePreToolUseMatcher('Bash');
+      const post = makePostToolUseMatcher();
+      registry.register('PreToolUse', pre);
+      registry.register('PostToolUse', post);
+      expect(registry.getMatchers('PreToolUse')).toEqual([pre]);
+      expect(registry.getMatchers('PostToolUse')).toEqual([post]);
+      expect(registry.getMatchers('Stop')).toEqual([]);
+    });
+
+    it('unregister removes the matcher from the registry', () => {
+      const registry = new HookRegistry();
+      const matcher = makePreToolUseMatcher();
+      const unregister = registry.register('PreToolUse', matcher);
+      expect(registry.getMatchers('PreToolUse')).toHaveLength(1);
+      unregister();
+      expect(registry.getMatchers('PreToolUse')).toHaveLength(0);
+    });
+
+    it('hasHookFor reflects registration state', () => {
+      const registry = new HookRegistry();
+      expect(registry.hasHookFor('PreToolUse')).toBe(false);
+      registry.register('PreToolUse', makePreToolUseMatcher());
+      expect(registry.hasHookFor('PreToolUse')).toBe(true);
+      expect(registry.hasHookFor('PostToolUse')).toBe(false);
+    });
+
+    it('supports multiple matchers on the same event', () => {
+      const registry = new HookRegistry();
+      const a = makePreToolUseMatcher('Bash');
+      const b = makePreToolUseMatcher('Edit');
+      registry.register('PreToolUse', a);
+      registry.register('PreToolUse', b);
+      const matchers = registry.getMatchers('PreToolUse');
+      expect(matchers).toHaveLength(2);
+      expect(matchers).toContain(a);
+      expect(matchers).toContain(b);
+    });
+
+    it('returns a fresh array on each getMatchers call', () => {
+      const registry = new HookRegistry();
+      registry.register('PreToolUse', makePreToolUseMatcher());
+      const first = registry.getMatchers('PreToolUse');
+      const second = registry.getMatchers('PreToolUse');
+      expect(first).not.toBe(second);
+      first.length = 0;
+      expect(registry.getMatchers('PreToolUse')).toHaveLength(1);
+    });
+  });
+
+  describe('session registration', () => {
+    it('scopes matchers to a single session', () => {
+      const registry = new HookRegistry();
+      const sessionA = makePreToolUseMatcher('Bash');
+      const sessionB = makePreToolUseMatcher('Edit');
+      registry.registerSession('run-a', 'PreToolUse', sessionA);
+      registry.registerSession('run-b', 'PreToolUse', sessionB);
+
+      expect(registry.getMatchers('PreToolUse', 'run-a')).toEqual([sessionA]);
+      expect(registry.getMatchers('PreToolUse', 'run-b')).toEqual([sessionB]);
+      expect(registry.getMatchers('PreToolUse')).toEqual([]);
+    });
+
+    it('merges global matchers in front of session matchers', () => {
+      const registry = new HookRegistry();
+      const global = makePreToolUseMatcher('global');
+      const session = makePreToolUseMatcher('session');
+      registry.register('PreToolUse', global);
+      registry.registerSession('run-a', 'PreToolUse', session);
+
+      const matchers = registry.getMatchers('PreToolUse', 'run-a');
+      expect(matchers).toEqual([global, session]);
+    });
+
+    it('clearSession drops only the given session', () => {
+      const registry = new HookRegistry();
+      const a = makePreToolUseMatcher('a');
+      const b = makePreToolUseMatcher('b');
+      registry.registerSession('run-a', 'PreToolUse', a);
+      registry.registerSession('run-b', 'PreToolUse', b);
+
+      registry.clearSession('run-a');
+      expect(registry.getMatchers('PreToolUse', 'run-a')).toEqual([]);
+      expect(registry.getMatchers('PreToolUse', 'run-b')).toEqual([b]);
+    });
+
+    it('removeMatcher removes from the correct scope', () => {
+      const registry = new HookRegistry();
+      const global = makePreToolUseMatcher('global');
+      const session = makePreToolUseMatcher('session');
+      registry.register('PreToolUse', global);
+      registry.registerSession('run-a', 'PreToolUse', session);
+
+      expect(registry.removeMatcher('PreToolUse', session, 'run-a')).toBe(true);
+      expect(registry.getMatchers('PreToolUse', 'run-a')).toEqual([global]);
+
+      expect(registry.removeMatcher('PreToolUse', global)).toBe(true);
+      expect(registry.getMatchers('PreToolUse', 'run-a')).toEqual([]);
+    });
+
+    it('removeMatcher returns false when the matcher is not found', () => {
+      const registry = new HookRegistry();
+      const orphan = makePreToolUseMatcher('orphan');
+      expect(registry.removeMatcher('PreToolUse', orphan)).toBe(false);
+      expect(registry.removeMatcher('PreToolUse', orphan, 'run-a')).toBe(false);
+    });
+
+    it('session unregister function removes only that matcher', () => {
+      const registry = new HookRegistry();
+      const a = makePreToolUseMatcher('a');
+      const b = makePreToolUseMatcher('b');
+      const unregisterA = registry.registerSession('run-a', 'PreToolUse', a);
+      registry.registerSession('run-a', 'PreToolUse', b);
+
+      unregisterA();
+      expect(registry.getMatchers('PreToolUse', 'run-a')).toEqual([b]);
+    });
+
+    it('hasHookFor honours the sessionId parameter', () => {
+      const registry = new HookRegistry();
+      registry.registerSession('run-a', 'PreToolUse', makePreToolUseMatcher());
+      expect(registry.hasHookFor('PreToolUse')).toBe(false);
+      expect(registry.hasHookFor('PreToolUse', 'run-a')).toBe(true);
+      expect(registry.hasHookFor('PreToolUse', 'run-b')).toBe(false);
+    });
+  });
+
+  describe('session isolation under parallel registration', () => {
+    it('does not leak matchers between sessions registered in parallel', async () => {
+      const registry = new HookRegistry();
+      const sessions = Array.from({ length: 50 }, (_, i) => `run-${i}`);
+      await Promise.all(
+        sessions.map(async (sid): Promise<void> => {
+          registry.registerSession(
+            sid,
+            'PreToolUse',
+            makePreToolUseMatcher(sid)
+          );
+        })
+      );
+
+      for (const sid of sessions) {
+        const matchers = registry.getMatchers('PreToolUse', sid);
+        expect(matchers).toHaveLength(1);
+        expect(matchers[0]?.pattern).toBe(sid);
+      }
+    });
+  });
+});

package/src/hooks/__tests__/compactHooks.test.ts

@@ -0,0 +1,214 @@
+// src/hooks/__tests__/compactHooks.test.ts
+import { HumanMessage, AIMessage } from '@langchain/core/messages';
+import { FakeListChatModel } from '@langchain/core/utils/testing';
+import { HookRegistry } from '../HookRegistry';
+import { Run } from '@/run';
+import * as providers from '@/llm/providers';
+import { ToolEndHandler, ModelEndHandler } from '@/events';
+import { createTokenCounter } from '@/utils/tokens';
+import { Providers, GraphEvents } from '@/common';
+import type * as t from '@/types';
+import type {
+  HookCallback,
+  PreCompactHookInput,
+  PreCompactHookOutput,
+  PostCompactHookInput,
+  PostCompactHookOutput,
+} from '../types';
+
+const SUMMARY_RESPONSE = '## Summary\nUser asked a question and got an answer.';
+
+const callerConfig = {
+  configurable: { thread_id: 'compact-test' },
+  streamMode: 'values' as const,
+  version: 'v2' as const,
+};
+
+let getChatModelClassSpy: jest.SpyInstance;
+const originalGetChatModelClass = providers.getChatModelClass;
+
+function mockSummarizationModel(): void {
+  getChatModelClassSpy = jest
+    .spyOn(providers, 'getChatModelClass')
+    .mockImplementation(((provider: Providers) => {
+      if (provider === Providers.OPENAI) {
+        return class extends FakeListChatModel {
+          constructor(
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            _options: any
+          ) {
+            super({ responses: [SUMMARY_RESPONSE] });
+          }
+          // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        } as any;
+      }
+      return originalGetChatModelClass(provider);
+    }) as typeof providers.getChatModelClass);
+}
+
+function buildConversation(): t.IState {
+  const messages: import('@langchain/core/messages').BaseMessage[] = [];
+  for (let i = 0; i < 20; i++) {
+    messages.push(new HumanMessage(`Question ${i}: ` + 'padding '.repeat(50)));
+    messages.push(new AIMessage(`Answer ${i}: ` + 'response '.repeat(50)));
+  }
+  return { messages };
+}
+
+async function createCompactingRun(
+  tokenCounter: t.TokenCounter,
+  hooks?: HookRegistry,
+  runId = 'compact-run'
+): Promise<Run<t.IState>> {
+  const conversation = buildConversation();
+  const indexTokenCountMap: Record<string, number> = {};
+  for (let i = 0; i < conversation.messages.length; i++) {
+    indexTokenCountMap[String(i)] = tokenCounter(conversation.messages[i]);
+  }
+  return Run.create<t.IState>({
+    runId,
+    graphConfig: {
+      type: 'standard',
+      llmConfig: {
+        provider: Providers.OPENAI,
+        streaming: true,
+        streamUsage: false,
+      },
+      instructions: 'Be concise.',
+      maxContextTokens: 200,
+      summarizationEnabled: true,
+      summarizationConfig: {
+        provider: Providers.OPENAI,
+      },
+    },
+    returnContent: true,
+    skipCleanup: true,
+    customHandlers: {
+      [GraphEvents.TOOL_END]: new ToolEndHandler(),
+      [GraphEvents.CHAT_MODEL_END]: new ModelEndHandler(),
+    },
+    hooks,
+    tokenCounter,
+    indexTokenCountMap,
+  });
+}
+
+describe('Compaction hook integration', () => {
+  jest.setTimeout(30_000);
+
+  let tokenCounter: t.TokenCounter;
+
+  beforeAll(async () => {
+    tokenCounter = await createTokenCounter();
+  });
+
+  beforeEach(() => {
+    mockSummarizationModel();
+  });
+
+  afterEach(() => {
+    getChatModelClassSpy.mockRestore();
+  });
+
+  describe('PreCompact', () => {
+    it('fires with messagesBeforeCount and trigger', async () => {
+      const registry = new HookRegistry();
+      let captured: PreCompactHookInput | undefined;
+      const hook: HookCallback<'PreCompact'> = async (
+        input
+      ): Promise<PreCompactHookOutput> => {
+        captured = input;
+        return {};
+      };
+      registry.register('PreCompact', { hooks: [hook] });
+
+      const run = await createCompactingRun(tokenCounter, registry);
+      run.Graph!.overrideTestModel(['Final answer after compaction.']);
+      const inputs = buildConversation();
+      await run.processStream(inputs, callerConfig);
+
+      expect(captured).toBeDefined();
+      expect(captured!.hook_event_name).toBe('PreCompact');
+      expect(captured!.messagesBeforeCount).toBeGreaterThan(0);
+      expect(captured!.runId).toBe('compact-run');
+      expect(captured!.threadId).toBe('compact-test');
+      expect(captured!.trigger).toBe('default');
+      expect(captured!.agentId).toBeDefined();
+    });
+  });
+
+  describe('PostCompact', () => {
+    it('fires with summary text after compaction', async () => {
+      const registry = new HookRegistry();
+      let captured: PostCompactHookInput | undefined;
+      const hook: HookCallback<'PostCompact'> = async (
+        input
+      ): Promise<PostCompactHookOutput> => {
+        captured = input;
+        return {};
+      };
+      registry.register('PostCompact', { hooks: [hook] });
+
+      const run = await createCompactingRun(tokenCounter, registry);
+      run.Graph!.overrideTestModel(['Final answer after compaction.']);
+      const inputs = buildConversation();
+      await run.processStream(inputs, callerConfig);
+
+      expect(captured).toBeDefined();
+      expect(captured!.hook_event_name).toBe('PostCompact');
+      expect(captured!.threadId).toBe('compact-test');
+      expect(captured!.summary).toBe(SUMMARY_RESPONSE);
+      expect(captured!.messagesAfterCount).toBe(0);
+      expect(captured!.agentId).toBeDefined();
+    });
+  });
+
+  describe('error resilience', () => {
+    it('throwing PreCompact hook does not crash compaction', async () => {
+      const registry = new HookRegistry();
+      const throwingHook: HookCallback<
+        'PreCompact'
+      > = async (): Promise<PreCompactHookOutput> => {
+        throw new Error('pre hook crash');
+      };
+      registry.register('PreCompact', { hooks: [throwingHook] });
+
+      const run = await createCompactingRun(tokenCounter, registry);
+      run.Graph!.overrideTestModel(['Answer after compaction.']);
+      const inputs = buildConversation();
+
+      await expect(
+        run.processStream(inputs, callerConfig)
+      ).resolves.not.toThrow();
+    });
+
+    it('throwing PostCompact hook does not crash compaction', async () => {
+      const registry = new HookRegistry();
+      const throwingHook: HookCallback<
+        'PostCompact'
+      > = async (): Promise<PostCompactHookOutput> => {
+        throw new Error('post hook crash');
+      };
+      registry.register('PostCompact', { hooks: [throwingHook] });
+
+      const run = await createCompactingRun(tokenCounter, registry);
+      run.Graph!.overrideTestModel(['Answer after compaction.']);
+      const inputs = buildConversation();
+
+      await expect(
+        run.processStream(inputs, callerConfig)
+      ).resolves.not.toThrow();
+    });
+  });
+
+  describe('no-hooks baseline', () => {
+    it('summarization works identically without hooks', async () => {
+      const run = await createCompactingRun(tokenCounter);
+      run.Graph!.overrideTestModel(['Answer.']);
+      const inputs = buildConversation();
+      const result = await run.processStream(inputs, callerConfig);
+
+      expect(result).toBeDefined();
+    });
+  });
+});