diff --git a/packages/harness/src/agent/agent-session-port.ts b/packages/harness/src/agent/agent-session-port.ts new file mode 100644 index 0000000..aef321b --- /dev/null +++ b/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; + abort(): Promise; + dispose(): void; +} + +export type AgentSessionPortFactory = (input: { + resolvedModel: ResolvedModelConnection; + workspaceRoot: string; + sessionId: string; + getActiveRunId: () => string | undefined; +}) => Promise; diff --git a/packages/harness/src/agent/pi-agent-session-adapter.ts b/packages/harness/src/agent/pi-agent-session-adapter.ts new file mode 100644 index 0000000..c880de2 --- /dev/null +++ b/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 { + return this.session.prompt(text, { images: options.images, source: "rpc" }); + } + + public abort(): Promise { + 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 | undefined { + for (const key of ["path", "filePath", "directory"]) { + if (typeof input[key] === "string") return input[key]; + } + return undefined; +} diff --git a/packages/harness/src/agent/pi-event-adapter.ts b/packages/harness/src/agent/pi-event-adapter.ts new file mode 100644 index 0000000..d8064ad --- /dev/null +++ b/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; + 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 { + 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): Promise { + 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): 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 { + 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); +} diff --git a/packages/harness/src/index.ts b/packages/harness/src/index.ts index aa1ad20..ccf4ca6 100644 --- a/packages/harness/src/index.ts +++ b/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"; diff --git a/packages/harness/src/sessions/harness-session.ts b/packages/harness/src/sessions/harness-session.ts new file mode 100644 index 0000000..bbe9ce3 --- /dev/null +++ b/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; + unsubscribe: () => void; + settled: boolean; +}; + +export type HarnessSessionOptions = { + sessionId: string; + workspaceRoot: string; + eventStore: EventStore; + modelStore: ModelConnectionStore; + attachmentStore: Pick; + approvalBroker: Pick; + createAgentPort: AgentSessionPortFactory; + createSkillInference?: (runId: string) => Pick; + runTimeoutMs?: number; +}; + +export class HarnessSession { + public readonly sessionId: string; + public readonly workspaceRoot: string; + readonly #eventStore: EventStore; + readonly #modelStore: ModelConnectionStore; + readonly #attachmentStore: Pick; + readonly #approvalBroker: Pick; + 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 { + 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 { + 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 { + 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"); +} diff --git a/packages/harness/src/sessions/in-memory-session-registry.ts b/packages/harness/src/sessions/in-memory-session-registry.ts new file mode 100644 index 0000000..e8dbf4b --- /dev/null +++ b/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 = { + sessionsRoot: string; + createSession(input: { sessionId: string; workspaceRoot: string }): TSession; +}; + +export class InMemorySessionRegistry +implements SessionRegistry { + readonly #sessionsRoot: string; + readonly #createSession: InMemorySessionRegistryOptions["createSession"]; + readonly #sessions = new Map(); + + public constructor(options: InMemorySessionRegistryOptions) { + this.#sessionsRoot = resolve(options.sessionsRoot); + this.#createSession = (input) => options.createSession(input); + } + + public async create(): Promise { + 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 { + 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 { + 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"); +} diff --git a/packages/harness/src/sessions/session-registry.ts b/packages/harness/src/sessions/session-registry.ts new file mode 100644 index 0000000..5d156af --- /dev/null +++ b/packages/harness/src/sessions/session-registry.ts @@ -0,0 +1,12 @@ +export type SessionRecord = { + sessionId: string; + workspaceRoot: string; + dispose(): Promise; +}; + +export interface SessionRegistry { + create(): Promise; + get(sessionId: string): TSession | undefined; + delete(sessionId: string): Promise; + disposeAll(): Promise; +} diff --git a/packages/harness/src/workspace/safe-workspace.ts b/packages/harness/src/workspace/safe-workspace.ts index fb0b246..5dc00d9 100644 --- a/packages/harness/src/workspace/safe-workspace.ts +++ b/packages/harness/src/workspace/safe-workspace.ts @@ -32,16 +32,21 @@ export class SafeWorkspace { } public async resolveReadable(inputPath: string): Promise { + 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 { 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; } diff --git a/packages/harness/tests/pi-event-adapter.test.ts b/packages/harness/tests/pi-event-adapter.test.ts new file mode 100644 index 0000000..4ddfe14 --- /dev/null +++ b/packages/harness/tests/pi-event-adapter.test.ts @@ -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) => { + observe(event: unknown): Promise; + }; + 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) => { + observe(event: unknown): Promise; + }; + 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"); + }); +}); diff --git a/packages/harness/tests/session-registry.test.ts b/packages/harness/tests/session-registry.test.ts new file mode 100644 index 0000000..f8371ca --- /dev/null +++ b/packages/harness/tests/session-registry.test.ts @@ -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(); + 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 { + this.promptCalls.push({ text, images: options.images }); + return new Promise((resolve, reject) => { + this.resolvePrompt = resolve; + this.rejectPrompt = reject; + }); + } + + abort(): Promise { + 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 { + 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) => { + 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) => { + startRun(message: { text: string; attachmentIds: string[] }): Promise<{ runId: string }>; + cancelRun(reason: string): Promise; + }; + 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) => { + 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) => { + create(): Promise<{ sessionId: string; workspaceRoot: string }>; + delete(sessionId: string): Promise; + get(sessionId: string): unknown; + }; + expect(Registry, "InMemorySessionRegistry must be exported").toBeDefined(); + const registry = new Registry({ + sessionsRoot: root, + createSession: (options: Record) => ({ ...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(); + }); +});