@yolk-sdk/harness 0.1.0-canary.77

package/src/outcome.ts

@@ -0,0 +1,316 @@
+import { Effect, Stream } from 'effect'
+import {
+  addAgentUsage,
+  zeroAgentUsage,
+  type AgentEvent,
+  type AgentMessage,
+  type AgentUsage,
+  type HitlRequest,
+  type HitlResponse
+} from '@yolk-sdk/agent/protocol'
+import {
+  collectModelTurnAttempt,
+  LLMError,
+  runModelTurn,
+  runToolBatch,
+  type AgentLoopError,
+  type ContextTransformer,
+  type LLMProvider,
+  type LoopConfig,
+  type ModelTurnConfig,
+  type ModelTurnResult,
+  type ToolBatchConfig,
+  type ToolExecutor
+} from '@yolk-sdk/agent/loop'
+import {
+  applyOverflowCompaction,
+  isOverflowCompactionAttemptCount
+} from '@yolk-sdk/agent/compaction'
+
+export type OverflowCompactionResult =
+  | {
+      readonly _tag: 'Compacted'
+      readonly messages: ReadonlyArray<AgentMessage>
+    }
+  | {
+      readonly _tag: 'Skipped'
+    }
+
+export type CompletedTurn = {
+  readonly _tag: 'Completed'
+  readonly needsContinuation: boolean
+} & ModelTurnResult
+
+export type ModelTurnOutcome =
+  | CompletedTurn
+  | { readonly _tag: 'Retry'; readonly error: AgentLoopError }
+  | ({ readonly _tag: 'Continue'; readonly error: AgentLoopError } & ModelTurnResult)
+  | { readonly _tag: 'RecoverFull'; readonly error: AgentLoopError }
+  | {
+      readonly _tag: 'Compacted'
+      readonly messages: ReadonlyArray<AgentMessage>
+      readonly overflowCompactionAttempt: number
+    }
+
+export type ToolBatchOutcome =
+  | CompletedTurn
+  | {
+      readonly _tag: 'AwaitingInput'
+      readonly requests: ReadonlyArray<HitlRequest>
+      readonly usage: AgentUsage
+    }
+
+export type StepOutcome = ModelTurnOutcome | ToolBatchOutcome
+
+const completedOutcome = (result: ModelTurnResult): CompletedTurn => ({
+  _tag: 'Completed',
+  needsContinuation: result.stopReason === 'tool_use',
+  assistantMessage: result.assistantMessage,
+  toolCalls: result.toolCalls,
+  usage: result.usage,
+  stopReason: result.stopReason
+})
+
+const modelTurnResult = (result: ModelTurnResult): ModelTurnResult => ({
+  assistantMessage: result.assistantMessage,
+  toolCalls: result.toolCalls,
+  usage: result.usage,
+  stopReason: result.stopReason
+})
+
+const invalidOverflowCompactionAttemptError = () =>
+  new LLMError({
+    cause: 'validation_error',
+    message: 'overflowCompactionAttempt must be a finite integer >= 0',
+    retryable: false
+  })
+
+const isAgentLoopError = (error: unknown): error is AgentLoopError => {
+  if (typeof error !== 'object' || error === null || !('_tag' in error)) {
+    return false
+  }
+  const tag = error._tag
+  return (
+    tag === 'LLMError' ||
+    tag === 'FauxExhaustedError' ||
+    tag === 'ToolError' ||
+    tag === 'ContextTransformError' ||
+    tag === 'AbortError'
+  )
+}
+
+const isOverflow = (error: AgentLoopError) =>
+  (error._tag === 'LLMError' || error._tag === 'ContextTransformError') &&
+  error.cause === 'context_overflow'
+
+const isMissingDone = (error: AgentLoopError) =>
+  error._tag === 'LLMError' && error.responseIssue === 'missing_done'
+
+const isRetryableLlm = (
+  error: AgentLoopError
+): error is Extract<AgentLoopError, { _tag: 'LLMError' }> =>
+  error._tag === 'LLMError' && error.retryable
+
+const classifyModelTurnFailure = <E2, R2>(input: {
+  readonly error: AgentLoopError
+  readonly collected: ModelTurnResult
+  readonly outputStarted: boolean
+  readonly messages: ReadonlyArray<AgentMessage>
+  readonly overflowCompactionAttempt: number
+  readonly compact?: (
+    messages: ReadonlyArray<AgentMessage>
+  ) => Effect.Effect<OverflowCompactionResult, E2, R2>
+}): Effect.Effect<ModelTurnOutcome, AgentLoopError | E2, R2> => {
+  if (isOverflow(input.error)) {
+    if (input.compact === undefined) {
+      return Effect.fail(input.error)
+    }
+
+    return applyOverflowCompaction({
+      compact: input.compact,
+      messages: input.messages,
+      attempt: input.overflowCompactionAttempt,
+      outputStarted: input.outputStarted
+    }).pipe(
+      Effect.flatMap(result =>
+        result._tag === 'Compacted'
+          ? Effect.succeed({
+              _tag: 'Compacted' as const,
+              messages: result.messages,
+              overflowCompactionAttempt: input.overflowCompactionAttempt + 1
+            })
+          : Effect.fail(input.error)
+      )
+    )
+  }
+
+  if (isMissingDone(input.error) && !input.outputStarted) {
+    return Effect.succeed({ _tag: 'RecoverFull', error: input.error })
+  }
+
+  if (isMissingDone(input.error) && input.outputStarted) {
+    return Effect.succeed({
+      _tag: 'Continue',
+      error: input.error,
+      ...modelTurnResult(input.collected)
+    })
+  }
+
+  if (isRetryableLlm(input.error) && !input.outputStarted) {
+    if (input.error.cause === 'invalid_response') {
+      return Effect.succeed({ _tag: 'RecoverFull', error: input.error })
+    }
+    return Effect.succeed({ _tag: 'Retry', error: input.error })
+  }
+
+  if (isRetryableLlm(input.error) && input.outputStarted) {
+    return Effect.succeed({
+      _tag: 'Continue',
+      error: input.error,
+      ...modelTurnResult(input.collected)
+    })
+  }
+
+  return Effect.fail(input.error)
+}
+
+export const attemptModelTurn = <E2 = never, R2 = never>(
+  config: ModelTurnConfig,
+  options?: {
+    readonly onEvent?: (event: AgentEvent) => Effect.Effect<void, E2, R2>
+    readonly initialUsage?: AgentUsage
+    readonly overflowCompactionAttempt?: number
+    readonly compact?: (
+      messages: ReadonlyArray<AgentMessage>
+    ) => Effect.Effect<OverflowCompactionResult, E2, R2>
+  }
+): Effect.Effect<
+  ModelTurnOutcome,
+  AgentLoopError | E2,
+  ContextTransformer | LLMProvider | LoopConfig | R2
+> =>
+  Effect.gen(function* () {
+    const overflowCompactionAttempt = options?.overflowCompactionAttempt ?? 0
+    if (!isOverflowCompactionAttemptCount(overflowCompactionAttempt)) {
+      return yield* Effect.fail(invalidOverflowCompactionAttemptError())
+    }
+
+    const outcome = yield* collectModelTurnAttempt(runModelTurn(config), {
+      onEvent: options?.onEvent,
+      initialUsage: options?.initialUsage
+    })
+
+    switch (outcome._tag) {
+      case 'Collected':
+        return completedOutcome(outcome.collection)
+      case 'SinkFailed':
+        return yield* Effect.fail(outcome.error)
+      case 'StreamFailed':
+        if (!isAgentLoopError(outcome.error)) {
+          return yield* Effect.fail(outcome.error)
+        }
+        return yield* classifyModelTurnFailure({
+          error: outcome.error,
+          collected: {
+            assistantMessage:
+              outcome.collection.assistantMessage ?? outcome.collection.partialAssistantMessage,
+            toolCalls: outcome.collection.toolCalls,
+            usage: outcome.collection.usage,
+            stopReason: outcome.collection.stopReason
+          },
+          outputStarted: outcome.collection.outputStarted,
+          messages: config.messages,
+          overflowCompactionAttempt,
+          compact: options?.compact
+        })
+    }
+  })
+
+export const attemptToolBatch = <E2 = never, R2 = never>(
+  config: ToolBatchConfig,
+  options?: {
+    readonly onEvent?: (event: AgentEvent) => Effect.Effect<void, E2, R2>
+  }
+): Effect.Effect<ToolBatchOutcome, AgentLoopError | E2, LoopConfig | ToolExecutor | R2> => {
+  const onEvent = options?.onEvent
+
+  return runToolBatch(config).pipe(
+    Stream.runFoldEffect(
+      (): {
+        requests: ReadonlyArray<HitlRequest>
+        usage: AgentUsage
+        toolCalls: ModelTurnResult['toolCalls']
+      } => ({
+        requests: [],
+        usage: config.usage ?? zeroAgentUsage,
+        toolCalls: []
+      }),
+      (acc, event) => {
+        const next =
+          event._tag === 'AgentAwaitingInput'
+            ? { ...acc, requests: event.requests, usage: event.usage }
+            : event._tag === 'UsageUpdate'
+              ? { ...acc, usage: addAgentUsage(acc.usage, event.usage) }
+              : event._tag === 'ToolExecutionCompleted' || event._tag === 'ToolExecutionAccepted'
+                ? { ...acc, toolCalls: [...acc.toolCalls, event.call] }
+                : acc
+        return onEvent === undefined ? Effect.succeed(next) : onEvent(event).pipe(Effect.as(next))
+      }
+    ),
+    Effect.map((result): ToolBatchOutcome => {
+      if (result.requests.length === 0) {
+        const needsContinuation = result.toolCalls.length > 0
+        return {
+          _tag: 'Completed',
+          needsContinuation,
+          assistantMessage: undefined,
+          toolCalls: result.toolCalls,
+          usage: result.usage,
+          stopReason: needsContinuation ? 'tool_use' : 'stop'
+        }
+      }
+      return { _tag: 'AwaitingInput', requests: result.requests, usage: result.usage }
+    })
+  )
+}
+
+export type HitlMatch =
+  | { readonly _tag: 'Match'; readonly requestId: string }
+  | { readonly _tag: 'Mismatch' }
+
+const hitlResponseMatchesRequest = (response: HitlResponse, request: HitlRequest) => {
+  switch (response._tag) {
+    case 'ToolApprovalResponse':
+      return (
+        request._tag === 'ToolApprovalRequest' &&
+        response.requestId === request.requestId &&
+        response.toolCallId === request.toolCallId
+      )
+    case 'QuestionResponse':
+      return (
+        request._tag === 'QuestionRequest' &&
+        response.requestId === request.requestId &&
+        response.toolCallId === request.toolCallId
+      )
+  }
+}
+
+export const matchHitlResponse = (
+  pending: ReadonlyArray<HitlRequest>,
+  response: HitlResponse
+): HitlMatch => {
+  const matched = pending.find(request => hitlResponseMatchesRequest(response, request))
+  return matched === undefined
+    ? { _tag: 'Mismatch' }
+    : { _tag: 'Match', requestId: matched.requestId }
+}
+
+export const resumeHitlIfMatched = <A, E, R>(input: {
+  readonly pending: ReadonlyArray<HitlRequest>
+  readonly response: HitlResponse
+  readonly resume: (requestId: string) => Effect.Effect<A, E, R>
+}): Effect.Effect<A | { readonly _tag: 'Mismatch' }, E, R> => {
+  const matched = matchHitlResponse(input.pending, input.response)
+  if (matched._tag === 'Mismatch') return Effect.succeed(matched)
+  return input.resume(matched.requestId)
+}

package/src/store.ts

@@ -0,0 +1,137 @@
+import { Context, Effect, Layer, Ref, Semaphore } from 'effect'
+
+export type RunStoreShape = {
+  readonly claim: (runId: string) => Effect.Effect<void>
+  readonly release: (runId: string) => Effect.Effect<void>
+  readonly isClaimed: (runId: string) => Effect.Effect<boolean>
+  readonly claimed: Effect.Effect<ReadonlySet<string>>
+  readonly incrementResumeCount: (runId: string) => Effect.Effect<number>
+  readonly resumeCount: (runId: string) => Effect.Effect<number>
+}
+
+export class RunStore extends Context.Service<RunStore, RunStoreShape>()(
+  '@yolk-sdk/harness/RunStore'
+) {}
+
+export type DurableRunStoreSnapshot = {
+  readonly claimed: ReadonlyArray<string>
+  readonly resumes: ReadonlyArray<readonly [string, number]>
+}
+
+type MemorySnapshot = {
+  readonly claimed: ReadonlySet<string>
+  readonly resumes: ReadonlyMap<string, number>
+}
+
+const emptySnapshot: DurableRunStoreSnapshot = {
+  claimed: [],
+  resumes: []
+}
+
+const snapshotFromMemory = (memory: MemorySnapshot): DurableRunStoreSnapshot => ({
+  claimed: [...memory.claimed],
+  resumes: [...memory.resumes.entries()]
+})
+
+const memoryFromSnapshot = (snapshot: DurableRunStoreSnapshot): MemorySnapshot => ({
+  claimed: new Set(snapshot.claimed),
+  resumes: new Map(snapshot.resumes)
+})
+
+const makeSnapshotRunStore = (options: {
+  readonly load: Effect.Effect<DurableRunStoreSnapshot | undefined>
+  readonly save: (snapshot: DurableRunStoreSnapshot) => Effect.Effect<void>
+}): Effect.Effect<RunStore['Service']> =>
+  Effect.gen(function* () {
+    const loaded = yield* options.load
+    const memory = yield* Ref.make(memoryFromSnapshot(loaded ?? emptySnapshot))
+    const lock = yield* Semaphore.make(1)
+    const mutate = <A, E, R>(effect: Effect.Effect<A, E, R>) =>
+      Effect.uninterruptibleMask(restore =>
+        restore(lock.take(1)).pipe(
+          Effect.flatMap(() => effect.pipe(Effect.ensuring(lock.release(1))))
+        )
+      )
+    const commit = (next: MemorySnapshot) =>
+      options.save(snapshotFromMemory(next)).pipe(Effect.flatMap(() => Ref.set(memory, next)))
+
+    return RunStore.of({
+      claim: runId =>
+        mutate(
+          Effect.gen(function* () {
+            const current = yield* Ref.get(memory)
+            yield* commit({
+              claimed: new Set(current.claimed).add(runId),
+              resumes: current.resumes
+            })
+          })
+        ),
+      release: runId =>
+        mutate(
+          Effect.gen(function* () {
+            const current = yield* Ref.get(memory)
+            const claimed = new Set(current.claimed)
+            claimed.delete(runId)
+            const resumes = new Map(current.resumes)
+            resumes.delete(runId)
+            yield* commit({ claimed, resumes })
+          })
+        ),
+      isClaimed: runId => Ref.get(memory).pipe(Effect.map(current => current.claimed.has(runId))),
+      claimed: Ref.get(memory).pipe(Effect.map(current => new Set(current.claimed))),
+      incrementResumeCount: runId =>
+        mutate(
+          Effect.gen(function* () {
+            const current = yield* Ref.get(memory)
+            const nextCount = (current.resumes.get(runId) ?? 0) + 1
+            const resumes = new Map(current.resumes)
+            resumes.set(runId, nextCount)
+            yield* commit({ claimed: current.claimed, resumes })
+            return nextCount
+          })
+        ),
+      resumeCount: runId =>
+        Ref.get(memory).pipe(Effect.map(current => current.resumes.get(runId) ?? 0))
+    })
+  })
+
+export const makeSnapshotRunStoreLayer = (options: {
+  readonly load: Effect.Effect<DurableRunStoreSnapshot | undefined>
+  readonly save: (snapshot: DurableRunStoreSnapshot) => Effect.Effect<void>
+}): Layer.Layer<RunStore> => Layer.effect(RunStore, makeSnapshotRunStore(options))
+
+export const makeInMemoryRunStoreLayer = (): Layer.Layer<RunStore> =>
+  Layer.effect(
+    RunStore,
+    Effect.gen(function* () {
+      const claimed = yield* Ref.make(new Set<string>())
+      const resumes = yield* Ref.make(new Map<string, number>())
+
+      return RunStore.of({
+        claim: runId => Ref.update(claimed, current => new Set(current).add(runId)),
+        release: runId =>
+          Effect.zip(
+            Ref.update(claimed, current => {
+              const next = new Set(current)
+              next.delete(runId)
+              return next
+            }),
+            Ref.update(resumes, current => {
+              const next = new Map(current)
+              next.delete(runId)
+              return next
+            })
+          ).pipe(Effect.asVoid),
+        isClaimed: runId => Ref.get(claimed).pipe(Effect.map(current => current.has(runId))),
+        claimed: Ref.get(claimed).pipe(Effect.map(current => new Set(current))),
+        incrementResumeCount: runId =>
+          Ref.modify(resumes, current => {
+            const nextCount = (current.get(runId) ?? 0) + 1
+            const next = new Map(current)
+            next.set(runId, nextCount)
+            return [nextCount, next] as const
+          }),
+        resumeCount: runId => Ref.get(resumes).pipe(Effect.map(current => current.get(runId) ?? 0))
+      })
+    })
+  )