feat(harness): manage Pi agent sessions

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

@@ -0,0 +1,16 @@
+import type { ImageContent } from "@earendil-works/pi-ai";
+import type { ResolvedModelConnection } from "../models/model-connection-store.js";
+
+export interface AgentSessionPort {
+  subscribe(listener: (event: unknown) => void): () => void;
+  prompt(text: string, options: { images: ImageContent[] }): Promise<void>;
+  abort(): Promise<void>;
+  dispose(): void;
+}
+
+export type AgentSessionPortFactory = (input: {
+  resolvedModel: ResolvedModelConnection;
+  workspaceRoot: string;
+  sessionId: string;
+  getActiveRunId: () => string | undefined;
+}) => Promise<AgentSessionPort>;

packages/harness/src/agent/pi-agent-session-adapter.ts

@@ -0,0 +1,137 @@
+import { randomUUID } from "node:crypto";
+import type { ImageContent } from "@earendil-works/pi-ai";
+import {
+  DefaultResourceLoader,
+  SessionManager,
+  SettingsManager,
+  createAgentSession,
+  type AgentSession,
+  type ExtensionFactory,
+  type ToolDefinition,
+} from "@earendil-works/pi-coding-agent";
+import type { ApprovalBroker } from "../approvals/approval-broker.js";
+import { HarnessError } from "../errors/harness-error.js";
+import { createPiModelRuntime } from "../models/pi-model-runtime-factory.js";
+import { decideToolPolicy } from "../tools/tool-policy.js";
+import { SafeWorkspace } from "../workspace/safe-workspace.js";
+import type { AgentSessionPort, AgentSessionPortFactory } from "./agent-session-port.js";
+
+export class PiAgentSessionAdapter implements AgentSessionPort {
+  public constructor(private readonly session: AgentSession) {}
+
+  public subscribe(listener: (event: unknown) => void): () => void {
+    return this.session.subscribe(listener);
+  }
+
+  public prompt(text: string, options: { images: ImageContent[] }): Promise<void> {
+    return this.session.prompt(text, { images: options.images, source: "rpc" });
+  }
+
+  public abort(): Promise<void> {
+    return this.session.abort();
+  }
+
+  public dispose(): void {
+    this.session.dispose();
+  }
+}
+
+export type CreatePiAgentSessionPortOptions = {
+  agentDir: string;
+  customTools?: ToolDefinition[];
+  approvalBroker?: ApprovalBroker;
+};
+
+export function createPiAgentSessionPortFactory(
+  options: CreatePiAgentSessionPortOptions,
+): AgentSessionPortFactory {
+  return async ({ resolvedModel, workspaceRoot, sessionId, getActiveRunId }) => {
+    const modelRuntime = await createPiModelRuntime(resolvedModel);
+    const model = modelRuntime.getModel(
+      resolvedModel.connection.providerId,
+      resolvedModel.connection.modelId,
+    );
+    if (model === undefined) {
+      throw new HarnessError("MODEL_NOT_FOUND", "Configured model is not available in the Pi runtime");
+    }
+    const settingsManager = SettingsManager.inMemory();
+    const workspace = await SafeWorkspace.create(workspaceRoot);
+    const extensionFactories: ExtensionFactory[] = [
+      createToolPolicyExtension({
+        workspace,
+        ...(options.approvalBroker === undefined ? {} : { approvalBroker: options.approvalBroker }),
+        sessionId,
+        getActiveRunId,
+      }),
+    ];
+    const resourceLoader = new DefaultResourceLoader({
+      cwd: workspace.root,
+      agentDir: options.agentDir,
+      settingsManager,
+      extensionFactories,
+      noExtensions: true,
+      noPromptTemplates: true,
+      noThemes: true,
+      noContextFiles: true,
+    });
+    await resourceLoader.reload();
+    const { session } = await createAgentSession({
+      cwd: workspace.root,
+      agentDir: options.agentDir,
+      modelRuntime,
+      model,
+      tools: ["read", "grep", "find", "ls", "mock_ads_metrics", "save_report_draft"],
+      customTools: options.customTools ?? [],
+      resourceLoader,
+      settingsManager,
+      sessionManager: SessionManager.inMemory(workspace.root),
+    });
+    return new PiAgentSessionAdapter(session);
+  };
+}
+
+function createToolPolicyExtension(input: {
+  workspace: SafeWorkspace;
+  approvalBroker?: ApprovalBroker;
+  sessionId: string;
+  getActiveRunId: () => string | undefined;
+}): ExtensionFactory {
+  return (pi) => {
+    pi.on("tool_call", async (event, context) => {
+      const decision = decideToolPolicy(event.toolName);
+      if (decision === "deny") return { block: true, reason: "Tool is not allowed" };
+      const path = readPathArgument(event.input);
+      if (path !== undefined) {
+        try {
+          await input.workspace.resolveExisting(path);
+        } catch {
+          return { block: true, reason: "Tool path is outside the safe workspace" };
+        }
+      }
+      if (decision === "allow") return undefined;
+      const runId = input.getActiveRunId();
+      if (input.approvalBroker === undefined || runId === undefined) {
+        return { block: true, reason: "Approval is unavailable" };
+      }
+      const resolution = await input.approvalBroker.request({
+        approvalId: `ap_${randomUUID()}`,
+        sessionId: input.sessionId,
+        runId,
+        toolCallId: event.toolCallId,
+        toolName: event.toolName,
+        arguments: event.input,
+        riskLevel: "medium",
+      }, context.signal);
+      return resolution.status === "approved"
+        ? undefined
+        : { block: true, reason: `Tool approval ${resolution.status}` };
+    });
+  };
+}
+
+function readPathArgument(input: Record<string, unknown>): string | undefined {
+  for (const key of ["path", "filePath", "directory"]) {
+    if (typeof input[key] === "string") return input[key];
+  }
+  return undefined;
+}

packages/harness/src/agent/pi-event-adapter.ts

@@ -0,0 +1,125 @@
+import type { EventStore } from "../events/event-store.js";
+import { redactAndLimit } from "../security/redact.js";
+
+type SkillInferencePort = {
+  onReadRequested(input: { toolCallId: string; path: string }): Promise<void>;
+  onReadCompleted(input: { toolCallId: string; success: boolean }): void;
+};
+
+export type PiEventAdapterOptions = {
+  sessionId: string;
+  runId: string;
+  eventStore: EventStore;
+  skillInference?: SkillInferencePort;
+  maxPayloadBytes?: number;
+};
+
+export class PiEventAdapter {
+  readonly #messageId: string;
+  readonly #sessionId: string;
+  readonly #runId: string;
+  readonly #eventStore: EventStore;
+  readonly #skillInference: SkillInferencePort | undefined;
+  readonly #maxPayloadBytes: number;
+  #assistantCompleted = false;
+
+  public constructor(options: PiEventAdapterOptions) {
+    this.#messageId = `msg_${options.runId}`;
+    this.#sessionId = options.sessionId;
+    this.#runId = options.runId;
+    this.#eventStore = options.eventStore;
+    this.#skillInference = options.skillInference;
+    this.#maxPayloadBytes = options.maxPayloadBytes ?? 16_384;
+  }
+
+  public async observe(value: unknown): Promise<void> {
+    if (!isRecord(value) || typeof value.type !== "string") return;
+    if (value.type === "message_update") {
+      const update = isRecord(value.assistantMessageEvent) ? value.assistantMessageEvent : undefined;
+      if (update?.type === "text_delta" && typeof update.delta === "string" && update.delta !== "") {
+        this.#eventStore.append(this.#sessionId, {
+          type: "assistant.delta",
+          runId: this.#runId,
+          payload: { messageId: this.#messageId, delta: update.delta },
+        });
+      }
+      return;
+    }
+    if (value.type === "message_end" && isRecord(value.message) && value.message.role === "assistant") {
+      if (!this.#assistantCompleted) {
+        this.#assistantCompleted = true;
+        this.#eventStore.append(this.#sessionId, {
+          type: "assistant.completed",
+          runId: this.#runId,
+          payload: { messageId: this.#messageId },
+        });
+      }
+      return;
+    }
+    if (value.type === "tool_execution_start") {
+      await this.#toolStarted(value);
+      return;
+    }
+    if (value.type === "tool_execution_end") this.#toolEnded(value);
+  }
+
+  async #toolStarted(event: Record<string, unknown>): Promise<void> {
+    if (typeof event.toolCallId !== "string" || typeof event.toolName !== "string") return;
+    const args = redactAndLimit(event.args, this.#maxPayloadBytes);
+    this.#eventStore.append(this.#sessionId, {
+      type: "tool.requested",
+      runId: this.#runId,
+      payload: { toolCallId: event.toolCallId, toolName: event.toolName, arguments: args },
+    });
+    if (
+      event.toolName === "read" &&
+      isRecord(event.args) &&
+      typeof event.args.path === "string"
+    ) {
+      await this.#skillInference?.onReadRequested({
+        toolCallId: event.toolCallId,
+        path: event.args.path,
+      });
+    }
+  }
+
+  #toolEnded(event: Record<string, unknown>): void {
+    if (typeof event.toolCallId !== "string" || typeof event.toolName !== "string") return;
+    const isError = event.isError === true;
+    if (event.toolName === "read") {
+      this.#skillInference?.onReadCompleted({ toolCallId: event.toolCallId, success: !isError });
+    }
+    if (isError) {
+      const safeResult = redactAndLimit(event.result, this.#maxPayloadBytes);
+      this.#eventStore.append(this.#sessionId, {
+        type: "tool.failed",
+        runId: this.#runId,
+        payload: {
+          toolCallId: event.toolCallId,
+          toolName: event.toolName,
+          error: safeErrorText(safeResult),
+        },
+      });
+      return;
+    }
+    this.#eventStore.append(this.#sessionId, {
+      type: "tool.completed",
+      runId: this.#runId,
+      payload: {
+        toolCallId: event.toolCallId,
+        toolName: event.toolName,
+        result: redactAndLimit(event.result, this.#maxPayloadBytes),
+      },
+    });
+  }
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null;
+}
+
+function safeErrorText(value: unknown): string {
+  if (typeof value === "string" && value.trim() !== "") return value.slice(0, 2_000);
+  const serialized = JSON.stringify(value);
+  return serialized === undefined || serialized === "" ? "Tool execution failed" : serialized.slice(0, 2_000);
+}

packages/harness/src/index.ts

@@ -1,5 +1,8 @@
 export * from "./approvals/approval-broker.js";
 export * from "./approvals/in-memory-approval-broker.js";
+export * from "./agent/agent-session-port.js";
+export * from "./agent/pi-agent-session-adapter.js";
+export * from "./agent/pi-event-adapter.js";
 export * from "./attachments/attachment-store.js";
 export * from "./attachments/file-attachment-store.js";
 export * from "./attachments/image-signatures.js";
@@ -11,6 +14,9 @@ export * from "./models/in-memory-model-connection-store.js";
 export * from "./models/model-connection-store.js";
 export * from "./models/pi-model-runtime-factory.js";
 export * from "./security/redact.js";
+export * from "./sessions/harness-session.js";
+export * from "./sessions/in-memory-session-registry.js";
+export * from "./sessions/session-registry.js";
 export * from "./skills/skill-catalog.js";
 export * from "./skills/skill-event-inference.js";
 export * from "./tools/mock-ads-metrics.js";

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

@@ -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");
+}

packages/harness/src/sessions/in-memory-session-registry.ts

@@ -0,0 +1,58 @@
+import { randomUUID } from "node:crypto";
+import { mkdir, realpath, rm } from "node:fs/promises";
+import { isAbsolute, join, relative, resolve, sep } from "node:path";
+import { HarnessError } from "../errors/harness-error.js";
+import type { SessionRecord, SessionRegistry } from "./session-registry.js";
+
+export type InMemorySessionRegistryOptions<TSession extends SessionRecord> = {
+  sessionsRoot: string;
+  createSession(input: { sessionId: string; workspaceRoot: string }): TSession;
+};
+
+export class InMemorySessionRegistry<TSession extends SessionRecord = SessionRecord>
+implements SessionRegistry<TSession> {
+  readonly #sessionsRoot: string;
+  readonly #createSession: InMemorySessionRegistryOptions<TSession>["createSession"];
+  readonly #sessions = new Map<string, TSession>();
+
+  public constructor(options: InMemorySessionRegistryOptions<TSession>) {
+    this.#sessionsRoot = resolve(options.sessionsRoot);
+    this.#createSession = (input) => options.createSession(input);
+  }
+
+  public async create(): Promise<TSession> {
+    const canonicalRoot = await realpath(this.#sessionsRoot);
+    const sessionId = `ses_${randomUUID()}`;
+    const workspaceRoot = join(canonicalRoot, sessionId);
+    assertContained(canonicalRoot, workspaceRoot);
+    await mkdir(workspaceRoot, { mode: 0o700 });
+    const session = this.#createSession({ sessionId, workspaceRoot });
+    this.#sessions.set(sessionId, session);
+    return session;
+  }
+
+  public get(sessionId: string): TSession | undefined {
+    return this.#sessions.get(sessionId);
+  }
+
+  public async delete(sessionId: string): Promise<boolean> {
+    const session = this.#sessions.get(sessionId);
+    if (session === undefined) return false;
+    this.#sessions.delete(sessionId);
+    await session.dispose();
+    const canonicalRoot = await realpath(this.#sessionsRoot);
+    assertContained(canonicalRoot, session.workspaceRoot);
+    await rm(session.workspaceRoot, { recursive: true, force: true });
+    return true;
+  }
+
+  public async disposeAll(): Promise<void> {
+    for (const sessionId of [...this.#sessions.keys()]) await this.delete(sessionId);
+  }
+}
+
+function assertContained(root: string, candidate: string): void {
+  const fromRoot = relative(root, candidate);
+  if (fromRoot !== "" && fromRoot !== ".." && !fromRoot.startsWith(`..${sep}`) && !isAbsolute(fromRoot)) return;
+  throw new HarnessError("UNSAFE_SESSION_PATH", "Session workspace must be a child of the sessions root");
+}

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

@@ -0,0 +1,12 @@
+export type SessionRecord = {
+  sessionId: string;
+  workspaceRoot: string;
+  dispose(): Promise<void>;
+};
+
+export interface SessionRegistry<TSession extends SessionRecord = SessionRecord> {
+  create(): Promise<TSession>;
+  get(sessionId: string): TSession | undefined;
+  delete(sessionId: string): Promise<boolean>;
+  disposeAll(): Promise<void>;
+}

packages/harness/src/workspace/safe-workspace.ts

@@ -32,16 +32,21 @@ export class SafeWorkspace {
   }
 
   public async resolveReadable(inputPath: string): Promise<string> {
+    const canonical = await this.resolveExisting(inputPath);
+    const metadata = await lstat(canonical);
+    if (!metadata.isFile()) {
+      throw new HarnessError("NOT_A_REGULAR_FILE", "Only regular files may be read");
+    }
+    return canonical;
+  }
+
+  public async resolveExisting(inputPath: string): Promise<string> {
     const candidate = isAbsolute(inputPath) ? resolve(inputPath) : resolve(this.root, inputPath);
     this.#assertContained(candidate);
     this.#assertUnprotected(candidate);
     await this.#assertNoSymlinkSegments(candidate);
     const canonical = await realpath(candidate);
     this.#assertContained(canonical);
-    const metadata = await lstat(canonical);
-    if (!metadata.isFile()) {
-      throw new HarnessError("NOT_A_REGULAR_FILE", "Only regular files may be read");
-    }
     return canonical;
   }