hermes-agent / apps /desktop /src /app /session /session-state-cache.ts

Add files using upload-large-folder tool

28ad89a verified about 2 months ago

5.27 kB

import type { ClientSessionState } from '../types'

export const DEFAULT_WARM_SESSION_TRANSCRIPT_COUNT = 24
export const DEFAULT_WARM_SESSION_TRANSCRIPT_BYTES = 32 * 1024 * 1024

interface SessionStateCacheLimits {
  maxBytes?: number
  maxCount?: number
}

interface SessionStateCacheCallbacks {
  isReferenced: (runtimeId: string, state: ClientSessionState) => boolean
  onEvict: (runtimeId: string, state: ClientSessionState) => void
  /** Optional liveness check for a cached snapshot's in-flight claims. A
   *  connection death mid-turn orphans snapshots: the respawned backend
   *  re-mints runtime ids, so their frozen busy/awaitingResponse flags never
   *  receive a settling publish (#95189) and would pin megabytes of warm
   *  transcript per reconnect cycle until restart. When wired, those flags
   *  only block eviction while the authoritative store still claims work for
   *  the same runtime id; without the probe they always block. */
  isAuthoritativelyActive?: (runtimeId: string, state: ClientSessionState) => boolean
}

function transcriptBytes(state: ClientSessionState): number {
  if (state.messages.length === 0) {
    return 0
  }

  // JS strings occupy two bytes per UTF-16 code unit. JSON also accounts for
  // ids, part tags, tool payloads, attachment metadata, and error text without
  // retaining a second serialized copy in the cache.
  return JSON.stringify(state.messages).length * 2
}

function hasDraftOrInFlightMessage(state: ClientSessionState): boolean {
  return state.messages.some(message => message.pending === true)
}

/**
 * Runtime state map whose settled, unreferenced transcripts form a weighted
 * LRU. Live/visible states and unsaved drafts are outside both limits.
 */
export class SessionStateCache extends Map<string, ClientSessionState> {
  readonly #callbacks: SessionStateCacheCallbacks
  readonly #maxBytes: number
  readonly #maxCount: number
  readonly #recency = new Map<string, number>()
  readonly #transcriptWeights = new WeakMap<ClientSessionState['messages'], number>()
  #clock = 0

  constructor(callbacks: SessionStateCacheCallbacks, limits: SessionStateCacheLimits = {}) {
    super()
    this.#callbacks = callbacks
    this.#maxBytes = limits.maxBytes ?? DEFAULT_WARM_SESSION_TRANSCRIPT_BYTES
    this.#maxCount = limits.maxCount ?? DEFAULT_WARM_SESSION_TRANSCRIPT_COUNT
  }

  override get(runtimeId: string): ClientSessionState | undefined {
    const state = super.get(runtimeId)

    if (state) {
      this.#touch(runtimeId)
    }

    return state
  }

  override set(runtimeId: string, state: ClientSessionState): this {
    super.set(runtimeId, state)
    this.#touch(runtimeId)

    return this
  }

  override delete(runtimeId: string): boolean {
    this.#recency.delete(runtimeId)

    return super.delete(runtimeId)
  }

  override clear(): void {
    this.#recency.clear()
    super.clear()
  }

  prune(): void {
    const candidates: Array<{ bytes: number; runtimeId: string; state: ClientSessionState; touched: number }> = []
    let bytes = 0

    for (const [runtimeId, state] of this.entries()) {
      if (!this.#isWarmSettled(runtimeId, state)) {
        continue
      }

      const weight = this.#weight(state)
      candidates.push({ bytes: weight, runtimeId, state, touched: this.#recency.get(runtimeId) ?? 0 })
      bytes += weight
    }

    let count = candidates.length

    if (count <= this.#maxCount && bytes <= this.#maxBytes) {
      return
    }

    candidates.sort((a, b) => a.touched - b.touched)

    for (const candidate of candidates) {
      if (count <= this.#maxCount && bytes <= this.#maxBytes) {
        break
      }

      // References and activity can change between insertion and pruning.
      const current = super.get(candidate.runtimeId)

      if (current !== candidate.state || !this.#isWarmSettled(candidate.runtimeId, current)) {
        continue
      }

      super.delete(candidate.runtimeId)
      this.#recency.delete(candidate.runtimeId)
      count -= 1
      bytes -= candidate.bytes
      this.#callbacks.onEvict(candidate.runtimeId, candidate.state)
    }
  }

  #isWarmSettled(runtimeId: string, state: ClientSessionState): boolean {
    // In-flight claims pin a transcript only while they are trustworthy: with
    // an authority probe wired, a frozen busy/awaitingResponse on an orphaned
    // snapshot stops blocking eviction (see callback docs). Without one, the
    // legacy behavior holds and the flags always block.
    if (
      (state.busy || state.awaitingResponse) &&
      this.#callbacks.isAuthoritativelyActive?.(runtimeId, state) !== false
    ) {
      return false
    }

    return (
      Boolean(state.storedSessionId) &&
      state.messages.length > 0 &&
      !state.needsInput &&
      !hasDraftOrInFlightMessage(state) &&
      !this.#callbacks.isReferenced(runtimeId, state)
    )
  }

  #weight(state: ClientSessionState): number {
    const cached = this.#transcriptWeights.get(state.messages)

    if (cached !== undefined) {
      return cached
    }

    const weight = transcriptBytes(state)
    this.#transcriptWeights.set(state.messages, weight)

    return weight
  }

  #touch(runtimeId: string): void {
    this.#clock += 1
    this.#recency.set(runtimeId, this.#clock)
  }
}

	import type { ClientSessionState } from '../types'

	export const DEFAULT_WARM_SESSION_TRANSCRIPT_COUNT = 24
	export const DEFAULT_WARM_SESSION_TRANSCRIPT_BYTES = 32 * 1024 * 1024

	interface SessionStateCacheLimits {
	maxBytes?: number
	maxCount?: number
	}

	interface SessionStateCacheCallbacks {
	isReferenced: (runtimeId: string, state: ClientSessionState) => boolean
	onEvict: (runtimeId: string, state: ClientSessionState) => void
	/** Optional liveness check for a cached snapshot's in-flight claims. A
	* connection death mid-turn orphans snapshots: the respawned backend
	* re-mints runtime ids, so their frozen busy/awaitingResponse flags never
	* receive a settling publish (#95189) and would pin megabytes of warm
	* transcript per reconnect cycle until restart. When wired, those flags
	* only block eviction while the authoritative store still claims work for
	* the same runtime id; without the probe they always block. */
	isAuthoritativelyActive?: (runtimeId: string, state: ClientSessionState) => boolean
	}

	function transcriptBytes(state: ClientSessionState): number {
	if (state.messages.length === 0) {
	return 0
	}

	// JS strings occupy two bytes per UTF-16 code unit. JSON also accounts for
	// ids, part tags, tool payloads, attachment metadata, and error text without
	// retaining a second serialized copy in the cache.
	return JSON.stringify(state.messages).length * 2
	}

	function hasDraftOrInFlightMessage(state: ClientSessionState): boolean {
	return state.messages.some(message => message.pending === true)
	}

	/**
	* Runtime state map whose settled, unreferenced transcripts form a weighted
	* LRU. Live/visible states and unsaved drafts are outside both limits.
	*/
	export class SessionStateCache extends Map<string, ClientSessionState> {
	readonly #callbacks: SessionStateCacheCallbacks
	readonly #maxBytes: number
	readonly #maxCount: number
	readonly #recency = new Map<string, number>()
	readonly #transcriptWeights = new WeakMap<ClientSessionState['messages'], number>()
	#clock = 0

	constructor(callbacks: SessionStateCacheCallbacks, limits: SessionStateCacheLimits = {}) {
	super()
	this.#callbacks = callbacks
	this.#maxBytes = limits.maxBytes ?? DEFAULT_WARM_SESSION_TRANSCRIPT_BYTES
	this.#maxCount = limits.maxCount ?? DEFAULT_WARM_SESSION_TRANSCRIPT_COUNT
	}

	override get(runtimeId: string): ClientSessionState \| undefined {
	const state = super.get(runtimeId)

	if (state) {
	this.#touch(runtimeId)
	}

	return state
	}

	override set(runtimeId: string, state: ClientSessionState): this {
	super.set(runtimeId, state)
	this.#touch(runtimeId)

	return this
	}

	override delete(runtimeId: string): boolean {
	this.#recency.delete(runtimeId)

	return super.delete(runtimeId)
	}

	override clear(): void {
	this.#recency.clear()
	super.clear()
	}

	prune(): void {
	const candidates: Array<{ bytes: number; runtimeId: string; state: ClientSessionState; touched: number }> = []
	let bytes = 0

	for (const [runtimeId, state] of this.entries()) {
	if (!this.#isWarmSettled(runtimeId, state)) {
	continue
	}

	const weight = this.#weight(state)
	candidates.push({ bytes: weight, runtimeId, state, touched: this.#recency.get(runtimeId) ?? 0 })
	bytes += weight
	}

	let count = candidates.length

	if (count <= this.#maxCount && bytes <= this.#maxBytes) {
	return
	}

	candidates.sort((a, b) => a.touched - b.touched)

	for (const candidate of candidates) {
	if (count <= this.#maxCount && bytes <= this.#maxBytes) {
	break
	}

	// References and activity can change between insertion and pruning.
	const current = super.get(candidate.runtimeId)

	if (current !== candidate.state \|\| !this.#isWarmSettled(candidate.runtimeId, current)) {
	continue
	}

	super.delete(candidate.runtimeId)
	this.#recency.delete(candidate.runtimeId)
	count -= 1
	bytes -= candidate.bytes
	this.#callbacks.onEvict(candidate.runtimeId, candidate.state)
	}
	}

	#isWarmSettled(runtimeId: string, state: ClientSessionState): boolean {
	// In-flight claims pin a transcript only while they are trustworthy: with
	// an authority probe wired, a frozen busy/awaitingResponse on an orphaned
	// snapshot stops blocking eviction (see callback docs). Without one, the
	// legacy behavior holds and the flags always block.
	if (
	(state.busy \|\| state.awaitingResponse) &&
	this.#callbacks.isAuthoritativelyActive?.(runtimeId, state) !== false
	) {
	return false
	}

	return (
	Boolean(state.storedSessionId) &&
	state.messages.length > 0 &&
	!state.needsInput &&
	!hasDraftOrInFlightMessage(state) &&
	!this.#callbacks.isReferenced(runtimeId, state)
	)
	}

	#weight(state: ClientSessionState): number {
	const cached = this.#transcriptWeights.get(state.messages)

	if (cached !== undefined) {
	return cached
	}

	const weight = transcriptBytes(state)
	this.#transcriptWeights.set(state.messages, weight)

	return weight
	}

	#touch(runtimeId: string): void {
	this.#clock += 1
	this.#recency.set(runtimeId, this.#clock)
	}
	}