feat(harness): manage Pi agent sessions
This commit is contained in:
@@ -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<void>;
|
||||||
|
abort(): Promise<void>;
|
||||||
|
dispose(): void;
|
||||||
|
}
|
||||||
|
|
||||||
|
export type AgentSessionPortFactory = (input: {
|
||||||
|
resolvedModel: ResolvedModelConnection;
|
||||||
|
workspaceRoot: string;
|
||||||
|
sessionId: string;
|
||||||
|
getActiveRunId: () => string | undefined;
|
||||||
|
}) => Promise<AgentSessionPort>;
|
||||||
@@ -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<void> {
|
||||||
|
return this.session.prompt(text, { images: options.images, source: "rpc" });
|
||||||
|
}
|
||||||
|
|
||||||
|
public abort(): Promise<void> {
|
||||||
|
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, unknown>): string | undefined {
|
||||||
|
for (const key of ["path", "filePath", "directory"]) {
|
||||||
|
if (typeof input[key] === "string") return input[key];
|
||||||
|
}
|
||||||
|
return undefined;
|
||||||
|
}
|
||||||
@@ -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<void>;
|
||||||
|
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<void> {
|
||||||
|
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<string, unknown>): Promise<void> {
|
||||||
|
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<string, unknown>): 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<string, unknown> {
|
||||||
|
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);
|
||||||
|
}
|
||||||
@@ -1,5 +1,8 @@
|
|||||||
export * from "./approvals/approval-broker.js";
|
export * from "./approvals/approval-broker.js";
|
||||||
export * from "./approvals/in-memory-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/attachment-store.js";
|
||||||
export * from "./attachments/file-attachment-store.js";
|
export * from "./attachments/file-attachment-store.js";
|
||||||
export * from "./attachments/image-signatures.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/model-connection-store.js";
|
||||||
export * from "./models/pi-model-runtime-factory.js";
|
export * from "./models/pi-model-runtime-factory.js";
|
||||||
export * from "./security/redact.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-catalog.js";
|
||||||
export * from "./skills/skill-event-inference.js";
|
export * from "./skills/skill-event-inference.js";
|
||||||
export * from "./tools/mock-ads-metrics.js";
|
export * from "./tools/mock-ads-metrics.js";
|
||||||
|
|||||||
@@ -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>;
|
||||||
|
}
|
||||||
@@ -32,16 +32,21 @@ export class SafeWorkspace {
|
|||||||
}
|
}
|
||||||
|
|
||||||
public async resolveReadable(inputPath: string): Promise<string> {
|
public async resolveReadable(inputPath: string): Promise<string> {
|
||||||
|
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<string> {
|
||||||
const candidate = isAbsolute(inputPath) ? resolve(inputPath) : resolve(this.root, inputPath);
|
const candidate = isAbsolute(inputPath) ? resolve(inputPath) : resolve(this.root, inputPath);
|
||||||
this.#assertContained(candidate);
|
this.#assertContained(candidate);
|
||||||
this.#assertUnprotected(candidate);
|
this.#assertUnprotected(candidate);
|
||||||
await this.#assertNoSymlinkSegments(candidate);
|
await this.#assertNoSymlinkSegments(candidate);
|
||||||
const canonical = await realpath(candidate);
|
const canonical = await realpath(candidate);
|
||||||
this.#assertContained(canonical);
|
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;
|
return canonical;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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<string, unknown>) => {
|
||||||
|
observe(event: unknown): Promise<void>;
|
||||||
|
};
|
||||||
|
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<string, unknown>) => {
|
||||||
|
observe(event: unknown): Promise<void>;
|
||||||
|
};
|
||||||
|
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");
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -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<Listener>();
|
||||||
|
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<void> {
|
||||||
|
this.promptCalls.push({ text, images: options.images });
|
||||||
|
return new Promise((resolve, reject) => {
|
||||||
|
this.resolvePrompt = resolve;
|
||||||
|
this.rejectPrompt = reject;
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
abort(): Promise<void> {
|
||||||
|
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<harness.InMemoryModelConnectionStore> {
|
||||||
|
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<string, unknown>) => {
|
||||||
|
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<string, unknown>) => {
|
||||||
|
startRun(message: { text: string; attachmentIds: string[] }): Promise<{ runId: string }>;
|
||||||
|
cancelRun(reason: string): Promise<boolean>;
|
||||||
|
};
|
||||||
|
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<string, unknown>) => {
|
||||||
|
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<string, unknown>) => {
|
||||||
|
create(): Promise<{ sessionId: string; workspaceRoot: string }>;
|
||||||
|
delete(sessionId: string): Promise<boolean>;
|
||||||
|
get(sessionId: string): unknown;
|
||||||
|
};
|
||||||
|
expect(Registry, "InMemorySessionRegistry must be exported").toBeDefined();
|
||||||
|
const registry = new Registry({
|
||||||
|
sessionsRoot: root,
|
||||||
|
createSession: (options: Record<string, unknown>) => ({ ...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();
|
||||||
|
});
|
||||||
|
});
|
||||||
Reference in New Issue
Block a user