|
| 1 | +import type { AgentMessage } from '@earendil-works/pi-agent-core'; |
| 2 | +import type { Api, Model } from '@earendil-works/pi-ai'; |
| 3 | +import type { PiBeforeLlmCallHookInput } from '@mavis/agent-core/pi-turn-runner'; |
| 4 | +import { describe, expect, it } from 'vitest'; |
| 5 | + |
| 6 | +import { validateCanonicalHistoryMessages } from '../../agent-host/history/canonical-history-validation.js'; |
| 7 | +import { ContextUsageAnchorState } from '../../compaction/execution/usage-anchor.js'; |
| 8 | +import { messagesCarrySessionId } from '../../agent-host/compaction/session-identity.js'; |
| 9 | +import { createSessionIdentityReminderHook } from './session-identity-reminder.js'; |
| 10 | + |
| 11 | +const SESSION_ID = 'ses_identity_1'; |
| 12 | +const MODEL = { |
| 13 | + id: 'test-model', |
| 14 | + name: 'test-model', |
| 15 | + api: 'anthropic-messages', |
| 16 | + provider: 'test-provider', |
| 17 | + baseUrl: 'https://example.invalid', |
| 18 | + reasoning: false, |
| 19 | + input: ['text'], |
| 20 | + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, |
| 21 | + contextWindow: 200_000, |
| 22 | + maxTokens: 8_192, |
| 23 | +} as Model<Api>; |
| 24 | + |
| 25 | +const hook = createSessionIdentityReminderHook(new ContextUsageAnchorState()); |
| 26 | + |
| 27 | +function message(value: Record<string, unknown>): AgentMessage { |
| 28 | + return value as unknown as AgentMessage; |
| 29 | +} |
| 30 | + |
| 31 | +function user(text: string, timestamp: number): AgentMessage { |
| 32 | + return message({ role: 'user', content: [{ type: 'text', text }], timestamp }); |
| 33 | +} |
| 34 | + |
| 35 | +const idBearingUser = user( |
| 36 | + `<system-reminder>\n YOUR SESSION ID: ${SESSION_ID}\n</system-reminder>\nfirst ask`, |
| 37 | + 1, |
| 38 | +); |
| 39 | +const checkpoint = message({ |
| 40 | + role: 'compactionSummary', |
| 41 | + summary: '## Goal\nKeep going', |
| 42 | + tokensBefore: 100, |
| 43 | + timestamp: 2, |
| 44 | +}); |
| 45 | +const currentUserWithoutId = user('<system-reminder>\n date: today\n</system-reminder>\nnext', 3); |
| 46 | + |
| 47 | +function hookInput( |
| 48 | + messages: AgentMessage[], |
| 49 | + phase: 'initial' | 'iteration' = 'iteration', |
| 50 | +): PiBeforeLlmCallHookInput { |
| 51 | + return { |
| 52 | + sessionId: SESSION_ID, |
| 53 | + turnId: 'turn-1', |
| 54 | + phase, |
| 55 | + messages, |
| 56 | + canonicalMessages: messages, |
| 57 | + model: MODEL, |
| 58 | + thinkingLevel: 'off', |
| 59 | + } as PiBeforeLlmCallHookInput; |
| 60 | +} |
| 61 | + |
| 62 | +describe('Session identity reminder hook', () => { |
| 63 | + it('stays silent while the model-visible context still shows the session ID', async () => { |
| 64 | + expect(await hook(hookInput([idBearingUser, currentUserWithoutId]))).toBeUndefined(); |
| 65 | + }); |
| 66 | + |
| 67 | + it('restores the ID before the current user when initial compaction summarized it away', async () => { |
| 68 | + const decision = await hook(hookInput([checkpoint, currentUserWithoutId], 'initial')); |
| 69 | + |
| 70 | + expect(decision).toMatchObject({ |
| 71 | + type: 'appendMessage', |
| 72 | + reason: 'session_identity_reminder', |
| 73 | + placement: 'before-current-user', |
| 74 | + message: { |
| 75 | + role: 'custom', |
| 76 | + customType: 'session_identity_reminder', |
| 77 | + display: false, |
| 78 | + content: expect.stringContaining(`YOUR SESSION ID: ${SESSION_ID}`), |
| 79 | + }, |
| 80 | + }); |
| 81 | + if (decision?.type !== 'appendMessage') throw new Error('Expected an append decision'); |
| 82 | + expect(() => validateCanonicalHistoryMessages([decision.message])).not.toThrow(); |
| 83 | + }); |
| 84 | + |
| 85 | + it('appends at the tail mid-turn and does not repeat once the marker is in context', async () => { |
| 86 | + const decision = await hook(hookInput([checkpoint, currentUserWithoutId])); |
| 87 | + if (decision?.type !== 'appendMessage') throw new Error('Expected an append decision'); |
| 88 | + expect(decision).not.toHaveProperty('placement'); |
| 89 | + |
| 90 | + expect( |
| 91 | + await hook(hookInput([checkpoint, currentUserWithoutId, decision.message])), |
| 92 | + ).toBeUndefined(); |
| 93 | + }); |
| 94 | +}); |
| 95 | + |
| 96 | +describe('messagesCarrySessionId', () => { |
| 97 | + it('only counts user messages or identity markers after the latest compaction', () => { |
| 98 | + expect(messagesCarrySessionId([idBearingUser, currentUserWithoutId], SESSION_ID)).toBe(true); |
| 99 | + expect( |
| 100 | + messagesCarrySessionId([idBearingUser, checkpoint, currentUserWithoutId], SESSION_ID), |
| 101 | + ).toBe(false); |
| 102 | + expect( |
| 103 | + messagesCarrySessionId( |
| 104 | + [ |
| 105 | + message({ |
| 106 | + role: 'assistant', |
| 107 | + content: [{ type: 'text', text: SESSION_ID }], |
| 108 | + timestamp: 4, |
| 109 | + }), |
| 110 | + ], |
| 111 | + SESSION_ID, |
| 112 | + ), |
| 113 | + ).toBe(false); |
| 114 | + expect(messagesCarrySessionId([idBearingUser], '')).toBe(false); |
| 115 | + }); |
| 116 | +}); |
0 commit comments