758 lines · 26.8 KB · TypeScript
RawDownload
import type {  EnvironmentId,  OrchestrationV2DomainEvent,  OrchestrationV2ThreadShell,  Project,  ThreadId,} from "@t3tools/contracts";import {  RelayApi,  type RelayAgentActivityPublishProofPayload,  type RelayAgentActivityState,} from "@t3tools/contracts/relay";import { projectThreadAwarenessV2 } from "@t3tools/shared/agentAwareness";import { makeDrainableWorker } from "@t3tools/shared/DrainableWorker";import { withRelayClientTracing } from "@t3tools/shared/relayTracing";import {  normalizeRelayIssuer,  RELAY_ACTIVITY_PUBLISH_TYP,  signRelayJwt,} from "@t3tools/shared/relayJwt";import * as Cause from "effect/Cause";import * as Context from "effect/Context";import * as Crypto from "effect/Crypto";import * as DateTime from "effect/DateTime";import * as Effect from "effect/Effect";import * as Fiber from "effect/Fiber";import * as Layer from "effect/Layer";import * as Option from "effect/Option";import * as Ref from "effect/Ref";import type * as Scope from "effect/Scope";import * as Stream from "effect/Stream";import * as FetchHttpClient from "effect/unstable/http/FetchHttpClient";import * as HttpClient from "effect/unstable/http/HttpClient";import * as HttpClientRequest from "effect/unstable/http/HttpClientRequest";import * as HttpApiClient from "effect/unstable/httpapi/HttpApiClient"; import * as ServerSecretStore from "../auth/ServerSecretStore.ts";import {  isAgentActivityPublishingEnabledValue,  PUBLISH_AGENT_ACTIVITY_SECRET,  RELAY_ENVIRONMENT_CREDENTIAL_SECRET,  RELAY_ISSUER_SECRET,  RELAY_URL_SECRET,} from "../cloud/config.ts";import { getOrCreateEnvironmentKeyPairFromSecretStore } from "../cloud/environmentKeys.ts";import * as ServerEnvironment from "../environment/ServerEnvironment.ts";import * as ThreadManagement from "../orchestration-v2/ThreadManagementService.ts";import * as ProjectService from "../project/ProjectService.ts";import { forkParked } from "../serverActivation.ts"; export class AgentAwarenessRelay extends Context.Service<  AgentAwarenessRelay,  {    readonly publishThread: (threadId: ThreadId) => Effect.Effect<void>;    readonly drain: Effect.Effect<void>;    readonly start: () => Effect.Effect<void, never, Scope.Scope>;  }>()("t3/relay/AgentAwarenessRelay") {} export function eventThreadId(event: OrchestrationV2DomainEvent): ThreadId {  return event.threadId;} export function shouldPublishAgentAwarenessEvent(  event: Pick<OrchestrationV2DomainEvent, "type">,): boolean {  // projectThreadAwarenessV2 reads thread metadata, run status, and pending requests.  // Message bodies and tool progress cannot change the published activity.  switch (event.type) {    case "thread.created":    case "thread.archived":    case "thread.unarchived":    case "thread.deleted":    case "thread.metadata-updated":    case "thread.model-selection-updated":    case "thread.provider-switched":    case "run.created":    case "run.updated":    case "runtime-request.updated":      return true;    case "thread.settled":    case "thread.unsettled":    case "thread.snoozed":    case "thread.unsnoozed":    case "thread.pinned":    case "thread.unpinned":    case "thread.pin-reordered":    case "thread.visited":    case "thread.marked-unread":    case "thread.runtime-mode-updated":    case "thread.interaction-mode-updated":    case "run-attempt.created":    case "run-attempt.updated":    case "node.updated":    case "subagent.updated":    case "provider-session.attached":    case "provider-session.updated":    case "provider-session.detached":    case "provider-thread.updated":    case "provider-turn.updated":    case "message.updated":    case "turn-item.updated":    case "plan.updated":    case "checkpoint-scope.created":    case "checkpoint.captured":    case "checkpoint.rollback-requested":    case "context-handoff.updated":    case "context-transfer.created":    case "context-transfer.updated":      return false;  }} export const makeAgentAwarenessPublishWorker = Effect.fnUntraced(function* <R>(  publish: (threadId: ThreadId) => Effect.Effect<void, never, R>,) {  const queued = new Set<ThreadId>();  const worker = yield* makeDrainableWorker((threadId: ThreadId) =>    Effect.sync(() => queued.delete(threadId)).pipe(Effect.andThen(publish(threadId))),  );  const enqueue = (threadId: ThreadId) =>    Effect.suspend(() => {      if (queued.has(threadId)) return Effect.void;      // Removing the ID when processing starts allows one new queue entry      // for updates received while the current snapshot is being published.      queued.add(threadId);      return worker.enqueue(threadId);    }).pipe(Effect.uninterruptible);  return { enqueue, drain: worker.drain };}); export function agentAwarenessPublishIdentity(state: RelayAgentActivityState | null): string {  if (state === null) {    return "null";  }  const { updatedAt: _updatedAt, ...meaningfulState } = state;  return JSON.stringify(meaningfulState);} export function isAgentActivityPublishingEnabled(value: string | null): boolean {  return isAgentActivityPublishingEnabledValue(value);} export function resolveAgentActivityPublishingStartupState(input: {  readonly relayConfigured: boolean;  readonly publishEnabled: boolean;}): "waiting-for-link" | "disabled" | "enabled" {  if (!input.relayConfigured) {    return "waiting-for-link";  }  return input.publishEnabled ? "enabled" : "disabled";} const RELAY_AGENT_ACTIVITY_DETAIL_MAX_LENGTH = 160;const REDACTED_RELAY_AGENT_FAILURE_DETAIL = "The agent run failed.";const RELAY_AGENT_ACTIVITY_RETRY_DELAYS_MS = [1_000, 2_000, 4_000, 8_000, 16_000] as const; export function sanitizeRelayAgentActivityState(  state: RelayAgentActivityState | null,): RelayAgentActivityState | null {  if (state === null) {    return null;  }  const { detail: _detail, ...rest } = state;  const detail = (state.phase === "failed" ? REDACTED_RELAY_AGENT_FAILURE_DETAIL : state.detail)    ?.trim()    .slice(0, RELAY_AGENT_ACTIVITY_DETAIL_MAX_LENGTH)    .trim();  return detail ? { ...rest, detail } : rest;} function relayEnvironmentClient(token: string) {  return HttpClient.mapRequest(HttpClientRequest.setHeader("authorization", `Bearer ${token}`));} function deliveryStats(  deliveries: ReadonlyArray<{    readonly ok: boolean;    readonly queued?: boolean | undefined;    readonly kind: string;    readonly apnsStatus?: number | null;    readonly apnsReason?: string | null;  }>,) {  let queued = 0;  let successful = 0;  let failed = 0;  const failedReasons: string[] = [];  const kinds = new Set<string>();   for (const delivery of deliveries) {    kinds.add(delivery.kind);    if (delivery.queued) {      queued += 1;      continue;    }    if (delivery.ok) {      successful += 1;      continue;    }    failed += 1;    failedReasons.push(`${delivery.apnsStatus ?? "transport"}:${delivery.apnsReason ?? "unknown"}`);  }   return {    total: deliveries.length,    queued,    successful,    failed,    kinds: [...kinds],    failedReasons,  };} export function signRelayAgentActivityPublishProof(input: {  readonly privateKey: string;  readonly payload: RelayAgentActivityPublishProofPayload;}) {  return signRelayJwt({    privateKey: input.privateKey,    typ: RELAY_ACTIVITY_PUBLISH_TYP,    payload: input.payload,  });} const makePublishProof = Effect.fn("makePublishProof")(function* (input: {  readonly privateKey: string;  readonly relayIssuer: string;  readonly environmentId: string;  readonly threadId: ThreadId;  readonly state: RelayAgentActivityState | null;  readonly jti: string;}) {  const now = yield* DateTime.now;  const expiresAt = DateTime.add(now, { minutes: 5 });  const payload = {    iss: `t3-env:${input.environmentId}`,    aud: normalizeRelayIssuer(input.relayIssuer),    sub: input.environmentId,    jti: input.jti,    iat: Math.floor(now.epochMilliseconds / 1_000),    exp: Math.floor(expiresAt.epochMilliseconds / 1_000),    environmentId: input.environmentId as RelayAgentActivityPublishProofPayload["environmentId"],    threadId: input.threadId,    state: input.state,  } satisfies RelayAgentActivityPublishProofPayload;  return yield* signRelayAgentActivityPublishProof({ privateKey: input.privateKey, payload });}); // Compact, log-safe view of the fields the awareness phase ladder reads.export function describeThreadShellForAwareness(  thread: Option.Option<OrchestrationV2ThreadShell>,): Record<string, unknown> {  if (Option.isNone(thread)) {    return { found: false };  }  const shell = thread.value;  return {    found: true,    status: shell.status,    activityRunStatus: shell.activityRunStatus ?? null,    activeRunId: shell.activeRunId ?? null,    latestRunId: shell.latestRunId ?? null,    pendingRuntimeRequestKind: shell.pendingRuntimeRequest?.kind ?? null,    hasActionableProposedPlan: shell.hasActionableProposedPlan,  };} export function resolveAgentAwarenessRelayPublishSnapshot(input: {  readonly environmentId: EnvironmentId;  readonly threadId: ThreadId;  readonly thread: Option.Option<OrchestrationV2ThreadShell>;  readonly project: Option.Option<Project>;}): {  readonly projectId: string | null;  readonly state: RelayAgentActivityState | null;  readonly reason: "snapshot" | "thread-not-found" | "project-not-found";} {  if (Option.isNone(input.thread)) {    return {      projectId: null,      state: null,      reason: "thread-not-found",    };  }  if (Option.isNone(input.project)) {    return {      projectId: input.thread.value.projectId,      state: null,      reason: "project-not-found",    };  }  return {    projectId: input.thread.value.projectId,    state: sanitizeRelayAgentActivityState(      projectThreadAwarenessV2({        environmentId: input.environmentId,        project: input.project.value,        thread: input.thread.value,      }),    ),    reason: "snapshot",  };} export function resolveAgentAwarenessRelayActiveThreadIds(input: {  readonly environmentId: EnvironmentId;  readonly projects: ReadonlyArray<Pick<Project, "id" | "title">>;  readonly threads: ReadonlyArray<OrchestrationV2ThreadShell>;}): ReadonlyArray<ThreadId> {  const projectById = new Map(input.projects.map((project) => [project.id, project]));  return input.threads    .filter((thread) => {      const project = projectById.get(thread.projectId);      if (!project) {        return false;      }      return (        projectThreadAwarenessV2({          environmentId: input.environmentId,          project,          thread,        }) !== null      );    })    .map((thread) => thread.id);} export const make = Effect.gen(function* () {  const secrets = yield* ServerSecretStore.ServerSecretStore;  const serverEnvironment = yield* ServerEnvironment.ServerEnvironment;  const threads = yield* ThreadManagement.ThreadManagementService;  const projects = yield* ProjectService.ProjectService;  const crypto = yield* Crypto.Crypto;  const scope = yield* Effect.scope;  const cloudLinkKeyPair = yield* getOrCreateEnvironmentKeyPairFromSecretStore(secrets);  const activeSnapshotPublishedRef = yield* Ref.make(false);  const publishedStateByThreadRef = yield* Ref.make(new Map<ThreadId, string>());   const readSecretString = (name: string) =>    secrets      .get(name)      .pipe(        Effect.map((bytes) =>          Option.isSome(bytes) ? new TextDecoder().decode(bytes.value) : null,        ),      );   const readRelayConfig = Effect.gen(function* () {    const [url, issuer, environmentCredential] = yield* Effect.all([      readSecretString(RELAY_URL_SECRET),      readSecretString(RELAY_ISSUER_SECRET),      readSecretString(RELAY_ENVIRONMENT_CREDENTIAL_SECRET),    ]);    return url && environmentCredential      ? { url, issuer: issuer ?? url, environmentCredential }      : null;  });   const readPublishAgentActivityEnabled = readSecretString(PUBLISH_AGENT_ACTIVITY_SECRET).pipe(    Effect.map(isAgentActivityPublishingEnabled),  );   const makeRelayClient = (relayConfig: {    readonly url: string;    readonly environmentCredential: string;  }) =>    HttpApiClient.make(RelayApi, {      baseUrl: relayConfig.url,      transformClient: relayEnvironmentClient(relayConfig.environmentCredential),    }).pipe(Effect.provide(FetchHttpClient.layer));   let cachedRelayClient: Effect.Success<ReturnType<typeof makeRelayClient>> | undefined;  let publishedRelayConfig: NonNullable<Effect.Success<typeof readRelayConfig>> | undefined;   // Deadlines for publishes that need confirmation (tombstones and  // first-state completions). The confirming publish is re-enqueued through  // the same drainable worker as every other publish, so a confirmed  // tombstone can never race an in-flight live update; a recovered state  // clears the deadline. Assigned after the worker exists.  const publishConfirmDeadlines = new Map<ThreadId, number>();  let schedulePublishConfirm: (threadId: ThreadId) => Effect.Effect<void> = () => Effect.void;  const publishRetries = new Map<    ThreadId,    { readonly attempts: number; timer: Fiber.Fiber<void> | undefined }  >();  const cancelPublishRetry = Effect.fnUntraced(function* (threadId: ThreadId) {    const retry = publishRetries.get(threadId);    publishRetries.delete(threadId);    if (retry?.timer !== undefined) yield* Fiber.interrupt(retry.timer);  });  const cancelPublishRetries = Effect.suspend(() =>    Effect.forEach([...publishRetries.keys()], cancelPublishRetry, { discard: true }),  );  let schedulePublishRetry: (threadId: ThreadId) => Effect.Effect<void> = () => Effect.void;  const resetPublishedConnection = Effect.gen(function* () {    publishedRelayConfig = undefined;    cachedRelayClient = undefined;    publishConfirmDeadlines.clear();    yield* Ref.set(publishedStateByThreadRef, new Map());  });   const publishThreadUnsafe = Effect.fn("publishThreadUnsafe")(function* (threadId: ThreadId) {    const publishAgentActivity = yield* readPublishAgentActivityEnabled;    if (!publishAgentActivity) {      yield* cancelPublishRetries;      yield* resetPublishedConnection;      yield* Effect.logDebug("agent activity publish skipped; publication disabled", {        threadId,      });      return;    }    const relayConfig = yield* readRelayConfig;    if (!relayConfig) {      yield* cancelPublishRetries;      yield* resetPublishedConnection;      yield* Effect.logDebug("agent activity publish skipped; relay link credentials unavailable", {        threadId,      });      return;    }    if (      publishedRelayConfig?.url !== relayConfig.url ||      publishedRelayConfig.issuer !== relayConfig.issuer ||      publishedRelayConfig.environmentCredential !== relayConfig.environmentCredential    ) {      yield* resetPublishedConnection;      publishedRelayConfig = relayConfig;    }    const environmentId = yield* serverEnvironment.getEnvironmentId;     const publishState = (input: {      readonly projectId: string | null;      readonly state: RelayAgentActivityState | null;      readonly reason: string;    }) =>      Effect.gen(function* () {        if (cachedRelayClient === undefined) {          cachedRelayClient = yield* makeRelayClient(relayConfig);        }        const relayClient = cachedRelayClient;        const proof = yield* makePublishProof({          privateKey: cloudLinkKeyPair.privateKey,          relayIssuer: relayConfig.issuer,          environmentId,          threadId,          state: input.state,          jti: yield* crypto.randomUUIDv4,        });         yield* Effect.logInfo("publishing agent activity for thread", {          environmentId,          threadId,          projectId: input.projectId,          statePhase: input.state?.phase ?? null,          hasState: input.state !== null,          reason: input.reason,        });         const response = yield* relayClient.server.publishAgentActivity({          params: {            environmentId,            threadId,          },          payload: {            state: input.state,            proof,          },        });         yield* Effect.logInfo("agent activity publish completed", {          environmentId,          threadId,          ok: response.ok,          deliveries: deliveryStats(response.deliveries),        });      });     // Per-thread shell read: this publish runs for every activity-relevant    // domain event, so materializing the full shell here would make the cost    // of one thread's activity proportional to how many threads exist.    const threadShell = yield* threads.getThreadShell(threadId);    const thread =      threadShell === null || threadShell.archivedAt !== null        ? Option.none<OrchestrationV2ThreadShell>()        : Option.some(threadShell);    const project = Option.isSome(thread)      ? yield* projects.getById(thread.value.projectId)      : Option.none<Project>();    const snapshot = resolveAgentAwarenessRelayPublishSnapshot({      environmentId,      threadId,      thread,      project,    });    const publishIdentity = agentAwarenessPublishIdentity(snapshot.state);    const publishedStateByThread = yield* Ref.get(publishedStateByThreadRef);    if (publishedStateByThread.get(threadId) === publishIdentity) {      // The projection is back at (or never left) the last published state, so      // any pending deferred confirmation is moot. Leaving the deadline in

Lines 1–500 of 758.