agent/packages/harness/tests/session-registry.test.ts

import { mkdir, mkdtemp, realpath, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import * as harness from "../src/index.js";

type Listener = (event: unknown) => void;
class FakePort {
  listeners = new Set<Listener>();
  promptCalls: Array<{ text: string; images: unknown[] }> = [];
  aborted = false;
  disposed = false;
  resolvePrompt!: () => void;
  rejectPrompt!: (error: Error) => void;

  subscribe(listener: Listener): () => void {
    this.listeners.add(listener);
    return () => this.listeners.delete(listener);
  }

  prompt(text: string, options: { images: unknown[] }): Promise<void> {
    this.promptCalls.push({ text, images: options.images });
    return new Promise((resolve, reject) => {
      this.resolvePrompt = resolve;
      this.rejectPrompt = reject;
    });
  }

  abort(): Promise<void> {
    this.aborted = true;
    this.resolvePrompt?.();
    return Promise.resolve();
  }

  dispose(): void {
    this.disposed = true;
  }
}

let root = "";

beforeEach(async () => {
  root = await mkdtemp(join(tmpdir(), "agent-studio-sessions-"));
});

afterEach(async () => {
  vi.useRealTimers();
  await rm(root, { recursive: true, force: true });
});

async function configuredModelStore(supportsImages = true): Promise<harness.InMemoryModelConnectionStore> {
  const store = new harness.InMemoryModelConnectionStore({
    validateUrl: (url) => Promise.resolve({ normalizedUrl: url, addresses: ["93.184.216.34"], redirectPolicy: "same-origin" }),
  });
  await store.configure("ses_1", {
    providerId: "training",
    api: "openai-responses",
    apiKey: "secret",
    baseUrl: "https://api.example.com/v1",
    modelId: "training-model",
    supportsImages,
  });
  return store;
}

describe("HarnessSession", () => {
  it("returns a run id before prompt completion and rejects overlap", async () => {
    const port = new FakePort();
    const events = new harness.InMemoryEventStore();
    const Session = Reflect.get(harness, "HarnessSession") as unknown as new (options: Record<string, unknown>) => {
      startRun(message: { text: string; attachmentIds: string[] }): Promise<{ runId: string }>;
    };
    expect(Session, "HarnessSession must be exported").toBeDefined();
    const session = new Session({
      sessionId: "ses_1",
      workspaceRoot: root,
      eventStore: events,
      modelStore: await configuredModelStore(),
      attachmentStore: { toPiImages: () => Promise.resolve([]), removeAll: () => Promise.resolve() },
      approvalBroker: { cancelForRun: () => 0, cancelForSession: () => 0 },
      createAgentPort: () => Promise.resolve(port),
    });

    const first = await session.startRun({ text: "first", attachmentIds: [] });
    expect(first.runId).toMatch(/^run_/);
    expect(port.promptCalls).toHaveLength(1);
    await expect(session.startRun({ text: "second", attachmentIds: [] })).rejects.toMatchObject({ code: "SESSION_BUSY" });
    port.resolvePrompt();
    await vi.waitFor(() => expect(events.listAfter("ses_1").events.at(-1)?.type).toBe("run.completed"));
  });

  it("distinguishes user cancellation from timeout", async () => {
    vi.useFakeTimers();
    const port = new FakePort();
    const events = new harness.InMemoryEventStore();
    const Session = Reflect.get(harness, "HarnessSession") as unknown as new (options: Record<string, unknown>) => {
      startRun(message: { text: string; attachmentIds: string[] }): Promise<{ runId: string }>;
      cancelRun(reason: string): Promise<boolean>;
    };
    const session = new Session({
      sessionId: "ses_1", workspaceRoot: root, eventStore: events,
      modelStore: await configuredModelStore(),
      attachmentStore: { toPiImages: () => Promise.resolve([]), removeAll: () => Promise.resolve() },
      approvalBroker: { cancelForRun: () => 0, cancelForSession: () => 0 },
      createAgentPort: () => Promise.resolve(port), runTimeoutMs: 50,
    });

    await session.startRun({ text: "cancel me", attachmentIds: [] });
    await session.cancelRun("user");
    expect(events.listAfter("ses_1").events.at(-1)).toMatchObject({ type: "run.cancelled", payload: { reason: "user" } });

    await session.startRun({ text: "time out", attachmentIds: [] });
    await vi.advanceTimersByTimeAsync(50);
    expect(events.listAfter("ses_1").events.at(-1)).toMatchObject({ type: "run.cancelled", payload: { reason: "timeout" } });
  });

  it("rejects images for a text-only model before launching the port", async () => {
    const createAgentPort = vi.fn(() => Promise.resolve(new FakePort()));
    const Session = Reflect.get(harness, "HarnessSession") as unknown as new (options: Record<string, unknown>) => {
      startRun(message: { text: string; attachmentIds: string[] }): Promise<{ runId: string }>;
    };
    const session = new Session({
      sessionId: "ses_1", workspaceRoot: root, eventStore: new harness.InMemoryEventStore(),
      modelStore: await configuredModelStore(false),
      attachmentStore: { toPiImages: () => Promise.resolve([{ type: "image", data: "AA==", mimeType: "image/png" }]), removeAll: () => Promise.resolve() },
      approvalBroker: { cancelForRun: () => 0, cancelForSession: () => 0 }, createAgentPort,
    });

    await expect(session.startRun({ text: "look", attachmentIds: ["att_1"] })).rejects.toMatchObject({ code: "MODEL_DOES_NOT_SUPPORT_IMAGES" });
    expect(createAgentPort).not.toHaveBeenCalled();
  });

  it("serializes async agent events before completing the run", async () => {
    const port = new FakePort();
    const events = new harness.InMemoryEventStore();
    let finishSkillLookup: (() => void) | undefined;
    const order: string[] = [];
    const session = new harness.HarnessSession({
      sessionId: "ses_1",
      workspaceRoot: root,
      eventStore: events,
      modelStore: await configuredModelStore(),
      attachmentStore: { toPiImages: () => Promise.resolve([]), removeAll: () => Promise.resolve() },
      approvalBroker: { cancelForRun: () => 0, cancelForSession: () => 0 },
      createAgentPort: () => Promise.resolve(port),
      createSkillInference: () => ({
        onReadRequested: () => new Promise<void>((resolve) => {
          finishSkillLookup = () => {
            order.push("requested");
            resolve();
          };
        }),
        onReadCompleted: () => { order.push("completed"); },
      }),
    });

    await session.startRun({ text: "load skill", attachmentIds: [] });
    for (const listener of port.listeners) {
      listener({ type: "tool_execution_start", toolCallId: "call_1", toolName: "read", args: { path: "/skill/SKILL.md" } });
      listener({ type: "tool_execution_end", toolCallId: "call_1", toolName: "read", result: "ok", isError: false });
    }
    port.resolvePrompt();
    await vi.waitFor(() => expect(finishSkillLookup).toBeDefined());
    expect(events.listAfter("ses_1").events.at(-1)?.type).not.toBe("run.completed");

    finishSkillLookup?.();
    await vi.waitFor(() => expect(events.listAfter("ses_1").events.at(-1)?.type).toBe("run.completed"));
    expect(order).toEqual(["requested", "completed"]);
  });
});

describe("InMemorySessionRegistry", () => {
  it("creates isolated workspaces and disposes a deleted session", async () => {
    await mkdir(root, { recursive: true });
    const Registry = Reflect.get(harness, "InMemorySessionRegistry") as unknown as new (options: Record<string, unknown>) => {
      create(): Promise<{ sessionId: string; workspaceRoot: string }>;
      delete(sessionId: string): Promise<boolean>;
      get(sessionId: string): unknown;
    };
    expect(Registry, "InMemorySessionRegistry must be exported").toBeDefined();
    const registry = new Registry({
      sessionsRoot: root,
      createSession: (options: Record<string, unknown>) => ({ ...options, dispose: () => Promise.resolve() }),
    });

    const created = await registry.create();
    expect(created.sessionId).toMatch(/^ses_/);
    expect(created.workspaceRoot.startsWith(await realpath(root))).toBe(true);
    await expect(registry.delete(created.sessionId)).resolves.toBe(true);
    await expect(registry.delete(created.sessionId)).resolves.toBe(false);
    expect(registry.get(created.sessionId)).toBeUndefined();
  });
});