diff --git a/README.md b/README.md index de7780ba..ca65d264 100644 --- a/README.md +++ b/README.md @@ -108,7 +108,7 @@ For account bundles, define `upstreams` instead of `upstream`. Tools, resource n Multi-upstream resource and prompt lists use opaque, bounded in-memory Miftah cursors so individual upstream cursors are not exposed. Re-list after a profile change or restart instead of retaining an old cursor. A standard `upstream` or a one-entry `upstreams` map preserves credential-free raw resource URIs, prompt names, and native upstream cursors for compatibility; URI and icon fields with userinfo, query values, or fragments are structurally redacted before they cross the boundary. A zero-entry bundle still omits resource and prompt capabilities. -`tooling.toolDiscoveryMode` defaults to `"permissive"`. In that mode, a bundled list returns only healthy upstream contributions, removes routes for failed upstreams, and retries incomplete tool discovery on later tool requests. If every upstream fails a capability list, Miftah returns `UPSTREAM_DISCOVERY_FAILED` instead of advertising an empty upstream surface. Set it to `"strict"` to reject any unavailable upstream; strict tool discovery additionally checks every configured profile for an identical client-visible tool contract. `miftah_health` reports the profile, upstream name, overall and process state, transition time, restart count, redacted error, and per-capability discovery state for each started upstream. See `examples/multi-upstream.miftah.json`. +`tooling.toolDiscoveryMode` defaults to `"permissive"`. In that mode, a bundled list returns only healthy upstream contributions, removes routes for failed upstreams, and retries incomplete tool discovery on later tool requests. If every upstream fails a capability list, Miftah returns `UPSTREAM_DISCOVERY_FAILED` instead of advertising an empty upstream surface. Set it to `"strict"` to reject any unavailable upstream; strict tool discovery additionally checks every configured profile for an identical client-visible tool contract. `miftah_health` reports the profile, upstream name, overall and process state, transition time, automatic restart count, redacted error, intentional stop reason, pending recovery time, and per-capability discovery state for each started upstream. See `examples/multi-upstream.miftah.json`. ## Routing and safety @@ -167,7 +167,7 @@ Use `miftah doctor` to inspect config and upstream readiness without printing pr ## Current boundaries -The current experimental code implements local STDIO and remote HTTP/SSE upstream clients, profile switching, hybrid routing rules, policies, namespaced tools/resources/prompts for account bundles, resilient healthy-upstream discovery, in-memory process/session caching, redacted JSONL audit logging, and a packageable CLI. Unsupported lifecycle tuning, persisted state, UI, routing plugins, profile matchers, and configurable tool namespaces are rejected with `UNSUPPORTED_CONFIG_OPTION` rather than silently ignored. +The current experimental code implements local STDIO and remote HTTP/SSE upstream clients, profile switching, hybrid routing rules, policies, namespaced tools/resources/prompts for account bundles, resilient healthy-upstream discovery, configurable local process lifecycle controls, in-memory process/session caching, redacted JSONL audit logging, and a packageable CLI. Local process controls cover startup and shutdown deadlines, optional idle cleanup, opt-in crash recovery with a bounded retry budget, and no-eviction profile-session capacity limits. Persisted state, UI, routing plugins, profile matchers, `process.startMode`, `process.cache`, and configurable tool namespaces are rejected with `UNSUPPORTED_CONFIG_OPTION` rather than silently ignored. ## License diff --git a/docs/architecture.md b/docs/architecture.md index de506e69..ba38827a 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -12,7 +12,7 @@ MCP client -> upstream STDIO server ``` -The public server is built with the official `@modelcontextprotocol/sdk` `Server` and `StdioServerTransport`. Each profile gets an SDK `Client` and `StdioClientTransport` on first use. A single-flight start map prevents concurrent calls from launching duplicate processes. Tool discovery is held in immutable, per-profile capability snapshots; concurrent callers share one discovery operation. Complete snapshots are cached, while a permissive partial snapshot remains callable but is retried by the next tool list or call request. A call captures its active profile before discovery, so a later profile switch cannot change that call's active-profile fallback. +The public server is built with the official `@modelcontextprotocol/sdk` `Server` and `StdioServerTransport`. Each profile/upstream pair gets an SDK `Client` and its configured transport on first use: local processes use `StdioClientTransport`, while remote upstreams use the matching SSE or streamable HTTP client transport. Per-profile start and restart single-flight state, attempt tokens, and lifecycle generations prevent concurrent calls or stale transport close events from creating duplicate children or replacing a newer session. Operations report begin/end activity so idle shutdown cannot close an in-flight request. Unexpected close is health `failed`; idle, explicit restart, wrapper shutdown, and timeout-enforced shutdown are intentional `stopped` states with a reason. Optional crash recovery uses bounded exponential backoff with jitter and a terminal retry budget; it retains its capacity reservation across the backoff window. A shared no-eviction limiter counts distinct profile bundles across named upstreams, so capacity pressure returns a typed refusal instead of terminating a live credential session. Tool discovery is held in immutable, per-profile capability snapshots; concurrent callers share one discovery operation. Complete snapshots are cached, while a permissive partial snapshot remains callable but is retried by the next tool list or call request. A call captures its active profile before discovery, so a later profile switch cannot change that call's active-profile fallback. The server advertises management tools plus tools discovered from the active profile. It advertises `tools.listChanged` and emits `notifications/tools/list_changed` after profile changes, restarts, and recovery that changes the public tool snapshot. Unknown names are rejected from the capability snapshot rather than guessed or forwarded. If a routing rule targets another profile, Miftah forwards only when that profile exposes the same name with an identical client-visible schema; otherwise it returns `TOOL_SCHEMA_MISMATCH`. In strict discovery mode, Miftah preflights every configured profile and rejects unavailable upstreams or any mismatched exposed tool contract before publishing a snapshot. diff --git a/docs/config.md b/docs/config.md index 320ff7d8..bcc18d4e 100644 --- a/docs/config.md +++ b/docs/config.md @@ -25,7 +25,7 @@ A standard `upstream` and a one-entry `upstreams` map retain credential-free raw - `"permissive"` returns tools, resources, and prompts from healthy upstreams when a peer cannot start or list a capability. Failed resource and prompt routes are removed rather than remaining callable. Incomplete tool snapshots are retried by later tool list or call requests. If no configured upstream succeeds, the operation fails with `UPSTREAM_DISCOVERY_FAILED`. - `"strict"` rejects a capability list if any configured upstream is unavailable. For tools, it discovers every configured profile and requires the complete exposed tool names and client-visible schemas to match; an availability failure returns `UPSTREAM_DISCOVERY_FAILED` and a mismatch returns `TOOL_SCHEMA_MISMATCH`. -Recovery that changes an aggregate surface emits the matching MCP list-change notification. `miftah_health` reports each started upstream with `upstreamName`, `profile`, overall `state`, `processState`, `lastTransition`, `restartCount`, redacted `error`, and `capabilities` for tools, resources, and prompts. +Recovery that changes an aggregate surface emits the matching MCP list-change notification. `miftah_health` reports each started upstream with `upstreamName`, `profile`, overall `state`, `processState`, `lastTransition`, `restartCount`, redacted `error`, and `capabilities` for tools, resources, and prompts. Intentional stops include `lastStopReason`; pending automatic recovery includes `nextRestartAt`; exhausted recovery includes `restartLimitReached`. Before loading secret sources or starting an upstream, Miftah validates that `defaultProfile`, profile policy names, routing-rule profiles, `security.lockToProfile`, and per-profile upstream override names all exist. Validation errors are human-readable and expose `MiftahError.details.diagnostics` for programs; each diagnostic includes a stable `code`, dotted `path`, `severity`, `message`, and remediation. @@ -50,6 +50,19 @@ Audit logging defaults to local JSONL when a path is configured. Arguments are e ## Runtime-supported controls -Miftah rejects settings without a runtime implementation with `UNSUPPORTED_CONFIG_OPTION` and the exact config path. The only supported routing mode is `"hybrid"`; use `routing.rules` and `routing.fallback` to control its behavior. Routing plugins, profile metadata and matchers, persistent state, UI settings, configurable management-tool namespaces, and lifecycle tuning other than `process.startupTimeoutMs` are not available yet. +Miftah rejects settings without a runtime implementation with `UNSUPPORTED_CONFIG_OPTION` and the exact config path. The only supported routing mode is `"hybrid"`; use `routing.rules` and `routing.fallback` to control its behavior. Routing plugins, profile metadata and matchers, persistent state, UI settings, and configurable management-tool namespaces are not available yet. + +### Process lifecycle + +`process` controls real local session behavior: + +- `startupTimeoutMs` is a positive integer and defaults to 30 seconds. A hung startup is terminated and reported as `UPSTREAM_START_FAILED`. +- `shutdownTimeoutMs` is a positive integer and defaults to 5 seconds. Miftah first closes STDIO gracefully, then force-terminates a child that exceeds the deadline or rejects close. Health records `lastStopReason: "shutdown-timeout"` for a deadline and `"shutdown-error"` for a rejected close rather than treating either cleanup as a crash. +- `idleTimeoutMs` is a positive integer. When omitted, sessions stay warm until an explicit restart or wrapper shutdown. When set, an inactive profile session closes after the timeout; in-flight upstream requests hold the session open. +- `restartOnCrash` defaults to `false`. When `true`, only an unexpected upstream loss schedules automatic recovery; explicit restart, idle expiry, and wrapper shutdown do not. +- `maxRestarts` is a non-negative integer and requires `restartOnCrash: true`. It defaults to 3 automatic attempts. Retry delay is bounded exponential backoff from 100 ms to 5 seconds with 20% jitter. A profile that keeps crashing exhausts its budget, reports `UPSTREAM_RESTART_LIMIT_EXCEEDED`, and rejects ordinary demand starts until `miftah_restart_profile` explicitly retries it. A recovered session must remain up for 30 seconds before its consecutive-crash budget resets. +- `maxConcurrentProfiles` is a positive integer. It limits distinct live profile bundles across all named upstreams. Miftah never evicts a live credential session to make room: a new profile instead receives `UPSTREAM_CONCURRENCY_LIMIT`. All upstreams for one profile share one slot, and idle, final crash, failed start, or intentional close releases it. + +`startMode` and `cache` remain unsupported because Miftah currently always lazily creates cached sessions. They are rejected rather than silently ignored. Secret and audit redaction are force-on protections. `security.redactSecrets` and `audit.redact` may be declared as `true`, but setting either to `false` is rejected. Audit output is always JSONL; `audit.format` therefore accepts only `"jsonl"`. diff --git a/src/cli/create-runtime.ts b/src/cli/create-runtime.ts index 7de16cbf..ff98c6b5 100644 --- a/src/cli/create-runtime.ts +++ b/src/cli/create-runtime.ts @@ -4,7 +4,12 @@ import { SecretResolver } from "../secrets/secret-resolver.js"; import { MultiUpstreamProcessManager } from "../upstream/multi-upstream-process-manager.js"; import { UpstreamProcessManager } from "../upstream/upstream-process-manager.js"; -/** Validates configuration before resolving secrets or constructing upstream runtime state. */ +/** + * Loads configuration, resolves its secrets, and constructs the runtime managers. + * + * @param configPath - Path to the configuration file. + * @returns The resolved configuration, upstream process manager, and profile manager. + */ export async function createRuntime(configPath: string) { const config = await loadConfig(configPath); const resolver = new SecretResolver({ @@ -12,21 +17,28 @@ export async function createRuntime(configPath: string) { allowPlaintextSecrets: config.secrets?.allowPlaintextSecrets ?? config.security?.allowPlaintextSecrets }); await resolver.load(); + const secretValues = new Set(); + const resolveMap = (values: Record | undefined): Record | undefined => { + if (!values) return values; + const resolved = resolver.resolveMapWithSecretValues(values); + for (const value of resolved.secretValues) secretValues.add(value); + return resolved.values; + }; const profiles = Object.fromEntries( Object.entries(config.profiles).map(([name, profile]) => [ name, { ...profile, - env: profile.env ? resolver.resolveMap(profile.env) : profile.env, - headers: profile.headers ? resolver.resolveMap(profile.headers) : profile.headers, + env: resolveMap(profile.env), + headers: resolveMap(profile.headers), upstreams: profile.upstreams ? Object.fromEntries( Object.entries(profile.upstreams).map(([upstreamName, override]) => [ upstreamName, { ...override, - env: override.env ? resolver.resolveMap(override.env) : override.env, - headers: override.headers ? resolver.resolveMap(override.headers) : override.headers + env: resolveMap(override.env), + headers: resolveMap(override.headers) } ]) ) @@ -40,8 +52,8 @@ export async function createRuntime(configPath: string) { name, { ...upstream, - env: upstream.env ? resolver.resolveMap(upstream.env) : undefined, - headers: upstream.headers ? resolver.resolveMap(upstream.headers) : undefined + env: resolveMap(upstream.env), + headers: resolveMap(upstream.headers) } ]) ) @@ -50,17 +62,14 @@ export async function createRuntime(configPath: string) { const upstream = resolvedConfig.upstream ? { ...resolvedConfig.upstream, - env: resolvedConfig.upstream.env ? resolver.resolveMap(resolvedConfig.upstream.env) : undefined, - headers: resolvedConfig.upstream.headers ? resolver.resolveMap(resolvedConfig.upstream.headers) : undefined + env: resolveMap(resolvedConfig.upstream.env), + headers: resolveMap(resolvedConfig.upstream.headers) } : undefined; + const managerOptions = { ...config.process, secretValues: [...secretValues] }; const manager = resolvedConfig.upstreams - ? new MultiUpstreamProcessManager(resolvedConfig, { - startupTimeoutMs: config.process?.startupTimeoutMs - }) - : new UpstreamProcessManager(upstream!, profiles, { - startupTimeoutMs: config.process?.startupTimeoutMs - }); + ? new MultiUpstreamProcessManager(resolvedConfig, managerOptions) + : new UpstreamProcessManager(upstream!, profiles, managerOptions); const profileManager = new ProfileManager(resolvedConfig, resolvedConfig.security); return { config: resolvedConfig, manager, profileManager }; } diff --git a/src/config/env-expand.ts b/src/config/env-expand.ts index bc23778f..2159ba3f 100644 --- a/src/config/env-expand.ts +++ b/src/config/env-expand.ts @@ -2,11 +2,20 @@ import { MiftahError } from "../utils/errors.js"; const environmentReference = /\$\{([A-Za-z_][A-Za-z0-9_]*)\}/g; -export function expandEnvironmentReferences( +/** + * Expands environment variable references in configuration values and collects their resolved values. + * + * @param values - Configuration values that may contain `${NAME}` references. + * @param environment - Environment variables used to resolve references. + * @returns The expanded values and the unique resolved environment values. + * @throws `MiftahError` if a referenced environment variable is undefined. + */ +export function expandEnvironmentReferencesWithSecretValues( values: Record, environment: NodeJS.ProcessEnv = process.env -): Record { - return Object.fromEntries( +): { values: Record; secretValues: string[] } { + const secretValues = new Set(); + const expandedValues = Object.fromEntries( Object.entries(values).map(([key, value]) => [ key, value.replace(environmentReference, (_, name: string) => { @@ -17,8 +26,24 @@ export function expandEnvironmentReferences( `SECRET_ENV_MISSING: profile environment variable '${name}' is not defined` ); } + secretValues.add(resolved); return resolved; }) ]) ); + return { values: expandedValues, secretValues: [...secretValues] }; +} + +/** + * Expands environment variable references in string values. + * + * @param values - Key-value pairs whose values may contain `${NAME}` references. + * @param environment - Environment variables used to resolve references. + * @returns The values with environment variable references replaced by their resolved values. + */ +export function expandEnvironmentReferences( + values: Record, + environment: NodeJS.ProcessEnv = process.env +): Record { + return expandEnvironmentReferencesWithSecretValues(values, environment).values; } diff --git a/src/config/schema.ts b/src/config/schema.ts index 263bbe9c..83403fa5 100644 --- a/src/config/schema.ts +++ b/src/config/schema.ts @@ -129,17 +129,26 @@ const securitySchema = z }) .strict(); -const publicProcessSchema = z.object({ startupTimeoutMs: z.number().int().positive().optional() }).strict(); +const publicProcessSchema = z + .object({ + startupTimeoutMs: z.number().int().positive().optional(), + shutdownTimeoutMs: z.number().int().positive().optional(), + idleTimeoutMs: z.number().int().positive().optional(), + restartOnCrash: z.boolean().optional(), + maxRestarts: z.number().int().nonnegative().optional(), + maxConcurrentProfiles: z.number().int().positive().optional() + }) + .strict(); const processSchema = z .object({ startMode: unsupportedOptionSchema, cache: unsupportedOptionSchema, - idleTimeoutMs: unsupportedOptionSchema, - restartOnCrash: unsupportedOptionSchema, - maxRestarts: unsupportedOptionSchema, startupTimeoutMs: z.number().int().positive().optional(), - shutdownTimeoutMs: unsupportedOptionSchema, - maxConcurrentProfiles: unsupportedOptionSchema + shutdownTimeoutMs: z.number().int().positive().optional(), + idleTimeoutMs: z.number().int().positive().optional(), + restartOnCrash: z.boolean().optional(), + maxRestarts: z.number().int().nonnegative().optional(), + maxConcurrentProfiles: z.number().int().positive().optional() }) .strict(); @@ -400,22 +409,23 @@ export const miftahConfigSchema = z rejectUnsupportedOption(["routing", "plugins"], "routing plugins are not implemented"); } - for (const option of [ - "startMode", - "cache", - "idleTimeoutMs", - "restartOnCrash", - "maxRestarts", - "shutdownTimeoutMs", - "maxConcurrentProfiles" - ] as const) { + for (const option of ["startMode", "cache"] as const) { if (value.process?.[option] !== undefined) { rejectUnsupportedOption( ["process", option], - `${option} is not implemented; only process.startupTimeoutMs currently changes runtime behavior` + `${option} is not implemented; Miftah always starts cached upstream sessions lazily` ); } } + if (value.process?.maxRestarts !== undefined && value.process.restartOnCrash !== true) { + addConfigIssue( + context, + "CONFIG_SCHEMA_INVALID", + ["process", "maxRestarts"], + "maxRestarts requires process.restartOnCrash to be true", + "Set process.restartOnCrash to true or remove process.maxRestarts." + ); + } if (value.security?.requireProfileSwitchConfirmation !== undefined) { rejectUnsupportedOption( diff --git a/src/config/types.ts b/src/config/types.ts index 40220314..0ffecebc 100644 --- a/src/config/types.ts +++ b/src/config/types.ts @@ -63,8 +63,14 @@ export interface SecurityConfig { lockToProfile?: string | null; } +/** Configures lifecycle behavior for profile-bound upstream processes. */ export interface ProcessConfig { startupTimeoutMs?: number; + shutdownTimeoutMs?: number; + idleTimeoutMs?: number; + restartOnCrash?: boolean; + maxRestarts?: number; + maxConcurrentProfiles?: number; } export interface AuditConfig { diff --git a/src/mcp/server/miftah-server.ts b/src/mcp/server/miftah-server.ts index 7d72f01f..29eeccc9 100644 --- a/src/mcp/server/miftah-server.ts +++ b/src/mcp/server/miftah-server.ts @@ -88,6 +88,7 @@ interface ResourcePromptProxyUnavailable { type ResourcePromptProxyAvailability = ResourcePromptProxyAvailable | ResourcePromptProxyUnavailable; +/** Hosts Miftah's MCP surface and coordinates profile routing, upstream discovery, and client notifications. */ export class MiftahServer { readonly server: Server; private readonly routing: RoutingEngine; @@ -98,7 +99,7 @@ export class MiftahServer { private readonly operationPipeline: OperationPipeline; private readonly resourcePromptRegistry?: ResourcePromptRegistry; private readonly invalidatedToolSnapshots = new Map(); - private readonly restartingProfiles = new Set(); + private readonly restartingProfiles = new Map>(); private readonly pendingResourceListChanges = new Set(); private readonly pendingPromptListChanges = new Set(); @@ -363,22 +364,7 @@ export class MiftahServer { } if (name === "miftah_restart_profile") { const profile = requiredString(args, "profile"); - this.restartingProfiles.add(profile); - try { - if (this.upstreams instanceof MultiUpstreamProcessManager) { - await this.upstreams.restartProfile(profile); - } else { - await this.upstreams.restart(profile); - } - } finally { - this.restartingProfiles.delete(profile); - this.toolRegistry.invalidate(profile); - this.resourcePromptRegistry?.invalidate(profile); - if (profile === this.profiles.current().activeProfile) { - await this.notifyToolListChanged(undefined, undefined); - await this.notifyResourcePromptListChanged(); - } - } + await this.restartUpstreamProfile(profile); return textResult("Profile restarted."); } if (name === "miftah_route_preview") { @@ -735,6 +721,39 @@ export class MiftahServer { } } + private restartUpstreamProfile(profile: string): Promise { + const existing = this.restartingProfiles.get(profile); + if (existing) return existing; + const restart = this.restartUpstreamProfileOnce(profile); + this.restartingProfiles.set(profile, restart); + void restart.then( + () => { + if (this.restartingProfiles.get(profile) === restart) this.restartingProfiles.delete(profile); + }, + () => { + if (this.restartingProfiles.get(profile) === restart) this.restartingProfiles.delete(profile); + } + ); + return restart; + } + + private async restartUpstreamProfileOnce(profile: string): Promise { + try { + if (this.upstreams instanceof MultiUpstreamProcessManager) { + await this.upstreams.restartProfile(profile); + } else { + await this.upstreams.restart(profile); + } + } finally { + this.toolRegistry.invalidate(profile); + this.resourcePromptRegistry?.invalidate(profile); + if (profile === this.profiles.current().activeProfile) { + await this.notifyToolListChanged(undefined, undefined); + await this.notifyResourcePromptListChanged(); + } + } + } + private handleUpstreamHealthChange(health: UpstreamHealth): void { if (health.processState !== "failed" || this.restartingProfiles.has(health.profile)) return; if (this.config.tooling?.toolDiscoveryMode === "strict") { diff --git a/src/secrets/secret-resolver.ts b/src/secrets/secret-resolver.ts index a96caf55..fc55f05f 100644 --- a/src/secrets/secret-resolver.ts +++ b/src/secrets/secret-resolver.ts @@ -2,12 +2,22 @@ import { parse } from "dotenv"; import { readFile as readFileAsync } from "node:fs/promises"; import { MiftahError } from "../utils/errors.js"; +const exactEnvironmentReferencePattern = /^\$\{([A-Za-z_][A-Za-z0-9_]*)\}$/; +const embeddedEnvironmentReferencePattern = /\$\{([A-Za-z_][A-Za-z0-9_]*)\}/g; + export interface SecretResolverOptions { environment?: NodeJS.ProcessEnv; envFiles?: string[]; allowPlaintextSecrets?: boolean; } +/** Contains resolved configuration values and every value sourced from a secret reference. */ +export interface ResolvedSecretMap { + values: Record; + secretValues: string[]; +} + +/** Resolves configured secret references without retaining provider-specific configuration in runtime objects. */ export class SecretResolver { private readonly environment: NodeJS.ProcessEnv; private readonly values: Record; @@ -38,19 +48,42 @@ export class SecretResolver { } resolveMap(values: Record): Record { - return Object.fromEntries(Object.entries(values).map(([key, value]) => [key, this.resolveValue(value)])); + return this.resolveMapWithSecretValues(values).values; + } + + /** Resolves a map while retaining secret-reference values for downstream diagnostic redaction. */ + resolveMapWithSecretValues(values: Record): ResolvedSecretMap { + const secretValues = new Set(); + const resolvedValues = Object.fromEntries( + Object.entries(values).map(([key, value]) => { + const resolved = this.resolveValueWithSecretValues(value); + for (const secretValue of resolved.secretValues) secretValues.add(secretValue); + return [key, resolved.value]; + }) + ); + return { values: resolvedValues, secretValues: [...secretValues] }; } resolveValue(value: string): string { - const environmentReference = value.match(/^\$\{([A-Za-z_][A-Za-z0-9_]*)\}$/); + return this.resolveValueWithSecretValues(value).value; + } + + private resolveValueWithSecretValues(value: string): { value: string; secretValues: string[] } { + const secretValues = new Set(); + const resolveReference = (name: string): string => { + const resolved = this.require(name); + secretValues.add(resolved); + return resolved; + }; + const environmentReference = value.match(exactEnvironmentReferencePattern); if (environmentReference) { - return this.require(environmentReference[1]!); + return { value: resolveReference(environmentReference[1]!), secretValues: [...secretValues] }; } if (value.startsWith("secretref:env://")) { - return this.require(value.slice("secretref:env://".length)); + return { value: resolveReference(value.slice("secretref:env://".length)), secretValues: [...secretValues] }; } if (value.startsWith("secretref:dotenv://")) { - return this.require(value.slice("secretref:dotenv://".length)); + return { value: resolveReference(value.slice("secretref:dotenv://".length)), secretValues: [...secretValues] }; } if (value.startsWith("secretref:plain://")) { if (this.options.allowPlaintextSecrets !== true) { @@ -59,7 +92,9 @@ export class SecretResolver { "SECRET_PROVIDER_FAILED: PLAINTEXT secret references are disabled" ); } - return value.slice("secretref:plain://".length); + const resolved = value.slice("secretref:plain://".length); + secretValues.add(resolved); + return { value: resolved, secretValues: [...secretValues] }; } if (value.startsWith("secretref:")) { throw new MiftahError( @@ -67,7 +102,10 @@ export class SecretResolver { `SECRET_PROVIDER_FAILED: unsupported secret provider in '${value.slice(0, value.indexOf("://") + 3)}'` ); } - return value.replace(/\$\{([A-Za-z_][A-Za-z0-9_]*)\}/g, (_, name: string) => this.require(name)); + return { + value: value.replace(embeddedEnvironmentReferencePattern, (_, name: string) => resolveReference(name)), + secretValues: [...secretValues] + }; } private require(name: string): string { diff --git a/src/upstream/multi-upstream-process-manager.ts b/src/upstream/multi-upstream-process-manager.ts index bdb532c1..38fdf92a 100644 --- a/src/upstream/multi-upstream-process-manager.ts +++ b/src/upstream/multi-upstream-process-manager.ts @@ -6,17 +6,21 @@ import { type UpstreamHealth, type UpstreamManagerOptions } from "./upstream-process-manager.js"; +import { ProfileSessionLimiter } from "./profile-session-limiter.js"; import { UpstreamSession } from "./upstream-session.js"; import { MiftahError } from "../utils/errors.js"; +/** Coordinates named upstream managers while sharing profile-capacity accounting across the bundle. */ export class MultiUpstreamProcessManager { private readonly managers: Record; private readonly healthListeners = new Set<(health: UpstreamHealth) => void>(); + private readonly limiter: ProfileSessionLimiter; constructor(config: MiftahConfig, options: UpstreamManagerOptions = {}) { + this.limiter = new ProfileSessionLimiter(options.maxConcurrentProfiles); this.managers = Object.fromEntries( Object.entries(config.upstreams ?? {}).map(([name, upstream]) => { - const manager = new UpstreamProcessManager(upstream, scopedProfiles(config.profiles, name), options, name); + const manager = new UpstreamProcessManager(upstream, scopedProfiles(config.profiles, name), options, name, this.limiter); manager.addHealthListener((health) => this.publishHealth(health)); return [name, manager]; }) diff --git a/src/upstream/profile-session-limiter.ts b/src/upstream/profile-session-limiter.ts new file mode 100644 index 00000000..e9fee544 --- /dev/null +++ b/src/upstream/profile-session-limiter.ts @@ -0,0 +1,34 @@ +import { MiftahError } from "../utils/errors.js"; + +/** + * Reserves a bounded number of distinct profile session bundles without evicting active processes. + */ +export class ProfileSessionLimiter { + private readonly ownersByProfile = new Map>(); + + constructor(private readonly maximumProfiles?: number) {} + + /** Acquires this owner's share of a profile slot, rejecting a new profile when the global limit is reached. */ + acquire(profile: string, owner: string): boolean { + const existingOwners = this.ownersByProfile.get(profile); + if (existingOwners?.has(owner)) return false; + if (existingOwners === undefined && this.maximumProfiles !== undefined && this.ownersByProfile.size >= this.maximumProfiles) { + throw new MiftahError( + "UPSTREAM_CONCURRENCY_LIMIT", + `UPSTREAM_CONCURRENCY_LIMIT: cannot start profile '${profile}' because ${this.maximumProfiles} profile session(s) are already active` + ); + } + const owners = existingOwners ?? new Set(); + owners.add(owner); + this.ownersByProfile.set(profile, owners); + return true; + } + + /** Releases an owner's share and frees the profile slot after its final upstream stops. */ + release(profile: string, owner: string): void { + const owners = this.ownersByProfile.get(profile); + if (!owners) return; + owners.delete(owner); + if (owners.size === 0) this.ownersByProfile.delete(profile); + } +} diff --git a/src/upstream/upstream-process-manager.ts b/src/upstream/upstream-process-manager.ts index 582e6fad..8ea12b06 100644 --- a/src/upstream/upstream-process-manager.ts +++ b/src/upstream/upstream-process-manager.ts @@ -1,18 +1,34 @@ +import { Client } from "@modelcontextprotocol/sdk/client/index.js"; import { getDefaultEnvironment, StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js"; import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"; import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js"; -import { Client } from "@modelcontextprotocol/sdk/client/index.js"; import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js"; -import type { Stream } from "node:stream"; import type { Tool } from "@modelcontextprotocol/sdk/types.js"; +import type { Stream } from "node:stream"; import type { ProfileConfig, UpstreamConfig } from "../config/types.js"; -import { expandEnvironmentReferences } from "../config/env-expand.js"; +import { expandEnvironmentReferencesWithSecretValues } from "../config/env-expand.js"; import { redactSecrets } from "../secrets/redact.js"; import { MiftahError } from "../utils/errors.js"; +import { ProfileSessionLimiter } from "./profile-session-limiter.js"; import { UpstreamSession } from "./upstream-session.js"; +const defaultStartupTimeoutMs = 30_000; +const defaultShutdownTimeoutMs = 5_000; +const defaultMaxRestarts = 3; +const initialRestartDelayMs = 100; +const maximumRestartDelayMs = 5_000; +const restartJitterFraction = 0.2; +const restartStabilityWindowMs = 30_000; +const credentialKeyPattern = /(token|secret|password|api[_-]?key|auth|private|credential|cookie)/i; + +/** Configures lifecycle behavior, capacity, and redacted diagnostics for an upstream manager. */ export interface UpstreamManagerOptions { startupTimeoutMs?: number; + shutdownTimeoutMs?: number; + idleTimeoutMs?: number; + restartOnCrash?: boolean; + maxRestarts?: number; + maxConcurrentProfiles?: number; secretValues?: readonly string[]; onStderr?: (profile: string, message: string) => void; } @@ -21,6 +37,9 @@ export type UpstreamCapability = "tools" | "resources" | "prompts"; export type UpstreamCapabilityState = "unknown" | "available" | "failed"; export type UpstreamProcessState = "stopped" | "starting" | "running" | "failed"; export type UpstreamState = UpstreamProcessState | "degraded"; +type ShutdownFailureReason = "shutdown-timeout" | "shutdown-error"; +/** Identifies an intentional reason a profile's upstream process stopped. */ +export type UpstreamStopReason = "idle" | "manual" | "restart" | "shutdown" | ShutdownFailureReason; export interface UpstreamCapabilityHealth { state: UpstreamCapabilityState; @@ -28,6 +47,7 @@ export interface UpstreamCapabilityHealth { error?: string; } +/** Describes the latest lifecycle and capability state for one profile/upstream pair. */ export interface UpstreamHealth { profile: string; upstreamName: string; @@ -36,56 +56,112 @@ export interface UpstreamHealth { state: UpstreamState; processState: UpstreamProcessState; lastTransition: string; + /** Number of automatic restart attempts, excluding manual restarts and idle wake-ups. */ restartCount: number; pid?: number | null; error?: string; + lastStopReason?: UpstreamStopReason; + nextRestartAt?: string; + restartLimitReached?: boolean; capabilities: Record; } +type StartSource = "demand" | "manual" | "automatic"; + +interface ManagedSession { + readonly session: UpstreamSession; + readonly transport: Transport; + readonly pid: number | null; + readonly token: number; + readonly generation: number; + inFlight: number; + closing: boolean; +} + +interface StartingAttempt { + readonly transport: Transport; + readonly generation: number; + pid: number | null; +} + +interface ScheduledRestart { + readonly generation: number; + readonly timer: NodeJS.Timeout; + readonly promise: Promise; + resolve(): void; +} + +type ResolvedOptions = Required< + Pick +> & + Omit; + +/** + * Owns one upstream process pool. Sessions are cached by profile and created only on demand. + */ export class UpstreamProcessManager { - private readonly sessions = new Map(); + private readonly sessions = new Map(); private readonly health = new Map(); private readonly starts = new Map>(); + private readonly startingAttempts = new Map(); + private readonly manualRestarts = new Map>(); + private readonly automaticRestarts = new Map(); + private readonly idleTimers = new Map(); + private readonly stabilityTimers = new Map(); + private readonly generations = new Map(); + private readonly startEpochs = new Map(); private readonly secretValuesSet = new Set(); - private readonly startAttempts = new Map(); + private readonly automaticRestartCounts = new Map(); + private readonly consecutiveRestartAttempts = new Map(); + private readonly restartExhausted = new Set(); private readonly processErrors = new Map(); private readonly healthListeners = new Set<(health: UpstreamHealth) => void>(); - private readonly options: Required> & - Omit; + private readonly options: ResolvedOptions; + private readonly limiter: ProfileSessionLimiter; + private nextToken = 0; + private closed = false; constructor( private readonly upstream: UpstreamConfig, private readonly profiles: Record, options: UpstreamManagerOptions = {}, - private readonly upstreamName = "default" + private readonly upstreamName = "default", + limiter?: ProfileSessionLimiter ) { this.options = { ...options, - startupTimeoutMs: options.startupTimeoutMs ?? 30_000 + startupTimeoutMs: options.startupTimeoutMs ?? defaultStartupTimeoutMs, + shutdownTimeoutMs: options.shutdownTimeoutMs ?? defaultShutdownTimeoutMs, + restartOnCrash: options.restartOnCrash ?? false, + maxRestarts: options.maxRestarts ?? defaultMaxRestarts }; + this.limiter = limiter ?? new ProfileSessionLimiter(options.maxConcurrentProfiles); } + /** Returns the live session or starts it after any scheduled recovery completes. */ async get(profile: string, _upstreamName?: string): Promise { void _upstreamName; - const current = this.sessions.get(profile); - if (current) return current; - const pending = this.starts.get(profile); - if (pending) return pending; - const start = this.start(profile); - this.starts.set(profile, start); - try { - return await start; - } finally { - this.starts.delete(profile); + this.assertOpen(); + const manualRestart = this.manualRestarts.get(profile); + if (manualRestart) return manualRestart; + const automaticRestart = this.automaticRestarts.get(profile); + if (automaticRestart) { + await automaticRestart.promise; + return this.get(profile); } + const current = this.sessions.get(profile); + if (current && !current.closing) return current.session; + this.assertRestartAvailable(profile); + return this.startOnce(profile, "demand"); } listHealth(): UpstreamHealth[] { return [...this.health.values()].map((health) => structuredClone(health)); } + /** Returns every configured or dynamically resolved value that must be redacted from upstream output. */ getSecretValues(): string[] { - return [...this.secretValuesSet]; + return [...new Set([...this.secretValuesSet, ...(this.options.secretValues ?? [])])]; } addHealthListener(listener: (health: UpstreamHealth) => void): () => void { @@ -109,29 +185,59 @@ export class UpstreamProcessManager { profile, capability, "failed", - redactSecrets(error instanceof Error ? error.message : String(error), this.options.secretValues ?? []) + this.redactProcessOutput(error instanceof Error ? error.message : String(error)) ); } + /** Manually restarts a profile and resets any exhausted automatic-recovery budget. */ async restart(profile: string, _upstreamName?: string): Promise { void _upstreamName; - await this.closeProfile(profile); - return this.get(profile); + this.assertOpen(); + const existing = this.manualRestarts.get(profile); + if (existing) return existing; + + this.cancelAutomaticRestart(profile, false); + this.resetRecoveryBudget(profile); + const restart = this.restartProfile(profile); + this.manualRestarts.set(profile, restart); + void restart.then( + () => { + if (this.manualRestarts.get(profile) === restart) this.manualRestarts.delete(profile); + }, + () => { + if (this.manualRestarts.get(profile) === restart) this.manualRestarts.delete(profile); + } + ); + return restart; } + /** Stops every known profile and waits for active startup and restart work to settle. */ async close(): Promise { - await Promise.all([...this.sessions.keys()].map((profile) => this.closeProfile(profile))); + if (this.closed) return; + this.closed = true; + const profiles = new Set([ + ...this.sessions.keys(), + ...this.starts.keys(), + ...this.manualRestarts.keys(), + ...this.automaticRestarts.keys(), + ...this.health.keys() + ]); + for (const profile of profiles) { + this.cancelAutomaticRestart(profile, true); + this.clearIdleTimer(profile); + this.clearStabilityTimer(profile); + } + await Promise.all(Array.from(profiles, (profile) => this.stopProfile(profile, "shutdown", true))); + await Promise.allSettled([...this.starts.values(), ...this.manualRestarts.values()]); } + /** Stops one profile without closing the manager itself. */ async closeProfile(profile: string): Promise { - const session = this.sessions.get(profile); - this.sessions.delete(profile); - if (session) { - await session.close(); - } - this.setProcessState(profile, "stopped", { resetCapabilities: true }); + this.cancelAutomaticRestart(profile, true); + await this.stopProfile(profile, "manual", true); } + /** Lists upstream tools and records a redacted capability outcome for health reporting. */ async listTools(profile: string, _upstreamName?: string): Promise { void _upstreamName; try { @@ -140,108 +246,462 @@ export class UpstreamProcessManager { return tools; } catch (error) { const failure = new MiftahError("UPSTREAM_TOOL_LIST_FAILED", `UPSTREAM_TOOL_LIST_FAILED: unable to list tools for '${profile}'`, { - cause: redactSecrets(error instanceof Error ? error.message : String(error), this.options.secretValues ?? []) + cause: this.redactProcessOutput(error instanceof Error ? error.message : String(error)) }); this.recordCapabilityFailure(profile, "tools", failure); throw failure; } } - private async start(profile: string): Promise { + private async restartProfile(profile: string): Promise { + await this.stopProfile(profile, "restart", false); + return this.startOnce(profile, "manual"); + } + + private startOnce(profile: string, source: StartSource): Promise { + this.assertOpen(); + const pending = this.starts.get(profile); + if (pending) return pending; + const start = this.start(profile, source); + this.starts.set(profile, start); + void start.then( + () => { + if (this.starts.get(profile) === start) this.starts.delete(profile); + }, + () => { + if (this.starts.get(profile) === start) this.starts.delete(profile); + } + ); + return start; + } + + private async start(profile: string, source: StartSource): Promise { const profileConfig = this.profiles[profile]; if (!profileConfig) { throw new MiftahError("PROFILE_NOT_FOUND", `PROFILE_NOT_FOUND: profile '${profile}' does not exist`); } if (this.upstream.transport === "stdio" && !this.upstream.command) { - throw new MiftahError( + throw new MiftahError("UPSTREAM_START_FAILED", "UPSTREAM_START_FAILED: stdio upstream requires a command"); + } + + const generation = this.generation(profile); + this.incrementStartEpoch(profile); + const token = ++this.nextToken; + let transport: Transport | undefined; + let stdioTransport: StdioClientTransport | undefined; + let startingAttempt: StartingAttempt | undefined; + let reserved = false; + + try { + reserved = this.limiter.acquire(profile, this.upstreamName); + this.setProcessState(profile, "starting", { resetCapabilities: true, pid: null }); + const { environment, headers } = this.resolveProfileOptions(profileConfig); + + if (this.upstream.transport === "stdio") { + stdioTransport = new StdioClientTransport({ + command: this.upstream.command!, + args: profileConfig.args ?? this.upstream.args ?? [], + env: { ...getDefaultEnvironment(), ...environment }, + ...(profileConfig.cwd ?? this.upstream.cwd ? { cwd: profileConfig.cwd ?? this.upstream.cwd } : {}), + stderr: "pipe" + }); + transport = stdioTransport; + this.attachStderr(profile, stdioTransport.stderr); + } else { + if (!this.upstream.url) { + throw new MiftahError("UPSTREAM_START_FAILED", "UPSTREAM_START_FAILED: remote upstream requires a url"); + } + const options = Object.keys(headers).length > 0 ? { requestInit: { headers } } : undefined; + transport = + this.upstream.transport === "sse" + ? new SSEClientTransport(new URL(this.upstream.url), options) + : new StreamableHTTPClientTransport(new URL(this.upstream.url), options); + } + + transport.onclose = () => this.handleTransportClosed(profile, token, generation); + const client = new Client({ name: "miftah", version: "0.1.1" }); + startingAttempt = { transport, generation, pid: null }; + this.startingAttempts.set(profile, startingAttempt); + const connection = client.connect(transport); + startingAttempt.pid = stdioTransport?.pid ?? null; + void connection.catch(() => undefined); + await withTimeout( + connection, + this.options.startupTimeoutMs, "UPSTREAM_START_FAILED", - "UPSTREAM_START_FAILED: stdio upstream requires a command" + `UPSTREAM_START_FAILED: startup timed out after ${this.options.startupTimeoutMs}ms` ); + + const pid = stdioTransport?.pid ?? null; + if (!this.isCurrent(profile, generation)) { + await this.terminateTransport(transport, pid); + throw new MiftahError("UPSTREAM_START_FAILED", `UPSTREAM_START_FAILED: startup for '${profile}' was cancelled`); + } + + const session = new UpstreamSession(profile, client, () => client.close(), { + begin: () => this.beginOperation(profile, token, entry), + end: () => this.endOperation(profile, token, entry) + }); + const entry: ManagedSession = { + session, + transport, + pid, + token, + generation, + inFlight: 0, + closing: false + }; + this.sessions.set(profile, entry); + if (this.startingAttempts.get(profile) === startingAttempt) this.startingAttempts.delete(profile); + this.setProcessState(profile, "running", { pid }); + this.scheduleIdleShutdown(profile, entry); + if (source === "automatic") this.scheduleStabilityWindow(profile, entry); + return session; + } catch (error) { + const pid = stdioTransport?.pid ?? null; + if (transport) await this.terminateTransport(transport, pid); + if (this.startingAttempts.get(profile) === startingAttempt) this.startingAttempts.delete(profile); + const current = this.isCurrent(profile, generation); + const failure = current + ? this.asStartFailure(profile, error) + : new MiftahError("UPSTREAM_START_FAILED", `UPSTREAM_START_FAILED: startup for '${profile}' was cancelled`); + if (current) { + this.setProcessState(profile, "failed", { error: failure.message, resetCapabilities: true, pid: null }); + } + if (source !== "automatic" || !this.canAutomaticallyRetry(profile)) { + this.limiter.release(profile, this.upstreamName); + } else if (!reserved) { + this.limiter.acquire(profile, this.upstreamName); + } + throw failure; } - this.startAttempts.set(profile, (this.startAttempts.get(profile) ?? 0) + 1); - this.setProcessState(profile, "starting", { resetCapabilities: true }); + } - const resolvedEnvironment = { - ...(this.upstream.env ? expandEnvironmentReferences(this.upstream.env) : {}), - ...(profileConfig.env ? expandEnvironmentReferences(profileConfig.env) : {}) + /** Resolves per-profile process settings and retains credential values for later diagnostic redaction. */ + private resolveProfileOptions(profile: ProfileConfig): { + environment: Record; + headers: Record; + } { + const upstreamEnvironment = this.upstream.env + ? expandEnvironmentReferencesWithSecretValues(this.upstream.env) + : undefined; + const profileEnvironment = profile.env ? expandEnvironmentReferencesWithSecretValues(profile.env) : undefined; + const upstreamHeaders = this.upstream.headers + ? expandEnvironmentReferencesWithSecretValues(this.upstream.headers) + : undefined; + const profileHeaders = profile.headers ? expandEnvironmentReferencesWithSecretValues(profile.headers) : undefined; + const environment = { + ...(upstreamEnvironment?.values ?? {}), + ...(profileEnvironment?.values ?? {}) }; - const resolvedHeaders = { - ...(this.upstream.headers ? expandEnvironmentReferences(this.upstream.headers) : {}), - ...(profileConfig.headers ? expandEnvironmentReferences(profileConfig.headers) : {}) + const headers = { + ...(upstreamHeaders?.values ?? {}), + ...(profileHeaders?.values ?? {}) }; - for (const [key, value] of Object.entries({ ...resolvedEnvironment, ...resolvedHeaders })) { - if (/(token|secret|password|api[_-]?key|auth|private|credential)/i.test(key) && value.length > 0) { + for (const value of [ + ...(upstreamEnvironment?.secretValues ?? []), + ...(profileEnvironment?.secretValues ?? []), + ...(upstreamHeaders?.secretValues ?? []), + ...(profileHeaders?.secretValues ?? []) + ]) { + this.secretValuesSet.add(value); + } + for (const [key, value] of Object.entries({ ...environment, ...headers })) { + if (credentialKeyPattern.test(key) && value.length > 0) { this.secretValuesSet.add(value); } } + return { environment, headers }; + } - let transport: Transport; - let pid: number | null = null; - let stderr: Stream | null = null; - if (this.upstream.transport === "stdio") { - const environment = { - ...getDefaultEnvironment(), - ...resolvedEnvironment - }; - const stdioTransport = new StdioClientTransport({ - command: this.upstream.command!, - args: profileConfig.args ?? this.upstream.args ?? [], - env: environment, - ...(profileConfig.cwd ?? this.upstream.cwd ? { cwd: profileConfig.cwd ?? this.upstream.cwd } : {}), - stderr: "pipe" - }); - transport = stdioTransport; - stderr = stdioTransport.stderr; + /** Emits process stderr only after applying static and dynamically resolved secret redaction. */ + private attachStderr(profile: string, stderr: Stream | null): void { + stderr?.on("data", (chunk: Buffer) => { + this.options.onStderr?.(profile, this.redactProcessOutput(chunk.toString("utf8"))); + }); + } + + /** Redacts process-originated text using static and dynamically resolved secret values. */ + private redactProcessOutput(value: string): string { + return redactSecrets(value, [...this.secretValuesSet, ...(this.options.secretValues ?? [])]); + } + + /** Handles an unexpected close only when it belongs to the current live session generation. */ + private handleTransportClosed(profile: string, token: number, generation: number): void { + const entry = this.sessions.get(profile); + if (!entry || entry.token !== token || entry.generation !== generation || entry.closing) return; + + this.sessions.delete(profile); + this.clearIdleTimer(profile); + this.clearStabilityTimer(profile); + if (!this.isCurrent(profile, generation)) return; + + this.setProcessState(profile, "failed", { + error: "UPSTREAM_START_FAILED: upstream process closed unexpectedly", + resetCapabilities: true, + pid: null + }); + if (this.options.restartOnCrash) { + this.scheduleAutomaticRestart(profile, generation); } else { - if (!this.upstream.url) { - throw new MiftahError("UPSTREAM_START_FAILED", "UPSTREAM_START_FAILED: remote upstream requires a url"); - } - const options = Object.keys(resolvedHeaders).length > 0 ? { requestInit: { headers: resolvedHeaders } } : undefined; - transport = - this.upstream.transport === "sse" - ? new SSEClientTransport(new URL(this.upstream.url), options) - : new StreamableHTTPClientTransport(new URL(this.upstream.url), options); + this.limiter.release(profile, this.upstreamName); } - stderr?.on("data", (chunk: Buffer) => { - this.options.onStderr?.( - profile, - redactSecrets(chunk.toString("utf8"), this.options.secretValues ?? []) - ); + } + + /** Schedules one bounded automatic retry while retaining the profile's capacity reservation. */ + private scheduleAutomaticRestart(profile: string, generation: number): void { + if (!this.isCurrent(profile, generation) || !this.options.restartOnCrash || this.manualRestarts.has(profile)) return; + if (this.automaticRestarts.has(profile)) return; + + const nextAttempt = (this.consecutiveRestartAttempts.get(profile) ?? 0) + 1; + if (nextAttempt > this.options.maxRestarts) { + this.restartExhausted.add(profile); + this.clearStabilityTimer(profile); + this.setProcessState(profile, "failed", { + error: `UPSTREAM_RESTART_LIMIT_EXCEEDED: automatic restart limit (${this.options.maxRestarts}) reached for '${profile}'`, + resetCapabilities: true, + pid: null, + restartLimitReached: true + }); + this.limiter.release(profile, this.upstreamName); + return; + } + + this.consecutiveRestartAttempts.set(profile, nextAttempt); + this.automaticRestartCounts.set(profile, (this.automaticRestartCounts.get(profile) ?? 0) + 1); + const delayMs = this.restartDelay(nextAttempt); + const nextRestartAt = new Date(Date.now() + delayMs).toISOString(); + this.setProcessState(profile, "failed", { + error: `UPSTREAM_START_FAILED: upstream process closed unexpectedly; automatic restart ${nextAttempt}/${this.options.maxRestarts} scheduled`, + resetCapabilities: true, + pid: null, + nextRestartAt }); - transport.onclose = () => { - if (this.sessions.get(profile)) { - this.sessions.delete(profile); - this.setProcessState(profile, "failed", { error: "upstream process closed", resetCapabilities: true }); - } + + let resolve!: () => void; + const promise = new Promise((done) => { + resolve = done; + }); + const scheduled: ScheduledRestart = { + generation, + timer: setTimeout(() => { + void this.runAutomaticRestart(profile, scheduled); + }, delayMs), + promise, + resolve }; - const client = new Client({ name: "miftah", version: "0.1.1" }); - const startPromise = client.connect(transport).then( - () => new UpstreamSession(profile, client, () => client.close()), - async (error: unknown) => { - await transport.close().catch(() => undefined); - throw new MiftahError( - "UPSTREAM_INIT_FAILED", - `UPSTREAM_INIT_FAILED: could not initialize profile '${profile}'`, - { cause: redactSecrets(error instanceof Error ? error.message : String(error), this.options.secretValues ?? []) } - ); + this.automaticRestarts.set(profile, scheduled); + } + + /** Runs the scheduled retry and chains another one only while this lifecycle generation remains current. */ + private async runAutomaticRestart(profile: string, scheduled: ScheduledRestart): Promise { + if (this.automaticRestarts.get(profile) !== scheduled) { + scheduled.resolve(); + return; + } + this.automaticRestarts.delete(profile); + try { + if (this.isCurrent(profile, scheduled.generation)) { + await this.startOnce(profile, "automatic"); + } + } catch { + if (this.isCurrent(profile, scheduled.generation)) { + this.scheduleAutomaticRestart(profile, scheduled.generation); } + } finally { + scheduled.resolve(); + } + } + + private restartDelay(attempt: number): number { + const baseDelay = Math.min(maximumRestartDelayMs, initialRestartDelayMs * 2 ** (attempt - 1)); + const jitter = baseDelay * restartJitterFraction; + return Math.max(1, Math.round(baseDelay - jitter + Math.random() * jitter * 2)); + } + + /** Clears the consecutive-crash budget only after an automatically recovered session remains stable. */ + private scheduleStabilityWindow(profile: string, entry: ManagedSession): void { + this.clearStabilityTimer(profile); + const timer = setTimeout(() => { + const current = this.sessions.get(profile); + if (!current || current.token !== entry.token || current.generation !== entry.generation || current.closing) return; + this.consecutiveRestartAttempts.delete(profile); + }, restartStabilityWindowMs); + this.stabilityTimers.set(profile, timer); + } + + private cancelAutomaticRestart(profile: string, releaseReservation: boolean): void { + const scheduled = this.automaticRestarts.get(profile); + if (scheduled) { + clearTimeout(scheduled.timer); + this.automaticRestarts.delete(profile); + scheduled.resolve(); + } + if (releaseReservation && !this.sessions.has(profile) && !this.starts.has(profile)) { + this.limiter.release(profile, this.upstreamName); + } + } + + private resetRecoveryBudget(profile: string): void { + this.restartExhausted.delete(profile); + this.consecutiveRestartAttempts.delete(profile); + this.clearStabilityTimer(profile); + } + + private canAutomaticallyRetry(profile: string): boolean { + return ( + this.options.restartOnCrash && + !this.restartExhausted.has(profile) && + (this.consecutiveRestartAttempts.get(profile) ?? 0) < this.options.maxRestarts + ); + } + + private assertRestartAvailable(profile: string): void { + if (!this.restartExhausted.has(profile)) return; + throw new MiftahError( + "UPSTREAM_RESTART_LIMIT_EXCEEDED", + `UPSTREAM_RESTART_LIMIT_EXCEEDED: automatic recovery for '${profile}' is exhausted; use miftah_restart_profile to retry` ); + } + + /** Stops a profile, cancels an in-progress start, and avoids stale lifecycle state after a replacement begins. */ + private async stopProfile(profile: string, reason: UpstreamStopReason, releaseReservation: boolean): Promise { + const startEpoch = this.startEpoch(profile); + const startingAttempt = this.startingAttempts.get(profile); + this.incrementGeneration(profile); + this.clearIdleTimer(profile); + this.clearStabilityTimer(profile); + const entry = this.sessions.get(profile); + this.sessions.delete(profile); + let shutdownFailure: ShutdownFailureReason | undefined; + if (entry) { + entry.closing = true; + shutdownFailure = await this.closeSession(entry); + } + if (startingAttempt) { + await this.terminateTransport(startingAttempt.transport, startingAttempt.pid); + } + if (this.startEpoch(profile) !== startEpoch) return; + if (releaseReservation) this.limiter.release(profile, this.upstreamName); + this.setProcessState(profile, "stopped", { + pid: null, + resetCapabilities: true, + lastStopReason: shutdownFailure ?? reason + }); + } + + /** Finalizes a session after a timeout or close error so lifecycle capacity is never stranded. */ + private async closeSession( + entry: ManagedSession + ): Promise { + const close = entry.session.close(); + void close.catch(() => undefined); try { - const session = await withTimeout(startPromise, this.options.startupTimeoutMs); - this.sessions.set(profile, session); - if ("pid" in transport && typeof transport.pid === "number") pid = transport.pid; - this.setProcessState(profile, "running", { pid }); - return session; + await withTimeout( + close, + this.options.shutdownTimeoutMs, + "UPSTREAM_SHUTDOWN_TIMEOUT", + `UPSTREAM_SHUTDOWN_TIMEOUT: shutdown timed out after ${this.options.shutdownTimeoutMs}ms` + ); + return undefined; } catch (error) { - await transport.close().catch(() => undefined); - this.setProcessState(profile, "failed", { - error: redactSecrets(error instanceof Error ? error.message : String(error), this.options.secretValues ?? []) - }); - throw error; + // Send SIGKILL before releasing capacity, but never let best-effort transport cleanup exceed the shutdown deadline. + void this.terminateTransport(entry.transport, entry.pid).catch(() => undefined); + return error instanceof MiftahError && error.code === "UPSTREAM_SHUTDOWN_TIMEOUT" + ? "shutdown-timeout" + : "shutdown-error"; + } + } + + private async terminateTransport(transport: Transport, pid: number | null): Promise { + if (pid !== null) { + try { + process.kill(pid, "SIGKILL"); + } catch (error) { + if (!(error instanceof Error) || !("code" in error) || error.code !== "ESRCH") throw error; + } + } + await transport.close().catch(() => undefined); + } + + /** Prevents an idle timer from closing a session after the activity callback has admitted an operation. */ + private beginOperation(profile: string, token: number, entry: ManagedSession | undefined): void { + const current = this.sessions.get(profile); + if (!entry || current !== entry || entry.token !== token || entry.closing) { + throw new MiftahError("UPSTREAM_CALL_FAILED", `UPSTREAM_CALL_FAILED: profile '${profile}' is no longer available`); + } + entry.inFlight += 1; + this.clearIdleTimer(profile); + } + + /** Re-arms idle shutdown only after the final in-flight operation completes. */ + private endOperation(profile: string, token: number, entry: ManagedSession | undefined): void { + const current = this.sessions.get(profile); + if (!entry || current !== entry || entry.token !== token) return; + entry.inFlight = Math.max(0, entry.inFlight - 1); + if (entry.inFlight === 0 && !entry.closing) this.scheduleIdleShutdown(profile, entry); + } + + private scheduleIdleShutdown(profile: string, entry: ManagedSession): void { + this.clearIdleTimer(profile); + if (this.options.idleTimeoutMs === undefined || entry.inFlight > 0 || entry.closing || this.closed) return; + const timer = setTimeout(() => { + void this.closeIdleSession(profile, entry); + }, this.options.idleTimeoutMs); + this.idleTimers.set(profile, timer); + } + + private async closeIdleSession(profile: string, entry: ManagedSession): Promise { + if (this.sessions.get(profile) !== entry || entry.inFlight > 0 || entry.closing || this.closed) return; + await this.stopProfile(profile, "idle", true); + } + + private clearIdleTimer(profile: string): void { + const timer = this.idleTimers.get(profile); + if (timer) clearTimeout(timer); + this.idleTimers.delete(profile); + } + + private clearStabilityTimer(profile: string): void { + const timer = this.stabilityTimers.get(profile); + if (timer) clearTimeout(timer); + this.stabilityTimers.delete(profile); + } + + private generation(profile: string): number { + return this.generations.get(profile) ?? 0; + } + + private incrementGeneration(profile: string): void { + this.generations.set(profile, this.generation(profile) + 1); + } + + private startEpoch(profile: string): number { + return this.startEpochs.get(profile) ?? 0; + } + + private incrementStartEpoch(profile: string): void { + this.startEpochs.set(profile, this.startEpoch(profile) + 1); + } + + private isCurrent(profile: string, generation: number): boolean { + return !this.closed && this.generation(profile) === generation; + } + + private assertOpen(): void { + if (this.closed) { + throw new MiftahError("UPSTREAM_START_FAILED", "UPSTREAM_START_FAILED: upstream manager is closed"); } } + private asStartFailure(profile: string, error: unknown): MiftahError { + if (error instanceof MiftahError) return error; + return new MiftahError("UPSTREAM_INIT_FAILED", `UPSTREAM_INIT_FAILED: could not initialize profile '${profile}'`, { + cause: this.redactProcessOutput(error instanceof Error ? error.message : String(error)) + }); + } + private recordCapability( profile: string, capability: UpstreamCapability, @@ -258,7 +718,10 @@ export class UpstreamProcessManager { this.publishHealth( this.healthEntry(profile, current.processState, { pid: current.pid, - capabilities + capabilities, + lastStopReason: current.lastStopReason, + nextRestartAt: current.nextRestartAt, + restartLimitReached: current.restartLimitReached }) ); } @@ -266,7 +729,14 @@ export class UpstreamProcessManager { private setProcessState( profile: string, processState: UpstreamProcessState, - options: { pid?: number | null; error?: string; resetCapabilities?: boolean } = {} + options: { + pid?: number | null; + error?: string; + resetCapabilities?: boolean; + lastStopReason?: UpstreamStopReason; + nextRestartAt?: string; + restartLimitReached?: boolean; + } = {} ): void { if (processState === "failed" && options.error !== undefined) { this.processErrors.set(profile, options.error); @@ -276,8 +746,11 @@ export class UpstreamProcessManager { const current = this.health.get(profile); this.publishHealth( this.healthEntry(profile, processState, { - pid: options.pid, - capabilities: options.resetCapabilities ? this.initialCapabilities() : current?.capabilities + pid: options.pid === undefined ? current?.pid : options.pid, + capabilities: options.resetCapabilities ? this.initialCapabilities() : current?.capabilities, + lastStopReason: processState === "stopped" ? options.lastStopReason : undefined, + nextRestartAt: processState === "failed" ? options.nextRestartAt : undefined, + restartLimitReached: processState === "failed" ? options.restartLimitReached : undefined }) ); } @@ -285,7 +758,13 @@ export class UpstreamProcessManager { private healthEntry( profile: string, processState: UpstreamProcessState, - options: { pid?: number | null; capabilities?: Record } = {} + options: { + pid?: number | null; + capabilities?: Record; + lastStopReason?: UpstreamStopReason; + nextRestartAt?: string; + restartLimitReached?: boolean; + } = {} ): UpstreamHealth { const capabilities = options.capabilities ?? this.initialCapabilities(); const capabilityFailure = Object.values(capabilities) @@ -301,9 +780,12 @@ export class UpstreamProcessManager { state, processState, lastTransition: new Date().toISOString(), - restartCount: Math.max(0, (this.startAttempts.get(profile) ?? 0) - 1), + restartCount: this.automaticRestartCounts.get(profile) ?? 0, ...(options.pid === undefined ? {} : { pid: options.pid }), ...(error === undefined ? {} : { error }), + ...(options.lastStopReason === undefined ? {} : { lastStopReason: options.lastStopReason }), + ...(options.nextRestartAt === undefined ? {} : { nextRestartAt: options.nextRestartAt }), + ...(options.restartLimitReached === undefined ? {} : { restartLimitReached: options.restartLimitReached }), capabilities }; } @@ -323,16 +805,27 @@ export class UpstreamProcessManager { } } -async function withTimeout(promise: Promise, timeoutMs: number): Promise { +/** + * Resolves with the promise value or fails when the timeout elapses. + * + * @param promise - The asynchronous operation to bound. + * @param timeoutMs - The maximum time to wait in milliseconds. + * @param code - The error code used when the timeout elapses. + * @param message - The error message used when the timeout elapses. + * @returns The resolved value of `promise`. + */ +async function withTimeout( + promise: Promise, + timeoutMs: number, + code: "UPSTREAM_START_FAILED" | "UPSTREAM_SHUTDOWN_TIMEOUT", + message: string +): Promise { let timer: NodeJS.Timeout | undefined; try { return await Promise.race([ promise, new Promise((_, reject) => { - timer = setTimeout( - () => reject(new MiftahError("UPSTREAM_START_FAILED", `UPSTREAM_START_FAILED: startup timed out after ${timeoutMs}ms`)), - timeoutMs - ); + timer = setTimeout(() => reject(new MiftahError(code, message)), timeoutMs); }) ]); } finally { diff --git a/src/upstream/upstream-session.ts b/src/upstream/upstream-session.ts index 60af40f9..271ca820 100644 --- a/src/upstream/upstream-session.ts +++ b/src/upstream/upstream-session.ts @@ -11,38 +11,55 @@ import type { ReadResourceRequest } from "@modelcontextprotocol/sdk/types.js"; +/** Lets the process manager bracket upstream work so idle shutdown cannot interrupt an active request. */ +export interface UpstreamSessionActivity { + begin(): void; + end(): void; +} + +/** Provides profile-bound MCP operations while reporting their activity to lifecycle management. */ export class UpstreamSession { constructor( readonly profile: string, private readonly client: Client, - private readonly closeTransport: () => Promise + private readonly closeTransport: () => Promise, + private readonly activity?: UpstreamSessionActivity ) {} listTools(): Promise { - return this.client.listTools(); + return this.request(() => this.client.listTools()); } callTool(params: CallToolRequest["params"]): Promise { - return this.client.callTool(params) as Promise; + return this.request(() => this.client.callTool(params) as Promise); } listResources(params?: ListResourcesRequest["params"]): Promise { - return this.client.listResources(params); + return this.request(() => this.client.listResources(params)); } readResource(params: ReadResourceRequest["params"]) { - return this.client.readResource(params); + return this.request(() => this.client.readResource(params)); } listPrompts(params?: ListPromptsRequest["params"]): Promise { - return this.client.listPrompts(params); + return this.request(() => this.client.listPrompts(params)); } getPrompt(params: GetPromptRequest["params"]) { - return this.client.getPrompt(params); + return this.request(() => this.client.getPrompt(params)); } async close(): Promise { await this.closeTransport(); } + + private async request(operation: () => Promise): Promise { + this.activity?.begin(); + try { + return await operation(); + } finally { + this.activity?.end(); + } + } } diff --git a/src/utils/errors.ts b/src/utils/errors.ts index edbfa305..18c56041 100644 --- a/src/utils/errors.ts +++ b/src/utils/errors.ts @@ -19,6 +19,9 @@ export type MiftahErrorCode = | "SECRET_PROVIDER_FAILED" | "UPSTREAM_START_FAILED" | "UPSTREAM_INIT_FAILED" + | "UPSTREAM_SHUTDOWN_TIMEOUT" + | "UPSTREAM_CONCURRENCY_LIMIT" + | "UPSTREAM_RESTART_LIMIT_EXCEEDED" | "UPSTREAM_TOOL_LIST_FAILED" | "UPSTREAM_DISCOVERY_FAILED" | "UPSTREAM_CALL_FAILED" diff --git a/tests/config-preflight.test.ts b/tests/config-preflight.test.ts index 404d8e7b..623e0a1d 100644 --- a/tests/config-preflight.test.ts +++ b/tests/config-preflight.test.ts @@ -1,9 +1,11 @@ -import { access, mkdtemp, writeFile } from "node:fs/promises"; +import { access, mkdtemp, rm, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { dirname, join } from "node:path"; import { fileURLToPath, pathToFileURL } from "node:url"; +import { setTimeout as delay } from "node:timers/promises"; import { describe, expect, it } from "vitest"; import { createRuntime } from "../src/cli/create-runtime.js"; +import { MiftahError } from "../src/utils/errors.js"; const fixture = join(dirname(fileURLToPath(import.meta.url)), "fixtures", "fake-upstream.mjs"); @@ -107,4 +109,89 @@ describe("configuration preflight", () => { await runtime.manager.close(); } }); + + it("redacts dotenv secrets assigned to noncredential upstream configuration keys", async () => { + const directory = await mkdtemp(join(tmpdir(), "miftah-runtime-redaction-")); + const configPath = join(directory, "miftah.json"); + const secretKey = `MIFTAH_TEST_COOKIE_${process.pid}_${Date.now()}`; + const secret = "runtime-cookie-secret"; + await writeFile(join(directory, ".env"), `${secretKey}=${secret}\n`); + await writeFile( + configPath, + JSON.stringify({ + version: "1", + name: "runtime-redaction", + defaultProfile: "default", + upstream: { + transport: "stdio", + command: process.execPath, + args: [fixture], + env: { + TEST_ERROR_MESSAGE: `secretref:dotenv://${secretKey}`, + TEST_FAIL_LIST_TOOLS: "true" + } + }, + profiles: { default: {} }, + secrets: { envFiles: [".env"] } + }) + ); + + const runtime = await createRuntime(configPath); + try { + let failure: unknown; + try { + await runtime.manager.listTools("default"); + } catch (error) { + failure = error; + } + if (!(failure instanceof MiftahError)) throw new Error("Expected a Miftah tool-list failure"); + const cause = failure.details?.cause; + if (typeof cause !== "string") throw new Error("Expected a redacted capability diagnostic cause"); + expect(cause).not.toContain(secret); + expect(cause).toContain("[REDACTED]"); + expect(runtime.manager.getSecretValues()).toContain(secret); + } finally { + await runtime.manager.close(); + await rm(directory, { recursive: true, force: true }); + } + }); + + it("wires supported lifecycle settings from configuration into the process manager", async () => { + const directory = await mkdtemp(join(tmpdir(), "miftah-lifecycle-runtime-")); + const configPath = join(directory, "miftah.json"); + await writeFile( + configPath, + JSON.stringify({ + version: "1", + name: "lifecycle-runtime", + defaultProfile: "default", + upstream: { + transport: "stdio", + command: process.execPath, + args: [fixture] + }, + profiles: { default: {} }, + process: { + startupTimeoutMs: 1_000, + shutdownTimeoutMs: 100, + idleTimeoutMs: 50, + restartOnCrash: true, + maxRestarts: 2, + maxConcurrentProfiles: 1 + } + }) + ); + + const runtime = await createRuntime(configPath); + try { + await runtime.manager.get("default"); + await delay(150); + expect(runtime.manager.listHealth()).toMatchObject([ + { profile: "default", processState: "stopped", lastStopReason: "idle" } + ]); + } finally { + await runtime.manager.close(); + await rm(directory, { recursive: true, force: true }); + } + }); }); diff --git a/tests/config-runtime-parity.test.ts b/tests/config-runtime-parity.test.ts index 7793cf24..4fce41b3 100644 --- a/tests/config-runtime-parity.test.ts +++ b/tests/config-runtime-parity.test.ts @@ -39,12 +39,6 @@ describe("config runtime parity", () => { it.each([ ["process.startMode", { process: { startMode: "lazy" } }], ["process.cache", { process: { cache: true } }], - ["process.idleTimeoutMs", { process: { idleTimeoutMs: 1 } }], - ["process.restartOnCrash", { process: { restartOnCrash: true } }], - ["process.restartOnCrash", { process: { restartOnCrash: "true" } }], - ["process.maxRestarts", { process: { maxRestarts: 1 } }], - ["process.shutdownTimeoutMs", { process: { shutdownTimeoutMs: 1 } }], - ["process.maxConcurrentProfiles", { process: { maxConcurrentProfiles: 1 } }], ["routing.mode", { routing: { mode: "active" } }], ["routing.mode", { routing: { mode: "rules" } }], ["routing.mode", { routing: { mode: true } }], @@ -153,6 +147,41 @@ describe("config runtime parity", () => { expect(config.audit?.redact).toBe(true); }); + it("accepts implemented lifecycle controls", () => { + const config = validateConfig( + baseConfig({ + process: { + startupTimeoutMs: 1_000, + shutdownTimeoutMs: 500, + idleTimeoutMs: 1_500, + restartOnCrash: true, + maxRestarts: 2, + maxConcurrentProfiles: 3 + } + }) + ); + + expect(config.process).toEqual({ + startupTimeoutMs: 1_000, + shutdownTimeoutMs: 500, + idleTimeoutMs: 1_500, + restartOnCrash: true, + maxRestarts: 2, + maxConcurrentProfiles: 3 + }); + }); + + it("requires automatic recovery when a restart limit is configured", () => { + const error = validationError( + baseConfig({ + process: { maxRestarts: 2 } + }) + ); + + expect(error.code).toBe("CONFIG_SCHEMA_INVALID"); + expect(error.message).toContain("process.maxRestarts"); + }); + it.each(["permissive", "strict"] as const)("accepts the %s discovery mode", (toolDiscoveryMode) => { const config = validateConfig( baseConfig({ diff --git a/tests/config-schema-contract.test.ts b/tests/config-schema-contract.test.ts index 2ab4f984..41b87dde 100644 --- a/tests/config-schema-contract.test.ts +++ b/tests/config-schema-contract.test.ts @@ -101,7 +101,14 @@ describe("published config schema", () => { expect(profileUpstreamOverride.properties).not.toHaveProperty("transport"); expect(profileUpstreamOverride.properties).not.toHaveProperty("command"); expect(profileUpstreamOverride.properties).not.toHaveProperty("url"); - expect(Object.keys(process ?? {})).toEqual(["startupTimeoutMs"]); + expect(Object.keys(process ?? {})).toEqual([ + "startupTimeoutMs", + "shutdownTimeoutMs", + "idleTimeoutMs", + "restartOnCrash", + "maxRestarts", + "maxConcurrentProfiles" + ]); expect(security).toMatchObject({ redactSecrets: { const: true } }); expect(security).not.toHaveProperty("requireProfileSwitchConfirmation"); expect(audit).toMatchObject({ format: { const: "jsonl" }, redact: { const: true } }); diff --git a/tests/config.test.ts b/tests/config.test.ts index 92889287..81dc08a5 100644 --- a/tests/config.test.ts +++ b/tests/config.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from "vitest"; import { validateConfig } from "../src/config/validate-config.js"; -import { expandEnvironmentReferences } from "../src/config/env-expand.js"; +import { expandEnvironmentReferences, expandEnvironmentReferencesWithSecretValues } from "../src/config/env-expand.js"; import { MiftahError } from "../src/utils/errors.js"; const policyNotFoundPattern = /POLICY_NOT_FOUND/u; @@ -25,6 +25,18 @@ describe("config foundation", () => { ).toEqual({ API_TOKEN: "secret-value", ACCOUNT: "work" }); }); + it("tracks values resolved from environment references regardless of their configuration key", () => { + expect( + expandEnvironmentReferencesWithSecretValues( + { COOKIE: "session=${SESSION_ID}", ACCOUNT: "work" }, + { SESSION_ID: "dynamic-session-secret" } + ) + ).toEqual({ + values: { COOKIE: "session=dynamic-session-secret", ACCOUNT: "work" }, + secretValues: ["dynamic-session-secret"] + }); + }); + it("rejects a config whose default profile does not exist", () => { expect(() => validateConfig({ diff --git a/tests/fixtures/fake-upstream.mjs b/tests/fixtures/fake-upstream.mjs index 1faf02e7..32a1805e 100644 --- a/tests/fixtures/fake-upstream.mjs +++ b/tests/fixtures/fake-upstream.mjs @@ -5,6 +5,7 @@ import { setTimeout as delay } from "node:timers/promises"; import { CallToolRequestSchema, GetPromptRequestSchema, + InitializeRequestSchema, ListPromptsRequestSchema, ListResourcesRequestSchema, ListToolsRequestSchema, @@ -34,6 +35,49 @@ const failOnRestartPath = process.env.TEST_FAIL_ON_RESTART_PATH; const failListResourcesPath = process.env.TEST_FAIL_LIST_RESOURCES_PATH; const failListPromptsPath = process.env.TEST_FAIL_LIST_PROMPTS_PATH; const crashOnCallToolPath = process.env.TEST_CRASH_ON_CALL_TOOL_PATH; +const crashAfterConnectPath = process.env.TEST_CRASH_AFTER_CONNECT_PATH; +const crashAfterConnectDelayMs = Number(process.env.TEST_CRASH_AFTER_CONNECT_DELAY_MS ?? "10"); +const startCountPath = process.env.TEST_START_COUNT_PATH; +const failInitialize = process.env.TEST_FAIL_INITIALIZE === "true"; +const stderrMessage = process.env.TEST_STDERR_MESSAGE; +const hangOnStartPath = process.env.TEST_HANG_ON_START_PATH; +const hangOnStartReadyPath = process.env.TEST_HANG_ON_START_READY_PATH; +const shutdownDelayMs = Number(process.env.TEST_SHUTDOWN_DELAY_MS ?? "0"); +const shutdownEndPath = process.env.TEST_SHUTDOWN_END_PATH; + +if (startCountPath) { + appendFileSync(startCountPath, "1\n"); +} +if (process.env.TEST_HANG_ON_START === "true" || (hangOnStartPath && existsSync(hangOnStartPath))) { + if (hangOnStartReadyPath) { + writeFileSync(hangOnStartReadyPath, "ready"); + } + if (hangOnStartPath) { + while (existsSync(hangOnStartPath)) { + await delay(5); + } + } else { + for (;;) { + await delay(1_000); + } + } +} +if (shutdownEndPath || shutdownDelayMs > 0) { + process.stdin.once("end", () => { + if (shutdownEndPath) { + writeFileSync(shutdownEndPath, "ended"); + } + if (shutdownDelayMs > 0) { + void delay(shutdownDelayMs).then(() => process.exit(0)); + } + }); +} +if (process.env.TEST_IGNORE_SIGTERM === "true") { + process.on("SIGTERM", () => undefined); +} +if (stderrMessage) { + process.stderr.write(`${stderrMessage}\n`); +} if (crashOnCallToolPath && existsSync(crashOnCallToolPath)) { throw new Error("test upstream configured to stay unavailable after an abrupt exit"); } @@ -72,6 +116,12 @@ const server = new Server( { capabilities: { tools: {}, resources: {}, prompts: {} } } ); +if (failInitialize) { + server.setRequestHandler(InitializeRequestSchema, async () => { + throw new Error(`test initialize failure: ${process.env.API_TOKEN}`); + }); +} + server.setRequestHandler(ListToolsRequestSchema, async () => { if (process.env.TEST_LIST_TOOLS_STARTED_PATH) { writeFileSync(process.env.TEST_LIST_TOOLS_STARTED_PATH, "started"); @@ -82,6 +132,9 @@ server.setRequestHandler(ListToolsRequestSchema, async () => { if (listToolsDelayMs > 0) { await delay(listToolsDelayMs); } + if (process.env.TEST_FAIL_LIST_TOOLS === "true") { + throw new Error(`test tool list failure: ${process.env.TEST_ERROR_MESSAGE ?? process.env.API_TOKEN}`); + } return { tools: [ { @@ -271,3 +324,6 @@ server.setRequestHandler(GetPromptRequestSchema, async () => { }); await server.connect(new StdioServerTransport()); +if (crashAfterConnectPath && existsSync(crashAfterConnectPath)) { + void delay(crashAfterConnectDelayMs).then(() => process.exit(1)); +} diff --git a/tests/upstream-manager.test.ts b/tests/upstream-manager.test.ts index 3a58b699..1c26a746 100644 --- a/tests/upstream-manager.test.ts +++ b/tests/upstream-manager.test.ts @@ -1,10 +1,38 @@ +import { existsSync } from "node:fs"; +import { mkdtemp, readFile, rm, unlink, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; import { dirname, join } from "node:path"; import { fileURLToPath } from "node:url"; -import { describe, expect, it } from "vitest"; +import { setTimeout as delay } from "node:timers/promises"; +import { describe, expect, it, vi } from "vitest"; +import { MultiUpstreamProcessManager } from "../src/upstream/multi-upstream-process-manager.js"; import { UpstreamProcessManager } from "../src/upstream/upstream-process-manager.js"; +import { MiftahError } from "../src/utils/errors.js"; const fixture = join(dirname(fileURLToPath(import.meta.url)), "fixtures", "fake-upstream.mjs"); +async function countStarts(path: string): Promise { + const contents = await readFile(path, "utf8"); + return contents.split("\n").filter(Boolean).length; +} + +async function waitFor( + read: () => Value | Promise, + matches: (value: Value) => boolean, + timeoutMs = 2_000 +): Promise { + const deadline = Date.now() + timeoutMs; + let value = await read(); + while (!matches(value)) { + if (Date.now() >= deadline) { + throw new Error(`Timed out after ${timeoutMs}ms waiting for lifecycle state; last value: ${JSON.stringify(value)}`); + } + await delay(10); + value = await read(); + } + return value; +} + describe("upstream process manager", () => { it("starts one cached upstream per profile and forwards MCP operations", async () => { const manager = new UpstreamProcessManager( @@ -52,4 +80,603 @@ describe("upstream process manager", () => { await manager.close(); } }); + + it("redacts dynamically resolved secrets from manager stderr and capability diagnostics", async () => { + const secret = "dynamic-profile-secret"; + const stderr: string[] = []; + const manager = new UpstreamProcessManager( + { + transport: "stdio", + command: process.execPath, + args: [fixture] + }, + { + work: { + env: { + API_TOKEN: secret, + TEST_STDERR_MESSAGE: `test stderr: ${secret}`, + TEST_FAIL_LIST_TOOLS: "true" + } + } + }, + { + startupTimeoutMs: 1_000, + onStderr: (_profile, message) => stderr.push(message) + } + ); + + try { + await manager.get("work"); + await waitFor(() => stderr.join("\n"), (output) => output.includes("test stderr")); + expect(stderr.join("\n")).not.toContain(secret); + expect(stderr.join("\n")).toContain("[REDACTED]"); + + let failure: unknown; + try { + await manager.listTools("work"); + } catch (error) { + failure = error; + } + if (!(failure instanceof MiftahError)) throw new Error("Expected a Miftah error from failed tool discovery"); + const cause = failure.details?.cause; + if (typeof cause !== "string") throw new Error("Expected a redacted diagnostic cause"); + expect(cause).not.toContain(secret); + expect(cause).toContain("[REDACTED]"); + } finally { + await manager.close(); + } + }); + + it("redacts dynamically resolved secrets from startup diagnostics", async () => { + const secret = "dynamic-startup-secret"; + const manager = new UpstreamProcessManager( + { + transport: "stdio", + command: process.execPath, + args: [fixture] + }, + { + work: { + env: { + API_TOKEN: secret, + TEST_FAIL_INITIALIZE: "true" + } + } + }, + { startupTimeoutMs: 1_000 } + ); + + try { + let failure: unknown; + try { + await manager.get("work"); + } catch (error) { + failure = error; + } + if (!(failure instanceof MiftahError)) throw new Error("Expected a Miftah startup error"); + const cause = failure.details?.cause; + if (typeof cause !== "string") throw new Error("Expected a redacted startup diagnostic cause"); + expect(cause).not.toContain(secret); + expect(cause).toContain("[REDACTED]"); + } finally { + await manager.close(); + } + }); + + it("shuts down an idle profile and starts a fresh process on its next use", async () => { + const directory = await mkdtemp(join(tmpdir(), "miftah-idle-")); + const startCountPath = join(directory, "starts"); + await writeFile(startCountPath, ""); + const manager = new UpstreamProcessManager( + { + transport: "stdio", + command: process.execPath, + args: [fixture], + env: { TEST_START_COUNT_PATH: startCountPath } + }, + { work: {} }, + { startupTimeoutMs: 1_000, idleTimeoutMs: 50 } + ); + + try { + await (await manager.get("work")).listTools(); + const stopped = await waitFor( + () => manager.listHealth().find((health) => health.profile === "work"), + (health) => health?.processState === "stopped" && health.lastStopReason === "idle" + ); + if (!stopped) throw new Error("Expected idle shutdown health"); + expect(stopped.restartCount).toBe(0); + expect(stopped.pid).toBeNull(); + expect(await countStarts(startCountPath)).toBe(1); + + await (await manager.get("work")).listTools(); + await waitFor(() => countStarts(startCountPath), (count) => count === 2); + } finally { + await manager.close(); + await rm(directory, { recursive: true, force: true }); + } + }); + + it("keeps a replacement session healthy and reserved while an idle predecessor is still closing", async () => { + const directory = await mkdtemp(join(tmpdir(), "miftah-idle-race-")); + const shutdownEndPath = join(directory, "stdin-ended"); + const manager = new UpstreamProcessManager( + { + transport: "stdio", + command: process.execPath, + args: [fixture], + env: { + TEST_SHUTDOWN_DELAY_MS: "500", + TEST_SHUTDOWN_END_PATH: shutdownEndPath, + TEST_LIST_TOOLS_DELAY_MS: "800" + } + }, + { work: {}, personal: {} }, + { startupTimeoutMs: 1_000, shutdownTimeoutMs: 1_000, idleTimeoutMs: 30, maxConcurrentProfiles: 1 } + ); + + try { + await manager.get("work"); + await waitFor(() => existsSync(shutdownEndPath), Boolean); + + const replacement = await manager.get("work"); + const inFlightList = replacement.listTools(); + await delay(600); + + expect(manager.listHealth()).toMatchObject([{ profile: "work", processState: "running" }]); + await expect(manager.get("personal")).rejects.toMatchObject({ code: "UPSTREAM_CONCURRENCY_LIMIT" }); + await expect(inFlightList).resolves.toMatchObject({ tools: expect.any(Array) }); + } finally { + await manager.close(); + await rm(directory, { recursive: true, force: true }); + } + }); + + it("does not idle-shutdown a profile while an upstream request is in flight", async () => { + const manager = new UpstreamProcessManager( + { + transport: "stdio", + command: process.execPath, + args: [fixture], + env: { TEST_LIST_TOOLS_DELAY_MS: "125" } + }, + { work: {} }, + { startupTimeoutMs: 1_000, idleTimeoutMs: 30 } + ); + + try { + const session = await manager.get("work"); + await expect(session.listTools()).resolves.toMatchObject({ tools: expect.any(Array) }); + expect(manager.listHealth()).toMatchObject([{ profile: "work", processState: "running" }]); + await waitFor( + () => manager.listHealth().find((health) => health.profile === "work"), + (health) => health?.processState === "stopped" && health.lastStopReason === "idle" + ); + } finally { + await manager.close(); + } + }); + + it("does not restart a crashed process unless automatic recovery is configured", async () => { + const directory = await mkdtemp(join(tmpdir(), "miftah-no-restart-")); + const crashPath = join(directory, "crash"); + const startCountPath = join(directory, "starts"); + await writeFile(startCountPath, ""); + const manager = new UpstreamProcessManager( + { + transport: "stdio", + command: process.execPath, + args: [fixture], + env: { + TEST_CRASH_ON_CALL_TOOL_PATH: crashPath, + TEST_START_COUNT_PATH: startCountPath + } + }, + { work: {} }, + { startupTimeoutMs: 1_000, restartOnCrash: false } + ); + + try { + const session = await manager.get("work"); + await writeFile(crashPath, "crash"); + await expect(session.callTool({ name: "whoami", arguments: {} })).rejects.toThrow(); + await waitFor( + () => manager.listHealth().find((health) => health.profile === "work"), + (health) => health?.processState === "failed" + ); + await delay(300); + expect(await countStarts(startCountPath)).toBe(1); + } finally { + await manager.close(); + await rm(directory, { recursive: true, force: true }); + } + }); + + it("single-flights concurrent starts and management restarts for the same profile", async () => { + const directory = await mkdtemp(join(tmpdir(), "miftah-single-flight-")); + const startCountPath = join(directory, "starts"); + const restartBlockPath = join(directory, "restart-block"); + const restartReadyPath = join(directory, "restart-ready"); + const restartReleasePath = join(directory, "restart-release"); + await writeFile(startCountPath, ""); + const manager = new UpstreamProcessManager( + { + transport: "stdio", + command: process.execPath, + args: [fixture], + env: { + TEST_START_COUNT_PATH: startCountPath, + TEST_BLOCK_ON_RESTART_PATH: restartBlockPath, + TEST_BLOCK_ON_RESTART_READY_PATH: restartReadyPath, + TEST_BLOCK_ON_RESTART_RELEASE_PATH: restartReleasePath + } + }, + { work: {} }, + { startupTimeoutMs: 1_000 } + ); + + try { + const initial = await Promise.all([manager.get("work"), manager.get("work"), manager.get("work")]); + expect(initial[0]).toBe(initial[1]); + expect(initial[1]).toBe(initial[2]); + expect(await countStarts(startCountPath)).toBe(1); + + const firstRestart = manager.restart("work"); + const secondRestart = manager.restart("work"); + await waitFor(() => existsSync(restartReadyPath), Boolean); + const joinedRestart = manager.get("work"); + await writeFile(restartReleasePath, "release"); + const [first, second, joined] = await Promise.all([firstRestart, secondRestart, joinedRestart]); + + expect(first).toBe(second); + expect(second).toBe(joined); + expect(await countStarts(startCountPath)).toBe(2); + } finally { + await manager.close(); + await rm(directory, { recursive: true, force: true }); + } + }); + + it("automatically recovers a crashed profile after a bounded backoff", async () => { + const directory = await mkdtemp(join(tmpdir(), "miftah-auto-restart-")); + const crashPath = join(directory, "crash"); + const startCountPath = join(directory, "starts"); + await writeFile(startCountPath, ""); + const manager = new UpstreamProcessManager( + { + transport: "stdio", + command: process.execPath, + args: [fixture], + env: { + TEST_CRASH_ON_CALL_TOOL_PATH: crashPath, + TEST_START_COUNT_PATH: startCountPath + } + }, + { work: {} }, + { startupTimeoutMs: 1_000, restartOnCrash: true, maxRestarts: 2 } + ); + + try { + const session = await manager.get("work"); + await writeFile(crashPath, "crash"); + await expect(session.callTool({ name: "whoami", arguments: {} })).rejects.toThrow(); + await unlink(crashPath); + + await waitFor(() => countStarts(startCountPath), (count) => count === 2); + const recovered = await waitFor( + () => manager.listHealth().find((health) => health.profile === "work"), + (health) => health?.processState === "running" + ); + if (!recovered) throw new Error("Expected recovered health"); + expect(recovered.restartCount).toBe(1); + await expect((await manager.get("work")).callTool({ name: "whoami", arguments: {} })).resolves.toMatchObject({ + content: [{ type: "text", text: "unknown" }] + }); + } finally { + await manager.close(); + await rm(directory, { recursive: true, force: true }); + } + }); + + it("stops automatic recovery when the configured restart budget is exhausted", async () => { + const directory = await mkdtemp(join(tmpdir(), "miftah-restart-limit-")); + const crashPath = join(directory, "crash"); + const startCountPath = join(directory, "starts"); + await Promise.all([writeFile(crashPath, "crash"), writeFile(startCountPath, "")]); + const manager = new UpstreamProcessManager( + { + transport: "stdio", + command: process.execPath, + args: [fixture], + env: { + TEST_CRASH_AFTER_CONNECT_PATH: crashPath, + TEST_CRASH_AFTER_CONNECT_DELAY_MS: "10", + TEST_START_COUNT_PATH: startCountPath + } + }, + { work: {} }, + { startupTimeoutMs: 1_000, restartOnCrash: true, maxRestarts: 2 } + ); + + try { + await manager.get("work"); + const exhausted = await waitFor( + () => manager.listHealth().find((health) => health.profile === "work"), + (health) => health?.restartLimitReached === true + ); + expect(exhausted).toMatchObject({ + processState: "failed", + restartCount: 2, + error: expect.stringContaining("UPSTREAM_RESTART_LIMIT_EXCEEDED") + }); + expect(await countStarts(startCountPath)).toBe(3); + await expect(manager.get("work")).rejects.toMatchObject({ code: "UPSTREAM_RESTART_LIMIT_EXCEEDED" }); + await unlink(crashPath); + const manuallyRestarted = await manager.restart("work"); + await expect(manuallyRestarted.callTool({ name: "whoami", arguments: {} })).resolves.toMatchObject({ + content: [{ type: "text", text: "unknown" }] + }); + expect(manager.listHealth()).toMatchObject([{ profile: "work", restartCount: 2, processState: "running" }]); + } finally { + await manager.close(); + await rm(directory, { recursive: true, force: true }); + } + }); + + it("bounds a hung startup and cleans up its child process", async () => { + const manager = new UpstreamProcessManager( + { + transport: "stdio", + command: process.execPath, + args: [fixture] + }, + { work: { env: { TEST_HANG_ON_START: "true" } } }, + { startupTimeoutMs: 200 } + ); + const startedAt = Date.now(); + + try { + const startup = manager.get("work"); + void startup.catch(() => undefined); + await expect(startup).rejects.toMatchObject({ code: "UPSTREAM_START_FAILED" }); + expect(Date.now() - startedAt).toBeLessThan(500); + expect(manager.listHealth()).toMatchObject([{ profile: "work", processState: "failed" }]); + } finally { + await manager.close(); + } + }); + + it("releases a profile capacity reservation after a failed startup", async () => { + const directory = await mkdtemp(join(tmpdir(), "miftah-failed-start-capacity-")); + const failurePath = join(directory, "fail"); + await writeFile(failurePath, "fail"); + const manager = new UpstreamProcessManager( + { + transport: "stdio", + command: process.execPath, + args: [fixture] + }, + { + work: { env: { TEST_CRASH_ON_CALL_TOOL_PATH: failurePath } }, + personal: {} + }, + { startupTimeoutMs: 1_000, maxConcurrentProfiles: 1 } + ); + + try { + await expect(manager.get("work")).rejects.toMatchObject({ code: "UPSTREAM_INIT_FAILED" }); + await expect((await manager.get("personal")).listTools()).resolves.toMatchObject({ tools: expect.any(Array) }); + } finally { + await manager.close(); + await rm(directory, { recursive: true, force: true }); + } + }); + + it("cancels a hanging startup when the manager shuts down", async () => { + const manager = new UpstreamProcessManager( + { + transport: "stdio", + command: process.execPath, + args: [fixture], + env: { TEST_HANG_ON_START: "true" } + }, + { work: {} }, + { startupTimeoutMs: 10_000 } + ); + + const startup = manager.get("work"); + await delay(50); + const startedAt = Date.now(); + await manager.close(); + + expect(Date.now() - startedAt).toBeLessThan(300); + await expect(startup).rejects.toMatchObject({ code: "UPSTREAM_START_FAILED" }); + }); + + it("restarts after cancelling an in-flight startup without reusing the cancelled attempt", async () => { + const directory = await mkdtemp(join(tmpdir(), "miftah-restart-startup-")); + const hangPath = join(directory, "hang"); + const hangReadyPath = join(directory, "hang-ready"); + const startCountPath = join(directory, "starts"); + await Promise.all([writeFile(hangPath, "hang"), writeFile(startCountPath, "")]); + const manager = new UpstreamProcessManager( + { + transport: "stdio", + command: process.execPath, + args: [fixture], + env: { + TEST_HANG_ON_START_PATH: hangPath, + TEST_HANG_ON_START_READY_PATH: hangReadyPath, + TEST_START_COUNT_PATH: startCountPath + } + }, + { work: {} }, + { startupTimeoutMs: 10_000 } + ); + + try { + const initialStartup = manager.get("work"); + void initialStartup.catch(() => undefined); + await waitFor(() => existsSync(hangReadyPath), Boolean); + await unlink(hangPath); + + const restarted = await manager.restart("work"); + await expect(restarted.listTools()).resolves.toMatchObject({ tools: expect.any(Array) }); + await expect(initialStartup).rejects.toMatchObject({ code: "UPSTREAM_START_FAILED" }); + expect(await countStarts(startCountPath)).toBe(2); + } finally { + await manager.close(); + await rm(directory, { recursive: true, force: true }); + } + }); + + it("forces a delayed shutdown to respect the configured timeout", async () => { + const directory = await mkdtemp(join(tmpdir(), "miftah-shutdown-")); + const shutdownEndPath = join(directory, "stdin-ended"); + const manager = new UpstreamProcessManager( + { + transport: "stdio", + command: process.execPath, + args: [fixture], + env: { + TEST_SHUTDOWN_DELAY_MS: "500", + TEST_SHUTDOWN_END_PATH: shutdownEndPath + } + }, + { work: {} }, + { startupTimeoutMs: 1_000, shutdownTimeoutMs: 50 } + ); + + try { + await manager.get("work"); + const startedAt = Date.now(); + await manager.close(); + expect(Date.now() - startedAt).toBeLessThan(300); + await expect(readFile(shutdownEndPath, "utf8")).resolves.toBe("ended"); + expect(manager.listHealth()).toMatchObject([ + { profile: "work", processState: "stopped", lastStopReason: "shutdown-timeout" } + ]); + } finally { + await manager.close(); + await rm(directory, { recursive: true, force: true }); + } + }); + + it("releases capacity after a session close rejects", async () => { + const manager = new UpstreamProcessManager( + { + transport: "stdio", + command: process.execPath, + args: [fixture] + }, + { work: {}, personal: {} }, + { startupTimeoutMs: 1_000, maxConcurrentProfiles: 1 } + ); + + try { + const work = await manager.get("work"); + vi.spyOn(work, "close").mockRejectedValueOnce(new Error("simulated close failure")); + + await expect(manager.closeProfile("work")).resolves.toBeUndefined(); + expect(manager.listHealth()).toMatchObject([ + { profile: "work", processState: "stopped", lastStopReason: "shutdown-error" } + ]); + await expect(manager.get("personal")).resolves.toBeDefined(); + } finally { + await manager.close(); + } + }); + + it("finalizes a timed-out close without waiting for its original promise", async () => { + const manager = new UpstreamProcessManager( + { + transport: "stdio", + command: process.execPath, + args: [fixture] + }, + { work: {}, personal: {} }, + { startupTimeoutMs: 1_000, shutdownTimeoutMs: 25, maxConcurrentProfiles: 1 } + ); + + try { + const work = await manager.get("work"); + vi.spyOn(work, "close").mockImplementation(() => new Promise(() => undefined)); + + const completion = await Promise.race([ + manager.closeProfile("work").then(() => "closed"), + delay(500).then(() => "timed-out") + ]); + expect(completion).toBe("closed"); + expect(manager.listHealth()).toMatchObject([ + { profile: "work", processState: "stopped", lastStopReason: "shutdown-timeout" } + ]); + await expect(manager.get("personal")).resolves.toBeDefined(); + } finally { + await manager.close(); + } + }); + + it("limits a multi-upstream bundle by distinct active profiles without evicting a live session", async () => { + const manager = new MultiUpstreamProcessManager( + { + version: "1", + name: "multi", + defaultProfile: "work", + upstreams: { + primary: { transport: "stdio", command: process.execPath, args: [fixture] }, + secondary: { transport: "stdio", command: process.execPath, args: [fixture] } + }, + profiles: { work: {}, personal: {} } + }, + { startupTimeoutMs: 1_000, maxConcurrentProfiles: 1 } + ); + + try { + await manager.get("work", "primary"); + await manager.get("work", "secondary"); + await expect(manager.get("personal", "primary")).rejects.toMatchObject({ code: "UPSTREAM_CONCURRENCY_LIMIT" }); + expect(manager.listHealth()).toEqual( + expect.arrayContaining([ + expect.objectContaining({ profile: "work", upstreamName: "primary", processState: "running" }), + expect.objectContaining({ profile: "work", upstreamName: "secondary", processState: "running" }) + ]) + ); + } finally { + await manager.close(); + } + }); + + it("cancels pending automatic recovery when the manager shuts down", async () => { + const directory = await mkdtemp(join(tmpdir(), "miftah-restart-cleanup-")); + const crashPath = join(directory, "crash"); + const startCountPath = join(directory, "starts"); + await Promise.all([writeFile(crashPath, "crash"), writeFile(startCountPath, "")]); + const manager = new UpstreamProcessManager( + { + transport: "stdio", + command: process.execPath, + args: [fixture], + env: { + TEST_CRASH_AFTER_CONNECT_PATH: crashPath, + TEST_START_COUNT_PATH: startCountPath + } + }, + { work: {} }, + { startupTimeoutMs: 1_000, restartOnCrash: true, maxRestarts: 3 } + ); + + try { + await manager.get("work"); + await waitFor( + () => manager.listHealth().find((health) => health.profile === "work"), + (health) => health?.processState === "failed" + ); + await manager.close(); + await delay(300); + expect(await countStarts(startCountPath)).toBe(1); + } finally { + await manager.close(); + await rm(directory, { recursive: true, force: true }); + } + }); });