288 lines · 11.0 KB · TypeScript
RawDownload
import {  MessageId,  ProviderSessionId,  ProviderThreadId,  ProviderTurnId,  RunAttemptId,  ThreadId,} from "@t3tools/contracts";import * as Context from "effect/Context";import * as Effect from "effect/Effect";import * as Layer from "effect/Layer";import * as Option from "effect/Option";import * as Schema from "effect/Schema"; import { ProjectionStoreV2 } from "./ProjectionStore.ts";import { ProviderSessionManagerV2 } from "./ProviderSessionManager.ts"; const yieldToRuntime = Effect.yieldNow.pipe(  Effect.andThen(    Effect.promise(      () =>        new Promise<void>((resolve) => {          setImmediate(resolve);        }),    ),  ),); export class ProviderTurnControlError extends Schema.TaggedErrorClass<ProviderTurnControlError>()(  "ProviderTurnControlError",  {    threadId: ThreadId,    operation: Schema.Literals(["interrupt", "restart", "steer"]),    providerTurnId: ProviderTurnId,    cause: Schema.optional(Schema.Defect()),  },) {} const isProviderTurnControlError = Schema.is(ProviderTurnControlError); export interface ProviderTurnControlServiceV2Shape {  readonly interrupt: (input: {    readonly threadId: ThreadId;    readonly providerSessionId: ProviderSessionId;    readonly providerThreadId: ProviderThreadId;    readonly providerTurnId: ProviderTurnId;  }) => Effect.Effect<void, ProviderTurnControlError>;  readonly steer: (input: {    readonly threadId: ThreadId;    readonly providerSessionId: ProviderSessionId;    readonly providerThreadId: ProviderThreadId;    readonly providerTurnId: ProviderTurnId;    readonly messageId: MessageId;  }) => Effect.Effect<void, ProviderTurnControlError>;  readonly interruptAndAwaitTerminal: (input: {    readonly threadId: ThreadId;    readonly providerSessionId: ProviderSessionId;    readonly replacementProviderSessionId?: ProviderSessionId;    readonly providerThreadId: ProviderThreadId;    readonly providerTurnId: ProviderTurnId;    readonly interruptedAttemptId: RunAttemptId;  }) => Effect.Effect<void, ProviderTurnControlError>;} export class ProviderTurnControlServiceV2 extends Context.Service<  ProviderTurnControlServiceV2,  ProviderTurnControlServiceV2Shape>()("t3/orchestration-v2/ProviderTurnControlService/ProviderTurnControlServiceV2") {} export const layer: Layer.Layer<  ProviderTurnControlServiceV2,  never,  ProjectionStoreV2 | ProviderSessionManagerV2> = Layer.effect(  ProviderTurnControlServiceV2,  Effect.gen(function* () {    const projections = yield* ProjectionStoreV2;    const sessions = yield* ProviderSessionManagerV2;     const load = (input: {      readonly threadId: ThreadId;      readonly providerSessionId: ProviderSessionId;      readonly replacementProviderSessionId?: ProviderSessionId;      readonly providerThreadId: ProviderThreadId;      readonly providerTurnId: ProviderTurnId;      readonly operation: "interrupt" | "restart" | "steer";      readonly messageId?: MessageId;    }) =>      Effect.gen(function* () {        const context = yield* projections.getProviderControlContext(input.threadId, input);        const { providerThread, providerTurn } = context;        const targetsRecordedSession =          providerThread?.providerSessionId === input.providerSessionId;        const targetsCommittedReplacement =          input.operation === "restart" &&          input.replacementProviderSessionId !== undefined &&          providerThread?.providerSessionId === input.replacementProviderSessionId;        if (          providerThread === undefined ||          providerTurn === undefined ||          (!targetsRecordedSession && !targetsCommittedReplacement) ||          providerTurn.providerThreadId !== providerThread.id        ) {          return yield* new ProviderTurnControlError({            threadId: input.threadId,            operation: input.operation,            providerTurnId: input.providerTurnId,            cause: "The recorded provider execution target is no longer valid.",          });        }        // A restart-session command commits the replacement binding before its        // process-bound effect runs. The old live runtime must still receive        // the interrupt, but only when the projection matches the exact        // replacement captured by that same durable effect.        const interruptProviderThread = targetsRecordedSession          ? providerThread          : { ...providerThread, providerSessionId: input.providerSessionId };        if (providerTurn.status !== "running") {          return {            context,            providerThread: interruptProviderThread,            providerTurn,            session: Option.none(),          };        }        const session = yield* sessions.get(input.providerSessionId);        if (Option.isNone(session)) {          // Interrupt/restart against a already-released session must not fail          // the durable effect (and retry 5x). The turn may still look running          // in projection until recovery/finalization; there is no live adapter          // to interrupt.          if (input.operation === "interrupt" || input.operation === "restart") {            yield* Effect.logWarning(              "Provider interrupt/restart found no live session; treating as already stopped",              {                threadId: input.threadId,                operation: input.operation,                providerSessionId: input.providerSessionId,                providerTurnId: input.providerTurnId,                providerTurnStatus: providerTurn.status,              },            );            return {              context,              providerThread: interruptProviderThread,              providerTurn,              session: Option.none(),            };          }          return yield* new ProviderTurnControlError({            threadId: input.threadId,            operation: input.operation,            providerTurnId: input.providerTurnId,            cause: `Provider session ${input.providerSessionId} is not active.`,          });        }        return { context, providerThread: interruptProviderThread, providerTurn, session };      });     return ProviderTurnControlServiceV2.of({      interrupt: (input) =>        Effect.gen(function* () {          const loaded = yield* load({ ...input, operation: "interrupt" });          if (Option.isNone(loaded.session)) return;          yield* loaded.session.value.interruptTurn({            providerThread: loaded.providerThread,            providerTurnId: loaded.providerTurn.id,            requestRuntimeRestart: true,          });        }).pipe(          Effect.mapError((cause) =>            isProviderTurnControlError(cause)              ? cause              : new ProviderTurnControlError({                  threadId: input.threadId,                  operation: "interrupt",                  providerTurnId: input.providerTurnId,                  cause,                }),          ),        ),      interruptAndAwaitTerminal: (input) =>        Effect.gen(function* () {          const loaded = yield* load({ ...input, operation: "restart" });          if (Option.isNone(loaded.session)) {            // No live adapter: nothing can emit a terminal provider-turn update            // from interrupt. Do not poll for projection terminalization or the            // restart effect stalls; let detach/start proceed.            if (loaded.providerTurn.status === "running") {              yield* Effect.logWarning(                "Provider restart interrupt skipped; no live session for a still-projected running turn",                {                  threadId: input.threadId,                  providerSessionId: input.providerSessionId,                  providerTurnId: input.providerTurnId,                },              );            }            return;          }           yield* loaded.session.value.interruptTurn({            providerThread: loaded.providerThread,            providerTurnId: loaded.providerTurn.id,          });           for (let remaining = 1_000; remaining > 0; remaining -= 1) {            const { providerTurn, attempt } = yield* projections.getProviderControlContext(              input.threadId,              {                providerThreadId: input.providerThreadId,                providerTurnId: input.providerTurnId,                attemptId: input.interruptedAttemptId,              },            );            if (              providerTurn !== undefined &&              providerTurn.status !== "running" &&              attempt !== undefined &&              attempt.status !== "running"            ) {              return;            }            // Provider terminal events are projected on a detached ingestion            // fiber. Yield through the Node event loop instead of sleeping on            // Effect's clock so deterministic runtimes cannot deadlock a            // command that is waiting for that projection.            yield* yieldToRuntime;          }          return yield* new ProviderTurnControlError({            threadId: input.threadId,            operation: "restart",            providerTurnId: input.providerTurnId,            cause: `Provider turn ${input.providerTurnId} did not terminalize before restart.`,          });        }).pipe(          Effect.mapError((cause) =>            isProviderTurnControlError(cause)              ? cause              : new ProviderTurnControlError({                  threadId: input.threadId,                  operation: "restart",                  providerTurnId: input.providerTurnId,                  cause,                }),          ),        ),      steer: (input) =>        Effect.gen(function* () {          const loaded = yield* load({ ...input, operation: "steer" });          if (Option.isNone(loaded.session)) return;          const { message, run } = loaded.context;          if (message === undefined || run === undefined) {            return yield* new ProviderTurnControlError({              threadId: input.threadId,              operation: "steer",              providerTurnId: input.providerTurnId,              cause: "The persisted steering message or target run is missing.",            });          }          yield* loaded.session.value.steerTurn({            threadId: input.threadId,            runId: run.id,            providerThread: loaded.providerThread,            providerTurnId: loaded.providerTurn.id,            message: {              messageId: message.id,              text: message.text,              attachments: message.attachments,              createdBy: message.createdBy,              creationSource: message.creationSource,            },          });        }).pipe(          Effect.mapError((cause) =>            isProviderTurnControlError(cause)              ? cause              : new ProviderTurnControlError({                  threadId: input.threadId,                  operation: "steer",                  providerTurnId: input.providerTurnId,                  cause,                }),          ),        ),    });  }),);