token-meter.spec.ts 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726
  1. import { describe, expect, expectTypeOf, it } from 'vitest'
  2. import { Context } from '@deepseek-ai/cordis'
  3. import { createUserMessage, CallId, createMessage } from '@deepseek-ai/dsh-llm'
  4. import type { ContentBlock, Message, TokenUsage } from '@deepseek-ai/dsh-llm'
  5. import SessionStore, { Session, SessionId, canonicalHeader } from '@deepseek-ai/dsh-session'
  6. import type { EpochHeader, SessionEvent } from '@deepseek-ai/dsh-session'
  7. import SessionProjectionRegistry from '@deepseek-ai/dsh-session-projection'
  8. import TokenMeter from '@deepseek-ai/dsh-token-meter'
  9. import type { TokenMeasurement, TokenMeterConfig } from '@deepseek-ai/dsh-token-meter'
  10. function header(model: string, extras: Omit<EpochHeader, 'config'> = {}): EpochHeader {
  11. return canonicalHeader({ config: { provider: 'mock', model }, ...extras })
  12. }
  13. function textMessage(text: string, role: Message['role'] = 'user'): Message {
  14. return createMessage({
  15. role,
  16. content: [{ type: 'text', text }],
  17. source: role === 'assistant'
  18. ? { kind: 'model', provider: 'mock', model: 'mock' }
  19. : { kind: 'user' },
  20. })
  21. }
  22. function appendHeader(session: Session, value: EpochHeader): void {
  23. session.append('request/header', { header: value, reason: 'initial' })
  24. }
  25. /** Inject malformed persisted history after the live append boundary for defensive replay tests. */
  26. function appendUnchecked(session: Session, event: SessionEvent): void {
  27. const log = (session as unknown as { log: SessionEvent[] }).log
  28. log.push(event)
  29. }
  30. interface SuccessfulCallOptions {
  31. turn?: number
  32. step?: number
  33. providerText?: string
  34. durableText?: string
  35. usage?: TokenUsage
  36. provenance?: 'exact' | 'empty' | 'absent'
  37. }
  38. function appendSuccessfulCall(
  39. session: Session,
  40. value: EpochHeader,
  41. options: SuccessfulCallOptions = {},
  42. ): void {
  43. const turn = options.turn ?? 1
  44. const step = options.step ?? 1
  45. const providerText = options.providerText ?? 'provider answer'
  46. const durableText = options.durableText ?? providerText
  47. const provenance = options.provenance ?? 'exact'
  48. session.append('step/start', { turn, step })
  49. appendHeader(session, value)
  50. const sources: number[] = []
  51. if (provenance === 'exact') {
  52. const chunks = [
  53. { type: 'block-start' as const, index: 0, blockType: 'text' as const },
  54. { type: 'text-delta' as const, index: 0, text: providerText },
  55. { type: 'block-end' as const, index: 0, block: { type: 'text' as const, text: providerText } },
  56. ...options.usage === undefined ? [] : [{ type: 'usage' as const, usage: options.usage }],
  57. { type: 'finish' as const, reason: { kind: 'stop' as const } },
  58. ]
  59. for (const chunk of chunks) {
  60. sources.push(session.append('assistant/chunk', { turn, step, chunk }).seq)
  61. }
  62. }
  63. const intent = provenance === 'absent'
  64. ? { surfaceOp: 'append' as const }
  65. : { surfaceOp: 'append' as const, sourceEventSeqs: provenance === 'empty' ? [] : sources }
  66. session.append('assistant/message', {
  67. turn,
  68. step,
  69. message: createMessage({
  70. role: 'assistant',
  71. content: durableText.length === 0 ? [] : [{ type: 'text', text: durableText }],
  72. source: {
  73. kind: 'model',
  74. ...{
  75. provider: value.config.provider,
  76. model: value.config.model,
  77. },
  78. },
  79. }),
  80. ...options.usage === undefined ? {} : { usage: options.usage },
  81. }, intent)
  82. session.append('step/end', { turn, step })
  83. }
  84. function meter(config: TokenMeterConfig = {}): TokenMeter {
  85. const ctx = new Context()
  86. // The registry is a required injection of the service (its three projection
  87. // units register in the constructor); mount it synchronously.
  88. new SessionProjectionRegistry(ctx)
  89. return new TokenMeter(ctx, config)
  90. }
  91. function expectSurfaceTotal(measurement: TokenMeasurement): void {
  92. expect(measurement.nodes.reduce((total, node) => total + node.tokens, 0))
  93. .toBe(measurement.surfaceTokens)
  94. }
  95. describe('TokenMeter configuration and registration', () => {
  96. it('exposes an empty public configuration type', () => {
  97. expectTypeOf<{}>().toExtend<TokenMeterConfig>()
  98. expectTypeOf<{ contextWindow: number }>().not.toExtend<TokenMeterConfig>()
  99. })
  100. it.each(['models', 'contextWindow', 'contextWidow'])(
  101. 'rejects stale or unknown top-level config key %s',
  102. (key) => {
  103. expect(() => meter({ [key]: {} } as unknown as TokenMeterConfig))
  104. .toThrow(`TokenMeterConfig: unknown key "${key}"`)
  105. },
  106. )
  107. it('registers and unregisters ctx.tokenMeter with its plugin fiber', async () => {
  108. const ctx = new Context()
  109. await ctx.plugin(SessionStore)
  110. await ctx.plugin(SessionProjectionRegistry)
  111. const fiber = await ctx.plugin(TokenMeter)
  112. expect(ctx.get('tokenMeter')).toBeInstanceOf(TokenMeter)
  113. await fiber.dispose()
  114. expect(ctx.get('tokenMeter')).toBeUndefined()
  115. })
  116. })
  117. describe('TokenMeter pricing', () => {
  118. it('prices every built-in content shape and merge-extended blocks with one fixed heuristic', () => {
  119. const service = meter()
  120. const blocks: ContentBlock[] = [
  121. { type: 'text', text: 'abcd' },
  122. { type: 'reasoning', text: 'ab' },
  123. { type: 'tool-call', id: CallId('c'), name: 'read', arguments: '{"x":1}' },
  124. {
  125. type: 'tool-result',
  126. toolCallId: CallId('c'),
  127. content: [{ type: 'text', text: 'xy' }],
  128. isError: false,
  129. },
  130. { type: 'future-block', payload: 'abcd' } as unknown as ContentBlock,
  131. ]
  132. const estimated = service.estimateMessage(createMessage({
  133. role: 'assistant', content: blocks,
  134. source: { kind: 'plugin', plugin: 'test' },
  135. }))
  136. expect(estimated).toBeGreaterThan(30)
  137. expect(service.estimateMessage(textMessage('abcd'))).toBe(9)
  138. })
  139. it('returns a detached deeply immutable empty measurement', () => {
  140. const service = meter()
  141. const session = Session.create(SessionId('empty'))
  142. const result = service.measure(session)
  143. expect(result).toEqual({
  144. logRevision: 0,
  145. baseline: { kind: 'none', tokens: 0 },
  146. surfaceDeltaTokens: 0,
  147. totalTokens: 0,
  148. surfaceTokens: 0,
  149. nodes: [],
  150. })
  151. expect(Object.isFrozen(result)).toBe(true)
  152. expect(Object.isFrozen(result.baseline)).toBe(true)
  153. expect(Object.isFrozen(result.nodes)).toBe(true)
  154. expectSurfaceTotal(result)
  155. expect(() => {
  156. ;(result as { totalTokens: number }).totalTokens = 1
  157. }).toThrow(TypeError)
  158. })
  159. it('keeps an earlier unified snapshot detached from later replay', () => {
  160. const service = meter()
  161. const session = Session.create(SessionId('detached'))
  162. session.append('user/message', createUserMessage({
  163. content: [{ type: 'text', text: 'first' }],
  164. source: { kind: 'user' },
  165. }), { surfaceOp: 'append' })
  166. const snapshot = service.measure(session)
  167. const snapshotCopy = structuredClone(snapshot)
  168. expect(Object.isFrozen(snapshot.nodes)).toBe(true)
  169. expect(Object.isFrozen(snapshot.nodes[0])).toBe(true)
  170. expectSurfaceTotal(snapshot)
  171. expect(() => {
  172. ;(snapshot.nodes as Array<{ seq: number; tokens: number }>).push({ seq: 99, tokens: 1 })
  173. }).toThrow(TypeError)
  174. expect(() => {
  175. ;(snapshot.nodes[0] as { seq: number; tokens: number }).tokens = 1
  176. }).toThrow(TypeError)
  177. session.append('user/message', createUserMessage({
  178. content: [{ type: 'text', text: 'second' }],
  179. source: { kind: 'user' },
  180. }), { surfaceOp: 'append' })
  181. const advanced = service.measure(session)
  182. expect(advanced.logRevision).toBe(2)
  183. expect(advanced.nodes).toHaveLength(2)
  184. expectSurfaceTotal(advanced)
  185. expect(snapshot).toEqual(snapshotCopy)
  186. expect(snapshot.logRevision).toBe(1)
  187. expect(snapshot.nodes).toHaveLength(1)
  188. })
  189. it('prices header, tools, and surface when no reusable usage exists', () => {
  190. const service = meter()
  191. const session = Session.create(SessionId('heuristic'))
  192. session.append('user/message', createUserMessage({
  193. content: [{ type: 'text', text: 'question' }],
  194. source: { kind: 'user' },
  195. }), { surfaceOp: 'append' })
  196. appendHeader(session, header('deepseek-v4-flash', {
  197. system: 'system',
  198. tools: [{ name: 'read', description: 'read', parameters: { type: 'object' } }],
  199. }))
  200. const result = service.measure(session)
  201. expect(result.baseline.kind).toBe('estimated')
  202. expect(result.totalTokens).toBeGreaterThan(result.surfaceTokens)
  203. expect(result.logRevision).toBe(session.events.length)
  204. expectSurfaceTotal(result)
  205. })
  206. it('keeps request-header overrides out of the returned surface', () => {
  207. const service = meter()
  208. const session = Session.create(SessionId('override-surface'))
  209. session.append('user/message', createUserMessage({
  210. content: [{ type: 'text', text: 'question' }],
  211. source: { kind: 'user' },
  212. }), { surfaceOp: 'append' })
  213. const logged = service.measure(session)
  214. const overridden = service.measure(session, header('another-model', {
  215. system: 'large override '.repeat(100),
  216. }))
  217. expect(overridden.totalTokens).toBeGreaterThan(logged.totalTokens)
  218. expect(overridden.surfaceTokens).toBe(logged.surfaceTokens)
  219. expect(overridden.nodes).toEqual(logged.nodes)
  220. expectSurfaceTotal(overridden)
  221. })
  222. })
  223. describe('replay anchors and surface folds', () => {
  224. const USAGE: TokenUsage = {
  225. inputTokens: 20,
  226. cacheReadTokens: 3,
  227. cacheWriteTokens: 4,
  228. outputTokens: 7,
  229. reasoningTokens: 6,
  230. }
  231. it('uses disjoint provider usage and signed durable-output rewrites', () => {
  232. const service = meter()
  233. const session = Session.create(SessionId('usage'))
  234. session.append('user/message', createUserMessage({
  235. content: [{ type: 'text', text: 'before' }],
  236. source: { kind: 'user' },
  237. }), { surfaceOp: 'append' })
  238. appendSuccessfulCall(session, header('deepseek-v4-flash'), {
  239. providerText: 'short',
  240. durableText: 'a much longer rewritten durable assistant answer',
  241. usage: USAGE,
  242. })
  243. const result = service.measure(session)
  244. expect(result.baseline).toMatchObject({ kind: 'usage', tokens: 34, usage: USAGE })
  245. expect(result.surfaceDeltaTokens).toBeGreaterThan(0)
  246. expect(result.totalTokens).toBe(34 + result.surfaceDeltaTokens)
  247. expect(() => {
  248. ;((result.baseline as { usage: { inputTokens: number } }).usage.inputTokens) = 1
  249. }).toThrow(TypeError)
  250. })
  251. it('selects a heuristic anchor when provider usage would undercut its scale', () => {
  252. const service = meter()
  253. const session = Session.create(SessionId('low-usage-anchor'))
  254. const system = 'system context'
  255. const requestHeader = header('deepseek-v4-flash', { system })
  256. appendSuccessfulCall(session, requestHeader, {
  257. providerText: 'abcd'.repeat(512),
  258. usage: { inputTokens: 20, outputTokens: 7 },
  259. })
  260. const anchored = service.measure(session)
  261. expect(anchored.baseline.kind).toBe('estimated')
  262. const assistant = anchored.nodes[0]!.seq
  263. session.append('user/message', createUserMessage({
  264. content: [{ type: 'text', text: 'short' }],
  265. source: { kind: 'plugin', plugin: 'test' },
  266. }), {
  267. surfaceOp: { op: 'replace', start: assistant, end: assistant },
  268. sourceEventSeqs: [assistant],
  269. })
  270. const shrunken = service.measure(session)
  271. expect(27 + shrunken.surfaceDeltaTokens).toBeLessThan(0)
  272. expect(shrunken.totalTokens).toBeGreaterThan(0)
  273. expect(shrunken.totalTokens).toBe(service.measure(
  274. session,
  275. header('different-model', { system }),
  276. ).totalTokens)
  277. })
  278. it('uses an estimated anchor when provider usage is absent', () => {
  279. const service = meter()
  280. const session = Session.create(SessionId('missing-usage'))
  281. appendSuccessfulCall(session, header('deepseek-v4-flash', { system: 's' }), {
  282. providerText: 'provider',
  283. durableText: 'rewritten',
  284. })
  285. const anchored = service.measure(session)
  286. expect(anchored.baseline.kind).toBe('estimated')
  287. expect(anchored.surfaceDeltaTokens).toBe(0)
  288. session.append('user/message', createUserMessage({
  289. content: [{ type: 'text', text: 'later' }],
  290. source: { kind: 'user' },
  291. }), { surfaceOp: 'append' })
  292. const advanced = service.measure(session)
  293. expect(advanced.surfaceDeltaTokens).toBeGreaterThan(0)
  294. })
  295. it('distinguishes an explicit empty source-event list from an absent legacy list', () => {
  296. const explicit = Session.create(SessionId('explicit-empty'))
  297. const legacy = Session.create(SessionId('legacy-absent'))
  298. appendSuccessfulCall(explicit, header('deepseek-v4-flash'), {
  299. durableText: 'listener injected text',
  300. providerText: '',
  301. usage: USAGE,
  302. provenance: 'empty',
  303. })
  304. appendSuccessfulCall(legacy, header('deepseek-v4-flash'), {
  305. durableText: 'listener injected text',
  306. providerText: '',
  307. usage: USAGE,
  308. provenance: 'absent',
  309. })
  310. const service = meter()
  311. expect(service.measure(explicit).surfaceDeltaTokens).toBeGreaterThan(0)
  312. expect(service.measure(legacy).surfaceDeltaTokens).toBe(0)
  313. })
  314. it('keeps only the latest successful request anchor across model switches', () => {
  315. const service = meter()
  316. const session = Session.create(SessionId('switch'))
  317. const alphaHeader = header('alpha', { system: 'same envelope' })
  318. appendSuccessfulCall(session, alphaHeader, { usage: USAGE, providerText: 'alpha' })
  319. expect(service.measure(session).baseline).toMatchObject({ kind: 'usage', tokens: 34 })
  320. appendSuccessfulCall(session, header('beta'), {
  321. turn: 1,
  322. step: 2,
  323. usage: { inputTokens: 100, outputTokens: 50 },
  324. providerText: 'beta response',
  325. })
  326. expect(service.measure(session).baseline).toMatchObject({ kind: 'usage', tokens: 150 })
  327. appendHeader(session, alphaHeader)
  328. const switchedBack = service.measure(session)
  329. expect(switchedBack.baseline.kind).toBe('estimated')
  330. expect(switchedBack.surfaceDeltaTokens).toBe(0)
  331. })
  332. it('invalidates usage for any canonical envelope change or explicit override', () => {
  333. const service = meter()
  334. const session = Session.create(SessionId('envelope'))
  335. const anchoredHeader = header('deepseek-v4-flash', { system: 'one' })
  336. appendSuccessfulCall(session, anchoredHeader, { usage: USAGE })
  337. expect(service.measure(session, { ...anchoredHeader, tools: [] }).baseline.kind).toBe('usage')
  338. expect(service.measure(session, header('deepseek-v4-flash', { system: 'two' })).baseline.kind)
  339. .toBe('estimated')
  340. expect(service.measure(session, header('deepseek-v4-pro', { system: 'one' })).baseline.kind)
  341. .toBe('estimated')
  342. expect(service.measure(session, {
  343. ...anchoredHeader,
  344. config: { ...anchoredHeader.config, temperature: 0.2 },
  345. }).baseline.kind).toBe('estimated')
  346. expect(service.measure(session, {
  347. ...anchoredHeader,
  348. tools: [{ name: 'read', description: 'read', parameters: { type: 'object' } }],
  349. }).baseline.kind).toBe('estimated')
  350. })
  351. it('folds the latest full header snapshot into the effective envelope', () => {
  352. const session = Session.create(SessionId('header-snapshot'))
  353. appendHeader(session, header('deepseek-v4-flash'))
  354. session.append('request/header', {
  355. header: header('deepseek-v4-pro'),
  356. reason: 'change',
  357. })
  358. const result = meter().measure(session)
  359. expect(result.baseline.kind).toBe('estimated')
  360. expect(result.logRevision).toBe(2)
  361. })
  362. it('replays seeded append and replace operations with signed deltas', () => {
  363. const service = meter()
  364. const original = Session.create(SessionId('surface-original'))
  365. appendSuccessfulCall(original, header('deepseek-v4-flash'), {
  366. usage: USAGE,
  367. providerText: 'long provider answer '.repeat(100),
  368. })
  369. original.append('user/message', createUserMessage({
  370. content: [{ type: 'text', text: 'new tail' }],
  371. source: { kind: 'user' },
  372. }), { surfaceOp: 'append' })
  373. const seeded = Session.create(SessionId('surface-seeded'), original.events)
  374. const before = service.measure(seeded)
  375. expect(before.nodes).toHaveLength(2)
  376. expect(before.surfaceDeltaTokens).toBeGreaterThan(0)
  377. expectSurfaceTotal(before)
  378. const first = seeded.surface.nodes[0]!
  379. seeded.append('user/message', createUserMessage({
  380. content: [{ type: 'text', text: 'replacement' }],
  381. source: { kind: 'plugin', plugin: 'test' },
  382. }), { surfaceOp: { op: 'replace', start: first, end: first }, sourceEventSeqs: [first] })
  383. const after = service.measure(seeded)
  384. expect(after.nodes).toHaveLength(2)
  385. expect(after.nodes[0]!.seq).toBe(seeded.events.length - 1)
  386. expect(after.logRevision).toBe(seeded.events.length)
  387. expect(Object.isFrozen(after.nodes)).toBe(true)
  388. expect(Object.isFrozen(after.nodes[0])).toBe(true)
  389. expect(after.surfaceDeltaTokens).toBeLessThan(0)
  390. expectSurfaceTotal(after)
  391. expect(before.nodes).toHaveLength(2)
  392. // The earlier snapshot still reports the log it measured: seed + boundary.
  393. expect(before.logRevision).toBe(original.events.length + 1)
  394. expect(before.surfaceDeltaTokens).toBeGreaterThan(0)
  395. })
  396. it('prices an empty assistant surface anchor as zero', () => {
  397. const session = Session.create(SessionId('empty-assistant'))
  398. appendSuccessfulCall(session, header('deepseek-v4-flash'), {
  399. providerText: '',
  400. durableText: '',
  401. provenance: 'empty',
  402. })
  403. const measurement = meter().measure(session)
  404. const assistant = session.events.find(event => event.type === 'assistant/message')!
  405. expect(measurement.nodes).toEqual([{ seq: assistant.seq, tokens: 0 }])
  406. expect(measurement.surfaceTokens).toBe(0)
  407. expectSurfaceTotal(measurement)
  408. })
  409. })
  410. describe('malformed replay and listener lifecycle', () => {
  411. function expectRepeatedFailure(service: TokenMeter, session: Session, pattern: RegExp): void {
  412. expect(() => service.measure(session)).toThrow(pattern)
  413. expect(() => service.measure(session)).toThrow(pattern)
  414. }
  415. it('rejects an assistant without its step boundary transactionally', () => {
  416. const session = Session.create(SessionId('bad-step'))
  417. appendHeader(session, header('deepseek-v4-flash'))
  418. session.append('assistant/message', {
  419. turn: 1,
  420. step: 1,
  421. message: createMessage({
  422. role: 'assistant',
  423. content: [{ type: 'text', text: 'bad' }],
  424. source: {
  425. kind: 'model',
  426. ...{ provider: 'mock', model: 'deepseek-v4-flash' },
  427. },
  428. }),
  429. }, { surfaceOp: 'append', sourceEventSeqs: [] })
  430. expectRepeatedFailure(meter(), session, /no matching step\/start/)
  431. })
  432. it('leaves the priced surface uncommitted when a later validation step rejects the event', () => {
  433. // A valid append plan whose anchor validation throws: only commit
  434. // ordering keeps the surface from double-counting across retries.
  435. const session = Session.create(SessionId('bad-step-surface'))
  436. appendHeader(session, header('deepseek-v4-flash'))
  437. session.append('assistant/message', {
  438. turn: 1,
  439. step: 1,
  440. message: createMessage({
  441. role: 'assistant',
  442. content: [{ type: 'text', text: 'planned but never committed' }],
  443. source: {
  444. kind: 'model',
  445. ...{ provider: 'mock', model: 'deepseek-v4-flash' },
  446. },
  447. }),
  448. }, { surfaceOp: 'append', sourceEventSeqs: [] })
  449. const service = meter()
  450. const states = (service as unknown as {
  451. states: WeakMap<Session, { surface: unknown[]; surfaceTokens: number }>
  452. }).states
  453. expectRepeatedFailure(service, session, /no matching step\/start/)
  454. const state = states.get(session)
  455. expect(state?.surface).toEqual([])
  456. expect(state?.surfaceTokens).toBe(0)
  457. })
  458. it('clears completed step boundaries and rejects overlapping or late step events', () => {
  459. const overlapping = Session.create(SessionId('overlapping-step'))
  460. overlapping.append('step/start', { turn: 1, step: 1 })
  461. overlapping.append('step/start', { turn: 1, step: 2 })
  462. expectRepeatedFailure(
  463. meter(),
  464. overlapping,
  465. /arrived before turn 1\/step 1 ended/,
  466. )
  467. const late = Session.create(SessionId('late-assistant'))
  468. late.append('step/start', { turn: 1, step: 1 })
  469. appendHeader(late, header('deepseek-v4-flash'))
  470. late.append('step/end', { turn: 1, step: 1 })
  471. late.append('assistant/message', {
  472. turn: 1,
  473. step: 1,
  474. message: createMessage({
  475. role: 'assistant',
  476. content: [],
  477. source: {
  478. kind: 'model',
  479. ...{ provider: 'mock', model: 'deepseek-v4-flash' },
  480. },
  481. }),
  482. }, { surfaceOp: 'append', sourceEventSeqs: [] })
  483. expectRepeatedFailure(
  484. meter(),
  485. late,
  486. /no matching step\/start/,
  487. )
  488. const mismatchedEnd = Session.create(SessionId('mismatched-end'))
  489. mismatchedEnd.append('step/start', { turn: 1, step: 1 })
  490. mismatchedEnd.append('step/end', { turn: 1, step: 2 })
  491. expectRepeatedFailure(
  492. meter(),
  493. mismatchedEnd,
  494. /step\/end .* no matching step\/start/,
  495. )
  496. })
  497. it('rejects invalid assistant source-event references', () => {
  498. const cases: Array<{
  499. name: string
  500. appendSource(session: Session): number[]
  501. pattern: RegExp
  502. }> = [
  503. {
  504. name: 'non-chunk',
  505. appendSource(session) {
  506. return [session.append('user/message', createUserMessage({
  507. content: [{ type: 'text', text: 'x' }],
  508. source: { kind: 'user' },
  509. }), { surfaceOp: 'append' }).seq]
  510. },
  511. pattern: /is not assistant\/chunk/,
  512. },
  513. {
  514. name: 'wrong-step',
  515. appendSource(session) {
  516. return [session.append('assistant/chunk', {
  517. turn: 1,
  518. step: 2,
  519. chunk: { type: 'finish', reason: { kind: 'stop' } },
  520. }).seq]
  521. },
  522. pattern: /belongs to another step/,
  523. },
  524. ]
  525. for (const testCase of cases) {
  526. const session = Session.create(SessionId(`bad-source-${testCase.name}`))
  527. session.append('step/start', { turn: 1, step: 1 })
  528. appendHeader(session, header('deepseek-v4-flash'))
  529. const sourceEventSeqs = testCase.appendSource(session)
  530. session.append('assistant/message', {
  531. turn: 1,
  532. step: 1,
  533. message: createMessage({
  534. role: 'assistant',
  535. content: [{ type: 'text', text: 'bad' }],
  536. source: {
  537. kind: 'model',
  538. ...{ provider: 'mock', model: 'deepseek-v4-flash' },
  539. },
  540. }),
  541. usage: { inputTokens: 1, outputTokens: 1 },
  542. }, { surfaceOp: 'append', sourceEventSeqs })
  543. expect(() => meter().measure(session)).toThrow(testCase.pattern)
  544. }
  545. })
  546. it('rejects repeated and non-earlier assistant source-event references', () => {
  547. const duplicate = Session.create(SessionId('duplicate-source'))
  548. duplicate.append('step/start', { turn: 1, step: 1 })
  549. appendHeader(duplicate, header('deepseek-v4-flash'))
  550. const source = duplicate.append('assistant/chunk', {
  551. turn: 1,
  552. step: 1,
  553. chunk: { type: 'finish', reason: { kind: 'stop' } },
  554. }).seq
  555. appendUnchecked(duplicate, {
  556. type: 'assistant/message',
  557. seq: duplicate.seq,
  558. time: 0,
  559. data: {
  560. turn: 1,
  561. step: 1,
  562. message: createMessage({
  563. role: 'assistant',
  564. content: [],
  565. source: {
  566. kind: 'model',
  567. ...{ provider: 'mock', model: 'deepseek-v4-flash' },
  568. },
  569. }),
  570. usage: { inputTokens: 1, outputTokens: 0 },
  571. },
  572. surfaceOp: 'append',
  573. sourceEventSeqs: [source, source],
  574. })
  575. expect(() => meter().measure(duplicate)).toThrow(/repeats source seq/)
  576. const future = Session.create(SessionId('future-source'))
  577. future.append('step/start', { turn: 1, step: 1 })
  578. appendHeader(future, header('deepseek-v4-flash'))
  579. appendUnchecked(future, {
  580. type: 'assistant/message',
  581. seq: future.seq,
  582. time: 0,
  583. data: {
  584. turn: 1,
  585. step: 1,
  586. message: createMessage({
  587. role: 'assistant',
  588. content: [],
  589. source: {
  590. kind: 'model',
  591. ...{ provider: 'mock', model: 'deepseek-v4-flash' },
  592. },
  593. }),
  594. usage: { inputTokens: 1, outputTokens: 0 },
  595. },
  596. surfaceOp: 'append',
  597. sourceEventSeqs: [99],
  598. })
  599. expect(() => meter().measure(future)).toThrow(/is not earlier/)
  600. })
  601. it('does not partially apply a malformed assistant replacement', () => {
  602. const session = Session.create(SessionId('transactional-replace'))
  603. session.append('user/message', createUserMessage({
  604. content: [{ type: 'text', text: 'head' }],
  605. source: { kind: 'user' },
  606. }), { surfaceOp: 'append' })
  607. appendHeader(session, header('deepseek-v4-flash'))
  608. const head = session.events[0]!.seq
  609. session.append('assistant/message', {
  610. turn: 1,
  611. step: 1,
  612. message: createMessage({
  613. role: 'assistant',
  614. content: [{ type: 'text', text: 'replacement' }],
  615. source: {
  616. kind: 'model',
  617. ...{ provider: 'mock', model: 'deepseek-v4-flash' },
  618. },
  619. }),
  620. }, { surfaceOp: { op: 'replace', start: head, end: head }, sourceEventSeqs: [head] })
  621. expectRepeatedFailure(
  622. meter(),
  623. session,
  624. /no matching step\/start/,
  625. )
  626. })
  627. it('rejects corrupt replacement ranges without advancing the replay cursor', () => {
  628. const session = Session.create(SessionId('bad-replace'))
  629. const head = session.append('user/message', createUserMessage({
  630. content: [{ type: 'text', text: 'head' }],
  631. source: { kind: 'user' },
  632. }), { surfaceOp: 'append' }).seq
  633. appendUnchecked(session, {
  634. type: 'user/message',
  635. seq: session.seq,
  636. time: 0,
  637. data: createUserMessage({
  638. content: [{ type: 'text', text: 'bad' }],
  639. source: { kind: 'user' },
  640. }),
  641. surfaceOp: { op: 'replace', start: 99, end: 99 },
  642. sourceEventSeqs: [head],
  643. })
  644. expectRepeatedFailure(meter(), session, /invalid current range/)
  645. })
  646. it('handles earlier-reader catch-up, eager observation, and service reload', async () => {
  647. const ctx = new Context()
  648. await ctx.plugin(SessionStore)
  649. await ctx.plugin(SessionProjectionRegistry)
  650. let activeMeter: TokenMeter | undefined
  651. const revisions: number[] = []
  652. ctx.on('session/event', (session) => {
  653. if (activeMeter !== undefined) revisions.push(activeMeter.measure(session).logRevision)
  654. })
  655. const firstFiber = await ctx.plugin(TokenMeter)
  656. activeMeter = ctx.tokenMeter
  657. const session = ctx.sessions.create(SessionId('listener-order'), { seed: [{
  658. type: 'turn/start',
  659. seq: 0,
  660. time: 1,
  661. data: { turn: 1 },
  662. }] })
  663. activeMeter.measure(session)
  664. session.append('user/message', createUserMessage({
  665. content: [{ type: 'text', text: 'one' }],
  666. source: { kind: 'user' },
  667. }), { surfaceOp: 'append' })
  668. // Seed, end-seed, then one live append. Only the last event published:
  669. // end-seed predates store attachment, like the seed.
  670. expect(revisions).toEqual([3])
  671. expect(activeMeter.measure(session).logRevision).toBe(3)
  672. await firstFiber.dispose()
  673. const secondFiber = await ctx.plugin(TokenMeter)
  674. activeMeter = ctx.tokenMeter
  675. expect(activeMeter.measure(session).logRevision).toBe(3)
  676. await secondFiber.dispose()
  677. })
  678. })