context-breakdown-projection.spec.ts 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525
  1. // contextBreakdown projection: heuristic system/tools/message composition,
  2. // plus the shared estimator's pricing branches.
  3. import { afterEach, describe, expect, it } from 'vitest'
  4. import { Context } from '@deepseek-ai/cordis'
  5. import { createMessage, createSystemMessage, createUserMessage } from '@deepseek-ai/dsh-llm'
  6. import type { ContentBlock, ToolSchema } from '@deepseek-ai/dsh-llm'
  7. import SessionStore, { SessionLogOffset, SessionSeq } from '@deepseek-ai/dsh-session'
  8. import type { Session, SessionEvent, SessionSeq as SessionSeqType } from '@deepseek-ai/dsh-session'
  9. import SessionProjectionRegistry from '@deepseek-ai/dsh-session-projection'
  10. import TokenMeter from '@deepseek-ai/dsh-token-meter'
  11. import type { ContextBreakdownProjection } from '@deepseek-ai/dsh-token-meter/client'
  12. import { CompactionId } from '@deepseek-ai/dsh-compaction'
  13. import { contextBreakdownProjectionDefinition } from '../src/breakdown-projection.ts'
  14. import {
  15. estimateContent,
  16. estimateMessage,
  17. estimateSystemMessage,
  18. estimateToolsTokens,
  19. } from '../src/estimate.ts'
  20. const contexts: Context[] = []
  21. afterEach(async () => {
  22. await Promise.all(contexts.splice(0).map(ctx => ctx.fiber.dispose()))
  23. })
  24. const CONFIG = { provider: 'test', model: 'test-model' }
  25. const TOOLS: ToolSchema[] = [{
  26. name: 'bash',
  27. description: 'run a command',
  28. parameters: { type: 'object', properties: {} },
  29. }]
  30. async function harness(): Promise<{ ctx: Context; session: Session }> {
  31. const ctx = new Context()
  32. contexts.push(ctx)
  33. await ctx.plugin(SessionStore)
  34. await ctx.plugin(SessionProjectionRegistry)
  35. await ctx.plugin(TokenMeter)
  36. return { ctx, session: ctx.sessions.create() }
  37. }
  38. const projected = (ctx: Context, session: Session): ContextBreakdownProjection => {
  39. const value = ctx.sessionProjections.snapshot(session).values.contextBreakdown
  40. if (value === undefined) throw new Error('contextBreakdown projection is not registered')
  41. return value
  42. }
  43. function appendUser(session: Session, text: string): SessionSeqType {
  44. return session.append('user/message', createUserMessage({
  45. content: [{ type: 'text', text }],
  46. source: { kind: 'user' },
  47. }), { surfaceOp: 'append' }).seq
  48. }
  49. const SYSTEM_PLUGIN = '@deepseek-ai/dsh-system-prompt'
  50. /** Append the rendered system prompt as surface node 0, the way the loop does before the first user message. */
  51. function appendSystem(session: Session, text: string): SessionSeqType {
  52. return session.append('system/message', {
  53. turn: 1,
  54. step: 1,
  55. message: createSystemMessage(text, SYSTEM_PLUGIN),
  56. }, { surfaceOp: 'append' }).seq
  57. }
  58. /** Replace the system node in place, the way the loop does when the rendered prompt changes. */
  59. function replaceSystem(session: Session, node: SessionSeqType, text: string): SessionSeqType {
  60. return session.append('system/message', {
  61. turn: 1,
  62. step: 1,
  63. message: createSystemMessage(text, SYSTEM_PLUGIN),
  64. }, { surfaceOp: { op: 'replace', startSeq: node, endSeq: node }, sourceEventSeqs: [node] }).seq
  65. }
  66. /**
  67. * Meter one upcoming replacement the way compaction-basic does: price the
  68. * replaced span from the measurement service's own nodes and log the
  69. * shadow-price event directly before the replace.
  70. */
  71. function appendSummaryMeter(ctx: Context, session: Session, start: SessionSeqType, end: SessionSeqType): void {
  72. const nodes = ctx.tokenMeter.measure(session).nodes
  73. const startIdx = nodes.findIndex(node => node.seq === start)
  74. const endIdx = nodes.findIndex(node => node.seq === end)
  75. const shadowed = nodes.slice(startIdx, endIdx + 1)
  76. session.append('compaction/summary', {
  77. compactionId: CompactionId('context-breakdown-summary'),
  78. summary: [{ type: 'text', text: 'summary' }],
  79. shadowedRange: { start, end },
  80. shadowedSeqs: shadowed.map(node => node.seq),
  81. shadowedTokenCount: shadowed.reduce((total, node) => total + node.tokens, 0),
  82. provider: 'mock',
  83. model: 'mock',
  84. })
  85. }
  86. describe('contextBreakdown session projection', () => {
  87. it('serves zeros for an empty log', async () => {
  88. const { ctx, session } = await harness()
  89. expect(projected(ctx, session)).toEqual({ systemTokens: 0, toolsTokens: 0, messageTokens: 0 })
  90. })
  91. it('prices the system node and the newest envelope last-wins and pushes no change for a restated envelope', async () => {
  92. const { ctx, session } = await harness()
  93. const systemNode = appendSystem(session, 'You are terse.')
  94. session.append('request/header', {
  95. header: { config: CONFIG, tools: TOOLS },
  96. reason: 'initial',
  97. })
  98. // 'You are terse.' prices to 8 (4 text + 4 role): the same figure the
  99. // request envelope's former system field priced to.
  100. expect(projected(ctx, session)).toEqual({
  101. systemTokens: 8,
  102. toolsTokens: estimateToolsTokens({ config: CONFIG, tools: TOOLS }),
  103. messageTokens: 0,
  104. })
  105. const changed: string[] = []
  106. ctx.sessionProjections.onChanged((_session, key) => { changed.push(key) })
  107. session.append('request/header', {
  108. header: { config: CONFIG, tools: TOOLS },
  109. reason: 'change',
  110. })
  111. session.append('session/end-seed', {})
  112. expect(changed).not.toContain('contextBreakdown')
  113. // A tool-less envelope prices the tools figure back to zero and leaves
  114. // the system node's figure alone.
  115. session.append('request/header', { header: { config: CONFIG }, reason: 'change' })
  116. expect(projected(ctx, session)).toEqual({ systemTokens: 8, toolsTokens: 0, messageTokens: 0 })
  117. // Replacing node 0 follows the new prompt; an empty prompt records none.
  118. const longer = replaceSystem(session, systemNode, 'You are terse and answer in one line.')
  119. expect(projected(ctx, session).systemTokens).toBe(Math.ceil('You are terse and answer in one line.'.length / 4) + 4)
  120. replaceSystem(session, longer, '')
  121. expect(projected(ctx, session)).toEqual({ systemTokens: 0, toolsTokens: 0, messageTokens: 0 })
  122. })
  123. it('keeps the system node out of the message figure across appends and a system replacement', async () => {
  124. const { ctx, session } = await harness()
  125. const systemNode = appendSystem(session, 'You are terse.')
  126. appendUser(session, 'abcd')
  127. expect(projected(ctx, session)).toMatchObject({ systemTokens: 8, messageTokens: 9 })
  128. replaceSystem(session, systemNode, 'You are verbose and thorough.')
  129. expect(projected(ctx, session)).toMatchObject({
  130. systemTokens: Math.ceil('You are verbose and thorough.'.length / 4) + 4,
  131. messageTokens: 9,
  132. })
  133. // The service prices the same system node identically, so the two
  134. // figures partition its surface total.
  135. const { systemTokens, messageTokens } = projected(ctx, session)
  136. expect(systemTokens + messageTokens).toBe(ctx.tokenMeter.measure(session).surfaceTokens)
  137. })
  138. it('moves a superseded in-history prompt into the message figure and subtracts it with a compaction', async () => {
  139. const { ctx, session } = await harness()
  140. const agree = (): ContextBreakdownProjection => {
  141. const projection = projected(ctx, session)
  142. expect(projection.systemTokens + projection.messageTokens).toBe(ctx.tokenMeter.measure(session).surfaceTokens)
  143. return projection
  144. }
  145. appendSystem(session, 'You are terse.')
  146. const question = appendUser(session, 'abcd')
  147. expect(agree()).toMatchObject({ systemTokens: 8, messageTokens: 9 })
  148. // An in-history route appends the changed prompt; node 0 stays model-visible history.
  149. const verbose = 'You are verbose and thorough.'
  150. const superseded = appendSystem(session, verbose)
  151. expect(agree()).toMatchObject({ systemTokens: Math.ceil(verbose.length / 4) + 4, messageTokens: 9 + 8 })
  152. const followUp = appendUser(session, 'efgh')
  153. appendSystem(session, 'You are terse once more.')
  154. expect(agree()).toMatchObject({
  155. systemTokens: Math.ceil('You are terse once more.'.length / 4) + 4,
  156. messageTokens: 9 + 8 + 9 + Math.ceil(verbose.length / 4) + 4,
  157. })
  158. // Compacting the span that holds the superseded mid-history prompt shrinks the message figure by it.
  159. appendSummaryMeter(ctx, session, question, followUp)
  160. const summary = createUserMessage({
  161. content: [{ type: 'text', text: 'summary' }],
  162. source: { kind: 'plugin', plugin: 'test' },
  163. })
  164. session.append('user/message', summary, {
  165. surfaceOp: { op: 'replace', startSeq: question, endSeq: followUp },
  166. sourceEventSeqs: [question, superseded, followUp],
  167. })
  168. expect(agree().messageTokens).toBe(8 + estimateMessage(summary))
  169. })
  170. it('restores the surviving head when compaction removes the newest large prompt', async () => {
  171. const { ctx, session } = await harness()
  172. appendSystem(session, 'head')
  173. const question = appendUser(session, 'abcd')
  174. const newest = appendSystem(session, 'x'.repeat(4000))
  175. appendSummaryMeter(ctx, session, question, newest)
  176. const summary = createUserMessage({ content: [{ type: 'text', text: 'summary' }], source: { kind: 'user' } })
  177. session.append('user/message', summary, {
  178. surfaceOp: { op: 'replace', startSeq: question, endSeq: newest },
  179. sourceEventSeqs: [question, newest],
  180. })
  181. expect(projected(ctx, session)).toEqual({ systemTokens: 5, toolsTokens: 0, messageTokens: estimateMessage(summary) })
  182. expect(projected(ctx, session).systemTokens + projected(ctx, session).messageTokens)
  183. .toBe(ctx.tokenMeter.measure(session).nodes.reduce((sum, node) => sum + node.heuristicTokens, 0))
  184. })
  185. it('sums surface appends and skips an empty-content assistant message', async () => {
  186. const { ctx, session } = await harness()
  187. appendUser(session, 'abcd')
  188. session.append('step/start', { turn: 1, step: 1 })
  189. session.append('assistant/message', {
  190. stream: [],
  191. turn: 1,
  192. step: 1,
  193. message: createMessage({
  194. role: 'assistant',
  195. content: [],
  196. source: { kind: 'model', provider: 'mock', model: 'mock' },
  197. }),
  198. usage: { inputTokens: 9, outputTokens: 0 },
  199. }, { surfaceOp: 'append' })
  200. session.append('step/end', { turn: 1, step: 1 })
  201. // 'abcd' prices to 9 (1 text + 4 block + 4 role); the usage-only assistant
  202. // message derives to no transcript entry and adds nothing.
  203. expect(projected(ctx, session).messageTokens).toBe(9)
  204. })
  205. it('shrinks the message figure when a metered replacement compacts the surface', async () => {
  206. const { ctx, session } = await harness()
  207. const first = appendUser(session, 'before compaction, a longer message')
  208. const second = appendUser(session, 'and a second entry')
  209. const summary = createUserMessage({
  210. content: [{ type: 'text', text: 'summary' }],
  211. source: { kind: 'plugin', plugin: 'test' },
  212. })
  213. appendSummaryMeter(ctx, session, first, second)
  214. session.append('user/message', summary, {
  215. surfaceOp: { op: 'replace', startSeq: first, endSeq: second },
  216. sourceEventSeqs: [first, second],
  217. })
  218. expect(projected(ctx, session).messageTokens).toBe(estimateMessage(summary))
  219. })
  220. it('keeps the message figure equal to the service result across appends and a compaction', async () => {
  221. const { ctx, session } = await harness()
  222. // The panel's composition rows and `measure()` answer the same question in
  223. // the same vocabulary; one shared fold is what makes that true.
  224. const agree = (): number => {
  225. const { systemTokens, messageTokens } = projected(ctx, session)
  226. expect(systemTokens + messageTokens).toBe(ctx.tokenMeter.measure(session).surfaceTokens)
  227. return messageTokens
  228. }
  229. appendSystem(session, 'You are terse.')
  230. session.append('request/header', {
  231. header: { config: CONFIG, tools: TOOLS },
  232. reason: 'initial',
  233. })
  234. expect(agree()).toBe(0)
  235. const question = appendUser(session, 'a first question, long enough to price above zero')
  236. session.append('step/start', { turn: 1, step: 1 })
  237. const answer = session.append('assistant/message', {
  238. stream: [],
  239. turn: 1,
  240. step: 1,
  241. message: createMessage({
  242. role: 'assistant',
  243. content: [{ type: 'text', text: 'a considered answer' }],
  244. source: { kind: 'model', provider: 'mock', model: 'mock' },
  245. }),
  246. usage: { inputTokens: 40, outputTokens: 7 },
  247. }, { surfaceOp: 'append' }).seq
  248. session.append('step/end', { turn: 1, step: 1 })
  249. const grown = agree()
  250. expect(grown).toBeGreaterThan(0)
  251. appendSummaryMeter(ctx, session, question, answer)
  252. // The armed shadow price must not move the published figure by itself.
  253. expect(agree()).toBe(grown)
  254. session.append('user/message', createUserMessage({
  255. content: [{ type: 'text', text: 'summary' }],
  256. source: { kind: 'plugin', plugin: 'test' },
  257. }), {
  258. surfaceOp: { op: 'replace', startSeq: question, endSeq: answer },
  259. sourceEventSeqs: [question, answer],
  260. })
  261. expect(agree()).toBeLessThan(grown)
  262. })
  263. it('prices unmetered replacements and rejects absent ranges without mutating prior state', async () => {
  264. const { session } = await harness()
  265. const first = appendUser(session, 'first message')
  266. const last = appendUser(session, 'last message')
  267. const definition = contextBreakdownProjectionDefinition
  268. const state = session.snapshotEvents().reduce(definition.apply, definition.init())
  269. const before = JSON.stringify(state)
  270. const replacement = session.append('user/message', createUserMessage({
  271. content: [{ type: 'text', text: 'summary' }], source: { kind: 'user' },
  272. }), { surfaceOp: { op: 'replace', startSeq: first, endSeq: last }, sourceEventSeqs: [first, last] })
  273. expect(definition.wire.view(definition.apply(state, replacement)).messageTokens).toBe(10)
  274. expect(JSON.stringify(state)).toBe(before)
  275. const invalid = { ...replacement, surfaceOp: { op: 'replace', startSeq: SessionSeq(999), endSeq: last } } as SessionEvent
  276. expect(() => definition.apply(state, invalid)).toThrow('invalid current range')
  277. expect(JSON.stringify(state)).toBe(before)
  278. })
  279. it('classifies the last nonempty system by position despite rewrites and extra source citations', async () => {
  280. const { ctx, session } = await harness()
  281. let head = appendSystem(session, 'head')
  282. let question = appendUser(session, 'question')
  283. const middle = appendSystem(session, 'middle prompt')
  284. const tail = appendSystem(session, 'last prompt in surface order')
  285. head = replaceSystem(session, head, 'head rewritten at a newer event seq')
  286. const agree = (text: string): void => {
  287. const view = projected(ctx, session)
  288. expect(view.systemTokens).toBe(estimateSystemMessage(createSystemMessage(text, SYSTEM_PLUGIN)))
  289. expect(view.messageTokens).toBeGreaterThanOrEqual(0)
  290. expect(view.systemTokens + view.messageTokens)
  291. .toBe(ctx.tokenMeter.measure(session).nodes.reduce((total, node) => total + node.heuristicTokens, 0))
  292. }
  293. agree('last prompt in surface order')
  294. question = session.append('user/message', createUserMessage({
  295. content: [{ type: 'text', text: 'rewritten question' }], source: { kind: 'user' },
  296. }), { surfaceOp: { op: 'replace', startSeq: question, endSeq: question }, sourceEventSeqs: [question] }).seq
  297. expect(question).toBeGreaterThan(middle)
  298. // Source-event references can cite a surviving prompt outside the replaced span.
  299. session.append('user/message', createUserMessage({
  300. content: [{ type: 'text', text: 'middle summary' }], source: { kind: 'user' },
  301. }), { surfaceOp: { op: 'replace', startSeq: question, endSeq: middle }, sourceEventSeqs: [question, middle, tail] })
  302. agree('last prompt in surface order')
  303. session.append('user/message', createUserMessage({
  304. content: [{ type: 'text', text: 'tail summary' }], source: { kind: 'user' },
  305. }), { surfaceOp: { op: 'replace', startSeq: tail, endSeq: tail }, sourceEventSeqs: [tail, head] })
  306. agree('head rewritten at a newer event seq')
  307. const repeated = appendSystem(session, 'head rewritten at a newer event seq')
  308. agree('head rewritten at a newer event seq')
  309. replaceSystem(session, repeated, '')
  310. agree('head rewritten at a newer event seq')
  311. replaceSystem(session, head, '')
  312. agree('')
  313. })
  314. it('ignores dormant empty system tails across per-node clearing and head fallback', async () => {
  315. const { ctx, session } = await harness()
  316. let head = appendSystem(session, 'head')
  317. appendUser(session, 'question')
  318. const middle = appendSystem(session, 'middle')
  319. const tail = appendSystem(session, 'tail')
  320. appendSystem(session, '')
  321. expect(projected(ctx, session).systemTokens).toBe(5)
  322. replaceSystem(session, tail, '')
  323. expect(projected(ctx, session).systemTokens).toBe(6)
  324. replaceSystem(session, middle, '')
  325. expect(projected(ctx, session)).toMatchObject({ systemTokens: 5, messageTokens: 10 })
  326. head = replaceSystem(session, head, 'fallback head')
  327. expect(projected(ctx, session)).toMatchObject({ systemTokens: 8, messageTokens: 10 })
  328. replaceSystem(session, head, '')
  329. expect(projected(ctx, session)).toMatchObject({ systemTokens: 0, messageTokens: 10 })
  330. })
  331. it('retains only compact current surface entries as history grows and compacts', async () => {
  332. const { ctx, session } = await harness()
  333. const first = appendUser(session, 'the first of many messages')
  334. for (let index = 0; index < 24; index += 1) appendUser(session, `message number ${index} with some text`)
  335. const last = appendUser(session, 'the last message before compaction')
  336. const state = () => {
  337. const current = ctx.sessionProjections.stateOf(session, 'contextBreakdown')
  338. if (current === undefined) throw new Error('registered context breakdown has no state')
  339. return current
  340. }
  341. expect(state().nodes).toHaveLength(26)
  342. expect(Object.keys(state().nodes[0]!).sort()).toEqual(['heuristicTokens', 'seq', 'system'])
  343. const shadowed = [...session.surface.nodes]
  344. session.append('user/message', createUserMessage({
  345. content: [{ type: 'text', text: 'summary' }], source: { kind: 'user' },
  346. }), { surfaceOp: { op: 'replace', startSeq: first, endSeq: last }, sourceEventSeqs: shadowed })
  347. expect(state().nodes).toHaveLength(1)
  348. expect(projected(ctx, session).messageTokens).toBe(10)
  349. })
  350. it('retains wire identity when a same-price rewrite changes only checkpoint positions', async () => {
  351. const { session } = await harness()
  352. const head = appendSystem(session, 'head')
  353. const definition = contextBreakdownProjectionDefinition
  354. const state = session.snapshotEvents().reduce(definition.apply, definition.init())
  355. replaceSystem(session, head, 'same')
  356. const next = definition.apply(state, session.snapshotEvents().at(-1)!)
  357. expect(next).not.toBe(state)
  358. expect(definition.wire.view(next)).toBe(definition.wire.view(state))
  359. expect(state.nodes[0]?.seq).toBe(head)
  360. })
  361. it('replays late registration, resumes a compact checkpoint, and discards scalar version 2', async () => {
  362. const ctx = new Context()
  363. contexts.push(ctx)
  364. await ctx.plugin(SessionStore)
  365. await ctx.plugin(SessionProjectionRegistry)
  366. const session = ctx.sessions.create()
  367. appendSystem(session, 'head')
  368. const first = appendUser(session, 'question')
  369. const last = appendSystem(session, 'x'.repeat(4000))
  370. await ctx.plugin(TokenMeter)
  371. expect(projected(ctx, session)).toMatchObject({ systemTokens: 1004, messageTokens: 15 })
  372. const checkpoint = JSON.parse(JSON.stringify(
  373. ctx.sessionProjections.checkpoint(session),
  374. )) as ReturnType<typeof ctx.sessionProjections.checkpoint>
  375. const row = checkpoint['contextBreakdown']!
  376. expect(row.ver).toBe(5)
  377. expect(ctx.sessionProjections.viewCheckpoint(checkpoint).contextBreakdown).toEqual(projected(ctx, session))
  378. const replacement = session.append('user/message', createUserMessage({
  379. content: [{ type: 'text', text: 'summary' }], source: { kind: 'user' },
  380. }), { surfaceOp: { op: 'replace', startSeq: first, endSeq: last }, sourceEventSeqs: [first, last] })
  381. const restored = ctx.sessionProjections.restore(
  382. checkpoint, [replacement], SessionLogOffset(replacement.seq), session.header, session.inheritedEventCount,
  383. )
  384. expect(restored.snapshot.values.contextBreakdown).toEqual({ systemTokens: 5, toolsTokens: 0, messageTokens: 10 })
  385. const stale = { ...checkpoint, contextBreakdown: { ...row, ver: 2, val: { systemTokens: 1004, toolsTokens: 0, messageTokens: -989 } } }
  386. expect(ctx.sessionProjections.viewCheckpoint(stale).contextBreakdown).toBeUndefined()
  387. expect(ctx.sessionProjections.restoreFloor(stale)).toBe(0)
  388. expect(() => ctx.sessionProjections.restore(
  389. stale, [replacement], SessionLogOffset(replacement.seq), session.header, session.inheritedEventCount,
  390. )).toThrow('re-read from seq 0')
  391. const replayed = ctx.sessionProjections.restore(
  392. stale, session.snapshotEvents(), SessionLogOffset(0), session.header, session.inheritedEventCount,
  393. )
  394. expect(replayed.snapshot.values.contextBreakdown).toEqual(projected(ctx, session))
  395. expect(replayed.checkpoint['contextBreakdown']?.ver).toBe(5)
  396. const invalid = {
  397. ...checkpoint,
  398. contextBreakdown: {
  399. ...row,
  400. val: { nodes: [{ seq: 0, heuristicTokens: -1, system: true }], breakdown: { systemTokens: 0, toolsTokens: 0, messageTokens: 0 } },
  401. },
  402. }
  403. expect(() => ctx.sessionProjections.restore(
  404. invalid, session.snapshotEvents(), SessionLogOffset(0), session.header, session.inheritedEventCount,
  405. )).toThrow()
  406. })
  407. it('discards lower-layer version-3 scalar caches and refolds the full surface', async () => {
  408. const { ctx, session } = await harness()
  409. try {
  410. appendSystem(session, 'You are terse.')
  411. session.append('request/header', { header: { config: CONFIG, tools: TOOLS }, reason: 'initial' })
  412. appendUser(session, 'abcd')
  413. const current = ctx.sessionProjections.checkpoint(session)
  414. const staleValue = { systemTokens: 0, toolsTokens: estimateToolsTokens({ config: CONFIG, tools: TOOLS }), messageTokens: 17 }
  415. const checkpoint = {
  416. ...current,
  417. contextBreakdown: { ver: 3, seq: SessionSeq(session.seq - 1), val: staleValue },
  418. }
  419. expect.soft(ctx.sessionProjections.viewCheckpoint(checkpoint)).not.toHaveProperty('contextBreakdown')
  420. expect.soft(ctx.sessionProjections.restoreFloor(checkpoint)).toBe(0)
  421. const restored = ctx.sessionProjections.restore(
  422. checkpoint, session.snapshotEvents(), SessionLogOffset(0), session.header, session.inheritedEventCount,
  423. )
  424. expect(restored.snapshot.values.contextBreakdown).toEqual({
  425. systemTokens: 8, toolsTokens: staleValue.toolsTokens, messageTokens: 9,
  426. })
  427. expect(restored.checkpoint).toEqual(current)
  428. expect(restored.checkpoint['contextBreakdown']?.ver).toBe(5)
  429. } finally {
  430. await ctx.fiber.dispose()
  431. }
  432. })
  433. it('restores from a JSON checkpoint and unregisters with the token-meter fiber', async () => {
  434. const ctx = new Context()
  435. contexts.push(ctx)
  436. await ctx.plugin(SessionStore)
  437. await ctx.plugin(SessionProjectionRegistry)
  438. const meterFiber = await ctx.plugin(TokenMeter)
  439. const session = ctx.sessions.create()
  440. appendSystem(session, 'You are terse.')
  441. appendUser(session, 'abcd')
  442. const checkpoint = JSON.parse(JSON.stringify(
  443. ctx.sessionProjections.checkpoint(session),
  444. )) as ReturnType<typeof ctx.sessionProjections.checkpoint>
  445. await meterFiber.dispose()
  446. expect(ctx.sessionProjections.snapshot(session).values).not.toHaveProperty('contextBreakdown')
  447. await ctx.plugin(TokenMeter)
  448. expect(ctx.sessionProjections.viewCheckpoint(checkpoint).contextBreakdown).toEqual({
  449. systemTokens: 8,
  450. toolsTokens: 0,
  451. messageTokens: 9,
  452. })
  453. })
  454. })
  455. describe('shared estimator', () => {
  456. it('prices every content-block shape under the fixed heuristic', () => {
  457. expect(estimateContent([{ type: 'text', text: 'abcd' }])).toBe(5)
  458. expect(estimateContent([{ type: 'reasoning', text: 'abcdefgh' }] as ContentBlock[])).toBe(6)
  459. expect(estimateContent([{ type: 'tool-call', id: 'c' as never, name: 'bash', arguments: '{"a":1}' }])).toBe(7)
  460. expect(estimateContent([{
  461. type: 'tool-result', toolCallId: 'c' as never,
  462. content: [{ type: 'text', text: 'abcd' }],
  463. }])).toBe(9)
  464. const unknown = { type: 'mystery', payload: 'abc' } as unknown as ContentBlock
  465. expect(estimateContent([unknown])).toBe(4 + Math.ceil(JSON.stringify(unknown).length / 4))
  466. })
  467. it('prices the system node without block overhead and an empty prompt to zero', () => {
  468. expect(estimateSystemMessage(createSystemMessage('', SYSTEM_PLUGIN))).toBe(0)
  469. expect(estimateSystemMessage(createSystemMessage('abcdefgh', SYSTEM_PLUGIN))).toBe(6)
  470. // estimateMessage routes the system role to the same figure.
  471. expect(estimateMessage(createSystemMessage('abcdefgh', SYSTEM_PLUGIN))).toBe(6)
  472. // A non-text block in a system message keeps a conservative JSON price.
  473. const image = { type: 'image', attachment: { attachmentId: 'a' } } as unknown as ContentBlock
  474. expect(estimateSystemMessage(createMessage({
  475. role: 'system',
  476. content: [{ type: 'text', text: 'abcd' }, image],
  477. source: { kind: 'plugin', plugin: SYSTEM_PLUGIN },
  478. }))).toBe(Math.ceil((4 + JSON.stringify(image).length) / 4) + 4)
  479. })
  480. it('prices the envelope tool schemas and absent tools to zero', () => {
  481. expect(estimateToolsTokens(undefined)).toBe(0)
  482. expect(estimateToolsTokens({ config: CONFIG, tools: [] })).toBe(0)
  483. expect(estimateToolsTokens({ config: CONFIG, tools: TOOLS }))
  484. .toBe(Math.ceil(JSON.stringify(TOOLS).length / 4) + 4)
  485. })
  486. })