feat(harness): manage Pi agent sessions

2026-09-07 16:20:44 +08:00
parent 1eaa85c09a
commit 64ca126047
10 changed files with 760 additions and 4 deletions
@@ -0,0 +1,192 @@
+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 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 }),
+    });
+    const unsubscribe = port.subscribe((event) => {
+      void adapter.observe(event);
+    });
+    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(() => 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");
+}