diff --git a/src/extensions/configResolver.ts b/src/extensions/configResolver.ts new file mode 100644 index 0000000..6021764 --- /dev/null +++ b/src/extensions/configResolver.ts @@ -0,0 +1,154 @@ +/** + * Extension configuration resolver. + * + * Encapsulates the layered config lookup (env var > SQLite > schema default > caller default) + * and the settings cache lifecycle. Extracted from extensionContext.ts so the resolution + * logic is independently testable and reusable. + * + * @module + */ + +import { schema } from "@src/db"; +import { eq } from "drizzle-orm"; +import type { BunSQLiteDatabase } from "drizzle-orm/bun-sqlite"; +import type { ConfigValue } from "./types"; + +/** + * Dependencies for creating a config resolver instance. + */ +export interface ConfigResolverDeps { + /** The extension name (used to derive env var prefix). */ + extensionName: string; + /** The shared Drizzle database instance. */ + database: BunSQLiteDatabase>; + /** The extension's settingsSchema (TypeBox TObject), if declared. */ + settingsSchema?: Record; +} + +/** + * Creates a scoped config resolver for a single extension. + * + * Provides a `get(key, default?)` method with layered precedence: + * 1. Environment variable `EXT_{NAME}_{KEY}` + * 2. SQLite persisted settings (cached in memory) + * 3. Schema default value + * 4. Caller-provided default + * + * Also exposes `invalidateCache()` for reacting to settings:changed events. + * + * @param deps - Extension identity, database, and optional schema + * @returns Object with `get` and `invalidateCache` methods + */ +export function createConfigResolver(deps: ConfigResolverDeps) { + const { extensionName, database, settingsSchema } = deps; + + /** Cached settings object from SQLite (null = not yet loaded). */ + let settingsCache: Record | null = null; + + /** + * Load persisted settings from SQLite into the cache. + * Returns the cached object (may be empty `{}`). + */ + function loadSettingsCache(): Record { + if (settingsCache !== null) return settingsCache; + try { + const row = database + .select({ config: schema.extensionSettings.config }) + .from(schema.extensionSettings) + .where(eq(schema.extensionSettings.name, extensionName)) + .get(); + settingsCache = row?.config ? (JSON.parse(row.config) as Record) : {}; + } catch { + settingsCache = {}; + } + return settingsCache; + } + + /** Invalidate the settings cache so the next read fetches from SQLite. */ + function invalidateCache(): void { + settingsCache = null; + } + + /** + * Read a configuration value for this extension by key. + * Precedence: env var > SQLite persisted value > schema default > caller default. + * + * Values are coerced from the raw env-var string: + * `"true"`/`"false"` -> boolean, numeric strings -> number, + * JSON-shaped strings -> parsed object/array, everything else -> string. + * + * @param key - The configuration key (UPPER_SNAKE_CASE). + * @param defaultValue - Returned when no source provides a value. + * @returns The resolved value, or `undefined`. + */ + function get(key: string, defaultValue?: ConfigValue): ConfigValue | undefined { + // 1. Check environment variable (highest precedence) + const envKey = `EXT_${extensionName.toUpperCase().replace(/-/g, "_")}_${key}`; + const val = process.env[envKey]; + if (typeof val !== "undefined") { + return coerceEnvValue(val, defaultValue); + } + + // 2. Check SQLite persisted settings + const camelKey = envKeyToCamelCase(key); + const cached = loadSettingsCache(); + if (camelKey in cached) { + return cached[camelKey] as ConfigValue; + } + + // 3. Check schema default + if (settingsSchema) { + const properties = (settingsSchema as Record).properties as + | Record> + | undefined; + if (properties?.[camelKey]?.default !== undefined) { + return properties[camelKey].default as ConfigValue; + } + } + + // 4. Caller-provided default + return defaultValue; + } + + return { get, invalidateCache }; +} + +// --------------------------------------------------------------------------- +// Internal helpers +// --------------------------------------------------------------------------- + +/** + * Coerce a raw env-var string into a typed ConfigValue. + * + * @param val - The raw string from process.env + * @param defaultValue - Fallback if JSON parsing fails + * @returns The coerced value + */ +function coerceEnvValue(val: string, defaultValue?: ConfigValue): ConfigValue | undefined { + try { + if (val.toLowerCase() === "true") return true; + if (val.toLowerCase() === "false") return false; + + const num = Number(val); + if (!Number.isNaN(num) && val.trim() !== "") return num; + + if (val.startsWith("{") || val.startsWith("[")) { + return JSON.parse(val) as Record | unknown[]; + } + } catch { + return defaultValue; + } + + return val; +} + +/** + * Convert an UPPER_SNAKE_CASE key (e.g. "MAX_PAYLOAD_SIZE") to camelCase + * (e.g. "maxPayloadSize") for matching against schema property names. + * + * @param key - The UPPER_SNAKE_CASE key + * @returns The camelCase equivalent + */ +function envKeyToCamelCase(key: string): string { + return key.toLowerCase().replace(/_([a-z0-9])/g, (_, c: string) => c.toUpperCase()); +} diff --git a/src/extensions/discovery.ts b/src/extensions/discovery.ts new file mode 100644 index 0000000..813f22b --- /dev/null +++ b/src/extensions/discovery.ts @@ -0,0 +1,164 @@ +/** + * Extension discovery and validation. + * + * Pure functions for scanning extension directories and validating that + * discovered modules conform to the {@link Extension} interface contract. + * Extracted from ExtensionRegistry to keep discovery logic testable in + * isolation without needing the full registry lifecycle. + * + * @module + */ + +import { Value } from "@sinclair/typebox/value"; +import { formatValidationErrors } from "@src/utils/validation"; +import createLogger from "logging"; +import { type Extension, ExtensionManifestSchema } from "./types"; + +const logger = createLogger("ExtensionRegistry"); + +/** + * Scan one or more extension directories for modules (subdirectories + * containing an `index.ts`). Supports both top-level extensions and + * extensions nested under `core/`. + * + * Each discovered module is dynamically imported and validated. Invalid + * modules are logged and skipped. + * + * @param extensionDirs - Directories to scan for extensions + * @returns Array of valid Extension objects + */ +export async function discoverExtensions(extensionDirs: string[]): Promise { + const extensions: Extension[] = []; + const patterns = ["*/index.ts", "core/*/index.ts"]; + + for (const dir of extensionDirs) { + try { + for (const pattern of patterns) { + const glob = new Bun.Glob(pattern); + for (const entry of glob.scanSync({ cwd: dir, absolute: false })) { + const modulePath = `${dir}/${entry}`; + const ext = await loadExtensionModule(modulePath); + if (ext) extensions.push(ext); + } + } + } catch { + logger.warn(`Extensions directory not found or unreadable: ${dir}`); + } + } + + return extensions; +} + +/** + * Dynamically import a single extension module and validate its exports. + * + * @param modulePath - Absolute path to the extension's index.ts + * @returns The validated Extension object, or null if import/validation failed + */ +export async function loadExtensionModule(modulePath: string): Promise { + try { + const mod = await import(modulePath); + const ext: Extension = mod.default ?? mod; + + if (!validateExtension(ext, modulePath)) { + return null; + } + + return ext; + } catch (err) { + logger.error(`Failed to import extension module at ${modulePath}:`, err); + return null; + } +} + +/** + * Validate that a module export satisfies the {@link Extension} interface. + * + * Checks: + * - TypeBox schema conformance for the manifest + * - `settingsSchema` shape (must be a TObject with `type: "object"` and `properties`) + * - `secretsSchema` for duplicate key names + * - `ui.navigation` for duplicate routes + * - Presence of `initialize()` and `shutdown()` lifecycle methods + * + * @param ext - The candidate object to validate + * @param modulePath - Path used for error messages + * @returns `true` if the object is a valid Extension + */ +export function validateExtension(ext: unknown, modulePath: string): ext is Extension { + if (!ext || typeof ext !== "object") { + logger.error(`Extension at ${modulePath}: export is not an object`); + return false; + } + + const candidate = ext as Record; + + // Validate manifest with TypeBox + if (!candidate.manifest || !Value.Check(ExtensionManifestSchema, candidate.manifest)) { + const errorDetail = candidate.manifest + ? formatValidationErrors(ExtensionManifestSchema, candidate.manifest) + : "missing manifest"; + logger.error(`Extension at ${modulePath}: invalid manifest - ${errorDetail}`); + return false; + } + + // Validate settingsSchema if present (must be a TObject with type "object" and properties) + const manifest = candidate.manifest as Record; + if (manifest.settingsSchema != null) { + const settingsSchema = manifest.settingsSchema as Record; + if ( + settingsSchema.type !== "object" || + typeof settingsSchema.properties !== "object" || + settingsSchema.properties === null + ) { + logger.error( + `Extension at ${modulePath}: settingsSchema must be a TypeBox Type.Object() (got type="${settingsSchema.type}")`, + ); + return false; + } + } + + // Validate secretsSchema for duplicate key names (TypeBox catches structure, this catches duplicates) + if (manifest.secretsSchema != null) { + const secretsSchema = manifest.secretsSchema as Array<{ key: string }>; + const keyNames = new Set(); + const duplicates: string[] = []; + for (const entry of secretsSchema) { + if (keyNames.has(entry.key)) { + duplicates.push(entry.key); + } + keyNames.add(entry.key); + } + if (duplicates.length > 0) { + logger.warn( + `Extension at ${modulePath}: secretsSchema has duplicate key names: ${duplicates.join(", ")} - skipping secrets schema`, + ); + manifest.secretsSchema = undefined; + } + } + + // Check for duplicate routes within the manifest's ui.navigation array + const ui = manifest.ui as { navigation?: Array<{ route: string }> } | undefined; + if (ui?.navigation && ui.navigation.length > 0) { + const routes = new Set(); + for (const entry of ui.navigation) { + if (routes.has(entry.route)) { + logger.error(`Extension at ${modulePath}: duplicate route "${entry.route}" in ui.navigation`); + return false; + } + routes.add(entry.route); + } + } + + if (typeof candidate.initialize !== "function") { + logger.error(`Extension at ${modulePath}: missing initialize() method`); + return false; + } + + if (typeof candidate.shutdown !== "function") { + logger.error(`Extension at ${modulePath}: missing shutdown() method`); + return false; + } + + return true; +} diff --git a/src/extensions/extensionContext.ts b/src/extensions/extensionContext.ts index 7d58899..f9703f7 100644 --- a/src/extensions/extensionContext.ts +++ b/src/extensions/extensionContext.ts @@ -8,7 +8,6 @@ import type { RouteRegistry } from "@ext/types"; import type { AgentTool } from "@mariozechner/pi-agent-core"; import type { WebSocketMessage } from "@shared/types"; -import { schema } from "@src/db"; import type { PushMessageFn } from "@src/push"; import type { JobInfo, JobProcessor, ManagedQueueOptions, ManagedQueuePort, QueueJob, QueueJobLogs } from "@src/queue"; import { ManagedQueue } from "@src/queue"; @@ -19,15 +18,14 @@ import type { SkillEntry } from "@src/tools/sandbox"; import { authenticatedFetch } from "@src/utils/fetch"; import { registerDynamicItemProvider as registerProviderFn } from "@src/web/dynamicItemProviders"; import type { FlowProducer } from "bunqueue/client"; -import { eq } from "drizzle-orm"; import type { BunSQLiteDatabase } from "drizzle-orm/bun-sqlite"; import createLogger from "logging"; +import { createConfigResolver } from "./configResolver"; import type { EventBus } from "./eventBus"; import type { LoadedExtension, RegisteredRoute, RegisteredStepType } from "./internalTypes"; import type { AgentEventContext, AgentProcessorResult, - ConfigValue, CoreQueueName, EventCallback, EventParam, @@ -185,44 +183,11 @@ export function createExtensionContext(deps: ExtensionContextDeps): { const stepTypes: RegisteredStepType[] = []; // ------------------------------------------------------------------------- - // Settings cache - holds parsed config JSON from SQLite in memory. - // Invalidated when a `settings:changed` event fires for this extension. + // Config resolution - delegated to the focused configResolver module. // ------------------------------------------------------------------------- - /** Cached settings object from SQLite (null = not yet loaded). */ - let settingsCache: Record | null = null; - - /** - * Load persisted settings from SQLite into the cache. - * Returns the cached object (may be empty `{}`). - */ - function loadSettingsCache(): Record { - if (settingsCache !== null) return settingsCache; - try { - const row = database - .select({ config: schema.extensionSettings.config }) - .from(schema.extensionSettings) - .where(eq(schema.extensionSettings.name, extensionName)) - .get(); - settingsCache = row?.config ? (JSON.parse(row.config) as Record) : {}; - } catch { - settingsCache = {}; - } - return settingsCache; - } - - /** Invalidate the settings cache so the next read fetches from SQLite. */ - function invalidateSettingsCache(): void { - settingsCache = null; - } - - /** - * Convert an UPPER_SNAKE_CASE key (e.g. "MAX_PAYLOAD_SIZE") to camelCase - * (e.g. "maxPayloadSize") for matching against schema property names. - */ - function envKeyToCamelCase(key: string): string { - return key.toLowerCase().replace(/_([a-z0-9])/g, (_, c: string) => c.toUpperCase()); - } + const configResolver = createConfigResolver({ extensionName, database, settingsSchema }); + const getConfig = configResolver.get; /** * Register an agent tool with the system. @@ -366,68 +331,6 @@ export function createExtensionContext(deps: ExtensionContextDeps): { eventBus.dispatch(event); } - /** - * Read a configuration value for this extension by key. - * Precedence: env var > SQLite persisted value > schema default > caller default. - * - * Values are coerced from the raw env-var string: - * `"true"`/`"false"` -> boolean, numeric strings -> number, - * JSON-shaped strings -> parsed object/array, everything else -> string. - * - * @param key - The configuration key (UPPER_SNAKE_CASE). - * @param defaultValue - Returned when no source provides a value. - * @returns The resolved value, or `undefined`. - */ - function getConfig(key: string, defaultValue?: ConfigValue): ConfigValue | undefined { - // 1. Check environment variable (highest precedence) - const envKey = `EXT_${extensionName.toUpperCase().replace(/-/g, "_")}_${key}`; - const val = process.env[envKey]; - if (typeof val !== "undefined") { - return coerceEnvValue(val, defaultValue); - } - - // 2. Check SQLite persisted settings - const camelKey = envKeyToCamelCase(key); - const cached = loadSettingsCache(); - if (camelKey in cached) { - return cached[camelKey] as ConfigValue; - } - - // 3. Check schema default - if (settingsSchema) { - const properties = (settingsSchema as Record).properties as - | Record> - | undefined; - if (properties?.[camelKey]?.default !== undefined) { - return properties[camelKey].default as ConfigValue; - } - } - - // 4. Caller-provided default - return defaultValue; - } - - /** - * Coerce a raw env-var string into a typed ConfigValue. - */ - function coerceEnvValue(val: string, defaultValue?: ConfigValue): ConfigValue | undefined { - try { - if (val.toLowerCase() === "true") return true; - if (val.toLowerCase() === "false") return false; - - const num = Number(val); - if (!Number.isNaN(num) && val.trim() !== "") return num; - - if (val.startsWith("{") || val.startsWith("[")) { - return JSON.parse(val) as Record | unknown[]; - } - } catch { - return defaultValue; - } - - return val; - } - /** * Broadcast a WebSocket message to all connected frontend clients. * @@ -622,7 +525,7 @@ export function createExtensionContext(deps: ExtensionContextDeps): { // Subscribe to settings:changed events to invalidate the cache eventBus.subscribe(extensionName, "settings:changed", (event) => { if ("extensionName" in event && event.extensionName === extensionName) { - invalidateSettingsCache(); + configResolver.invalidateCache(); } }); diff --git a/src/extensions/lifecycle.ts b/src/extensions/lifecycle.ts new file mode 100644 index 0000000..4d6256d --- /dev/null +++ b/src/extensions/lifecycle.ts @@ -0,0 +1,244 @@ +/** + * Extension lifecycle operations. + * + * Extracted from ExtensionRegistry to consolidate the repeated patterns of + * "create context, call initialize, handle partial cleanup on failure" and + * "shutdown, tear down registrations". The registry delegates to these + * functions while retaining ownership of state (loaded list, global sets). + * + * @module + */ + +import type { WebSocketMessage } from "@shared/types"; +import type { ManagedQueuePort } from "@src/queue"; +import createLogger from "logging"; +import type { EventBus } from "./eventBus"; +import type { ExtensionContextDeps } from "./extensionContext"; +import { createExtensionContext } from "./extensionContext"; +import type { LoadedExtension } from "./internalTypes"; +import type { Extension } from "./types"; + +const logger = createLogger("ExtensionRegistry"); + +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- + +/** Mutable entry in the registry's loaded list (passed by reference). */ +export type LoadedEntry = { name: string } & LoadedExtension; + +/** Subset of registry state needed by lifecycle operations. */ +export interface LifecycleState { + /** Global tool name set — prevents duplicates across extensions and core. */ + toolNameSet: Set; + /** Global route key set — "METHOD:/full/path". */ + routeKeySet: Set; + /** Global step type name set — prevents duplicates across extensions. */ + stepTypeNameSet: Set; + /** Route prefixes for unloaded extensions — checked by the web server route guard. */ + disabledRoutePrefixes: Set; + /** The shared event bus instance. */ + eventBus: EventBus; +} + +/** Dependencies needed to build an ExtensionContext during activation. */ +export interface ActivationDeps { + /** Builds the full dependency set for createExtensionContext. */ + buildContextDeps: (extensionName: string, ext: Extension) => ExtensionContextDeps; + /** Broadcast a WebSocket message to all connected clients. */ + broadcastFn: (message: WebSocketMessage) => void; + /** Callback invoked when an extension creates a queue during loading. */ + onQueueCreated?: (queue: ManagedQueuePort) => void; +} + +// --------------------------------------------------------------------------- +// Cleanup helper (shared by deactivate, shutdownAll, and init-failure paths) +// --------------------------------------------------------------------------- + +/** + * Tears down all registrations for a loaded extension entry. + * + * Removes tools, step types, and route keys from the global sets, + * unsubscribes events, and closes queues. Does NOT call `shutdown()` + * on the extension itself — callers handle that separately. + * + * @param entry - The loaded extension entry to clean up + * @param state - The shared registry state holding the global sets + */ +export async function cleanupRegistrations(entry: LoadedEntry, state: LifecycleState): Promise { + for (const tool of entry.tools) { + state.toolNameSet.delete(tool.name); + } + for (const st of entry.stepTypes) { + state.stepTypeNameSet.delete(st.type); + } + for (const route of entry.routes) { + state.routeKeySet.delete(`${route.method}:${route.fullPath}`); + } + state.eventBus.unsubscribeAll(entry.name); + + for (const q of entry.queues) { + try { + await q.close(); + } catch (err) { + logger.error(`Error closing queue for extension "${entry.name}":`, err); + } + } +} + +// --------------------------------------------------------------------------- +// Activate +// --------------------------------------------------------------------------- + +/** + * Activate a suspended extension: creates a fresh ExtensionContext, calls + * `initialize()`, wires registrations, and transitions to active state. + * + * This function consolidates the duplicated pattern previously found in + * both `initializeExtension` and `activate` on the registry class. + * + * @param entry - The loaded extension entry (mutated in-place on success) + * @param state - The shared registry state + * @param deps - Dependencies for building the extension context + * @throws If `initialize()` fails (partial registrations are cleaned up first) + */ +export async function activateExtension( + entry: LoadedEntry, + state: LifecycleState, + deps: ActivationDeps, +): Promise { + // No-op if already active + if (entry.state === "active") return; + + const ext = entry.extension; + const name = entry.name; + + const contextDeps = deps.buildContextDeps(name, ext); + const { context, loaded } = createExtensionContext(contextDeps); + + try { + await ext.initialize(context); + } catch (err) { + // Partial registration cleanup on failed initialize + await cleanupPartialRegistrations(loaded, state, name); + entry.error = err instanceof Error ? err.message : String(err); + throw err; + } + + // Update the loaded entry with new registrations + entry.tools = loaded.tools; + entry.routes = loaded.routes; + entry.queues = loaded.queues; + entry.stepTypes = loaded.stepTypes; + entry.state = "active"; + entry.error = null; + + // Remove from disabled route prefixes + state.disabledRoutePrefixes.delete(`/ext/${name}`); + + // Notify monitor about any queues created + if (deps.onQueueCreated) { + for (const mq of loaded.queues) { + deps.onQueueCreated(mq); + } + } + + // Broadcast lifecycle event + deps.broadcastFn({ + type: "extension_lifecycle", + action: "activated", + name, + version: ext.manifest.version, + }); + + logger.info(`Activated extension "${name}" v${ext.manifest.version}`); +} + +// --------------------------------------------------------------------------- +// Deactivate +// --------------------------------------------------------------------------- + +/** + * Deactivate a loaded extension: calls `shutdown()`, tears down all + * registrations (tools, routes, queues, events), and transitions to + * suspended state. The extension remains in the loaded list for + * re-activation later. + * + * No-op if the extension is already suspended. + * + * @param entry - The loaded extension entry (mutated in-place) + * @param state - The shared registry state + * @param broadcastFn - Optional function to broadcast lifecycle events + */ +export async function deactivateExtension( + entry: LoadedEntry, + state: LifecycleState, + broadcastFn?: (message: WebSocketMessage) => void, +): Promise { + if (entry.state === "suspended") return; + + const name = entry.name; + + // Shutdown the extension (errors are logged but don't prevent cleanup) + try { + await entry.extension.shutdown(); + logger.debug(`Shut down extension "${name}"`); + } catch (err) { + logger.error(`Error shutting down extension "${name}":`, err); + } + + // Clean up all registrations + await cleanupRegistrations(entry, state); + state.disabledRoutePrefixes.add(`/ext/${name}`); + + // Clear registrations and transition to suspended + entry.tools = []; + entry.routes = []; + entry.queues = []; + entry.stepTypes = []; + entry.state = "suspended"; + + // Broadcast lifecycle event + if (broadcastFn) { + broadcastFn({ + type: "extension_lifecycle", + action: "deactivated", + name, + version: entry.extension.manifest.version, + }); + } + + logger.info(`Deactivated extension "${name}"`); +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +/** + * Clean up partial registrations from a failed `initialize()` call. + * Called when the extension's initialize method throws before completing. + */ +async function cleanupPartialRegistrations( + loaded: Omit, + state: LifecycleState, + name: string, +): Promise { + for (const tool of loaded.tools) { + state.toolNameSet.delete(tool.name); + } + for (const st of loaded.stepTypes) { + state.stepTypeNameSet.delete(st.type); + } + for (const route of loaded.routes) { + state.routeKeySet.delete(`${route.method}:${route.fullPath}`); + } + state.eventBus.unsubscribeAll(name); + for (const q of loaded.queues) { + try { + await q.close(); + } catch { + // Ignore cleanup errors + } + } +} diff --git a/src/extensions/registry.ts b/src/extensions/registry.ts index 2a94832..943110b 100644 --- a/src/extensions/registry.ts +++ b/src/extensions/registry.ts @@ -7,7 +7,6 @@ import type { FSWatcher } from "node:fs"; import { watch as fsWatch } from "node:fs"; import type { AgentTool } from "@mariozechner/pi-agent-core"; import type { ExtensionInfo, WebSocketMessage } from "@shared/types"; -import { Value } from "@sinclair/typebox/value"; import { PROJECT_DIR, serverOrigin } from "@src/config"; import { getDb, schema } from "@src/db"; import type { PushMessageFn } from "@src/push"; @@ -20,25 +19,26 @@ import { loadSkillScripts as loadSkillScriptsFn, } from "@src/skills/loader"; import type { SkillEntry } from "@src/tools/sandbox"; -import { formatValidationErrors } from "@src/utils/validation"; import { enrichSchemaWithDynamicItems } from "@src/web/dynamicItemProviders"; import { FlowProducer } from "bunqueue/client"; import { eq } from "drizzle-orm"; import type { BunSQLiteDatabase } from "drizzle-orm/bun-sqlite"; import createLogger from "logging"; import { resolveDependencyOrder } from "./dependencyResolver"; +import { discoverExtensions as discoverExtensionsFn, validateExtension as validateExtensionFn } from "./discovery"; import { EventBus } from "./eventBus"; -import { createExtensionContext } from "./extensionContext"; +import type { ExtensionContextDeps } from "./extensionContext"; import { ExternalDependencyResolver } from "./externalDependencyResolver"; import type { LoadedExtension } from "./internalTypes"; -import type { RouteRegistry } from "./types"; import { - type AgentProcessorResult, - type CoreQueueName, - type Extension, - ExtensionManifestSchema, - type RunAgentOptions, -} from "./types"; + type ActivationDeps, + activateExtension, + cleanupRegistrations, + deactivateExtension, + type LifecycleState, + type LoadedEntry, +} from "./lifecycle"; +import type { AgentProcessorResult, CoreQueueName, Extension, RouteRegistry, RunAgentOptions } from "./types"; const logger = createLogger("ExtensionRegistry"); @@ -360,124 +360,15 @@ export class ExtensionRegistry { * and extensions nested under core/. */ private async discoverExtensions(): Promise { - const extensions: Extension[] = []; - const dirs = this.extensionDirs; - const patterns = ["*/index.ts", "core/*/index.ts"]; - - for (const dir of dirs) { - try { - for (const pattern of patterns) { - const glob = new Bun.Glob(pattern); - for (const entry of glob.scanSync({ cwd: dir, absolute: false })) { - const modulePath = `${dir}/${entry}`; - const ext = await this.loadExtensionModule(modulePath); - if (ext) extensions.push(ext); - } - } - } catch { - logger.warn(`Extensions directory not found or unreadable: ${dir}`); - } - } - - return extensions; - } - - /** - * Dynamically import a single extension module and validate its exports. - */ - private async loadExtensionModule(modulePath: string): Promise { - try { - const mod = await import(modulePath); - const ext: Extension = mod.default ?? mod; - - if (!this.validateExtension(ext, modulePath)) { - return null; - } - - return ext; - } catch (err) { - logger.error(`Failed to import extension module at ${modulePath}:`, err); - return null; - } + return discoverExtensionsFn(this.extensionDirs); } /** * Validate that a module export satisfies the Extension interface. - * Checks TypeBox schema conformance, settingsSchema shape, duplicate - * routes within ui.navigation, and presence of lifecycle methods. + * Delegates to the standalone {@link validateExtensionFn} function. */ validateExtension(ext: unknown, modulePath: string): ext is Extension { - if (!ext || typeof ext !== "object") { - logger.error(`Extension at ${modulePath}: export is not an object`); - return false; - } - - const candidate = ext as Record; - - // Validate manifest with TypeBox - if (!candidate.manifest || !Value.Check(ExtensionManifestSchema, candidate.manifest)) { - const errorDetail = candidate.manifest - ? formatValidationErrors(ExtensionManifestSchema, candidate.manifest) - : "missing manifest"; - logger.error(`Extension at ${modulePath}: invalid manifest - ${errorDetail}`); - return false; - } - - // Validate settingsSchema if present (must be a TObject with type "object" and properties) - const manifest = candidate.manifest as Record; - if (manifest.settingsSchema != null) { - const schema = manifest.settingsSchema as Record; - if (schema.type !== "object" || typeof schema.properties !== "object" || schema.properties === null) { - logger.error( - `Extension at ${modulePath}: settingsSchema must be a TypeBox Type.Object() (got type="${schema.type}")`, - ); - return false; - } - } - - // Validate secretsSchema for duplicate key names (TypeBox catches structure, this catches duplicates) - if (manifest.secretsSchema != null) { - const secretsSchema = manifest.secretsSchema as Array<{ key: string }>; - const keyNames = new Set(); - const duplicates: string[] = []; - for (const entry of secretsSchema) { - if (keyNames.has(entry.key)) { - duplicates.push(entry.key); - } - keyNames.add(entry.key); - } - if (duplicates.length > 0) { - logger.warn( - `Extension at ${modulePath}: secretsSchema has duplicate key names: ${duplicates.join(", ")} - skipping secrets schema`, - ); - manifest.secretsSchema = undefined; - } - } - - // Check for duplicate routes within the manifest's ui.navigation array - const ui = manifest.ui as { navigation?: Array<{ route: string }> } | undefined; - if (ui?.navigation && ui.navigation.length > 0) { - const routes = new Set(); - for (const entry of ui.navigation) { - if (routes.has(entry.route)) { - logger.error(`Extension at ${modulePath}: duplicate route "${entry.route}" in ui.navigation`); - return false; - } - routes.add(entry.route); - } - } - - if (typeof candidate.initialize !== "function") { - logger.error(`Extension at ${modulePath}: missing initialize() method`); - return false; - } - - if (typeof candidate.shutdown !== "function") { - logger.error(`Extension at ${modulePath}: missing shutdown() method`); - return false; - } - - return true; + return validateExtensionFn(ext, modulePath); } /** @@ -513,89 +404,36 @@ export class ExtensionRegistry { */ private async initializeExtension(ext: Extension, modulePath?: string): Promise { const name = ext.manifest.name; - const deps = this.initDeps!; - const { context, loaded } = createExtensionContext({ - extensionName: name, - workDir: this.workDir, - dataDir: this.dataDir, - extensionsDir: this.builtinExtensionsDir, - toolNameSet: this.toolNameSet, - routeKeySet: this.routeKeySet, - stepTypeNameSet: this.stepTypeNameSet, - eventBus: this.eventBus, - broadcastFn: deps.broadcastFn, - flowProducer: this.flowProducer, - resolveSkillFn: (n) => this.resolveSkill(n), - database: deps.database, - getCoreQueueFn: this.getCoreQueueFn, - getExtensionQueuesFn: () => this.getRegisteredQueues(), - runAgentFn: deps.runAgentFn, - sessionStore: deps.sessionStore, - pushMessageFn: deps.pushMessageFn, - isExtensionEnabledFn: (n) => this.isExtensionEnabled(n), - secretVault: deps.secretVault, - routeRegistry: deps.routeRegistry, - rescanSkillsFn: () => this.discoverAndLoadSkills(), - getSkillNamesFn: () => this.getSkillNames(), - getStepHandlerFn: (type) => this.getRegisteredStepTypes().find((st) => st.type === type)?.handler, - settingsSchema: ext.manifest.settingsSchema as Record | undefined, - }); + // Add to loaded list as suspended first, then activate + const entry: LoadedEntry = { + name, + extension: ext, + modulePath, + tools: [], + routes: [], + queues: [], + stepTypes: [], + state: "suspended", + }; + this.loaded.push(entry); try { - await ext.initialize(context); - - this.loaded.push({ name, extension: ext, modulePath, ...loaded, state: "active" }); + await activateExtension(entry, this.getLifecycleState(), this.getActivationDeps()); // Build a concise summary of what the extension registered const parts: string[] = []; - if (loaded.tools.length > 0) parts.push(`${loaded.tools.length} tool(s)`); - if (loaded.routes.length > 0) parts.push(`${loaded.routes.length} route(s)`); - if (loaded.queues.length > 0) parts.push(`${loaded.queues.length} queue(s)`); + if (entry.tools.length > 0) parts.push(`${entry.tools.length} tool(s)`); + if (entry.routes.length > 0) parts.push(`${entry.routes.length} route(s)`); + if (entry.queues.length > 0) parts.push(`${entry.queues.length} queue(s)`); const summary = parts.length > 0 ? ` (${parts.join(", ")})` : ""; logger.info(`Initialized extension "${name}" v${ext.manifest.version}${summary}`); - - // Notify caller about any queues created by this extension - if (deps.onQueueCreated) { - for (const mq of loaded.queues) { - deps.onQueueCreated(mq); - } - } } catch (err) { logger.error(`Failed to initialize extension "${name}":`, err); - - // Clean up any partial registrations from the failed context - for (const tool of loaded.tools) { - this.toolNameSet.delete(tool.name); + // Entry remains in loaded list as suspended with error (set by activateExtension) + if (!entry.error) { + entry.error = err instanceof Error ? err.message : String(err); } - for (const st of loaded.stepTypes) { - this.stepTypeNameSet.delete(st.type); - } - for (const route of loaded.routes) { - this.routeKeySet.delete(`${route.method}:${route.fullPath}`); - } - this.eventBus.unsubscribeAll(name); - for (const q of loaded.queues) { - try { - await q.close(); - } catch { - // Ignore cleanup errors - } - } - - // Still add to loaded list as suspended so it remains visible in the UI - const errorMsg = err instanceof Error ? err.message : String(err); - this.loaded.push({ - name, - extension: ext, - modulePath, - tools: [], - routes: [], - queues: [], - stepTypes: [], - state: "suspended", - error: errorMsg, - }); } } @@ -822,65 +660,7 @@ export class ExtensionRegistry { throw new Error(`Cannot deactivate extension "${name}": not found in loaded list`); } - // No-op if already suspended - if (entry.state === "suspended") { - return; - } - - // Shutdown the extension (errors are logged but don't prevent cleanup) - try { - await entry.extension.shutdown(); - logger.debug(`Shut down extension "${name}"`); - } catch (err) { - logger.error(`Error shutting down extension "${name}":`, err); - } - - // Clean up tools from the global set - for (const tool of entry.tools) { - this.toolNameSet.delete(tool.name); - } - - // Clean up step types from the global set - for (const st of entry.stepTypes) { - this.stepTypeNameSet.delete(st.type); - } - - // Clean up route keys and add prefix to disabled set - for (const route of entry.routes) { - this.routeKeySet.delete(`${route.method}:${route.fullPath}`); - } - this.disabledRoutePrefixes.add(`/ext/${name}`); - - // Clean up event subscriptions - this.eventBus.unsubscribeAll(name); - - // Close queues - for (const q of entry.queues) { - try { - await q.close(); - } catch (err) { - logger.error(`Error closing queue for extension "${name}":`, err); - } - } - - // Clear registrations and transition to suspended - entry.tools = []; - entry.routes = []; - entry.queues = []; - entry.stepTypes = []; - entry.state = "suspended"; - - // Broadcast lifecycle event - if (this.initDeps) { - this.initDeps.broadcastFn({ - type: "extension_lifecycle", - action: "deactivated", - name, - version: entry.extension.manifest.version, - }); - } - - logger.info(`Deactivated extension "${name}"`); + await deactivateExtension(entry, this.getLifecycleState(), this.initDeps?.broadcastFn); } /** @@ -898,99 +678,11 @@ export class ExtensionRegistry { throw new Error(`Cannot activate extension "${name}": not found in loaded list`); } - // No-op if already active - if (entry.state === "active") { - return; - } - - const deps = this.initDeps; - if (!deps) { + if (!this.initDeps) { throw new Error(`Cannot activate extension "${name}": registry not initialized`); } - const ext = entry.extension; - - // Create a fresh context - const { context, loaded } = createExtensionContext({ - extensionName: name, - workDir: this.workDir, - dataDir: this.dataDir, - extensionsDir: this.builtinExtensionsDir, - toolNameSet: this.toolNameSet, - routeKeySet: this.routeKeySet, - stepTypeNameSet: this.stepTypeNameSet, - eventBus: this.eventBus, - broadcastFn: deps.broadcastFn, - flowProducer: this.flowProducer, - resolveSkillFn: (n) => this.resolveSkill(n), - database: deps.database, - getCoreQueueFn: this.getCoreQueueFn, - getExtensionQueuesFn: () => this.getRegisteredQueues(), - runAgentFn: deps.runAgentFn, - sessionStore: deps.sessionStore, - pushMessageFn: deps.pushMessageFn, - isExtensionEnabledFn: (n) => this.isExtensionEnabled(n), - secretVault: deps.secretVault, - routeRegistry: deps.routeRegistry, - rescanSkillsFn: () => this.discoverAndLoadSkills(), - getSkillNamesFn: () => this.getSkillNames(), - getStepHandlerFn: (type) => this.getRegisteredStepTypes().find((st) => st.type === type)?.handler, - settingsSchema: ext.manifest.settingsSchema as Record | undefined, - }); - - try { - await ext.initialize(context); - } catch (err) { - // Partial registration cleanup on failed initialize - for (const tool of loaded.tools) { - this.toolNameSet.delete(tool.name); - } - for (const st of loaded.stepTypes) { - this.stepTypeNameSet.delete(st.type); - } - for (const route of loaded.routes) { - this.routeKeySet.delete(`${route.method}:${route.fullPath}`); - } - this.eventBus.unsubscribeAll(name); - for (const q of loaded.queues) { - try { - await q.close(); - } catch { - // Ignore cleanup errors - } - } - // Store error on the entry so it's visible in the UI - entry.error = err instanceof Error ? err.message : String(err); - throw err; - } - - // Update the loaded entry with new registrations - entry.tools = loaded.tools; - entry.routes = loaded.routes; - entry.queues = loaded.queues; - entry.stepTypes = loaded.stepTypes; - entry.state = "active"; - entry.error = null; - - // Remove from disabled route prefixes - this.disabledRoutePrefixes.delete(`/ext/${name}`); - - // Notify monitor about any queues created - if (deps.onQueueCreated) { - for (const mq of loaded.queues) { - deps.onQueueCreated(mq); - } - } - - // Broadcast lifecycle event - deps.broadcastFn({ - type: "extension_lifecycle", - action: "activated", - name, - version: ext.manifest.version, - }); - - logger.info(`Activated extension "${name}" v${ext.manifest.version}`); + await activateExtension(entry, this.getLifecycleState(), this.getActivationDeps()); } /** @@ -1074,6 +766,7 @@ export class ExtensionRegistry { // Reverse init order const reversed = [...this.loaded].reverse(); + const state = this.getLifecycleState(); for (const entry of reversed) { try { @@ -1083,32 +776,7 @@ export class ExtensionRegistry { logger.error(`Error shutting down extension "${entry.name}":`, err); } - // Clean up tools from the global set - for (const tool of entry.tools) { - this.toolNameSet.delete(tool.name); - } - - // Clean up step types from the global set - for (const st of entry.stepTypes) { - this.stepTypeNameSet.delete(st.type); - } - - // Clean up route keys from the global set - for (const route of entry.routes) { - this.routeKeySet.delete(`${route.method}:${route.fullPath}`); - } - - // Clean up event subscriptions - this.eventBus.unsubscribeAll(entry.name); - - // Close queues (handles both worker and queue internally) - for (const q of entry.queues) { - try { - await q.close(); - } catch (err) { - logger.error(`Error closing queue for extension "${entry.name}":`, err); - } - } + await cleanupRegistrations(entry, state); } this.loaded = []; @@ -1117,6 +785,56 @@ export class ExtensionRegistry { // Close the shared FlowProducer await this.flowProducer.close(); } + + // --------------------------------------------------------------------------- + // Private helpers for lifecycle delegation + // --------------------------------------------------------------------------- + + /** Returns the subset of registry state needed by lifecycle operations. */ + private getLifecycleState(): LifecycleState { + return { + toolNameSet: this.toolNameSet, + routeKeySet: this.routeKeySet, + stepTypeNameSet: this.stepTypeNameSet, + disabledRoutePrefixes: this.disabledRoutePrefixes, + eventBus: this.eventBus, + }; + } + + /** Builds the activation dependencies for the lifecycle module. */ + private getActivationDeps(): ActivationDeps { + const deps = this.initDeps!; + return { + buildContextDeps: (extensionName: string, ext: Extension): ExtensionContextDeps => ({ + extensionName, + workDir: this.workDir, + dataDir: this.dataDir, + extensionsDir: this.builtinExtensionsDir, + toolNameSet: this.toolNameSet, + routeKeySet: this.routeKeySet, + stepTypeNameSet: this.stepTypeNameSet, + eventBus: this.eventBus, + broadcastFn: deps.broadcastFn, + flowProducer: this.flowProducer, + resolveSkillFn: (n) => this.resolveSkill(n), + database: deps.database, + getCoreQueueFn: this.getCoreQueueFn, + getExtensionQueuesFn: () => this.getRegisteredQueues(), + runAgentFn: deps.runAgentFn, + sessionStore: deps.sessionStore, + pushMessageFn: deps.pushMessageFn, + isExtensionEnabledFn: (n) => this.isExtensionEnabled(n), + secretVault: deps.secretVault, + routeRegistry: deps.routeRegistry, + rescanSkillsFn: () => this.discoverAndLoadSkills(), + getSkillNamesFn: () => this.getSkillNames(), + getStepHandlerFn: (type) => this.getRegisteredStepTypes().find((st) => st.type === type)?.handler, + settingsSchema: ext.manifest.settingsSchema as Record | undefined, + }), + broadcastFn: deps.broadcastFn, + onQueueCreated: deps.onQueueCreated, + }; + } } // --------------------------------------------------------------------------- diff --git a/src/jobs/agentQueue.ts b/src/jobs/agentQueue.ts index d9ff0c0..cba968d 100644 --- a/src/jobs/agentQueue.ts +++ b/src/jobs/agentQueue.ts @@ -1,17 +1,15 @@ /** * Agent queue - processes agent prompt jobs (spell-check, telegram messages, etc.) - * via the {@link ManagedQueue} abstraction. + * via the generic {@link createJobQueue} factory. */ -import type { AgentEventContext, EventBus } from "@src/extensions"; -import type { ManagedQueuePort, QueueJob } from "@src/queue"; -import { ManagedQueue } from "@src/queue"; -import type { AgentProcessorConfig, AgentProcessorResult } from "./agentProcessor"; -import { runAgent } from "./agentProcessor"; -import { AGENT_QUEUE_DEFAULTS } from "./defaults"; +import type { AgentEventContext } from "@src/extensions"; +import type { ManagedQueuePort } from "@src/queue"; +import type { BaseAgentJob, JobQueueDeps } from "./jobQueueFactory"; +import { createJobQueue } from "./jobQueueFactory"; /** Payload for agent prompt jobs. */ -export interface AgentJob { +export interface AgentJob extends BaseAgentJob { /** Optional event context for routing responses (e.g. telegram chat ID). */ context?: AgentEventContext; /** Optional system prompt override. */ @@ -21,14 +19,7 @@ export interface AgentJob { } /** Dependencies required to create the agent queue. */ -export interface AgentQueueDeps { - /** Builds an {@link AgentProcessorConfig} for each job at processing time (sessionId is merged from job data). */ - buildProcessor: ( - job: QueueJob, - ) => Omit | Promise>; - /** Getter for the event bus (resolved at job processing time). */ - getEventBus: () => EventBus; -} +export type AgentQueueDeps = JobQueueDeps; /** * Creates a {@link ManagedQueue} for agent prompt jobs. @@ -42,20 +33,5 @@ export interface AgentQueueDeps { * @returns The managed agent queue */ export function createAgentQueue(deps: AgentQueueDeps): ManagedQueuePort { - const { buildProcessor, getEventBus } = deps; - - return new ManagedQueue( - "agents", - async (job: QueueJob) => { - const config = await buildProcessor(job); - - return runAgent(job, { - ...config, - sessionId: job.data.sessionId, - eventBus: config.eventBus ?? getEventBus(), - context: config.context ?? job.data?.context, - }); - }, - AGENT_QUEUE_DEFAULTS, - ); + return createJobQueue("agents", deps); } diff --git a/src/jobs/chatQueue.ts b/src/jobs/chatQueue.ts index 010688b..820848b 100644 --- a/src/jobs/chatQueue.ts +++ b/src/jobs/chatQueue.ts @@ -1,16 +1,15 @@ /** - * Chat queue - processes conversational chat jobs with a dedicated system prompt. + * Chat queue - processes conversational chat jobs with a dedicated system prompt + * via the generic {@link createJobQueue} factory. */ -import type { AgentEventContext, EventBus } from "@src/extensions"; -import type { ManagedQueuePort, QueueJob } from "@src/queue"; -import { ManagedQueue } from "@src/queue"; -import type { AgentProcessorConfig, AgentProcessorResult } from "./agentProcessor"; -import { runAgent } from "./agentProcessor"; -import { AGENT_QUEUE_DEFAULTS } from "./defaults"; +import type { AgentEventContext } from "@src/extensions"; +import type { ManagedQueuePort } from "@src/queue"; +import type { BaseAgentJob, JobQueueDeps } from "./jobQueueFactory"; +import { createJobQueue } from "./jobQueueFactory"; /** Payload for chat jobs. */ -export interface ChatJob { +export interface ChatJob extends BaseAgentJob { /** Event context for routing responses (e.g. chat ID for frontend). */ context?: AgentEventContext; /** Session ID for conversation context (callers must append user message before enqueuing). */ @@ -18,14 +17,7 @@ export interface ChatJob { } /** Dependencies required to create the chat queue. */ -export interface ChatQueueDeps { - /** Builds an {@link AgentProcessorConfig} for each job at processing time (sessionId is merged from job data). */ - buildProcessor: ( - job: QueueJob, - ) => Omit | Promise>; - /** Getter for the event bus (resolved at job processing time). */ - getEventBus: () => EventBus; -} +export type ChatQueueDeps = JobQueueDeps; /** * Creates a {@link ManagedQueue} for chat jobs. @@ -38,20 +30,5 @@ export interface ChatQueueDeps { * @returns The managed chat queue */ export function createChatQueue(deps: ChatQueueDeps): ManagedQueuePort { - const { buildProcessor, getEventBus } = deps; - - return new ManagedQueue( - "chat", - async (job: QueueJob) => { - const config = await buildProcessor(job); - - return runAgent(job, { - ...config, - sessionId: job.data.sessionId, - eventBus: config.eventBus ?? getEventBus(), - context: config.context ?? job.data?.context, - }); - }, - AGENT_QUEUE_DEFAULTS, - ); + return createJobQueue("chat", deps); } diff --git a/src/jobs/index.ts b/src/jobs/index.ts index 770eaec..ee1dc2c 100644 --- a/src/jobs/index.ts +++ b/src/jobs/index.ts @@ -6,4 +6,6 @@ export { abortJob } from "./cancellation.ts"; export type { ChatJob, ChatQueueDeps } from "./chatQueue.ts"; export { createChatQueue } from "./chatQueue.ts"; export { AGENT_QUEUE_DEFAULTS } from "./defaults.ts"; +export type { BaseAgentJob, JobQueueDeps } from "./jobQueueFactory.ts"; +export { createJobQueue } from "./jobQueueFactory.ts"; export { buildAgentSystemPrompt, buildChatSystemPrompt } from "./systemPrompts.ts"; diff --git a/src/jobs/jobQueueFactory.ts b/src/jobs/jobQueueFactory.ts new file mode 100644 index 0000000..168b7de --- /dev/null +++ b/src/jobs/jobQueueFactory.ts @@ -0,0 +1,76 @@ +/** + * Generic job queue factory for agent-based processors. + * + * Both the agent queue and chat queue share the same processing pattern: + * build a processor config, merge the session ID from the job payload, + * fall back to the event bus and context from the job data, then run + * the agent. This module extracts that shared logic into a single + * generic factory parameterized by the job payload type. + * + * @module + */ + +import type { AgentEventContext, EventBus } from "@src/extensions"; +import type { ManagedQueuePort, QueueJob } from "@src/queue"; +import { ManagedQueue } from "@src/queue"; +import type { AgentProcessorConfig, AgentProcessorResult } from "./agentProcessor"; +import { runAgent } from "./agentProcessor"; +import { AGENT_QUEUE_DEFAULTS } from "./defaults"; + +/** + * Base constraint for job payloads processed by agent queues. + * All agent-based job types must carry a session ID and an optional + * event routing context. + */ +export interface BaseAgentJob { + /** Optional event context for routing responses (e.g. chat ID, telegram chat ID). */ + context?: AgentEventContext; + /** Session ID for conversation context (callers must append user message before enqueuing). */ + sessionId: string; +} + +/** + * Dependencies required to create an agent-based job queue. + * + * @typeParam T - The job payload type (must extend {@link BaseAgentJob}) + */ +export interface JobQueueDeps { + /** Builds an {@link AgentProcessorConfig} for each job at processing time (sessionId is merged from job data). */ + buildProcessor: ( + job: QueueJob, + ) => Omit | Promise>; + /** Getter for the event bus (resolved at job processing time). */ + getEventBus: () => EventBus; +} + +/** + * Creates a {@link ManagedQueue} for agent-based jobs. + * + * The processor resolves dependencies lazily via getter functions at job + * processing time, so extensions loaded after queue creation are still visible. + * Event dispatching is handled centrally by `runAgent` via `config.eventBus` + * and `config.context`. + * + * @typeParam T - The job payload type (must extend {@link BaseAgentJob}) + * @param name - Queue name (e.g. "agents", "chat") + * @param deps - Lazy getters for processor config and event bus + * @returns The managed queue instance + */ +export function createJobQueue(name: string, deps: JobQueueDeps): ManagedQueuePort { + const { buildProcessor, getEventBus } = deps; + + return new ManagedQueue( + name, + async (job: QueueJob) => { + const config = await buildProcessor(job); + + return runAgent(job, { + ...config, + sessionId: job.data.sessionId, + eventBus: config.eventBus ?? getEventBus(), + context: config.context ?? job.data?.context, + }); + }, + AGENT_QUEUE_DEFAULTS, + ); +} diff --git a/src/secrets/vault.ts b/src/secrets/vault.ts index c4afa29..4e9e890 100644 --- a/src/secrets/vault.ts +++ b/src/secrets/vault.ts @@ -329,11 +329,11 @@ export class SecretVault { * @param consumer - The consumer identity performing the deletion (for audit) * @returns True if the secret was deleted, false if it was not found */ - async remove(scope: string, key: string, consumer: string): Promise { + remove(scope: string, key: string, consumer: string): boolean { const secretName = `${scope}/${key}`; // Check existence before delete (drizzle .run() returns void) - const exists = await this.has(scope, key); + const exists = this.has(scope, key); if (!exists) { return false; } @@ -360,7 +360,7 @@ export class SecretVault { * @param key - The secret key to check * @returns True if the secret exists */ - async has(scope: string, key: string): Promise { + has(scope: string, key: string): boolean { const row = this.db .select({ count: sql`COUNT(*)` }) .from(secretsVault) @@ -595,7 +595,7 @@ export class SecretVault { * @param key - The secret key to remove * @returns True if the secret was deleted, false if not found */ - async removeGlobal(key: string): Promise { + removeGlobal(key: string): boolean { return this.remove(SecretVault.GLOBAL_SCOPE, key, "admin:web"); } diff --git a/src/session/sessionStore.ts b/src/session/sessionStore.ts index 5f8db45..fcd9952 100644 --- a/src/session/sessionStore.ts +++ b/src/session/sessionStore.ts @@ -154,6 +154,42 @@ function omitMetaFields(meta: Record): Record return result; } +/** + * Accumulates token usage counters on a session row within a transaction. + * + * Increments the running totals for input, output, cache read, cache write, + * and total tokens, and updates `lastInputTokens` with the current value. + * + * @param tx - The active database transaction + * @param sessionId - The session to update + * @param usage - The raw usage object from the assistant message + */ +function updateTokenTotals( + tx: Parameters["transaction"]>[0]>[0], + sessionId: string, + usage: unknown, +): void { + const u = usage as { + input?: number; + output?: number; + cacheRead?: number; + cacheWrite?: number; + totalTokens?: number; + }; + tx.update(schema.sessions) + .set({ + updatedAt: Date.now(), + totalInputTokens: sql`${schema.sessions.totalInputTokens} + ${u.input ?? 0}`, + totalOutputTokens: sql`${schema.sessions.totalOutputTokens} + ${u.output ?? 0}`, + totalCacheReadTokens: sql`${schema.sessions.totalCacheReadTokens} + ${u.cacheRead ?? 0}`, + totalCacheWriteTokens: sql`${schema.sessions.totalCacheWriteTokens} + ${u.cacheWrite ?? 0}`, + totalTokens: sql`${schema.sessions.totalTokens} + ${u.totalTokens ?? 0}`, + lastInputTokens: u.input ?? 0, + }) + .where(eq(schema.sessions.id, sessionId)) + .run(); +} + /** * SQLite-backed implementation of {@link SessionStorePort}. * @@ -371,25 +407,7 @@ export class SessionStore implements SessionStorePort { // Update session timestamp and token totals for assistant messages if (msg.role === "assistant" && usage) { - const u = usage as { - input?: number; - output?: number; - cacheRead?: number; - cacheWrite?: number; - totalTokens?: number; - }; - tx.update(schema.sessions) - .set({ - updatedAt: Date.now(), - totalInputTokens: sql`${schema.sessions.totalInputTokens} + ${u.input ?? 0}`, - totalOutputTokens: sql`${schema.sessions.totalOutputTokens} + ${u.output ?? 0}`, - totalCacheReadTokens: sql`${schema.sessions.totalCacheReadTokens} + ${u.cacheRead ?? 0}`, - totalCacheWriteTokens: sql`${schema.sessions.totalCacheWriteTokens} + ${u.cacheWrite ?? 0}`, - totalTokens: sql`${schema.sessions.totalTokens} + ${u.totalTokens ?? 0}`, - lastInputTokens: u.input ?? 0, - }) - .where(eq(schema.sessions.id, sessionId)) - .run(); + updateTokenTotals(tx, sessionId, usage); } else { tx.update(schema.sessions).set({ updatedAt: Date.now() }).where(eq(schema.sessions.id, sessionId)).run(); } diff --git a/src/tools/file.ts b/src/tools/file.ts index 2c0254f..df391fa 100644 --- a/src/tools/file.ts +++ b/src/tools/file.ts @@ -1,9 +1,6 @@ -import assert from "node:assert"; -import { default as path } from "node:path"; import type { AgentTool, AgentToolResult, AgentToolUpdateCallback } from "@mariozechner/pi-agent-core"; import { type Static, Type } from "@sinclair/typebox"; -import { WORK_DIR } from "@src/config"; -import { mainLogger as log, shellLogger as shellLog } from "@src/utils/logger"; +import { shellLogger as shellLog } from "@src/utils/logger"; import type { Bash } from "just-bash"; // --------------------------------------------------------------------------- @@ -72,29 +69,6 @@ const ExecuteCommandParams = Type.Object({ command: Type.String({ description: "Shell command to execute" }), }); -// --------------------------------------------------------------------------- -// Host path helper (used by OCR, not by sandbox file tools) -// --------------------------------------------------------------------------- - -/** - * Validate that an absolute path resolves within the work directory. - * Used by non-sandbox code (e.g. OCR) that operates on host paths directly. - * - * @param absolutePath - Already-resolved absolute path - * @returns The validated absolute path - * @throws If the path escapes the work directory - */ -export function assertInsideWorkDir(absolutePath: string): string { - assert(WORK_DIR && WORK_DIR.length > 0); - - const resolved = path.resolve(absolutePath); - if (!resolved.startsWith(WORK_DIR)) { - log.error(`Access denied: ${absolutePath}`); - throw new Error(`Access denied: ${absolutePath}`); - } - return resolved; -} - // --------------------------------------------------------------------------- // Tool names - exported so callers can filter / reference them by name // ---------------------------------------------------------------------------