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,50 @@
+import { describe, expect, it, vi } from "vitest";
+import * as harness from "../src/index.js";
+
+describe("PiEventAdapter", () => {
+  it("maps assistant deltas and completion with one stable message id", async () => {
+    const store = new harness.InMemoryEventStore();
+    const Adapter = Reflect.get(harness, "PiEventAdapter") as unknown as new (options: Record<string, unknown>) => {
+      observe(event: unknown): Promise<void>;
+    };
+    expect(Adapter, "PiEventAdapter must be exported").toBeDefined();
+    const adapter = new Adapter({ sessionId: "ses_1", runId: "run_1", eventStore: store });
+
+    await adapter.observe({ type: "message_update", message: { role: "assistant" }, assistantMessageEvent: { type: "text_delta", delta: "你好" } });
+    await adapter.observe({ type: "message_end", message: { role: "assistant" } });
+
+    expect(store.listAfter("ses_1").events.map(({ type, payload }) => ({ type, payload }))).toEqual([
+      { type: "assistant.delta", payload: { messageId: "msg_run_1", delta: "你好" } },
+      { type: "assistant.completed", payload: { messageId: "msg_run_1" } },
+    ]);
+  });
+
+  it("maps tool lifecycle, redacts payloads, and correlates successful reads to Skill inference", async () => {
+    const store = new harness.InMemoryEventStore();
+    const requested = vi.fn(() => Promise.resolve());
+    const completed = vi.fn();
+    const Adapter = Reflect.get(harness, "PiEventAdapter") as unknown as new (options: Record<string, unknown>) => {
+      observe(event: unknown): Promise<void>;
+    };
+    const adapter = new Adapter({
+      sessionId: "ses_1",
+      runId: "run_1",
+      eventStore: store,
+      skillInference: { onReadRequested: requested, onReadCompleted: completed },
+    });
+
+    await adapter.observe({ type: "tool_execution_start", toolCallId: "call_1", toolName: "read", args: { path: "/skill/SKILL.md", token: "hidden" } });
+    await adapter.observe({ type: "tool_execution_end", toolCallId: "call_1", toolName: "read", result: { content: "ok", apiKey: "hidden" }, isError: false });
+    await adapter.observe({ type: "tool_execution_end", toolCallId: "call_2", toolName: "grep", result: "failed", isError: true });
+    await expect(adapter.observe({ type: "future_event", anything: true })).resolves.toBeUndefined();
+
+    expect(requested).toHaveBeenCalledWith({ toolCallId: "call_1", path: "/skill/SKILL.md" });
+    expect(completed).toHaveBeenCalledWith({ toolCallId: "call_1", success: true });
+    expect(store.listAfter("ses_1").events.map((event) => event.type)).toEqual([
+      "tool.requested",
+      "tool.completed",
+      "tool.failed",
+    ]);
+    expect(JSON.stringify(store.listAfter("ses_1"))).not.toContain("hidden");
+  });
+});
@@ -0,0 +1,155 @@
+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();
+  });
+});
+
+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();
+  });
+});