Files
sw-ads-agent/packages/harness/src/session.ts
T

380 lines
23 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, ReportArtifact, 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;
reportGenerated: (toolCallId: string, artifact: ReportArtifact, markdown: string) => void;
toolFailed: (toolCallId: string, toolName: string, error: string) => void;
}
export interface AgentSessionFactory { create(input: AgentSessionFactoryInput): Promise<AgentSessionPort> }
class PiAgentSessionPort implements AgentSessionPort {
readonly modelId: string;
readonly supportsImages: boolean;
constructor(private readonly session: AgentSession, private readonly completeRun?: () => Promise<string | undefined> | string | undefined) {
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);
// Pi reports provider failures through its message stream and resolves prompt().
// Convert a terminal model error into Harness run.failed instead of run.completed.
if (this.session.agent.state.errorMessage) throw new Error("Model provider run failed");
const completionError = await this.completeRun?.();
if (completionError) throw new Error(completionError);
}
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);
let swadsWorkflow: { runId: string; rendered: boolean; queryCount: number; finalTextSeen: boolean; fatal: boolean; explained: boolean } | undefined;
const extension: InlineExtension = { name: "agent-studio-safety", hidden: true, factory: (pi) => {
const pendingSkills = new Map<string, SkillIndexEntry>();
pi.on("before_agent_start", (event) => {
const run = input.getRun();
swadsWorkflow = run && event.prompt.includes('<skill name="swads-daily-report"') ? { runId: run.runId, rendered: false, queryCount: 0, finalTextSeen: false, fatal: false, explained: false } : undefined;
});
pi.on("message_end", (event) => {
const run = input.getRun();
if (!run || swadsWorkflow?.runId !== run.runId || event.message.role !== "assistant") return;
if (swadsWorkflow.queryCount > 0 && assistantText(event.message).trim()) swadsWorkflow.finalTextSeen = true;
});
pi.on("turn_end", (event) => {
const run = input.getRun(); const workflow = swadsWorkflow;
if (!run || workflow?.runId !== run.runId || workflow.rendered || event.message.role !== "assistant") return;
if (workflow.fatal && assistantText(event.message).trim()) workflow.explained = true;
});
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 === "render_swads_daily_report" && swadsWorkflow?.runId === run.runId && !swadsWorkflow.finalTextSeen) return { block: true, reason: "Final Markdown analysis is required before rendering the SW Ads report" };
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();
if (run && swadsWorkflow?.runId === run.runId) {
const toolInput = objectValue(event.input);
if (event.toolName === "swads_cli" && !event.isError && toolInput.command === "metrics.query") swadsWorkflow.queryCount++;
if (event.toolName === "swads_cli" && event.isError && ["capabilities", "whoami", "metrics.catalog", "metrics.query"].includes(String(toolInput.command))) swadsWorkflow.fatal = true;
if (event.toolName === "render_swads_daily_report") {
if (event.isError) swadsWorkflow.fatal = true;
else swadsWorkflow.rendered = true;
}
}
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, async () => {
const run = input.getRun(); const workflow = swadsWorkflow;
if (!run || workflow?.runId !== run.runId || workflow.rendered || workflow.explained) return undefined;
if (!workflow.fatal && workflow.queryCount >= 4) {
const toolCallId = `swads-fallback-${randomUUID()}`;
input.toolRequested(toolCallId, "render_swads_daily_report", { status: "server-fallback" });
try {
const result = await swads.renderCollected(run.signal);
workflow.rendered = true;
input.reportGenerated(toolCallId, result.artifact, result.markdown);
return undefined;
} catch (error) {
input.toolFailed(toolCallId, "render_swads_daily_report", error instanceof HarnessError ? error.code : "SWADS_REPORT_FAILED");
}
}
return "SWADS_REPORT_INCOMPLETE";
});
}
}
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 = 300_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) })));
this.touchRunTimeout(session);
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) } }); } },
reportGenerated: (toolCallId, artifact, markdown) => { const run = session.currentRun; if (run) { this.events.append(session.sessionId, { type: "tool.completed", runId: run.runId, payload: { toolCallId, toolName: "render_swads_daily_report", result: artifact } }); this.events.append(session.sessionId, { type: "report.generated", runId: run.runId, payload: { assistantMessageId: run.assistantMessageId, artifact, markdown } }); } },
toolFailed: (toolCallId, toolName, error) => { const run = session.currentRun; if (run) this.events.append(session.sessionId, { type: "tool.failed", runId: run.runId, payload: { toolCallId, toolName, error } }); },
});
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;
this.touchRunTimeout(session);
// 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;
}
}
private touchRunTimeout(session: InternalSession): void {
const run = session.currentRun;
if (!run || !session.busy || run.controller.signal.aborted) return;
if (session.timeout) clearTimeout(session.timeout);
const runId = run.runId;
session.timeout = setTimeout(() => {
if (session.currentRun?.runId === runId) void this.cancel(session.sessionId, "timeout");
}, this.runTimeoutMs);
}
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 objectValue(value: unknown): Record<string, unknown> { return value && typeof value === "object" && !Array.isArray(value) ? value as Record<string, unknown> : {}; }
function assistantText(message: unknown): string {
const content = objectValue(message).content;
return Array.isArray(content) ? content.filter((part) => objectValue(part).type === "text").map((part) => String(objectValue(part).text ?? "")).join("\n") : "";
}
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;
}