session-prefetch.ts 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100
  1. const key = (directory: string, sessionID: string) => `${directory}\n${sessionID}`
  2. export const SESSION_PREFETCH_TTL = 15_000
  3. type Meta = {
  4. limit: number
  5. cursor?: string
  6. complete: boolean
  7. at: number
  8. }
  9. export function shouldSkipSessionPrefetch(input: { message: boolean; info?: Meta; chunk: number; now?: number }) {
  10. if (input.message) {
  11. if (!input.info) return true
  12. if (input.info.complete) return true
  13. if (input.info.limit > input.chunk) return true
  14. } else {
  15. if (!input.info) return false
  16. }
  17. return (input.now ?? Date.now()) - input.info.at < SESSION_PREFETCH_TTL
  18. }
  19. const cache = new Map<string, Meta>()
  20. const inflight = new Map<string, Promise<Meta | undefined>>()
  21. const rev = new Map<string, number>()
  22. const version = (id: string) => rev.get(id) ?? 0
  23. export function getSessionPrefetch(directory: string, sessionID: string) {
  24. return cache.get(key(directory, sessionID))
  25. }
  26. export function getSessionPrefetchPromise(directory: string, sessionID: string) {
  27. return inflight.get(key(directory, sessionID))
  28. }
  29. export function clearSessionPrefetchInflight() {
  30. inflight.clear()
  31. }
  32. export function isSessionPrefetchCurrent(directory: string, sessionID: string, value: number) {
  33. return version(key(directory, sessionID)) === value
  34. }
  35. export function runSessionPrefetch(input: {
  36. directory: string
  37. sessionID: string
  38. task: (value: number) => Promise<Meta | undefined>
  39. }) {
  40. const id = key(input.directory, input.sessionID)
  41. const pending = inflight.get(id)
  42. if (pending) return pending
  43. const value = version(id)
  44. const promise = input.task(value).finally(() => {
  45. if (inflight.get(id) === promise) inflight.delete(id)
  46. })
  47. inflight.set(id, promise)
  48. return promise
  49. }
  50. export function setSessionPrefetch(input: {
  51. directory: string
  52. sessionID: string
  53. limit: number
  54. cursor?: string
  55. complete: boolean
  56. at?: number
  57. }) {
  58. cache.set(key(input.directory, input.sessionID), {
  59. limit: input.limit,
  60. cursor: input.cursor,
  61. complete: input.complete,
  62. at: input.at ?? Date.now(),
  63. })
  64. }
  65. export function clearSessionPrefetch(directory: string, sessionIDs: Iterable<string>) {
  66. for (const sessionID of sessionIDs) {
  67. if (!sessionID) continue
  68. const id = key(directory, sessionID)
  69. rev.set(id, version(id) + 1)
  70. cache.delete(id)
  71. inflight.delete(id)
  72. }
  73. }
  74. export function clearSessionPrefetchDirectory(directory: string) {
  75. const prefix = `${directory}\n`
  76. const keys = new Set([...cache.keys(), ...inflight.keys()])
  77. for (const id of keys) {
  78. if (!id.startsWith(prefix)) continue
  79. rev.set(id, version(id) + 1)
  80. cache.delete(id)
  81. inflight.delete(id)
  82. }
  83. }