agent/packages/harness/src/sessions/harness-session.ts

import { createHash, randomUUID } from "node:crypto";
import { MessageInputSchema, type MessageInput } from "@agent-studio/shared";
import type { ApprovalBroker } from "../approvals/approval-broker.js";
import type { AttachmentStore } from "../attachments/attachment-store.js";
import type { AgentSessionPort, AgentSessionPortFactory } from "../agent/agent-session-port.js";
import { PiEventAdapter } from "../agent/pi-event-adapter.js";
import type { SkillEventInference } from "../skills/skill-event-inference.js";
import type { EventStore } from "../events/event-store.js";
import { HarnessError } from "../errors/harness-error.js";
import type { ModelConnectionStore, ResolvedModelConnection } from "../models/model-connection-store.js";
import { SafeWorkspace } from "../workspace/safe-workspace.js";

export type RunCancellationReason = "user" | "timeout" | "session_deleted" | "shutdown";

type ActiveRun = {
  runId: string;
  messageId: string;
  timer: ReturnType<typeof setTimeout>;
  unsubscribe: () => void;
  settled: boolean;
};

export type HarnessSessionOptions = {
  sessionId: string;
  workspaceRoot: string;
  eventStore: EventStore;
  modelStore: ModelConnectionStore;
  attachmentStore: Pick<AttachmentStore, "toPiImages" | "removeAll">;
  approvalBroker: Pick<ApprovalBroker, "cancelForRun" | "cancelForSession">;
  createAgentPort: AgentSessionPortFactory;
  createSkillInference?: (runId: string) => Pick<SkillEventInference, "onReadRequested" | "onReadCompleted">;
  runTimeoutMs?: number;
};

export class HarnessSession {
  public readonly sessionId: string;
  public readonly workspaceRoot: string;
  readonly #eventStore: EventStore;
  readonly #modelStore: ModelConnectionStore;
  readonly #attachmentStore: Pick<AttachmentStore, "toPiImages" | "removeAll">;
  readonly #approvalBroker: Pick<ApprovalBroker, "cancelForRun" | "cancelForSession">;
  readonly #createAgentPort: AgentSessionPortFactory;
  readonly #createSkillInference: HarnessSessionOptions["createSkillInference"];
  readonly #runTimeoutMs: number;
  #port: AgentSessionPort | undefined;
  #modelFingerprint: string | undefined;
  #active: ActiveRun | undefined;
  #disposed = false;

  public constructor(options: HarnessSessionOptions) {
    this.sessionId = options.sessionId;
    this.workspaceRoot = options.workspaceRoot;
    this.#eventStore = options.eventStore;
    this.#modelStore = options.modelStore;
    this.#attachmentStore = options.attachmentStore;
    this.#approvalBroker = options.approvalBroker;
    this.#createAgentPort = options.createAgentPort;
    this.#createSkillInference = options.createSkillInference;
    this.#runTimeoutMs = options.runTimeoutMs ?? 120_000;
  }

  public get isBusy(): boolean {
    return this.#active !== undefined;
  }

  public async startRun(input: MessageInput): Promise<{ runId: string }> {
    if (this.#disposed) throw new HarnessError("SESSION_DISPOSED", "Session has been disposed");
    if (this.#active !== undefined) throw new HarnessError("SESSION_BUSY", "Session already has an active run");
    const message = MessageInputSchema.parse(input);
    const resolvedModel = this.#modelStore.resolveConnection(this.sessionId);
    if (message.attachmentIds.length > 0 && !resolvedModel.connection.supportsImages) {
      throw new HarnessError("MODEL_DOES_NOT_SUPPORT_IMAGES", "Configured model does not support image input");
    }
    const images = await this.#attachmentStore.toPiImages(this.sessionId, message.attachmentIds);
    const port = await this.#ensurePort(resolvedModel);
    const runId = `run_${randomUUID()}`;
    const messageId = `msg_${runId}`;
    const skillInference = this.#createSkillInference?.(runId);
    const adapter = new PiEventAdapter({
      sessionId: this.sessionId,
      runId,
      eventStore: this.#eventStore,
      ...(skillInference === undefined ? {} : { skillInference }),
    });
    let eventQueue = Promise.resolve();
    let eventError: unknown;
    const unsubscribe = port.subscribe((event) => {
      eventQueue = eventQueue.then(async () => {
        if (eventError !== undefined) return;
        try {
          await adapter.observe(event);
        } catch (error) {
          eventError = error;
        }
      });
    });
    const active: ActiveRun = {
      runId,
      messageId,
      unsubscribe,
      settled: false,
      timer: setTimeout(() => {
        void this.cancelRun("timeout");
      }, this.#runTimeoutMs),
    };
    this.#active = active;
    this.#eventStore.append(this.sessionId, {
      type: "user.message",
      runId,
      payload: { messageId: `usr_${runId}`, text: message.text, attachmentIds: message.attachmentIds },
    });
    this.#eventStore.append(this.sessionId, {
      type: "run.started",
      runId,
      payload: { modelId: resolvedModel.connection.modelId },
    });

    void Promise.resolve()
      .then(() => port.prompt(message.text, { images }))
      .then(async () => {
        await eventQueue;
        if (eventError !== undefined) throw eventError;
        this.#complete(active);
      })
      .catch((error: unknown) => this.#fail(active, error));
    return { runId };
  }

  public async cancelRun(reason: RunCancellationReason): Promise<boolean> {
    const active = this.#active;
    if (active === undefined || active.settled) return false;
    active.settled = true;
    clearTimeout(active.timer);
    this.#approvalBroker.cancelForRun(active.runId);
    active.unsubscribe();
    await this.#port?.abort().catch(() => undefined);
    if (this.#active === active) this.#active = undefined;
    this.#eventStore.append(this.sessionId, {
      type: "run.cancelled",
      runId: active.runId,
      payload: { reason },
    });
    return true;
  }

  public async dispose(): Promise<void> {
    if (this.#disposed) return;
    this.#disposed = true;
    await this.cancelRun("session_deleted");
    this.#approvalBroker.cancelForSession(this.sessionId);
    this.#port?.dispose();
    this.#port = undefined;
    await this.#attachmentStore.removeAll(this.sessionId);
    const workspace = await SafeWorkspace.create(this.workspaceRoot);
    await workspace.remove();
    this.#eventStore.append(this.sessionId, {
      type: "session.ended",
      payload: { reason: "deleted" },
    });
  }

  async #ensurePort(resolvedModel: ResolvedModelConnection): Promise<AgentSessionPort> {
    const fingerprint = modelFingerprint(resolvedModel);
    if (this.#port !== undefined && this.#modelFingerprint === fingerprint) return this.#port;
    this.#port?.dispose();
    this.#port = await this.#createAgentPort({
      resolvedModel,
      workspaceRoot: this.workspaceRoot,
      sessionId: this.sessionId,
      getActiveRunId: () => this.#active?.runId,
    });
    this.#modelFingerprint = fingerprint;
    return this.#port;
  }

  #complete(active: ActiveRun): void {
    if (active.settled || this.#active !== active) return;
    active.settled = true;
    clearTimeout(active.timer);
    active.unsubscribe();
    this.#active = undefined;
    this.#eventStore.append(this.sessionId, {
      type: "run.completed",
      runId: active.runId,
      payload: { messageId: active.messageId },
    });
  }

  #fail(active: ActiveRun, error: unknown): void {
    if (active.settled || this.#active !== active) return;
    active.settled = true;
    clearTimeout(active.timer);
    active.unsubscribe();
    this.#approvalBroker.cancelForRun(active.runId);
    this.#active = undefined;
    this.#eventStore.append(this.sessionId, {
      type: "run.failed",
      runId: active.runId,
      payload: {
        code: error instanceof HarnessError ? error.code : "AGENT_RUN_FAILED",
        message: error instanceof Error ? error.message.slice(0, 2_000) : "Agent run failed",
      },
    });
  }
}

function modelFingerprint(resolved: ResolvedModelConnection): string {
  return createHash("sha256").update(JSON.stringify(resolved)).digest("hex");
}