test-session-bootstrap.mjs 9.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224
  1. import assert from 'node:assert/strict';
  2. import fs from 'node:fs';
  3. import { pathToFileURL } from 'node:url';
  4. const [, , inputPath] = process.argv;
  5. assert.ok(inputPath, 'pass the plugin module path');
  6. const pluginURL = pathToFileURL(fs.realpathSync(inputPath));
  7. const marker = '<EXTREMELY_IMPORTANT>\nYou have superpowers.';
  8. let generation = 0;
  9. function reply(flavor, session) {
  10. return flavor === 'v1' ? { data: session } : session;
  11. }
  12. function makeEvent(flavor, sessionID) {
  13. const text = { type: 'text', text: 'Execute the assigned task' };
  14. return {
  15. sessionID,
  16. messages: [flavor === 'v1'
  17. ? { info: { role: 'user', sessionID }, parts: [text] }
  18. : { role: 'user', content: [text] }],
  19. };
  20. }
  21. function bootstrapCount(event) {
  22. return event.messages.flatMap((message) => message.parts ?? message.content ?? []).filter(
  23. (part) => part.type === 'text' && part.text.startsWith(marker)
  24. ).length;
  25. }
  26. async function makeHarness(flavor, fetchSession) {
  27. const mod = await import(`${pluginURL.href}?session-test=${++generation}`);
  28. const lookups = [];
  29. const registered = [];
  30. const get = async (id) => {
  31. lookups.push(id);
  32. return fetchSession(id, lookups.length);
  33. };
  34. let invoke;
  35. if (flavor === 'v1') {
  36. const hooks = await mod.SuperpowersPlugin({
  37. client: { session: { get: ({ path: { id } }) => get(id) } },
  38. directory: '.',
  39. });
  40. invoke = (event) => hooks['experimental.chat.messages.transform']({}, event);
  41. } else {
  42. await mod.default.setup({
  43. skill: { transform: async (transform) => transform({ add: (skill) => registered.push(skill) }) },
  44. session: {
  45. get: ({ sessionID }) => get(sessionID),
  46. hook: async (name, callback) => { if (name === 'context') invoke = callback; },
  47. },
  48. });
  49. }
  50. assert.equal(typeof invoke, 'function');
  51. return { invoke, lookups, registered };
  52. }
  53. for (const flavor of ['v1', 'v2']) {
  54. for (const [kind, extra, expected] of [
  55. ['root', {}, 1],
  56. ['child', { parentID: 'parent' }, 0],
  57. ['fork', { fork: { sessionID: 'origin' } }, 1],
  58. ]) {
  59. const id = `${flavor}-${kind}`;
  60. const h = await makeHarness(flavor, () => reply(flavor, { id, ...extra }));
  61. const event = makeEvent(flavor, id);
  62. await h.invoke(event);
  63. assert.equal(bootstrapCount(event), expected, `${id}: first request`);
  64. await h.invoke(event);
  65. assert.equal(bootstrapCount(event), expected, `${id}: repeated event`);
  66. const fresh = makeEvent(flavor, id);
  67. await h.invoke(fresh);
  68. assert.equal(bootstrapCount(fresh), expected, `${id}: fresh request`);
  69. assert.deepEqual(h.lookups, [id], `${id}: cache successful classification`);
  70. if (flavor === 'v2' && kind === 'child') {
  71. assert.ok(h.registered.some((skill) => skill.id === 'brainstorming'));
  72. }
  73. }
  74. const failures = [
  75. ['throws', () => { throw new Error('temporary lookup failure'); }],
  76. ['missing', () => undefined],
  77. ['null', () => null],
  78. ['empty', () => reply(flavor, {})],
  79. ['wrong-id', () => reply(flavor, { id: 'different-session' })],
  80. ['invalid-parent', (id) => reply(flavor, { id, parentID: 42 })],
  81. ];
  82. if (flavor === 'v1') {
  83. failures.push(['resolved-http-error', () => ({
  84. data: undefined,
  85. error: { name: 'UnknownError', data: { message: 'temporary 503' } },
  86. response: { ok: false, status: 503 },
  87. })]);
  88. }
  89. for (const [kind, firstResult] of failures) {
  90. const id = `${flavor}-${kind}`;
  91. const h = await makeHarness(flavor, (sessionID, call) => call === 1
  92. ? firstResult(sessionID)
  93. : reply(flavor, { id: sessionID, parentID: 'parent' }));
  94. const counts = [];
  95. for (let step = 0; step < 2; step++) {
  96. const event = makeEvent(flavor, id);
  97. await h.invoke(event);
  98. counts.push(bootstrapCount(event));
  99. }
  100. assert.deepEqual(counts, [1, 0], `${id}: recover on the next request`);
  101. assert.deepEqual(h.lookups, [id, id], `${id}: never cache the failure`);
  102. }
  103. const isolated = await makeHarness(flavor, (id) => reply(flavor,
  104. id === 'child-session' ? { id, parentID: 'parent' } : { id }));
  105. for (const [id, expected] of [['root-session', 1], ['child-session', 0], ['root-session', 1], ['child-session', 0]]) {
  106. const event = makeEvent(flavor, id);
  107. await isolated.invoke(event);
  108. assert.equal(bootstrapCount(event), expected);
  109. }
  110. assert.deepEqual(isolated.lookups, ['root-session', 'child-session']);
  111. const bounded = await makeHarness(flavor, (id) => reply(flavor, { id, parentID: 'parent' }));
  112. for (let index = 0; index <= 512; index++) {
  113. const event = makeEvent(flavor, `eviction-${index}`);
  114. await bounded.invoke(event);
  115. assert.equal(bootstrapCount(event), 0);
  116. }
  117. const evicted = makeEvent(flavor, 'eviction-0');
  118. await bounded.invoke(evicted);
  119. assert.equal(bootstrapCount(evicted), 0);
  120. assert.equal(bounded.lookups.filter((id) => id === 'eviction-0').length, 2);
  121. const restarted = await makeHarness(flavor, (id) => reply(flavor, { id, parentID: 'parent' }));
  122. const afterRestart = makeEvent(flavor, 'eviction-0');
  123. await restarted.invoke(afterRestart);
  124. assert.equal(bootstrapCount(afterRestart), 0);
  125. assert.deepEqual(restarted.lookups, ['eviction-0']);
  126. const unknown = await makeHarness(flavor, () => { throw new Error('must not look up a missing ID'); });
  127. const noID = makeEvent(flavor, undefined);
  128. await unknown.invoke(noID);
  129. assert.equal(bootstrapCount(noID), 1);
  130. assert.deepEqual(unknown.lookups, []);
  131. }
  132. function compactedEvent(sessionID) {
  133. return {
  134. sessionID,
  135. system: [],
  136. messages: [{
  137. role: 'assistant',
  138. content: [{ type: 'compaction', provider: 'fixture', encrypted: 'opaque-checkpoint' }],
  139. }],
  140. };
  141. }
  142. const compactedRoot = await makeHarness('v2', (id) => ({ id }));
  143. const rootEvent = compactedEvent('compacted-root');
  144. const checkpoint = structuredClone(rootEvent.messages[0]);
  145. await compactedRoot.invoke(rootEvent);
  146. assert.equal(bootstrapCount(rootEvent), 1);
  147. assert.deepEqual(rootEvent.messages[0], checkpoint);
  148. assert.equal(rootEvent.messages.length, 2);
  149. assert.equal(rootEvent.messages[1].role, 'user');
  150. assert.deepEqual(rootEvent.system, []);
  151. await compactedRoot.invoke(rootEvent);
  152. assert.equal(bootstrapCount(rootEvent), 1);
  153. assert.equal(rootEvent.messages.length, 2);
  154. const freshRootEvent = compactedEvent('compacted-root');
  155. await compactedRoot.invoke(freshRootEvent);
  156. assert.equal(bootstrapCount(freshRootEvent), 1);
  157. assert.deepEqual(compactedRoot.lookups, ['compacted-root']);
  158. const compactedChild = await makeHarness('v2', (id) => ({ id, parentID: 'parent' }));
  159. const childEvent = compactedEvent('compacted-child');
  160. const originalChild = structuredClone(childEvent);
  161. await compactedChild.invoke(childEvent);
  162. assert.equal(bootstrapCount(childEvent), 0);
  163. assert.deepEqual(childEvent, originalChild);
  164. assert.deepEqual(compactedChild.lookups, ['compacted-child']);
  165. const retryChild = await makeHarness('v2', (id, call) => {
  166. if (call === 1) throw new Error('temporary lookup failure');
  167. return { id, parentID: 'parent' };
  168. });
  169. const unknownChild = compactedEvent('retry-compacted-child');
  170. await retryChild.invoke(unknownChild);
  171. assert.equal(bootstrapCount(unknownChild), 1);
  172. const recoveredChild = compactedEvent('retry-compacted-child');
  173. await retryChild.invoke(recoveredChild);
  174. assert.equal(bootstrapCount(recoveredChild), 0);
  175. assert.equal(recoveredChild.messages.length, 1);
  176. assert.equal(retryChild.lookups.length, 2);
  177. const newPromptAfterCheckpoint = compactedEvent('new-prompt-after-checkpoint-root');
  178. newPromptAfterCheckpoint.messages.push({ role: 'user', content: [{ type: 'text', text: 'Continue' }] });
  179. await compactedRoot.invoke(newPromptAfterCheckpoint);
  180. assert.equal(bootstrapCount(newPromptAfterCheckpoint), 1);
  181. assert.equal(newPromptAfterCheckpoint.messages.length, 2);
  182. assert.equal(newPromptAfterCheckpoint.messages[1].content.length, 2);
  183. const retainedUser = compactedEvent('retained-user-root');
  184. retainedUser.messages.unshift({ role: 'user', content: [{ type: 'text', text: 'Keep going' }] });
  185. const retainedCheckpoint = structuredClone(retainedUser.messages[1]);
  186. await compactedRoot.invoke(retainedUser);
  187. assert.equal(bootstrapCount(retainedUser), 1);
  188. assert.equal(retainedUser.messages.length, 2);
  189. assert.equal(retainedUser.messages[0].content.length, 2);
  190. assert.ok(retainedUser.messages[0].content[0].text.startsWith(marker));
  191. assert.equal(retainedUser.messages[0].content[1].text, 'Keep going');
  192. assert.deepEqual(retainedUser.messages[1], retainedCheckpoint);
  193. await compactedRoot.invoke(retainedUser);
  194. assert.equal(bootstrapCount(retainedUser), 1);
  195. assert.equal(retainedUser.messages.length, 2);
  196. const retainedUserChild = compactedEvent('retained-user-child');
  197. retainedUserChild.messages.unshift({ role: 'user', content: [{ type: 'text', text: 'Keep going' }] });
  198. const originalRetainedUserChild = structuredClone(retainedUserChild);
  199. await compactedChild.invoke(retainedUserChild);
  200. assert.equal(bootstrapCount(retainedUserChild), 0);
  201. assert.deepEqual(retainedUserChild, originalRetainedUserChild);
  202. const empty = { sessionID: 'empty', messages: [] };
  203. await compactedRoot.invoke(empty);
  204. assert.deepEqual(empty.messages, []);
  205. console.log('Session classification, recovery and cache lifetime passed');