feat: initial commit
This commit is contained in:
Jeffrey Wu
@@ -0,0 +1,303 @@
|
||||
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, 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 { 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,
|
||||
) {}
|
||||
|
||||
async create(input: AgentSessionFactoryInput): Promise<AgentSessionPort> {
|
||||
const resolved = await this.models.createRuntime(input.sessionId);
|
||||
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"]);
|
||||
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, 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"],
|
||||
customTools: [createMockAdsMetricsTool(this.fixturePath), createSaveReportDraftTool(input.workspace)],
|
||||
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 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);
|
||||
void this.executeRun(session, message.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 } }); },
|
||||
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 } }); },
|
||||
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;
|
||||
}
|
||||
Reference in New Issue
Block a user