feat(harness): manage Pi agent sessions
This commit is contained in:
Wenxin Duan
@@ -0,0 +1,192 @@
|
||||
import { createHash, randomUUID } from "node:crypto";
|
||||
import { MessageInputSchema, type MessageInput } from "@agent-studio/shared";
|
||||
import type { ApprovalBroker } from "../approvals/approval-broker.js";
|
||||
import type { AttachmentStore } from "../attachments/attachment-store.js";
|
||||
import type { AgentSessionPort, AgentSessionPortFactory } from "../agent/agent-session-port.js";
|
||||
import { PiEventAdapter } from "../agent/pi-event-adapter.js";
|
||||
import type { SkillEventInference } from "../skills/skill-event-inference.js";
|
||||
import type { EventStore } from "../events/event-store.js";
|
||||
import { HarnessError } from "../errors/harness-error.js";
|
||||
import type { ModelConnectionStore, ResolvedModelConnection } from "../models/model-connection-store.js";
|
||||
import { SafeWorkspace } from "../workspace/safe-workspace.js";
|
||||
|
||||
export type RunCancellationReason = "user" | "timeout" | "session_deleted" | "shutdown";
|
||||
|
||||
type ActiveRun = {
|
||||
runId: string;
|
||||
messageId: string;
|
||||
timer: ReturnType<typeof setTimeout>;
|
||||
unsubscribe: () => void;
|
||||
settled: boolean;
|
||||
};
|
||||
|
||||
export type HarnessSessionOptions = {
|
||||
sessionId: string;
|
||||
workspaceRoot: string;
|
||||
eventStore: EventStore;
|
||||
modelStore: ModelConnectionStore;
|
||||
attachmentStore: Pick<AttachmentStore, "toPiImages" | "removeAll">;
|
||||
approvalBroker: Pick<ApprovalBroker, "cancelForRun" | "cancelForSession">;
|
||||
createAgentPort: AgentSessionPortFactory;
|
||||
createSkillInference?: (runId: string) => Pick<SkillEventInference, "onReadRequested" | "onReadCompleted">;
|
||||
runTimeoutMs?: number;
|
||||
};
|
||||
|
||||
export class HarnessSession {
|
||||
public readonly sessionId: string;
|
||||
public readonly workspaceRoot: string;
|
||||
readonly #eventStore: EventStore;
|
||||
readonly #modelStore: ModelConnectionStore;
|
||||
readonly #attachmentStore: Pick<AttachmentStore, "toPiImages" | "removeAll">;
|
||||
readonly #approvalBroker: Pick<ApprovalBroker, "cancelForRun" | "cancelForSession">;
|
||||
readonly #createAgentPort: AgentSessionPortFactory;
|
||||
readonly #createSkillInference: HarnessSessionOptions["createSkillInference"];
|
||||
readonly #runTimeoutMs: number;
|
||||
#port: AgentSessionPort | undefined;
|
||||
#modelFingerprint: string | undefined;
|
||||
#active: ActiveRun | undefined;
|
||||
#disposed = false;
|
||||
|
||||
public constructor(options: HarnessSessionOptions) {
|
||||
this.sessionId = options.sessionId;
|
||||
this.workspaceRoot = options.workspaceRoot;
|
||||
this.#eventStore = options.eventStore;
|
||||
this.#modelStore = options.modelStore;
|
||||
this.#attachmentStore = options.attachmentStore;
|
||||
this.#approvalBroker = options.approvalBroker;
|
||||
this.#createAgentPort = options.createAgentPort;
|
||||
this.#createSkillInference = options.createSkillInference;
|
||||
this.#runTimeoutMs = options.runTimeoutMs ?? 120_000;
|
||||
}
|
||||
|
||||
public async startRun(input: MessageInput): Promise<{ runId: string }> {
|
||||
if (this.#disposed) throw new HarnessError("SESSION_DISPOSED", "Session has been disposed");
|
||||
if (this.#active !== undefined) throw new HarnessError("SESSION_BUSY", "Session already has an active run");
|
||||
const message = MessageInputSchema.parse(input);
|
||||
const resolvedModel = this.#modelStore.resolveConnection(this.sessionId);
|
||||
if (message.attachmentIds.length > 0 && !resolvedModel.connection.supportsImages) {
|
||||
throw new HarnessError("MODEL_DOES_NOT_SUPPORT_IMAGES", "Configured model does not support image input");
|
||||
}
|
||||
const images = await this.#attachmentStore.toPiImages(this.sessionId, message.attachmentIds);
|
||||
const port = await this.#ensurePort(resolvedModel);
|
||||
const runId = `run_${randomUUID()}`;
|
||||
const messageId = `msg_${runId}`;
|
||||
const skillInference = this.#createSkillInference?.(runId);
|
||||
const adapter = new PiEventAdapter({
|
||||
sessionId: this.sessionId,
|
||||
runId,
|
||||
eventStore: this.#eventStore,
|
||||
...(skillInference === undefined ? {} : { skillInference }),
|
||||
});
|
||||
const unsubscribe = port.subscribe((event) => {
|
||||
void adapter.observe(event);
|
||||
});
|
||||
const active: ActiveRun = {
|
||||
runId,
|
||||
messageId,
|
||||
unsubscribe,
|
||||
settled: false,
|
||||
timer: setTimeout(() => {
|
||||
void this.cancelRun("timeout");
|
||||
}, this.#runTimeoutMs),
|
||||
};
|
||||
this.#active = active;
|
||||
this.#eventStore.append(this.sessionId, {
|
||||
type: "user.message",
|
||||
runId,
|
||||
payload: { messageId: `usr_${runId}`, text: message.text, attachmentIds: message.attachmentIds },
|
||||
});
|
||||
this.#eventStore.append(this.sessionId, {
|
||||
type: "run.started",
|
||||
runId,
|
||||
payload: { modelId: resolvedModel.connection.modelId },
|
||||
});
|
||||
|
||||
void Promise.resolve()
|
||||
.then(() => port.prompt(message.text, { images }))
|
||||
.then(() => this.#complete(active))
|
||||
.catch((error: unknown) => this.#fail(active, error));
|
||||
return { runId };
|
||||
}
|
||||
|
||||
public async cancelRun(reason: RunCancellationReason): Promise<boolean> {
|
||||
const active = this.#active;
|
||||
if (active === undefined || active.settled) return false;
|
||||
active.settled = true;
|
||||
clearTimeout(active.timer);
|
||||
this.#approvalBroker.cancelForRun(active.runId);
|
||||
active.unsubscribe();
|
||||
await this.#port?.abort().catch(() => undefined);
|
||||
if (this.#active === active) this.#active = undefined;
|
||||
this.#eventStore.append(this.sessionId, {
|
||||
type: "run.cancelled",
|
||||
runId: active.runId,
|
||||
payload: { reason },
|
||||
});
|
||||
return true;
|
||||
}
|
||||
|
||||
public async dispose(): Promise<void> {
|
||||
if (this.#disposed) return;
|
||||
this.#disposed = true;
|
||||
await this.cancelRun("session_deleted");
|
||||
this.#approvalBroker.cancelForSession(this.sessionId);
|
||||
this.#port?.dispose();
|
||||
this.#port = undefined;
|
||||
await this.#attachmentStore.removeAll(this.sessionId);
|
||||
const workspace = await SafeWorkspace.create(this.workspaceRoot);
|
||||
await workspace.remove();
|
||||
this.#eventStore.append(this.sessionId, {
|
||||
type: "session.ended",
|
||||
payload: { reason: "deleted" },
|
||||
});
|
||||
}
|
||||
|
||||
async #ensurePort(resolvedModel: ResolvedModelConnection): Promise<AgentSessionPort> {
|
||||
const fingerprint = modelFingerprint(resolvedModel);
|
||||
if (this.#port !== undefined && this.#modelFingerprint === fingerprint) return this.#port;
|
||||
this.#port?.dispose();
|
||||
this.#port = await this.#createAgentPort({
|
||||
resolvedModel,
|
||||
workspaceRoot: this.workspaceRoot,
|
||||
sessionId: this.sessionId,
|
||||
getActiveRunId: () => this.#active?.runId,
|
||||
});
|
||||
this.#modelFingerprint = fingerprint;
|
||||
return this.#port;
|
||||
}
|
||||
|
||||
#complete(active: ActiveRun): void {
|
||||
if (active.settled || this.#active !== active) return;
|
||||
active.settled = true;
|
||||
clearTimeout(active.timer);
|
||||
active.unsubscribe();
|
||||
this.#active = undefined;
|
||||
this.#eventStore.append(this.sessionId, {
|
||||
type: "run.completed",
|
||||
runId: active.runId,
|
||||
payload: { messageId: active.messageId },
|
||||
});
|
||||
}
|
||||
|
||||
#fail(active: ActiveRun, error: unknown): void {
|
||||
if (active.settled || this.#active !== active) return;
|
||||
active.settled = true;
|
||||
clearTimeout(active.timer);
|
||||
active.unsubscribe();
|
||||
this.#approvalBroker.cancelForRun(active.runId);
|
||||
this.#active = undefined;
|
||||
this.#eventStore.append(this.sessionId, {
|
||||
type: "run.failed",
|
||||
runId: active.runId,
|
||||
payload: {
|
||||
code: error instanceof HarnessError ? error.code : "AGENT_RUN_FAILED",
|
||||
message: error instanceof Error ? error.message.slice(0, 2_000) : "Agent run failed",
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
function modelFingerprint(resolved: ResolvedModelConnection): string {
|
||||
return createHash("sha256").update(JSON.stringify(resolved)).digest("hex");
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { mkdir, realpath, rm } from "node:fs/promises";
|
||||
import { isAbsolute, join, relative, resolve, sep } from "node:path";
|
||||
import { HarnessError } from "../errors/harness-error.js";
|
||||
import type { SessionRecord, SessionRegistry } from "./session-registry.js";
|
||||
|
||||
export type InMemorySessionRegistryOptions<TSession extends SessionRecord> = {
|
||||
sessionsRoot: string;
|
||||
createSession(input: { sessionId: string; workspaceRoot: string }): TSession;
|
||||
};
|
||||
|
||||
export class InMemorySessionRegistry<TSession extends SessionRecord = SessionRecord>
|
||||
implements SessionRegistry<TSession> {
|
||||
readonly #sessionsRoot: string;
|
||||
readonly #createSession: InMemorySessionRegistryOptions<TSession>["createSession"];
|
||||
readonly #sessions = new Map<string, TSession>();
|
||||
|
||||
public constructor(options: InMemorySessionRegistryOptions<TSession>) {
|
||||
this.#sessionsRoot = resolve(options.sessionsRoot);
|
||||
this.#createSession = (input) => options.createSession(input);
|
||||
}
|
||||
|
||||
public async create(): Promise<TSession> {
|
||||
const canonicalRoot = await realpath(this.#sessionsRoot);
|
||||
const sessionId = `ses_${randomUUID()}`;
|
||||
const workspaceRoot = join(canonicalRoot, sessionId);
|
||||
assertContained(canonicalRoot, workspaceRoot);
|
||||
await mkdir(workspaceRoot, { mode: 0o700 });
|
||||
const session = this.#createSession({ sessionId, workspaceRoot });
|
||||
this.#sessions.set(sessionId, session);
|
||||
return session;
|
||||
}
|
||||
|
||||
public get(sessionId: string): TSession | undefined {
|
||||
return this.#sessions.get(sessionId);
|
||||
}
|
||||
|
||||
public async delete(sessionId: string): Promise<boolean> {
|
||||
const session = this.#sessions.get(sessionId);
|
||||
if (session === undefined) return false;
|
||||
this.#sessions.delete(sessionId);
|
||||
await session.dispose();
|
||||
const canonicalRoot = await realpath(this.#sessionsRoot);
|
||||
assertContained(canonicalRoot, session.workspaceRoot);
|
||||
await rm(session.workspaceRoot, { recursive: true, force: true });
|
||||
return true;
|
||||
}
|
||||
|
||||
public async disposeAll(): Promise<void> {
|
||||
for (const sessionId of [...this.#sessions.keys()]) await this.delete(sessionId);
|
||||
}
|
||||
}
|
||||
|
||||
function assertContained(root: string, candidate: string): void {
|
||||
const fromRoot = relative(root, candidate);
|
||||
if (fromRoot !== "" && fromRoot !== ".." && !fromRoot.startsWith(`..${sep}`) && !isAbsolute(fromRoot)) return;
|
||||
throw new HarnessError("UNSAFE_SESSION_PATH", "Session workspace must be a child of the sessions root");
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
export type SessionRecord = {
|
||||
sessionId: string;
|
||||
workspaceRoot: string;
|
||||
dispose(): Promise<void>;
|
||||
};
|
||||
|
||||
export interface SessionRegistry<TSession extends SessionRecord = SessionRecord> {
|
||||
create(): Promise<TSession>;
|
||||
get(sessionId: string): TSession | undefined;
|
||||
delete(sessionId: string): Promise<boolean>;
|
||||
disposeAll(): Promise<void>;
|
||||
}
|
||||
Reference in New Issue
Block a user