pi-harness-runtime 0.5.0 → 0.6.3-beta.13

package/packages/session/src/policy.ts

@@ -0,0 +1,363 @@
+/**
+ * Session Manager - Policy Engine
+ *
+ * Session policy enforcement and budget tracking.
+ */
+
+import type { PolicyEngineConfig, PolicyState, TokenUsage } from "./types.js";
+
+// ─── Rate Limiter ──────────────────────────────────────────────────────────
+
+interface RateLimitEntry {
+	count: number;
+	resetAt: number;
+}
+
+export class RateLimiter {
+	private entries: Map<string, RateLimitEntry> = new Map();
+	private readonly requestsPerMinute: number;
+
+	constructor(requestsPerMinute: number) {
+		this.requestsPerMinute = requestsPerMinute;
+	}
+
+	/**
+	 * Check if request is allowed
+	 */
+	canMakeRequest(key: string): boolean {
+		const now = Date.now();
+		const entry = this.entries.get(key);
+
+		if (!entry || now >= entry.resetAt) {
+			return true;
+		}
+
+		return entry.count < this.requestsPerMinute;
+	}
+
+	/**
+	 * Record a request
+	 */
+	recordRequest(key: string): void {
+		const now = Date.now();
+		const entry = this.entries.get(key);
+
+		if (!entry || now >= entry.resetAt) {
+			// Start new window
+			this.entries.set(key, {
+				count: 1,
+				resetAt: now + 60000, // 1 minute
+			});
+		} else {
+			entry.count++;
+		}
+	}
+
+	/**
+	 * Get remaining requests
+	 */
+	getRemaining(key: string): number {
+		const entry = this.entries.get(key);
+		if (!entry || Date.now() >= entry.resetAt) {
+			return this.requestsPerMinute;
+		}
+
+		return Math.max(0, this.requestsPerMinute - entry.count);
+	}
+
+	/**
+	 * Clear expired entries
+	 */
+	clearExpired(): void {
+		const now = Date.now();
+		for (const [key, entry] of this.entries) {
+			if (now >= entry.resetAt) {
+				this.entries.delete(key);
+			}
+		}
+	}
+}
+
+// ─── Budget Tracker ────────────────────────────────────────────────────────
+
+interface BudgetEntry {
+	spent: number;
+	limit: number;
+	resetAt: number;
+}
+
+export class BudgetTracker {
+	private entries: Map<string, BudgetEntry> = new Map();
+
+	/**
+	 * Set budget for a key
+	 */
+	setBudget(key: string, limit: number, resetAtMs: number): void {
+		this.entries.set(key, {
+			spent: 0,
+			limit,
+			resetAt: Date.now() + resetAtMs,
+		});
+	}
+
+	/**
+	 * Check if budget allows spending
+	 */
+	canSpend(key: string, amount: number): boolean {
+		const entry = this.entries.get(key);
+		if (!entry) return true;
+
+		// Check if reset time passed
+		if (Date.now() >= entry.resetAt) {
+			return true;
+		}
+
+		return entry.spent + amount <= entry.limit;
+	}
+
+	/**
+	 * Record spending
+	 */
+	recordSpend(key: string, amount: number): void {
+		const entry = this.entries.get(key);
+		if (!entry) return;
+
+		// Check if reset time passed
+		if (Date.now() >= entry.resetAt) {
+			entry.spent = amount;
+			entry.resetAt = Date.now() + (entry.resetAt - Date.now());
+		} else {
+			entry.spent += amount;
+		}
+	}
+
+	/**
+	 * Get remaining budget
+	 */
+	getRemaining(key: string): number | null {
+		const entry = this.entries.get(key);
+		if (!entry) return null;
+
+		if (Date.now() >= entry.resetAt) {
+			return entry.limit;
+		}
+
+		return Math.max(0, entry.limit - entry.spent);
+	}
+
+	/**
+	 * Get reset time
+	 */
+	getResetAt(key: string): number | null {
+		const entry = this.entries.get(key);
+		if (!entry) return null;
+
+		return entry.resetAt;
+	}
+}
+
+// ─── Policy Engine ─────────────────────────────────────────────────────────
+
+export class PolicyEngine {
+	private readonly config: PolicyEngineConfig;
+	private readonly rateLimiter: RateLimiter;
+	private readonly budgetTracker: BudgetTracker;
+	private sessionStates: Map<string, PolicyState> = new Map();
+
+	constructor(config: PolicyEngineConfig = {}) {
+		this.config = {
+			maxRequestsPerMinute: config.maxRequestsPerMinute ?? 60,
+			maxTokensPerDay: config.maxTokensPerDay ?? 1000000,
+			maxCostPerSession: config.maxCostPerSession ?? 100,
+			maxConcurrentSessions: config.maxConcurrentSessions ?? 10,
+			sessionBudget: config.sessionBudget ?? 50,
+		};
+
+		this.rateLimiter = new RateLimiter(this.config.maxRequestsPerMinute!);
+		this.budgetTracker = new BudgetTracker();
+	}
+
+	/**
+	 * Get policy state for a session
+	 */
+	getPolicyState(sessionId: string): PolicyState {
+		let state = this.sessionStates.get(sessionId);
+		if (!state) {
+			state = {
+				rateLimitRemaining: this.config.maxRequestsPerMinute ?? 60,
+				suspended: false,
+			};
+			this.sessionStates.set(sessionId, state);
+		}
+		return state;
+	}
+
+	/**
+	 * Check if action is allowed
+	 */
+	canProceed(
+		sessionId: string,
+		action: "message" | "tool_call" | "session_create",
+	): boolean {
+		const state = this.getPolicyState(sessionId);
+
+		// Check if suspended
+		if (state.suspended) {
+			return false;
+		}
+
+		// Check rate limit
+		if (!this.rateLimiter.canMakeRequest(sessionId)) {
+			return false;
+		}
+
+		// Check budget
+		if (state.budgetRemaining !== undefined && state.budgetRemaining <= 0) {
+			return false;
+		}
+
+		return true;
+	}
+
+	/**
+	 * Record an action
+	 */
+	recordAction(sessionId: string, action: "message" | "tool_call"): void {
+		// Record rate limit
+		this.rateLimiter.recordRequest(sessionId);
+
+		// Update state
+		const state = this.getPolicyState(sessionId);
+		state.rateLimitRemaining = this.rateLimiter.getRemaining(sessionId);
+
+		// Update rate limit reset time
+		const resetAt = Date.now() + 60000;
+		state.rateLimitResetAt = new Date(resetAt).toISOString();
+	}
+
+	/**
+	 * Record token usage
+	 */
+	recordTokenUsage(sessionId: string, usage: TokenUsage): void {
+		const state = this.getPolicyState(sessionId);
+
+		if (this.config.sessionBudget) {
+			// Convert cost to cents for tracking
+			const costInCents = Math.round(usage.totalCost * 100);
+			this.budgetTracker.recordSpend(sessionId, costInCents);
+
+			state.budgetRemaining = this.budgetTracker.getRemaining(sessionId)
+				? this.budgetTracker.getRemaining(sessionId)! / 100
+				: undefined;
+
+			const resetAt = this.budgetTracker.getResetAt(sessionId);
+			if (resetAt) {
+				state.budgetResetAt = new Date(resetAt).toISOString();
+			}
+		}
+	}
+
+	/**
+	 * Suspend a session
+	 */
+	suspend(sessionId: string, reason: string): void {
+		const state = this.getPolicyState(sessionId);
+		state.suspended = true;
+		state.suspensionReason = reason;
+	}
+
+	/**
+	 * Resume a session
+	 */
+	resume(sessionId: string): void {
+		const state = this.getPolicyState(sessionId);
+		state.suspended = false;
+		state.suspensionReason = undefined;
+	}
+
+	/**
+	 * Reset session policy state
+	 */
+	resetState(sessionId: string): void {
+		this.sessionStates.delete(sessionId);
+	}
+
+	/**
+	 * Set session budget
+	 */
+	setSessionBudget(sessionId: string, budget: number, periodMs: number): void {
+		const budgetInCents = Math.round(budget * 100);
+		this.budgetTracker.setBudget(sessionId, budgetInCents, periodMs);
+
+		const state = this.getPolicyState(sessionId);
+		state.budgetRemaining = budget;
+		state.budgetResetAt = new Date(Date.now() + periodMs).toISOString();
+	}
+
+	/**
+	 * Get violation type
+	 */
+	getViolationType(sessionId: string): string | null {
+		const state = this.getPolicyState(sessionId);
+
+		if (state.suspended) {
+			return `suspended: ${state.suspensionReason}`;
+		}
+
+		if (state.rateLimitRemaining <= 0) {
+			return "rate_limit_exceeded";
+		}
+
+		if (state.budgetRemaining !== undefined && state.budgetRemaining <= 0) {
+			return "budget_exceeded";
+		}
+
+		return null;
+	}
+
+	/**
+	 * Check if session exceeds cost limit
+	 */
+	exceedsCostLimit(sessionId: string, additionalCost: number): boolean {
+		const state = this.getPolicyState(sessionId);
+
+		if (this.config.maxCostPerSession && state.budgetRemaining !== undefined) {
+			return state.budgetRemaining - additionalCost < 0;
+		}
+
+		return false;
+	}
+
+	/**
+	 * Get rate limit info
+	 */
+	getRateLimitInfo(
+		sessionId: string,
+	): { remaining: number; resetAt: string } | null {
+		const remaining = this.rateLimiter.getRemaining(sessionId);
+		if (remaining === this.config.maxRequestsPerMinute) {
+			return null;
+		}
+
+		return {
+			remaining,
+			resetAt: new Date(Date.now() + 60000).toISOString(),
+		};
+	}
+
+	/**
+	 * Clean up expired state
+	 */
+	cleanup(): void {
+		this.rateLimiter.clearExpired();
+	}
+}
+
+// ─── Factory Function ────────────────────────────────────────────────────
+
+/**
+ * Create a policy engine
+ */
+export function createPolicyEngine(config?: PolicyEngineConfig): PolicyEngine {
+	return new PolicyEngine(config);
+}

package/packages/session/src/store.ts

@@ -0,0 +1,348 @@
+/**
+ * Session Manager - Store
+ *
+ * Session persistence to disk with indexing.
+ */
+
+import {
+	readdir,
+	readFile,
+	writeFile,
+	mkdir,
+	rename,
+	unlink,
+} from "node:fs/promises";
+import { dirname, join } from "node:path";
+import type {
+	Session,
+	SessionContext,
+	SessionStoreConfig,
+	SessionStatus,
+} from "./types.js";
+
+// ─── Session Index ────────────────────────────────────────────────────────
+
+/**
+ * Session index entry
+ */
+interface SessionIndexEntry {
+	id: string;
+	userId: string;
+	status: SessionStatus;
+	createdAt: string;
+	updatedAt: string;
+	expiresAt?: string;
+	messageCount: number;
+	totalTokens: number;
+}
+
+// ─── Session Store ────────────────────────────────────────────────────────
+
+export class SessionStore {
+	private readonly rootDir: string;
+	private readonly autoSave: boolean;
+	private readonly autoSaveIntervalMs: number;
+	private readonly maxSessions?: number;
+
+	/**
+	 * Get max sessions limit
+	 */
+	getMaxSessions(): number | undefined {
+		return this.maxSessions;
+	}
+	private indices: Map<string, SessionIndexEntry> = new Map();
+	private autoSaveTimers: Map<string, NodeJS.Timeout> = new Map();
+
+	constructor(config: SessionStoreConfig) {
+		this.rootDir = config.rootDir;
+		this.autoSave = config.autoSave;
+		this.autoSaveIntervalMs = config.autoSaveIntervalMs;
+		this.maxSessions = config.maxSessions;
+	}
+
+	/**
+	 * Ensure directory exists
+	 */
+	private async ensureDir(dir: string): Promise<void> {
+		await mkdir(dir, { recursive: true });
+	}
+
+	/**
+	 * Get session directory
+	 */
+	private sessionDir(sessionId: string): string {
+		return join(this.rootDir, "sessions", sessionId.slice(0, 2), sessionId);
+	}
+
+	/**
+	 * Get session file path
+	 */
+	private sessionPath(sessionId: string): string {
+		return join(this.sessionDir(sessionId), "session.json");
+	}
+
+	/**
+	 * Get index file path
+	 */
+	private indexPath(): string {
+		return join(this.rootDir, "sessions", "index.json");
+	}
+
+	/**
+	 * Load or create index
+	 */
+	private async loadIndex(): Promise<Map<string, SessionIndexEntry>> {
+		const path = this.indexPath();
+		try {
+			const content = await readFile(path, "utf-8");
+			const entries = JSON.parse(content) as SessionIndexEntry[];
+			return new Map(entries.map((e) => [e.id, e]));
+		} catch {
+			return new Map();
+		}
+	}
+
+	/**
+	 * Save index
+	 */
+	private async saveIndex(): Promise<void> {
+		const entries = Array.from(this.indices.values());
+		await this.ensureDir(dirname(this.indexPath()));
+		await writeFile(
+			this.indexPath(),
+			JSON.stringify(entries, null, 2),
+			"utf-8",
+		);
+	}
+
+	/**
+	 * Save session context to disk
+	 */
+	private async saveToDisk(context: SessionContext): Promise<void> {
+		const path = this.sessionPath(context.id);
+		await this.ensureDir(dirname(path));
+		const tmp = `${path}.tmp`;
+		await writeFile(tmp, JSON.stringify(context, null, 2), "utf-8");
+		await rename(tmp, path);
+	}
+
+	/**
+	 * Load session context from disk
+	 */
+	private async loadFromDisk(
+		sessionId: string,
+	): Promise<SessionContext | null> {
+		const path = this.sessionPath(sessionId);
+		try {
+			const content = await readFile(path, "utf-8");
+			return JSON.parse(content) as SessionContext;
+		} catch {
+			return null;
+		}
+	}
+
+	/**
+	 * Create or update session index entry
+	 */
+	private updateIndexEntry(context: SessionContext): void {
+		this.indices.set(context.id, {
+			id: context.id,
+			userId: context.userId,
+			status: context.status,
+			createdAt: context.createdAt,
+			updatedAt: context.updatedAt,
+			expiresAt: context.expiresAt,
+			messageCount: context.messages.length,
+			totalTokens: context.tokenUsage.totalTokens,
+		});
+	}
+
+	/**
+	 * Initialize the store
+	 */
+	async initialize(): Promise<void> {
+		await this.ensureDir(this.rootDir);
+		this.indices = await this.loadIndex();
+	}
+
+	/**
+	 * Save session context
+	 */
+	async save(context: SessionContext): Promise<void> {
+		this.updateIndexEntry(context);
+		await this.saveToDisk(context);
+
+		if (this.autoSave) {
+			this.scheduleAutoSave(context.id);
+		}
+
+		// Save index periodically
+		await this.saveIndex();
+	}
+
+	/**
+	 * Schedule auto-save for a session
+	 */
+	private scheduleAutoSave(sessionId: string): void {
+		// Clear existing timer
+		const existing = this.autoSaveTimers.get(sessionId);
+		if (existing) {
+			clearTimeout(existing);
+		}
+
+		// Schedule new save
+		const timer = setTimeout(async () => {
+			this.autoSaveTimers.delete(sessionId);
+			const context = await this.loadFromDisk(sessionId);
+			if (context) {
+				await this.saveToDisk(context);
+			}
+		}, this.autoSaveIntervalMs);
+
+		this.autoSaveTimers.set(sessionId, timer);
+	}
+
+	/**
+	 * Load session context
+	 */
+	async load(sessionId: string): Promise<SessionContext | null> {
+		return this.loadFromDisk(sessionId);
+	}
+
+	/**
+	 * Delete session
+	 */
+	async delete(sessionId: string): Promise<void> {
+		// Cancel auto-save
+		const timer = this.autoSaveTimers.get(sessionId);
+		if (timer) {
+			clearTimeout(timer);
+			this.autoSaveTimers.delete(sessionId);
+		}
+
+		// Remove from index
+		this.indices.delete(sessionId);
+
+		// Delete files
+		const dir = this.sessionDir(sessionId);
+		try {
+			const files = await readdir(dir);
+			await Promise.all(files.map((f) => unlink(join(dir, f))));
+		} catch {
+			// Ignore if directory doesn't exist
+		}
+
+		// Update index
+		await this.saveIndex();
+	}
+
+	/**
+	 * List sessions for a user
+	 */
+	async listByUser(userId: string): Promise<Session[]> {
+		const entries = Array.from(this.indices.values()).filter(
+			(e) => e.userId === userId,
+		);
+
+		return entries.map((e) => ({
+			id: e.id,
+			userId: e.userId,
+			status: e.status,
+			createdAt: e.createdAt,
+			updatedAt: e.updatedAt,
+			lastActivityAt: e.updatedAt,
+			expiresAt: e.expiresAt,
+		}));
+	}
+
+	/**
+	 * List all sessions
+	 */
+	async listAll(): Promise<Session[]> {
+		const entries = Array.from(this.indices.values());
+
+		return entries.map((e) => ({
+			id: e.id,
+			userId: e.userId,
+			status: e.status,
+			createdAt: e.createdAt,
+			updatedAt: e.updatedAt,
+			lastActivityAt: e.updatedAt,
+			expiresAt: e.expiresAt,
+		}));
+	}
+
+	/**
+	 * Get session summary without loading full context
+	 */
+	async getSummary(sessionId: string): Promise<Session | null> {
+		const entry = this.indices.get(sessionId);
+		if (!entry) return null;
+
+		return {
+			id: entry.id,
+			userId: entry.userId,
+			status: entry.status,
+			createdAt: entry.createdAt,
+			updatedAt: entry.updatedAt,
+			lastActivityAt: entry.updatedAt,
+			expiresAt: entry.expiresAt,
+		};
+	}
+
+	/**
+	 * Check if session exists
+	 */
+	async exists(sessionId: string): Promise<boolean> {
+		return this.indices.has(sessionId);
+	}
+
+	/**
+	 * Get session count
+	 */
+	async count(): Promise<number> {
+		return this.indices.size;
+	}
+
+	/**
+	 * Get sessions by status
+	 */
+	async listByStatus(status: SessionStatus): Promise<Session[]> {
+		const entries = Array.from(this.indices.values()).filter(
+			(e) => e.status === status,
+		);
+
+		return entries.map((e) => ({
+			id: e.id,
+			userId: e.userId,
+			status: e.status,
+			createdAt: e.createdAt,
+			updatedAt: e.updatedAt,
+			lastActivityAt: e.updatedAt,
+			expiresAt: e.expiresAt,
+		}));
+	}
+
+	/**
+	 * Close store and cleanup
+	 */
+	async close(): Promise<void> {
+		// Clear all auto-save timers
+		for (const timer of this.autoSaveTimers.values()) {
+			clearTimeout(timer);
+		}
+		this.autoSaveTimers.clear();
+
+		// Save final index
+		await this.saveIndex();
+	}
+}
+
+// ─── Factory Function ────────────────────────────────────────────────────
+
+/**
+ * Create a session store
+ */
+export function createSessionStore(config: SessionStoreConfig): SessionStore {
+	return new SessionStore(config);
+}