@lota-sdk/core 0.4.8 → 0.4.10

package/src/runtime/chat-run-orchestration.ts

@@ -1,25 +1,108 @@
-const COMPACTION_POLL_INTERVAL_MS = 200
+import { Cause, Context, Schema, Duration, Effect, Latch, Layer } from 'effect'
+
+import { effectTryPromise } from '../effect/helpers'
+import { nowEpochMillis } from '../utils/date-time'
+
+const COMPACTION_WAIT_REFRESH_MS = 1_000
 const COMPACTION_MAX_WAIT_MS = 120_000
 
-export async function waitForCompactionIfNeeded<TEntity>(params: {
+class WaitForCompactionError extends Schema.TaggedErrorClass<WaitForCompactionError>()('WaitForCompactionError', {
+  entityId: Schema.String,
+  entityLabel: Schema.String,
+  message: Schema.String,
+  cause: Schema.optional(Schema.Defect),
+}) {}
+
+function toWaitForCompactionError(params: {
   entityId: string
   entityLabel: string
-  loadEntity: () => Promise<TEntity>
-  isCompacting: (entity: TEntity) => boolean
-}): Promise<TEntity> {
-  let entity = await params.loadEntity()
-  if (!params.isCompacting(entity)) return entity
-
-  const deadline = Date.now() + COMPACTION_MAX_WAIT_MS
-  while (params.isCompacting(entity)) {
-    if (Date.now() > deadline) {
-      throw new Error(
-        `${params.entityLabel} ${params.entityId} compaction did not complete within ${COMPACTION_MAX_WAIT_MS}ms`,
-      )
-    }
-    await Bun.sleep(COMPACTION_POLL_INTERVAL_MS)
-    entity = await params.loadEntity()
-  }
+  cause: unknown
+}): WaitForCompactionError {
+  return new WaitForCompactionError({
+    entityId: params.entityId,
+    entityLabel: params.entityLabel,
+    message: params.cause instanceof Error ? params.cause.message : String(params.cause),
+    cause: params.cause,
+  })
+}
 
-  return entity
+interface CompactionCoordination {
+  readonly signal: (entityId: string, compacting: boolean) => Effect.Effect<void>
+  readonly waitIfNeeded: <TEntity>(params: {
+    entityId: string
+    entityLabel: string
+    loadEntity: () => PromiseLike<TEntity> | Effect.Effect<TEntity, unknown>
+    isCompacting: (entity: TEntity) => boolean
+  }) => Effect.Effect<TEntity, WaitForCompactionError>
 }
+
+export class CompactionCoordinationTag extends Context.Service<CompactionCoordinationTag, CompactionCoordination>()(
+  '@lota-sdk/core/CompactionCoordination',
+) {}
+
+export const CompactionCoordinationLive = Layer.effect(
+  CompactionCoordinationTag,
+  Effect.sync((): CompactionCoordination => {
+    const latches = new Map<string, Latch.Latch>()
+
+    function getLatch(entityId: string): Latch.Latch {
+      let latch = latches.get(entityId)
+      if (!latch) {
+        latch = Latch.makeUnsafe(true)
+        latches.set(entityId, latch)
+      }
+      return latch
+    }
+
+    return {
+      signal: (entityId, compacting) =>
+        Effect.sync(() => {
+          const latch = getLatch(entityId)
+          if (compacting) {
+            Latch.closeUnsafe(latch)
+          } else {
+            Latch.openUnsafe(latch)
+          }
+        }),
+
+      waitIfNeeded: Effect.fn('CompactionCoordination.waitIfNeeded')(function* <TEntity>(params: {
+        entityId: string
+        entityLabel: string
+        loadEntity: () => PromiseLike<TEntity> | Effect.Effect<TEntity, unknown>
+        isCompacting: (entity: TEntity) => boolean
+      }) {
+        const deadline = nowEpochMillis() + COMPACTION_MAX_WAIT_MS
+        const latch = getLatch(params.entityId)
+        let entity = yield* effectTryPromise(
+          () => params.loadEntity(),
+          (cause) => toWaitForCompactionError({ entityId: params.entityId, entityLabel: params.entityLabel, cause }),
+        )
+
+        while (params.isCompacting(entity)) {
+          Latch.closeUnsafe(latch)
+          if (nowEpochMillis() > deadline) {
+            return yield* new WaitForCompactionError({
+              entityId: params.entityId,
+              entityLabel: params.entityLabel,
+              message: `${params.entityLabel} ${params.entityId} compaction did not complete within ${COMPACTION_MAX_WAIT_MS}ms`,
+            })
+          }
+
+          const remainingMs = Math.max(0, deadline - nowEpochMillis())
+          const refreshWindowMs = Math.min(COMPACTION_WAIT_REFRESH_MS, remainingMs)
+          yield* latch.await.pipe(
+            Effect.timeout(Duration.millis(refreshWindowMs)),
+            Effect.catchIf(Cause.isTimeoutError, () => Effect.void),
+          )
+          entity = yield* effectTryPromise(
+            () => params.loadEntity(),
+            (cause) => toWaitForCompactionError({ entityId: params.entityId, entityLabel: params.entityLabel, cause }),
+          )
+        }
+
+        Latch.openUnsafe(latch)
+        return entity
+      }),
+    }
+  }),
+)

package/src/runtime/chat-run-registry.ts

@@ -1,5 +1,9 @@
+import { Effect, Fiber, FiberMap } from 'effect'
+
 export class ChatRunRegistry {
-  private controllers = new Map<string, AbortController>()
+  private readonly controllers = new Map<string, AbortController>()
+
+  constructor(private readonly trackedRuns: FiberMap.FiberMap<string>) {}
 
   has(runId: string): boolean {
     return this.controllers.has(runId)
@@ -16,9 +20,39 @@ export class ChatRunRegistry {
   stop(runId: string, reason?: unknown): boolean {
     const controller = this.controllers.get(runId)
     if (!controller) return false
-
     this.controllers.delete(runId)
     controller.abort(reason)
+    void Effect.runFork(FiberMap.remove(this.trackedRuns, runId))
     return true
   }
+
+  trackRunEffect<A, E, R>(
+    runId: string,
+    controller: AbortController,
+    effect: Effect.Effect<A, E, R>,
+  ): Effect.Effect<A, E, R> {
+    return Effect.gen(
+      function* (this: ChatRunRegistry) {
+        this.controllers.set(runId, controller)
+        const fiber = yield* FiberMap.run(
+          this.trackedRuns,
+          runId,
+          effect.pipe(
+            Effect.ensuring(
+              Effect.sync(() => {
+                if (this.controllers.get(runId) === controller) {
+                  this.controllers.delete(runId)
+                }
+              }),
+            ),
+          ),
+        )
+        return yield* Fiber.join(fiber)
+      }.bind(this),
+    )
+  }
+
+  stopEffect(runId: string, reason?: unknown): Effect.Effect<boolean> {
+    return Effect.sync(() => this.stop(runId, reason))
+  }
 }

package/src/runtime/context-compaction/context-compaction-runtime.ts

@@ -0,0 +1,107 @@
+import { Schema, Effect } from 'effect'
+
+import { createContextCompactionAgent } from '../../system-agents/context-compaction.agent'
+import type { GenerateHelperStructuredParams, GenerateHelperTextParams } from '../helper-model'
+import {
+  buildContextCompactionPrompt,
+  buildMemoryBlockCompactionPrompt,
+  ContextCompactionOutputSchema,
+  createContextCompactionRuntime,
+  parseCompactionOutput,
+} from './context-compaction'
+import type { ContextCompactionRunnerParams } from './context-compaction'
+import {
+  COMPACTION_CHUNK_MAX_CHARS,
+  CONTEXT_COMPACTION_INCLUDED_TOOL_NAMES,
+  CONTEXT_COMPACTION_INCLUDED_TOOL_PREFIXES,
+  CONTEXT_COMPACTION_THRESHOLD_RATIO,
+  CONTEXT_OUTPUT_RESERVE_TOKENS,
+  CONTEXT_SAFETY_MARGIN_TOKENS,
+  SUMMARY_ROLLUP_MAX_TOKENS,
+} from './context-compaction-constants'
+
+const CONTEXT_COMPACTION_MAX_OUTPUT_TOKENS = 512
+
+interface HelperModelRuntime {
+  generateHelperStructured<T>(params: GenerateHelperStructuredParams<T>): Promise<T>
+  generateHelperText(params: GenerateHelperTextParams): Promise<string>
+}
+
+interface CreateContextCompactionRuntimeDeps {
+  helperModelRuntime: HelperModelRuntime
+  now?: () => number
+  randomId?: () => string
+}
+
+class ContextCompactionRuntimeError extends Schema.TaggedErrorClass<ContextCompactionRuntimeError>()(
+  'ContextCompactionRuntimeError',
+  { message: Schema.String, cause: Schema.Defect },
+) {}
+
+function tryContextCompactionPromise<A>(
+  message: string,
+  thunk: () => PromiseLike<A>,
+): Effect.Effect<A, ContextCompactionRuntimeError> {
+  return Effect.tryPromise({
+    try: () => Promise.resolve(thunk()),
+    catch: (cause) => new ContextCompactionRuntimeError({ message, cause }),
+  })
+}
+
+function runContextCompacter(helperModelRuntime: HelperModelRuntime, params: ContextCompactionRunnerParams) {
+  return Effect.runPromise(
+    tryContextCompactionPromise('Failed to compact runtime context.', () =>
+      helperModelRuntime.generateHelperStructured({
+        tag: 'context-compaction',
+        createAgent: createContextCompactionAgent,
+        messages: [
+          {
+            role: 'user',
+            content: buildContextCompactionPrompt({
+              previousSummary: params.previousSummary,
+              transcript: params.transcript,
+            }),
+          },
+        ],
+        schema: ContextCompactionOutputSchema,
+        maxOutputTokens: 8_000,
+      }),
+    ).pipe(Effect.flatMap(parseCompactionOutput)),
+  )
+}
+
+export function createWiredContextCompactionRuntime(deps: CreateContextCompactionRuntimeDeps) {
+  const { helperModelRuntime } = deps
+
+  const runtime = createContextCompactionRuntime({
+    runCompacter: (params) => runContextCompacter(helperModelRuntime, params),
+    now: deps.now,
+    randomId: deps.randomId,
+    thresholdRatio: CONTEXT_COMPACTION_THRESHOLD_RATIO,
+    outputReserveTokens: CONTEXT_OUTPUT_RESERVE_TOKENS,
+    safetyMarginTokens: CONTEXT_SAFETY_MARGIN_TOKENS,
+    compactionChunkMaxChars: COMPACTION_CHUNK_MAX_CHARS,
+    summaryRollupMaxTokens: SUMMARY_ROLLUP_MAX_TOKENS,
+    includedToolNames: CONTEXT_COMPACTION_INCLUDED_TOOL_NAMES,
+    includedToolPrefixes: CONTEXT_COMPACTION_INCLUDED_TOOL_PREFIXES,
+  })
+
+  function compactMemoryBlockSummary(params: { previousSummary: string; newEntriesText: string }): Promise<string> {
+    const previousSummary = params.previousSummary.trim()
+    const newEntriesText = params.newEntriesText.trim()
+    if (!previousSummary && !newEntriesText) return Promise.resolve('')
+
+    return Effect.runPromise(
+      tryContextCompactionPromise('Failed to compact memory block summary.', () =>
+        helperModelRuntime.generateHelperText({
+          tag: 'memory-block-compaction',
+          createAgent: createContextCompactionAgent,
+          messages: [{ role: 'user', content: buildMemoryBlockCompactionPrompt({ previousSummary, newEntriesText }) }],
+          maxOutputTokens: CONTEXT_COMPACTION_MAX_OUTPUT_TOKENS,
+        }),
+      ),
+    )
+  }
+
+  return { ...runtime, compactMemoryBlockSummary }
+}

package/src/runtime/{context-compaction.ts → context-compaction/context-compaction.ts}

@@ -1,7 +1,16 @@
 import type { ChatMessage } from '@lota-sdk/shared'
+import { Duration, Effect, Schedule, Schema } from 'effect'
 import { z } from 'zod'
 
-import { CHARS_PER_TOKEN_ESTIMATE, compactWhitespace, readRecord, readString, stringifyUnknown } from '../utils/string'
+import { iterateEffect } from '../../effect/helpers'
+import { nowEpochMillis } from '../../utils/date-time'
+import {
+  CHARS_PER_TOKEN_ESTIMATE,
+  compactWhitespace,
+  readRecord,
+  readString,
+  stringifyUnknown,
+} from '../../utils/string'
 import {
   COMPACTION_CHUNK_MAX_CHARS,
   CONTEXT_COMPACTION_INCLUDED_TOOL_NAMES,
@@ -51,6 +60,18 @@ export interface ContextCompactionRunnerParams {
   transcript: string
 }
 
+class CompactionError extends Schema.TaggedErrorClass<CompactionError>()('CompactionError', {
+  message: Schema.String,
+  cause: Schema.optional(Schema.Defect),
+}) {}
+
+class CompactionParseError extends Schema.TaggedErrorClass<CompactionParseError>()('CompactionParseError', {
+  message: Schema.String,
+  cause: Schema.optional(Schema.Defect),
+}) {}
+
+const COMPACTION_RUNNER_RETRY_OPTIONS = { times: 2, schedule: Schedule.exponential(Duration.millis(500), 2) } as const
+
 export interface ContextCompactionPromptParams {
   previousSummary: string
   transcript: string
@@ -216,8 +237,8 @@ export function buildMemoryBlockCompactionPrompt(params: MemoryBlockCompactionPr
 export function createContextCompactionRuntime(
   options: CreateContextCompactionRuntimeOptions,
 ): ContextCompactionRuntime {
-  const now = options.now ?? (() => Date.now())
-  const randomId = options.randomId ?? (() => crypto.randomUUID())
+  const now = options.now ?? nowEpochMillis
+  const randomId = options.randomId ?? (() => Bun.randomUUIDv7())
   const thresholdRatio = options.thresholdRatio ?? CONTEXT_COMPACTION_THRESHOLD_RATIO
   const outputReserveTokens = options.outputReserveTokens ?? CONTEXT_OUTPUT_RESERVE_TOKENS
   const safetyMarginTokens = options.safetyMarginTokens ?? CONTEXT_SAFETY_MARGIN_TOKENS
@@ -312,114 +333,160 @@ export function createContextCompactionRuntime(
     return { estimatedTokens, threshold, shouldCompact: estimatedTokens >= threshold }
   }
 
-  const compactContextMessages = async (params: {
+  const compactContextMessagesEffect = Effect.fn('ContextCompaction.compactContextMessages')(function* (params: {
     previousSummary: string
     newMessages: ContextMessage[]
-  }): Promise<CompactionOutput> => {
+  }) {
     const chunks = splitByCharBudget(params.newMessages, compactionChunkMaxChars)
-    let summary = normalizeSummary(params.previousSummary)
-
-    for (const chunk of chunks) {
-      const transcript = toCompactionTranscript(chunk)
-      const output = await options.runCompacter({ previousSummary: summary, chunk, transcript })
-      summary = normalizeSummary(output.summary)
-    }
+    const initialSummary = normalizeSummary(params.previousSummary)
+
+    const finalSummary = yield* iterateEffect<{ summary: string; index: number }, CompactionError, never>(
+      { summary: initialSummary, index: 0 },
+      {
+        while: (state) => state.index < chunks.length,
+        body: (state) =>
+          Effect.gen(function* () {
+            const chunk = chunks[state.index]
+            const transcript = toCompactionTranscript(chunk)
+            const output = yield* Effect.tryPromise({
+              try: () => options.runCompacter({ previousSummary: state.summary, chunk, transcript }),
+              catch: (error: unknown) => new CompactionError({ message: String(error), cause: error }),
+            }).pipe(Effect.retry(COMPACTION_RUNNER_RETRY_OPTIONS))
+            return { summary: normalizeSummary(output.summary), index: state.index + 1 }
+          }),
+      },
+    )
+
+    return { summary: finalSummary.summary }
+  })
 
-    return { summary }
-  }
+  const rollupSummaryIfOversizedEffect = (summary: string) =>
+    Effect.gen(function* () {
+      if (estimateTokens(summary) <= summaryRollupMaxTokens) {
+        return summary
+      }
 
-  const rollupSummaryIfOversized = async (summary: string): Promise<string> => {
-    if (estimateTokens(summary) <= summaryRollupMaxTokens) {
-      return summary
-    }
+      const output = yield* compactContextMessagesEffect({
+        previousSummary: '',
+        newMessages: [{ role: 'assistant', text: summary, sourceMessageId: 'summary-rollup' }],
+      })
 
-    const output = await compactContextMessages({
-      previousSummary: '',
-      newMessages: [{ role: 'assistant', text: summary, sourceMessageId: 'summary-rollup' }],
+      return normalizeSummary(output.summary)
     })
 
-    return normalizeSummary(output.summary)
+  type CompactionLoopState = {
+    summaryText: string
+    remainingMessages: ChatMessage[]
+    compactedMessages: ChatMessage[]
+    lastCompactedMessageId: string | undefined
+    estimatedTokens: number
+    done: boolean
   }
 
-  const compactHistory = async (params: CompactHistoryParams): Promise<CompactHistoryResult> => {
-    let summaryText = normalizeSummary(params.summaryText)
-    let remainingMessages = [...params.liveMessages]
-    let compactedMessages: ChatMessage[] = []
-    let lastCompactedMessageId: string | undefined
-    const summaryPayload = buildSyntheticSummaryPayload(summaryText)
-    const initialPayload = JSON.stringify([...(summaryPayload ? [summaryPayload] : []), ...remainingMessages])
-    const inputChars = initialPayload.length
-
-    const buildEarlyExitResult = (estimatedTokens: number): CompactHistoryResult => {
-      const exitSummaryPayload = buildSyntheticSummaryPayload(summaryText)
-      const outputPayload = JSON.stringify([...(exitSummaryPayload ? [exitSummaryPayload] : []), ...remainingMessages])
-      return {
-        compacted: compactedMessages.length > 0,
-        summaryText,
-        ...(lastCompactedMessageId ? { lastCompactedMessageId } : {}),
-        compactedMessages,
-        compactedMessageCount: compactedMessages.length,
-        remainingMessageCount: remainingMessages.length,
-        estimatedTokens,
-        inputChars,
-        outputChars: outputPayload.length,
-      }
-    }
-
-    for (;;) {
-      const assessment = shouldCompactHistory({
-        summaryText,
-        liveMessages: remainingMessages,
-        contextSize: params.contextSize,
-      })
-
-      if (!assessment.shouldCompact) {
-        return buildEarlyExitResult(assessment.estimatedTokens)
-      }
-
-      const boundary = Math.max(0, remainingMessages.length - params.tailMessageCount)
-      if (boundary <= 0) {
-        return buildEarlyExitResult(assessment.estimatedTokens)
+  const compactHistoryEffect = (params: CompactHistoryParams) =>
+    Effect.gen(function* () {
+      const initialSummaryText = normalizeSummary(params.summaryText)
+      const initialRemaining = [...params.liveMessages]
+      const summaryPayload = buildSyntheticSummaryPayload(initialSummaryText)
+      const initialPayload = stringifyUnknown([...(summaryPayload ? [summaryPayload] : []), ...initialRemaining]) ?? ''
+      const inputChars = initialPayload.length
+
+      const buildExitResult = (state: CompactionLoopState): CompactHistoryResult => {
+        const exitSummaryPayload = buildSyntheticSummaryPayload(state.summaryText)
+        const outputPayload =
+          stringifyUnknown([...(exitSummaryPayload ? [exitSummaryPayload] : []), ...state.remainingMessages]) ?? ''
+        return {
+          compacted: state.compactedMessages.length > 0,
+          summaryText: state.summaryText,
+          ...(state.lastCompactedMessageId ? { lastCompactedMessageId: state.lastCompactedMessageId } : {}),
+          compactedMessages: state.compactedMessages,
+          compactedMessageCount: state.compactedMessages.length,
+          remainingMessageCount: state.remainingMessages.length,
+          estimatedTokens: state.estimatedTokens,
+          inputChars,
+          outputChars: outputPayload.length,
+        }
       }
 
-      const candidatePrefix = remainingMessages.slice(0, boundary)
-      const messagesToCompact = candidatePrefix.filter((message) => !readIsCompacted(message))
-      const contextMessages = messagesToCompact
-        .map(toContextMessageFromChatMessage)
-        .filter((message) => compactWhitespace(message.text).length > 0)
-      const sourceText = toCompactionTranscript(contextMessages)
-
-      if (!compactWhitespace(sourceText)) {
-        return buildEarlyExitResult(assessment.estimatedTokens)
+      const initialState: CompactionLoopState = {
+        summaryText: initialSummaryText,
+        remainingMessages: initialRemaining,
+        compactedMessages: [],
+        lastCompactedMessageId: undefined,
+        estimatedTokens: 0,
+        done: false,
       }
 
-      let nextSummary = normalizeSummary(
-        (await compactContextMessages({ previousSummary: summaryText, newMessages: contextMessages })).summary,
-      )
-      nextSummary = await rollupSummaryIfOversized(nextSummary)
-
-      if (nextSummary.length >= sourceText.length) {
-        throw new Error('Compaction summary is not shorter than compacted source')
-      }
+      const finalState = yield* iterateEffect<CompactionLoopState, CompactionError, never>(initialState, {
+        while: (state) => !state.done,
+        body: (state) =>
+          Effect.gen(function* () {
+            const assessment = shouldCompactHistory({
+              summaryText: state.summaryText,
+              liveMessages: state.remainingMessages,
+              contextSize: params.contextSize,
+            })
+
+            if (!assessment.shouldCompact) {
+              return { ...state, estimatedTokens: assessment.estimatedTokens, done: true }
+            }
+
+            const boundary = Math.max(0, state.remainingMessages.length - params.tailMessageCount)
+            if (boundary <= 0) {
+              return { ...state, estimatedTokens: assessment.estimatedTokens, done: true }
+            }
+
+            const candidatePrefix = state.remainingMessages.slice(0, boundary)
+            const messagesToCompact = candidatePrefix.filter((message) => !readIsCompacted(message))
+            const contextMessages = messagesToCompact
+              .map(toContextMessageFromChatMessage)
+              .filter((message) => compactWhitespace(message.text).length > 0)
+            const sourceText = toCompactionTranscript(contextMessages)
+
+            if (!compactWhitespace(sourceText)) {
+              return { ...state, estimatedTokens: assessment.estimatedTokens, done: true }
+            }
+
+            const compacted = yield* compactContextMessagesEffect({
+              previousSummary: state.summaryText,
+              newMessages: contextMessages,
+            })
+            const rolledSummary = yield* rollupSummaryIfOversizedEffect(normalizeSummary(compacted.summary))
+
+            if (rolledSummary.length >= sourceText.length) {
+              return yield* new CompactionError({ message: 'Compaction summary is not shorter than compacted source' })
+            }
+
+            return {
+              summaryText: rolledSummary,
+              remainingMessages: state.remainingMessages.slice(boundary),
+              compactedMessages: [
+                ...state.compactedMessages,
+                ...candidatePrefix.map((message) => markMessageCompacted(message, now)),
+              ],
+              lastCompactedMessageId: candidatePrefix.at(-1)?.id ?? state.lastCompactedMessageId,
+              estimatedTokens: assessment.estimatedTokens,
+              done: false,
+            }
+          }),
+      })
 
-      summaryText = nextSummary
-      compactedMessages = [
-        ...compactedMessages,
-        ...candidatePrefix.map((message) => markMessageCompacted(message, now)),
-      ]
-      lastCompactedMessageId = candidatePrefix.at(-1)?.id ?? lastCompactedMessageId
-      remainingMessages = remainingMessages.slice(boundary)
+      return buildExitResult(finalState)
+    })
 
-      if (remainingMessages.length <= params.tailMessageCount) {
-        continue
-      }
-    }
+  return {
+    createSummaryMessage,
+    prependSummaryMessage,
+    estimateThreshold,
+    shouldCompactHistory,
+    compactHistory: (params) => Effect.runPromise(compactHistoryEffect(params)),
   }
-
-  return { createSummaryMessage, prependSummaryMessage, estimateThreshold, shouldCompactHistory, compactHistory }
 }
 
-export function parseCompactionOutput(value: unknown): CompactionOutput {
-  const parsed = ContextCompactionOutputSchema.parse(value)
-  return { summary: parsed.summary }
+export function parseCompactionOutput(value: unknown): Effect.Effect<CompactionOutput, CompactionParseError> {
+  const parsed = ContextCompactionOutputSchema.safeParse(value)
+  if (!parsed.success) {
+    return Effect.fail(new CompactionParseError({ message: 'Failed to parse compaction output', cause: parsed.error }))
+  }
+  return Effect.succeed({ summary: parsed.data.summary })
 }