314 lines
18 KiB
TypeScript
314 lines
18 KiB
TypeScript
import { randomUUID } from "node:crypto";
|
|
import { mkdir, realpath, rm } from "node:fs/promises";
|
|
import path from "node:path";
|
|
import type { ImageContent } from "@earendil-works/pi-ai";
|
|
import {
|
|
SessionManager,
|
|
SettingsManager,
|
|
createAgentSession,
|
|
isToolCallEventType,
|
|
type AgentSession,
|
|
type AgentSessionEvent,
|
|
type InlineExtension,
|
|
} from "@earendil-works/pi-coding-agent";
|
|
import type { ModelConnectionInput, SendMessageInput } from "@agent-studio/shared";
|
|
import { InMemoryApprovalBroker, type ApprovalBroker } from "./approval.js";
|
|
import type { FileAttachmentStore } from "./attachments.js";
|
|
import { HarnessError, redactValue, redactServerSecret, safeErrorMessage } from "./errors.js";
|
|
import type { EventStore } from "./event-store.js";
|
|
import type { ModelConnectionStore } from "./model-store.js";
|
|
import { mapPiEvent } from "./pi-event-adapter.js";
|
|
import { SafeWorkspace } from "./safe-workspace.js";
|
|
import type { SkillCatalog } from "./skills.js";
|
|
import type { SkillIndexEntry } from "./skills.js";
|
|
import { SwadsGateway } from "./swads-gateway.js";
|
|
import { SwadsReportTools } from "./swads-tools.js";
|
|
import { ReportStore } from "./report-store.js";
|
|
import { createMockAdsMetricsTool, createSaveReportDraftTool } from "./tools.js";
|
|
|
|
export interface AgentSessionPort {
|
|
readonly modelId: string;
|
|
readonly supportsImages: boolean;
|
|
prompt(text: string, options?: { images?: unknown[] }): Promise<void>;
|
|
abort(): Promise<void>;
|
|
subscribe(listener: (event: unknown) => void): () => void;
|
|
dispose(): void;
|
|
}
|
|
|
|
export interface AgentSessionFactoryInput {
|
|
sessionId: string;
|
|
cwd: string;
|
|
workspace: SafeWorkspace;
|
|
getRun: () => { runId: string; signal: AbortSignal; assistantMessageId: string } | undefined;
|
|
skillRequested: (skill: SkillIndexEntry, toolCallId: string) => void;
|
|
skillLoaded: (skill: SkillIndexEntry, toolCallId: string) => void;
|
|
toolRequested: (toolCallId: string, toolName: string, argumentsValue: unknown) => void;
|
|
}
|
|
export interface AgentSessionFactory { create(input: AgentSessionFactoryInput): Promise<AgentSessionPort> }
|
|
|
|
class PiAgentSessionPort implements AgentSessionPort {
|
|
readonly modelId: string;
|
|
readonly supportsImages: boolean;
|
|
constructor(private readonly session: AgentSession) {
|
|
if (!session.model) throw new Error("Pi AgentSession was created without an explicit model");
|
|
this.modelId = session.model.id;
|
|
this.supportsImages = session.model.input.includes("image");
|
|
}
|
|
async prompt(text: string, options?: { images?: unknown[] }): Promise<void> {
|
|
const images = options?.images as ImageContent[] | undefined;
|
|
await this.session.prompt(text, images ? { images } : undefined);
|
|
}
|
|
abort(): Promise<void> { return this.session.abort(); }
|
|
subscribe(listener: (event: unknown) => void): () => void { return this.session.subscribe(listener); }
|
|
dispose(): void { this.session.dispose(); }
|
|
}
|
|
|
|
export class PiAgentSessionFactory implements AgentSessionFactory {
|
|
constructor(
|
|
private readonly fixturePath: string,
|
|
private readonly catalog: SkillCatalog,
|
|
private readonly models: ModelConnectionStore,
|
|
private readonly approvals: ApprovalBroker,
|
|
private readonly gateway = new SwadsGateway(),
|
|
private readonly reports = new ReportStore(path.resolve(fixturePath, "../../workspace/reports"), path.resolve(fixturePath, "../../.agents/skills/swads-daily-report/assets/report.css")),
|
|
) {}
|
|
|
|
async create(input: AgentSessionFactoryInput): Promise<AgentSessionPort> {
|
|
const resolved = await this.models.createRuntime(input.sessionId);
|
|
const swads = new SwadsReportTools(this.gateway, this.reports, input.getRun);
|
|
const extension: InlineExtension = { name: "agent-studio-safety", hidden: true, factory: (pi) => {
|
|
const pendingSkills = new Map<string, SkillIndexEntry>();
|
|
pi.on("tool_call", async (event) => {
|
|
const run = input.getRun();
|
|
if (!run) return { block: true, reason: "No active Harness run" };
|
|
const known = new Set(["read", "grep", "find", "ls", "mock_ads_metrics", "save_report_draft", "swads_cli", "render_swads_daily_report"]);
|
|
if (!known.has(event.toolName)) {
|
|
input.toolRequested(event.toolCallId, event.toolName, publicToolArguments(event.input, input.cwd));
|
|
return { block: true, reason: "Tool is not allowed by the Harness policy", terminate: true };
|
|
}
|
|
let matchedSkill: SkillIndexEntry | undefined;
|
|
if (isToolCallEventType("read", event)) {
|
|
const skill = await this.catalog.matchReadPath(event.input.path, input.cwd);
|
|
if (skill) {
|
|
matchedSkill = skill;
|
|
pendingSkills.set(event.toolCallId, skill);
|
|
input.skillRequested(skill, event.toolCallId);
|
|
} else await input.workspace.assertReadPath(event.input.path);
|
|
} else if (["grep", "find", "ls"].includes(event.toolName)) {
|
|
const maybePath = (event.input as Record<string, unknown>).path;
|
|
if (typeof maybePath === "string") await input.workspace.assertReadPath(maybePath);
|
|
}
|
|
input.toolRequested(event.toolCallId, event.toolName, event.toolName === "render_swads_daily_report" ? { status: "requested" } : publicToolArguments(event.input, input.cwd, matchedSkill));
|
|
if (event.toolName === "save_report_draft") {
|
|
const decision = await this.approvals.request({ sessionId: input.sessionId, runId: run.runId, toolCallId: event.toolCallId, toolName: event.toolName, arguments: event.input, riskLevel: "medium" }, run.signal);
|
|
if (decision.status !== "approved") return { block: true, reason: `Human approval ${decision.status}` };
|
|
}
|
|
return undefined;
|
|
});
|
|
pi.on("tool_result", async (event) => {
|
|
const run = input.getRun();
|
|
const skill = pendingSkills.get(event.toolCallId);
|
|
if (run && skill && !event.isError && event.toolName === "read") {
|
|
input.skillLoaded(skill, event.toolCallId);
|
|
pendingSkills.delete(event.toolCallId);
|
|
}
|
|
});
|
|
} };
|
|
const loader = await this.catalog.createSessionLoader(input.cwd, path.join(input.cwd, ".pi-agent"), [extension]);
|
|
const { session } = await createAgentSession({
|
|
cwd: input.cwd,
|
|
agentDir: path.join(input.cwd, ".pi-agent"),
|
|
modelRuntime: resolved.runtime,
|
|
model: resolved.model,
|
|
tools: ["read", "grep", "find", "ls", "mock_ads_metrics", "save_report_draft", "swads_cli", "render_swads_daily_report"],
|
|
customTools: [createMockAdsMetricsTool(this.fixturePath), createSaveReportDraftTool(input.workspace), ...swads.definitions()],
|
|
resourceLoader: loader,
|
|
sessionManager: SessionManager.inMemory(input.cwd),
|
|
settingsManager: SettingsManager.inMemory(),
|
|
});
|
|
return new PiAgentSessionPort(session);
|
|
}
|
|
}
|
|
|
|
export interface HarnessSession {
|
|
sessionId: string;
|
|
cwd: string;
|
|
workspace: SafeWorkspace;
|
|
busy: boolean;
|
|
currentRun: { runId: string; controller: AbortController; assistantMessageId: string; cancelReason?: "user" | "timeout" | "session_deleted" | "shutdown" } | undefined;
|
|
}
|
|
export interface CreateSessionInput { id?: string }
|
|
export interface SessionRegistry {
|
|
create(input?: CreateSessionInput): Promise<HarnessSession>;
|
|
get(sessionId: string): HarnessSession | undefined;
|
|
delete(sessionId: string): Promise<void>;
|
|
}
|
|
|
|
interface InternalSession extends HarnessSession { agent: AgentSessionPort | undefined; unsubscribe: (() => void) | undefined; timeout: ReturnType<typeof setTimeout> | undefined; disposed: boolean; requestedToolCalls: Set<string> }
|
|
|
|
export class InMemorySessionRegistry implements SessionRegistry {
|
|
private readonly sessions = new Map<string, InternalSession>();
|
|
readonly approvals: ApprovalBroker;
|
|
constructor(
|
|
private readonly sessionsRoot: string,
|
|
private readonly events: EventStore,
|
|
private readonly attachments: FileAttachmentStore,
|
|
private readonly models: ModelConnectionStore,
|
|
private readonly catalog: SkillCatalog,
|
|
private factory: AgentSessionFactory | undefined,
|
|
approvalTimeoutMs = 60_000,
|
|
private readonly runTimeoutMs = 120_000,
|
|
) {
|
|
this.approvals = new InMemoryApprovalBroker(approvalTimeoutMs, (resolution) => {
|
|
this.events.append(resolution.request.sessionId, { type: "approval.resolved", runId: resolution.request.runId, payload: { approvalId: resolution.request.approvalId, toolCallId: resolution.request.toolCallId, status: resolution.decision.status, resolvedAt: resolution.decision.resolvedAt } });
|
|
}, (request) => {
|
|
this.events.append(request.sessionId, { type: "approval.required", runId: request.runId, payload: request });
|
|
});
|
|
}
|
|
setFactory(factory: AgentSessionFactory): void { this.factory = factory; }
|
|
|
|
async create(input: CreateSessionInput = {}): Promise<HarnessSession> {
|
|
const sessionId = input.id ?? `ses_${randomUUID()}`;
|
|
const cwd = path.join(this.sessionsRoot, sessionId);
|
|
await mkdir(cwd, { recursive: false, mode: 0o700 });
|
|
const workspace = await SafeWorkspace.create(cwd);
|
|
const session: InternalSession = { sessionId, cwd: workspace.root, workspace, busy: false, currentRun: undefined, agent: undefined, unsubscribe: undefined, timeout: undefined, disposed: false, requestedToolCalls: new Set() };
|
|
this.sessions.set(sessionId, session);
|
|
this.events.append(sessionId, { type: "session.started", payload: { cwd: `workspace/sessions/${sessionId}` } });
|
|
return session;
|
|
}
|
|
get(id: string): HarnessSession | undefined { return this.sessions.get(id); }
|
|
private internal(id: string): InternalSession { const session = this.sessions.get(id); if (!session) throw new HarnessError("SESSION_NOT_FOUND", "Session not found", 404); return session; }
|
|
|
|
async configureModel(id: string, input: ModelConnectionInput) {
|
|
const session = this.internal(id);
|
|
if (session.busy) throw new HarnessError("SESSION_BUSY", "Cannot change model configuration while a run is active", 409);
|
|
await this.disposeAgent(session);
|
|
const value = await this.models.configure(id, input);
|
|
this.events.append(id, { type: "model.configured", payload: { providerId: input.providerId, modelId: input.modelId, api: input.api, supportsImages: input.supportsImages } });
|
|
return value;
|
|
}
|
|
async clearCredentials(id: string): Promise<void> { const session = this.internal(id); if (session.busy) throw new HarnessError("SESSION_BUSY", "Cannot clear credentials while a run is active", 409); await this.disposeAgent(session); await this.models.clearCredentials(id); }
|
|
|
|
async send(id: string, message: SendMessageInput): Promise<{ runId: string }> {
|
|
const session = this.internal(id);
|
|
if (session.busy) throw new HarnessError("SESSION_BUSY", "A run is already active for this session", 409);
|
|
const configuredModel = this.models.getRedacted(id);
|
|
if (!configuredModel.configured) throw new HarnessError("MODEL_NOT_CONFIGURED", "No model connection is configured", 422);
|
|
if (!configuredModel.keyConfigured) throw new HarnessError("MODEL_CREDENTIAL_NOT_CONFIGURED", "The selected provider has no API key", 422);
|
|
const maxImages = Number(process.env.MAX_IMAGES_PER_MESSAGE ?? 4);
|
|
if (message.attachmentIds.length > maxImages) throw new HarnessError("TOO_MANY_ATTACHMENTS", `A message can include at most ${maxImages} images`, 413);
|
|
const selected = message.attachmentIds.map((attachmentId) => {
|
|
const item = this.attachments.getStored(id, attachmentId);
|
|
if (!item) throw new HarnessError("ATTACHMENT_NOT_FOUND", "Attachment not found for this session", 404);
|
|
return item;
|
|
});
|
|
const expanded = await this.catalog.expandCommand(message.text);
|
|
const runId = `run_${randomUUID()}`;
|
|
const assistantMessageId = `msg_${randomUUID()}`;
|
|
const controller = new AbortController();
|
|
session.currentRun = { runId, controller, assistantMessageId };
|
|
session.busy = true;
|
|
try { await this.ensureAgent(session); } catch (error) { session.busy = false; session.currentRun = undefined; this.events.append(id, { type: "run.failed", runId, payload: { code: "AGENT_SESSION_CREATE_FAILED", message: safeErrorMessage(error) } }); throw error; }
|
|
if (selected.length > 0 && !session.agent?.supportsImages) {
|
|
session.busy = false;
|
|
session.currentRun = undefined;
|
|
throw new HarnessError("MODEL_DOES_NOT_SUPPORT_IMAGES", "The selected model does not support image input", 422);
|
|
}
|
|
const agent = session.agent;
|
|
if (!agent) throw new HarnessError("AGENT_SESSION_CREATE_FAILED", "Agent session was not created", 500);
|
|
this.events.append(id, { type: "user.message", runId, payload: { messageId: `msg_${randomUUID()}`, text: message.text, attachmentIds: message.attachmentIds } });
|
|
this.events.append(id, { type: "run.started", runId, payload: { modelId: agent.modelId } });
|
|
const images: ImageContent[] = await Promise.all(selected.map(async (item) => ({ type: "image" as const, mimeType: item.mediaType, data: await this.attachments.readBase64(id, item.attachmentId) })));
|
|
session.timeout = setTimeout(() => void this.cancel(id, "timeout"), this.runTimeoutMs);
|
|
if (expanded.skill) {
|
|
for (const type of ["skill.requested", "skill.loaded"] as const) this.events.append(id, { type, runId, payload: { name: expanded.skill.name, filePath: expanded.skill.filePath, source: "command" } });
|
|
}
|
|
void this.executeRun(session, redactServerSecret(expanded.text), images);
|
|
return { runId };
|
|
}
|
|
|
|
private async ensureAgent(session: InternalSession): Promise<void> {
|
|
if (session.agent) return;
|
|
if (!this.factory) throw new Error("Agent session factory is not configured");
|
|
session.agent = await this.factory.create({
|
|
sessionId: session.sessionId,
|
|
cwd: session.cwd,
|
|
workspace: session.workspace,
|
|
getRun: () => session.currentRun ? { runId: session.currentRun.runId, signal: session.currentRun.controller.signal, assistantMessageId: session.currentRun.assistantMessageId } : undefined,
|
|
skillRequested: (skill, toolCallId) => { const run = session.currentRun; if (run) this.events.append(session.sessionId, { type: "skill.requested", runId: run.runId, payload: { name: skill.name, filePath: skill.filePath, toolCallId, source: "read" } }); },
|
|
skillLoaded: (skill, toolCallId) => { const run = session.currentRun; if (run) this.events.append(session.sessionId, { type: "skill.loaded", runId: run.runId, payload: { name: skill.name, filePath: skill.filePath, toolCallId, source: "read" } }); },
|
|
toolRequested: (toolCallId, toolName, argumentsValue) => { const run = session.currentRun; if (run && !session.requestedToolCalls.has(toolCallId)) { session.requestedToolCalls.add(toolCallId); this.events.append(session.sessionId, { type: "tool.requested", runId: run.runId, payload: { toolCallId, toolName, arguments: redactValue(argumentsValue) } }); } },
|
|
});
|
|
session.unsubscribe = session.agent.subscribe((event) => this.onPiEvent(session, event as AgentSessionEvent));
|
|
}
|
|
private onPiEvent(session: InternalSession, event: AgentSessionEvent): void {
|
|
const run = session.currentRun;
|
|
if (!run) return;
|
|
// The blocking Extension preflight is the single source for tool.requested.
|
|
// Pi's later execution-start event has the same toolCallId and would create
|
|
// duplicate assistant-ui Tool parts (and can expose pre-normalized paths).
|
|
if (event.type === "tool_execution_start") return;
|
|
for (const mapped of mapPiEvent(event, { runId: run.runId, assistantMessageId: run.assistantMessageId })) this.events.append(session.sessionId, mapped);
|
|
}
|
|
private async executeRun(session: InternalSession, text: string, images: ImageContent[]): Promise<void> {
|
|
const run = session.currentRun;
|
|
if (!run || !session.agent) return;
|
|
try {
|
|
await session.agent.prompt(text, images.length ? { images } : undefined);
|
|
if (!run.controller.signal.aborted) this.events.append(session.sessionId, { type: "run.completed", runId: run.runId, payload: {} });
|
|
} catch (error) {
|
|
if (!run.controller.signal.aborted) this.events.append(session.sessionId, { type: "run.failed", runId: run.runId, payload: { code: "AGENT_RUN_FAILED", message: safeErrorMessage(error) } });
|
|
} finally {
|
|
if (session.timeout) clearTimeout(session.timeout);
|
|
session.timeout = undefined;
|
|
session.busy = false;
|
|
if (session.currentRun?.runId === run.runId) session.currentRun = undefined;
|
|
}
|
|
}
|
|
async cancel(id: string, reason: "user" | "timeout" | "session_deleted" | "shutdown" = "user"): Promise<void> {
|
|
const session = this.internal(id);
|
|
const run = session.currentRun;
|
|
if (!run || !session.busy || run.controller.signal.aborted) return;
|
|
run.cancelReason = reason;
|
|
run.controller.abort(reason);
|
|
this.approvals.cancelForSession(id);
|
|
await session.agent?.abort().catch(() => undefined);
|
|
this.events.append(id, { type: "run.cancelled", runId: run.runId, payload: { reason } });
|
|
}
|
|
async delete(id: string, endReason: "deleted" | "shutdown" = "deleted"): Promise<void> {
|
|
const session = this.internal(id);
|
|
if (session.busy) await this.cancel(id, endReason === "shutdown" ? "shutdown" : "session_deleted");
|
|
this.approvals.cancelForSession(id);
|
|
this.events.append(id, { type: "session.ended", payload: { reason: endReason } });
|
|
await this.disposeAgent(session);
|
|
await this.attachments.removeAll(id);
|
|
await this.models.remove(id);
|
|
const root = await realpath(this.sessionsRoot);
|
|
const target = await realpath(session.cwd);
|
|
const relative = path.relative(root, target);
|
|
if (!relative || relative.startsWith("..") || path.isAbsolute(relative)) throw new HarnessError("UNSAFE_SESSION_DELETE", "Refusing unsafe session directory deletion", 500);
|
|
await rm(target, { recursive: true, force: true });
|
|
session.disposed = true;
|
|
this.sessions.delete(id);
|
|
}
|
|
async shutdown(): Promise<void> { for (const id of [...this.sessions.keys()]) await this.delete(id, "shutdown"); }
|
|
private async disposeAgent(session: InternalSession): Promise<void> { session.unsubscribe?.(); session.unsubscribe = undefined; session.agent?.dispose(); session.agent = undefined; }
|
|
}
|
|
|
|
function publicToolArguments(value: unknown, cwd: string, skill?: SkillIndexEntry): unknown {
|
|
const redacted = redactValue(value);
|
|
if (!redacted || typeof redacted !== "object" || Array.isArray(redacted)) return redacted;
|
|
const result = { ...(redacted as Record<string, unknown>) };
|
|
for (const key of ["path", "cwd", "filePath"]) {
|
|
const candidate = result[key];
|
|
if (typeof candidate !== "string" || !path.isAbsolute(candidate)) continue;
|
|
if (skill && key === "path") result[key] = skill.filePath;
|
|
else {
|
|
const relative = path.relative(cwd, candidate);
|
|
result[key] = relative && !relative.startsWith("..") && !path.isAbsolute(relative) ? relative : "<protected-path>";
|
|
}
|
|
}
|
|
return result;
|
|
}
|