diff --git a/CHANGELOG.md b/CHANGELOG.md index c63974f1..2052d885 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ ### Added +- **Explicit runtime lifecycle ownership (PR #226).** `Tramai` and `SovereignTramai` are now `AutoCloseable` and own exactly one lazily-created runtime (one engine) shared by every `create()`/`runtime()` call — previously every `create()` leaked an unreachable engine. Closing is idempotent and concurrency-safe; after close, `create()`/`runtime()` and old proxies fail fast with a fixed `IllegalStateException` before any provider work. `TramaiEngine.close()` cancels once and awaits engine-hierarchy termination (self-close safe), and terminates in-flight suspend invocations; the caller continuation is always resumed exactly once. Spring closes the shared runtime via `destroyMethod = "close"`, so multiple `@AiService` beans share one owned engine. TramAI closes only resources it creates; externally supplied providers/stores/clients/observers remain caller-owned. API surface addition is additive: `Tramai`/`SovereignTramai` gain `close()`; all constructor descriptors remain byte-identical to 0.5.0 (note: adding the `AutoCloseable` supertype is source-compatible but affects compiled negative-`instanceof` checks). Epic 1.3 Runtime Lifecycle Ownership is complete. + - **Safe persistence failure boundaries (PR #225).** Persistence stores expose fixed, cause-free failure text; raw paths, SQL, and payloads flow only to `PersistenceFailureDiagnosticObserver`; worker observers receive safe failures; existing exception and store ABI is preserved by the binary fixture. Epic 1.2 Safe Error Boundaries is complete. - **Safe provider and built-in workflow-step failure boundaries (PRs #222, #223).** Provider HTTP rejections and built-in HTTP, shell, MCP, Codex, and Hermes workflow failures expose fixed cause-free public exceptions with typed failure codes. Original failure detail is retained only by an explicitly configured, fail-open diagnostic observer; public workflow events omit URLs, commands, raw tool names, and failure reasons. Existing public exception constructor descriptors remain compatible with 0.5.0 clients. diff --git a/docs/ROADMAP-0.6.0.md b/docs/ROADMAP-0.6.0.md index 1050f4af..46a083b5 100644 --- a/docs/ROADMAP-0.6.0.md +++ b/docs/ROADMAP-0.6.0.md @@ -356,6 +356,8 @@ This phase is intentionally completed before large decomposition work. **Goal:** Ensure no convenience API creates an uncloseable runtime. +> **Status:** ✅ Complete — PR #226. `Tramai`/`SovereignTramai` are `AutoCloseable` and own one runtime/engine; `close()` cancels and joins engine-owned work (blocking calls, suspend invocations, streaming collections) via an internal lifecycle job; caller-supplied `job`/`scope` are never cancelled/joined; Spring closes via `destroyMethod`; resource ownership documented. + ### Tasks 1. Make `Tramai` own one engine or one runtime session rather than constructing an unreachable engine per `create()` call. @@ -373,6 +375,16 @@ This phase is intentionally completed before large decomposition work. - Repeated creation does not accidentally create independent hidden engines. - Spring context shutdown leaves no TramAI-owned jobs or hooks active. +### Leak-test evidence matrix (roadmap task 6) + +| Requirement | Existing proof | +|---|---| +| Engine jobs | PR #226 lifecycle tests: `close() cancels and joins` blocking/suspend/streaming engine-owned work (`blocking invocation in long suspension is cancelled and joined by close`, `self close from owned coroutine does not deadlock`, `self close from streaming owned coroutine does not deadlock`, `stream start racing close never hangs the collector`, `close does not deadlock when caller supplied its own job and scope`) | +| Worker jobs | Existing worker shutdown/cancellation tests in tramai-orchestration (`TramaiWorkerTest`, lease-drain and shutdown coverage from Epics 1.1/1.2) | +| Subprocesses | PR #216/#221 cancellation contract (`SubprocessCancellationContractTest` in tramai-orchestration) | +| HTTP response streams | Provider-level InputStream cleanup tests in tramai-openai `OpenAiProviderTest`: `stream closes response body after done marker`, `stream closes response body after malformed chunk`, `stream closes response body when collector stops after first token`, `mid stream io failure is retryable sanitized and observed`; plus #226 engine streaming lifecycle tests (`streaming collection suspended indefinitely is cancelled and cleaned up by close`, `mid-collection close terminates an in-flight stream`) and the springboot example E2E smoke test | +| Shutdown hooks | `TramaiWorkerTest`: `close deregisters the JVM shutdown hook and retains no reference` — proves the registered hook is absent from `Runtime` after close (`removeShutdownHook` returns false) and the worker retains no `Thread` reference; plus Spring `destroyMethod` close + context-shutdown tests and `repeated close is harmless` idempotency test in tramai-standalone | + --- ## Epic 1.4: HTTP network-boundary correctness diff --git a/docs/modules/tramai-engine.md b/docs/modules/tramai-engine.md index 6077fc0c..306e1612 100644 --- a/docs/modules/tramai-engine.md +++ b/docs/modules/tramai-engine.md @@ -718,7 +718,7 @@ The test suite covers **28 behavioral scenarios** across 2 test files: |------|------|-------------| | `TramaiEngine` | Class | Main engine: creates AI-backed proxies from annotated interfaces | | `TramaiEngine.create()` | Method | Returns a JVM proxy implementing the given service type | -| `TramaiEngine.close()` | Method | Cancels the engine's coroutine job hierarchy | +| `TramaiEngine.close()` | Method | Cancels and joins all engine-owned work: blocking calls, suspend invocations, and streaming collections. The caller-supplied `job`/`scope` constructor parameters are never cancelled or joined | | `RetryPolicySettings` | Data class | Retry delay computation: max `Retry-After` cap, jitter ratio | | `CircuitBreakerSettings` | Data class | Per-provider circuit breaker: enabled, threshold, open duration | | `TokenBudgetSettings` | Data class | Token budgets: hard per-attempt, hard/soft per-operation | diff --git a/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt b/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt index 927a26ae..f685d8d6 100644 --- a/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt +++ b/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt @@ -90,8 +90,10 @@ import dev.tramai.core.exception.ToolInvalidInputException import dev.tramai.core.model.* import kotlinx.coroutines.CancellationException import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.CoroutineExceptionHandler import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.Job +import kotlinx.coroutines.asContextElement import kotlinx.coroutines.SupervisorJob import kotlinx.coroutines.TimeoutCancellationException import kotlinx.coroutines.currentCoroutineContext @@ -102,6 +104,7 @@ import kotlinx.coroutines.flow.collect import kotlinx.coroutines.flow.flow import kotlinx.coroutines.launch import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.withContext import kotlinx.coroutines.withTimeout import java.lang.reflect.InvocationHandler import java.lang.reflect.Method @@ -171,6 +174,44 @@ class TramaiEngine( ?: LegacyPermissivePolicyEngine private val isLegacyFallback: Boolean = policyEngine == null private val resumeOperationRegistry: ResumeOperationRegistry = ResumeOperationRegistry() + private val closed = java.util.concurrent.atomic.AtomicBoolean(false) + private val engineThreadMarker = ThreadLocal() + /** + * Internally owned lifecycle job and scope. The engine's OWN work (blocking + * calls, streaming collections) parents here — never to the caller-supplied + * [job]/[scope] constructor parameters, which remain for ABI compatibility + * only. close() cancels and joins [lifecycleJob], so it can prove that + * engine-initiated work has terminated, regardless of where the caller's + * job lives (and without risking the caller-job join deadlock). + */ + private val lifecycleJob: Job = SupervisorJob() + // Every engine-owned child (streaming collection, future lifecycle tasks) + // carries the engine-thread marker: close() called from ANY engine-owned + // coroutine (provider/interceptor/observer re-entering close) skips the + // join and cannot self-deadlock. Encoding ownership once at the scope + // level is stronger than decorating individual launches. + private val lifecycleScope: CoroutineScope = CoroutineScope( + lifecycleJob + Dispatchers.Default + engineThreadMarker.asContextElement(true) + CoroutineExceptionHandler { _, error -> + // Engine-owned background work can outlive its caller (e.g. a + // streaming collection abandoned mid-flight). Its failure is + // already surfaced to the caller's continuation when one exists. + // Orphaned failures must not crash the process or leak onto a + // shared global handler. Log FIXED safe metadata only — never the + // raw throwable: the failure may carry externally supplied + // exception messages (PII), which the safe-error-boundary work + // (Epic 1.2) keeps out of normal logs. + System.getLogger("dev.tramai.engine.TramaiEngine").log( + System.Logger.Level.WARNING, + "Engine-owned coroutine failed after close or abandonment (type: ${error::class.qualifiedName})", + ) + }, + ) + /** + * Suspend-invocation jobs launched for caller continuations. They are + * children of the CALLER's job (so parent cancellation propagates), but the + * engine tracks them so close() terminates in-flight work it owns. + */ + private val activeInvocationJobs = java.util.concurrent.ConcurrentHashMap.newKeySet() /** * Creates an engine backed by a single provider. @@ -384,6 +425,7 @@ class TramaiEngine( * Creates a proxy implementation for the given Tramai service interface. */ fun create(serviceType: KClass): T { + check(!closed.get()) { "Tramai runtime is closed" } val definition = ServiceDefinition.create( serviceType = serviceType, toolRegistry = toolRegistry, @@ -405,6 +447,11 @@ class TramaiEngine( chatMemory = chatMemory, conversationIdProvider = conversationIdProvider, scope = scope, + lifecycleJob = lifecycleJob, + lifecycleScope = lifecycleScope, + isClosed = closed, + engineThreadMarker = engineThreadMarker, + activeInvocationJobs = activeInvocationJobs, serviceDefinition = definition, policyEngine = resolvedPolicyEngine, migrationWarningGuard = migrationWarningGuard, @@ -447,6 +494,7 @@ class TramaiEngine( * Conflicting registration (same key, different digest) fails closed. */ fun registerService(serviceType: KClass<*>) { + check(!closed.get()) { "Tramai runtime is closed" } val definition = ServiceDefinition.create( serviceType = serviceType, toolRegistry = toolRegistry, @@ -468,6 +516,11 @@ class TramaiEngine( chatMemory = chatMemory, conversationIdProvider = conversationIdProvider, scope = scope, + lifecycleJob = lifecycleJob, + lifecycleScope = lifecycleScope, + isClosed = closed, + engineThreadMarker = engineThreadMarker, + activeInvocationJobs = activeInvocationJobs, serviceDefinition = definition, policyEngine = resolvedPolicyEngine, migrationWarningGuard = migrationWarningGuard, @@ -506,6 +559,7 @@ class TramaiEngine( * @throws dev.tramai.core.exception.ApprovalAuthorizationException on store-level failures. */ suspend fun resumeApproval(command: ResumeApprovalCommand): Any? { + check(!closed.get()) { "Tramai runtime is closed" } // P1-2: Check continuation status BEFORE loading metadata // (post-completion cleanup removes metadata, but continuation is authoritative) val store = approvalContinuationStore @@ -531,10 +585,43 @@ class TramaiEngine( resumeApproval(command) as R /** - * Cancels the engine-owned coroutine job hierarchy. + * Cancels all engine-initiated work and, except from one of the engine's + * own coroutines, waits for it to terminate. The caller-supplied [job] and + * [scope] constructor parameters are NEVER cancelled or joined here — the + * engine owns its own [lifecycleJob], so closing cannot deadlock a caller + * that passed its current job. Dependencies supplied by callers are not + * closed. */ override fun close() { - job.cancel() + if (closed.compareAndSet(false, true)) { + lifecycleJob.cancel() + // Suspend invocations are children of their CALLER's job, not the + // engine scope job; cancel them explicitly so close() owns them. + // Synchronized with the launch+add in invokeSuspend: either the + // launch completed first (its job is in the set and gets cancelled + // here) or close() won and the launch's in-lock re-check rejects it. + val tracked = synchronized(activeInvocationJobs) { + activeInvocationJobs.toList() + } + tracked.forEach { it.cancel() } + if (engineThreadMarker.get() != true) { + // Wait for the engine-owned hierarchy AND every tracked + // invocation: cancellation is a request, not termination — + // cleanup (e.g. NonCancellable finally blocks) must complete + // before close() returns. Invocation jobs run on the engine's + // own dispatcher (lifecycleScope's Dispatchers.Default; the + // caller's ContinuationInterceptor is stripped at launch so a + // single-threaded caller loop can't be blocked by close()), + // so joining is safe as long as close() is not called from a + // coroutine dispatched on that same engine dispatcher + // (documented caller constraint, matching the self-close + // marker guard below). + runBlocking { + lifecycleJob.join() + tracked.forEach { it.join() } + } + } + } } } @@ -576,6 +663,11 @@ internal class TramaiInvocationHandler( private val chatMemory: ChatMemory?, private val conversationIdProvider: ConversationIdProvider, private val scope: CoroutineScope, + private val lifecycleJob: Job, + private val lifecycleScope: CoroutineScope, + private val isClosed: java.util.concurrent.atomic.AtomicBoolean = java.util.concurrent.atomic.AtomicBoolean(false), + private val engineThreadMarker: ThreadLocal = ThreadLocal(), + private val activeInvocationJobs: MutableSet = java.util.concurrent.ConcurrentHashMap.newKeySet(), private val serviceDefinition: ServiceDefinition, policyEngine: PolicyEngine, private val migrationWarningGuard: java.util.concurrent.atomic.AtomicBoolean, @@ -639,6 +731,8 @@ internal class TramaiInvocationHandler( return handleObjectMethod(proxy, method, args.orEmpty()) } + check(!isClosed.get()) { "Tramai runtime is closed" } + val operation = serviceDefinition.operations[method] ?: throw ConfigurationException("No operation metadata registered for ${method.name}") @@ -646,9 +740,21 @@ internal class TramaiInvocationHandler( return if (operation.isSuspend) { invokeSuspend(operation, args.orEmpty(), conversationId) } else { - runBlocking { + // Run the blocking call as a child of the engine's OWN lifecycle + // job (not the caller-supplied job/scope): close() cancels and + // joins lifecycleJob, so it can terminate a blocking provider + // that is still executing when the engine is closed. The thread + // marker marks this coroutine as engine-owned so a blocking call + // that itself invokes close() skips the join (avoiding a + // self-deadlock on lifecycleJob). + val result = runBlocking(lifecycleJob + engineThreadMarker.asContextElement(true)) { execute(operation, args.orEmpty().toList(), conversationId) } + // The engine may have closed while this blocking call was in + // flight. Never deliver a result computed against a closed engine: + // the caller sees the fixed lifecycle error instead. + check(!isClosed.get()) { "Tramai runtime is closed" } + result } } @@ -663,10 +769,53 @@ internal class TramaiInvocationHandler( ?: throw ConfigurationException("Suspend invocation for ${operation.method.name} is missing its continuation") val callArguments = args.dropLast(1) - scope.launch(continuation.context) { - runCatching { execute(operation, callArguments, conversationId) } - .onSuccess { continuation.resumeWith(Result.success(it)) } - .onFailure { continuation.resumeWith(Result.failure(it)) } + // Launch as a child of the CALLER's job (continuation.context, with the + // caller's Job element retained) so parent cancellation propagates + // synchronously into the in-flight invocation (validated by the + // ToolSafeFailureContract / StructuredOutputFailureBoundary + // parent-cancellation tests). The invocation RUNS on the engine's own + // dispatcher (lifecycleScope), NOT the caller's: if it ran on the + // caller's single-threaded dispatcher, close() joining it could + // deadlock when that thread is blocked inside close(). Engine close() + // owns the work: the launch+add is synchronized with close()'s cancel + // snapshot, the closed flag is re-checked INSIDE the lock, and close() + // cancels AND joins every tracked invocation. Exactly-once resume: the + // block records the outcome BEFORE resuming the continuation, and + // invokeOnCompletion resumes with a cancellation when the block never + // ran (job cancelled pre-start by close()) — otherwise the caller's + // suspension would freeze forever. + val resumed = java.util.concurrent.atomic.AtomicReference?>(null) + val launched = synchronized(activeInvocationJobs) { + check(!isClosed.get()) { "Tramai runtime is closed" } + val job = lifecycleScope.launch( + continuation.context.minusKey(kotlin.coroutines.ContinuationInterceptor) + + engineThreadMarker.asContextElement(true), + ) { + var result = runCatching { execute(operation, callArguments, conversationId) } + // Never deliver a success computed against a closed engine: the + // engine may have closed while the invocation was in flight. + // The caller sees the fixed lifecycle error instead (mirrors + // the blocking path). + if (isClosed.get() && result.isSuccess) { + result = Result.failure(IllegalStateException("Tramai runtime is closed")) + } + resumed.set(result) + continuation.resumeWith(result) + } + activeInvocationJobs += job + job + } + launched.invokeOnCompletion { cause -> + activeInvocationJobs -= launched + if (resumed.get() == null) { + // Block never ran (e.g. cancelled before the dispatcher started + // it): resume so the caller's suspension does not freeze. + continuation.resumeWith( + Result.failure( + cause as? CancellationException ?: CancellationException("Engine closed", cause), + ), + ) + } } return COROUTINE_SUSPENDED } @@ -709,74 +858,127 @@ internal class TramaiInvocationHandler( ?: (emptyList() to initialMessages) return flow { - val correlationId = java.util.UUID.randomUUID().toString() - enforceBeforeProviderResolution(operation, correlationId, securityContext) - val candidates = providerRegistry.resolveCandidates(operation.operation) - var lastFailure: Throwable? = null - var lastCircuitOpen: CircuitBreakerOpenException? = null - val attemptCounter = AttemptCounter() - - for ((routeIndex, route) in candidates.withIndex()) { - val circuitOpen = handleCircuitBreakerOpenRoute( - route = route, - nextRoute = candidates.getOrNull(routeIndex + 1), - correlationId = correlationId, - securityContext = securityContext, - ) - if (circuitOpen != null) { - lastCircuitOpen = circuitOpen - continue - } - - when ( - val result = executeStreamingRoute( - StreamingExecutionRoute( - operation = operation, - route = route, - routeIndex = routeIndex, - attempt = attemptCounter.next(), - tokenBudgetTracker = tokenBudgetTracker, - memoryMessages = effectiveMessages, - historySize = history.size, - conversationId = conversationId, - emitChunk = { emit(it) }, - ), - correlationId = correlationId, - securityContext = securityContext, - arguments = arguments, - ) - ) { - is StreamingRouteResult.Completed -> { - if (chatMemory != null && conversationId != null) { - val assistantMessage = Message( - role = MessageRole.ASSISTANT, - content = result.fullText, - ) - val turnMessages = effectiveMessages - .drop(history.size) - .filter { it.role != MessageRole.SYSTEM } - chatMemory.add(conversationId, turnMessages + assistantMessage) - } - return@flow - } - is StreamingRouteResult.StartupFailure -> { - enforceStreamingFallbackAfterFailure( - error = result.error, + check(!isClosed.get()) { "Tramai runtime is closed" } + // The provider collection runs as a child of the engine's OWN + // lifecycle job (lifecycleScope), NOT the collector's job: close() + // cancels lifecycleJob, which cancels an in-flight collection + // (including the provider stream's cleanup), and close() joins + // lifecycleJob — so the collection has terminated before close() + // returns. Chunks are bridged to the caller's emit through a + // RENDEZVOUS channel: emit must stay in the collector's coroutine + // (SafeCollector invariant), and a rendezvous keeps Flow + // backpressure semantics — a slow caller blocks the provider + // instead of letting it race ahead into unbounded buffering. + val chunks = kotlinx.coroutines.channels.Channel(kotlinx.coroutines.channels.Channel.RENDEZVOUS) + val collectFailure = java.util.concurrent.atomic.AtomicReference(null) + val collectJob = lifecycleScope.launch { + try { + val correlationId = java.util.UUID.randomUUID().toString() + enforceBeforeProviderResolution(operation, correlationId, securityContext) + val candidates = providerRegistry.resolveCandidates(operation.operation) + var lastFailure: Throwable? = null + var lastCircuitOpen: CircuitBreakerOpenException? = null + val attemptCounter = AttemptCounter() + + for ((routeIndex, route) in candidates.withIndex()) { + val circuitOpen = handleCircuitBreakerOpenRoute( route = route, nextRoute = candidates.getOrNull(routeIndex + 1), correlationId = correlationId, securityContext = securityContext, ) - lastFailure = result.error - } - is StreamingRouteResult.TerminalError -> { - emit(result.errorChunk) - return@flow + if (circuitOpen != null) { + lastCircuitOpen = circuitOpen + continue + } + + when ( + val result = executeStreamingRoute( + StreamingExecutionRoute( + operation = operation, + route = route, + routeIndex = routeIndex, + attempt = attemptCounter.next(), + tokenBudgetTracker = tokenBudgetTracker, + memoryMessages = effectiveMessages, + historySize = history.size, + conversationId = conversationId, + emitChunk = { chunks.send(it) }, + ), + correlationId = correlationId, + securityContext = securityContext, + arguments = arguments, + ) + ) { + is StreamingRouteResult.Completed -> { + if (chatMemory != null && conversationId != null) { + val assistantMessage = Message( + role = MessageRole.ASSISTANT, + content = result.fullText, + ) + val turnMessages = effectiveMessages + .drop(history.size) + .filter { it.role != MessageRole.SYSTEM } + chatMemory.add(conversationId, turnMessages + assistantMessage) + } + return@launch + } + is StreamingRouteResult.StartupFailure -> { + enforceStreamingFallbackAfterFailure( + error = result.error, + route = route, + nextRoute = candidates.getOrNull(routeIndex + 1), + correlationId = correlationId, + securityContext = securityContext, + ) + lastFailure = result.error + } + is StreamingRouteResult.TerminalError -> { + chunks.send(result.errorChunk) + return@launch + } + } } + + chunks.send(noAvailableStreamingRouteChunk(operation, lastFailure, lastCircuitOpen)) + } catch (e: kotlinx.coroutines.CancellationException) { + // The engine closed (or the collector stopped): terminate + // the collection job normally; the invokeOnCompletion + // below closes the channel with the cancellation cause. + throw e + } catch (failure: Throwable) { + // Surface the failure to the collector WITHOUT rethrowing + // it here: rethrowing would let an arbitrary (possibly + // sensitive, externally supplied) throwable reach the + // lifecycle scope's CoroutineExceptionHandler and the + // normal logger. The collector rethrows it after the + // channel drains. + collectFailure.set(failure) } } - - emit(noAvailableStreamingRouteChunk(operation, lastFailure, lastCircuitOpen)) + // Channel termination depends on JOB completion, not on the + // coroutine body having started: if close() cancels lifecycleJob + // after the flow's open check but before this launch's body runs, + // the body's finally never executes — but invokeOnCompletion still + // fires, closing the channel so the collector terminates instead + // of hanging forever on receive. + collectJob.invokeOnCompletion { cause -> chunks.close(cause) } + try { + for (chunk in chunks) { + check(!isClosed.get()) { "Tramai runtime is closed" } + emit(chunk) + } + // The collection job may have failed (e.g. provider does not + // support streaming, or a route error aborted the loop): the + // channel closes either way, so rethrow the job's failure here + // instead of silently completing the flow. + collectFailure.get()?.let { throw it } + } finally { + // If the caller stops collecting (or the engine closed), the + // engine-owned collection job must not keep running. + collectJob.cancel() + collectJob.join() + } } } diff --git a/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt b/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt index 58131cc2..357ea649 100644 --- a/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt +++ b/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt @@ -1,5 +1,6 @@ package dev.tramai.engine +import dev.tramai.core.approval.ApprovalToken import dev.tramai.core.annotations.AiService import dev.tramai.core.annotations.AiRange import dev.tramai.core.annotations.ConversationId @@ -58,13 +59,18 @@ import dev.tramai.security.audit.toCanonicalJson import dev.tramai.structured.JacksonStructuredOutputHandler import kotlinx.coroutines.async import kotlinx.coroutines.awaitAll +import kotlinx.coroutines.NonCancellable +import kotlinx.coroutines.withContext import kotlinx.coroutines.awaitCancellation +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.Job import kotlinx.coroutines.delay import kotlinx.coroutines.flow.Flow import kotlinx.coroutines.flow.flow import kotlinx.coroutines.flow.take import kotlinx.coroutines.flow.toList import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.withTimeout import org.assertj.core.api.Assertions.assertThat import org.assertj.core.api.Assertions.assertThatThrownBy import org.junit.jupiter.api.AfterAll @@ -120,6 +126,110 @@ class TramaiEngineTest { .contains("invoice-123") } + @Test + fun `old proxy fails after close before provider executes`() { + val provider = RecordingProvider { ModelResponse(content = "unused") } + val engine = TramaiEngine(provider) + val service = engine.create() + + engine.close() + + assertThatThrownBy { runBlocking { service.analyze("invoice-123") } } + .isInstanceOf(IllegalStateException::class.java) + .hasMessage("Tramai runtime is closed") + assertThat(provider.requests).isEmpty() + } + + @Test + fun `in flight suspend invocation terminates on close`() = runBlocking { + val started = CompletableDeferred() + val provider = RecordingProvider { + started.complete(Unit) + awaitCancellation() + } + val engine = TramaiEngine(provider) + val service = engine.create() + val call = async { runCatching { service.analyze("invoice-123") } } + + started.await() + engine.close() + + assertThat(call.await().exceptionOrNull()).isInstanceOf(kotlinx.coroutines.CancellationException::class.java) + } + + @Test + fun `self close from owned coroutine does not deadlock`() = runBlocking { + lateinit var engine: TramaiEngine + val provider = RecordingProvider { + engine.close() + ModelResponse(content = "unreachable") + } + engine = TramaiEngine(provider) + val service = engine.create() + + val result = withTimeout(2_000) { runCatching { service.analyze("invoice-123") } } + + assertThat(result.exceptionOrNull()).isInstanceOf(kotlinx.coroutines.CancellationException::class.java) + } + + @Test + fun `close racing a fast suspend invocation never leaves work against a closed engine`() = runBlocking { + repeat(100) { + val providerEntered = CompletableDeferred() + val providerStartedAt = java.util.concurrent.atomic.AtomicLong(-1) + val provider = RecordingProvider { + providerStartedAt.set(System.nanoTime()) + providerEntered.complete(Unit) + ModelResponse(content = "ok") + } + val engine = TramaiEngine(provider) + val service = engine.create() + val closeCompletedAt = java.util.concurrent.atomic.AtomicLong(-1) + val closeDone = CompletableDeferred() + val outcome = CompletableDeferred() + // Close immediately while the invocation is launched: either the + // launch won (tracked + cancelled/joined) or close won (launch + // rejected with the closed error). A TOCTOU miss would run the + // provider AFTER close() returned — the leak this test guards. + val closer = Thread { + engine.close() + closeCompletedAt.set(System.nanoTime()) + closeDone.complete(Unit) + } + closer.start() + try { + service.analyze("invoice-1") + outcome.complete(null) + } catch (t: Throwable) { + outcome.complete(t) + } + // Suspend-await, never Thread.join: the continuation may resume on + // a Default worker (engine-owned invocation), and a blocking join + // there would deadlock against close() joining that same job. + closeDone.await() + closer.join() + val terminal = outcome.await() + val providerStarted = providerStartedAt.get() + val closeCompleted = closeCompletedAt.get() + if (terminal == null) { + // Success is only legal if the provider ran BEFORE close completed. + assertThat(providerStarted) + .describedAs("iteration $it: provider must not run after close") + .isNotNegative() + assertThat(providerStarted).isLessThan(closeCompleted) + } else { + // Close won the race (ISE via the in-lock re-check) or the + // invocation was cancelled in-flight / pre-start (CE via the + // exactly-once resume) — never a provider result delivered + // after close. + assertThat(terminal).satisfiesAnyOf( + { assertThat(it).isInstanceOf(IllegalStateException::class.java).hasMessage("Tramai runtime is closed") }, + { assertThat(it).isInstanceOf(kotlinx.coroutines.CancellationException::class.java) }, + ) + } + } + } + @Test fun `renders classified document payloads into prompts without wrapper metadata`() { val provider = RecordingProvider { ModelResponse(content = "hardcoded response") } @@ -143,6 +253,339 @@ class TramaiEngineTest { .doesNotContain("classification=") } + @Test + fun `blocking invocation racing close never delivers a result from a closed engine`() { + val provider = RecordingProvider { ModelResponse(content = "summary") } + val engine = TramaiEngine(provider) + val service = engine.create() + + val closer = Thread { engine.close() } + closer.start() + val outcome = try { + val result = service.summarize("raw input") + "result:$result" + } catch (t: Throwable) { + "error:${t.javaClass.simpleName}:${t.message}" + } + closer.join() + + // Either the call completed before close won the race (result), or it + // was rejected / cancelled by close (fixed lifecycle error, or a + // cancellation because the engine-owned blocking call was terminated + // mid-flight — the P1-1A guarantee). It must NEVER be a summary + // delivered from an engine that was already closed. + assertThat(outcome) + .describedAs("outcome: $outcome") + .satisfiesAnyOf( + { assertThat(it).startsWith("result:") }, + { assertThat(it).isEqualTo("error:IllegalStateException:Tramai runtime is closed") }, + { assertThat(it).startsWith("error:JobCancellationException") }, + { assertThat(it).startsWith("error:CancellationException") }, + ) + } + + @Test + fun `registerService and resumeApproval fail fast on a closed engine`() { + val engine = TramaiEngine(RecordingProvider { ModelResponse(content = "unused") }) + + engine.close() + + assertThatThrownBy { engine.registerService(SuspendAnalyzer::class) } + .isInstanceOf(IllegalStateException::class.java) + .hasMessage("Tramai runtime is closed") + // Guard fires before any store access, so placeholder command values suffice. + assertThatThrownBy { + runBlocking { + engine.resumeApproval( + ResumeApprovalCommand( + approvalId = "any", + approvalExpectedVersion = 0L, + continuationExpectedVersion = 0L, + presentedToken = ApprovalToken.parsePresented("token"), + resumedBy = "test", + ), + ) + } + } + .isInstanceOf(IllegalStateException::class.java) + .hasMessage("Tramai runtime is closed") + } + + @Test + fun `streaming flow collected after close fails before provider executes`() { + val provider = NamedStreamingProvider("p") { + flow { emit(StreamChunk.Token("unused")) } + } + val registry = ProviderRegistry.builder() + .provider("p", provider) + .model("claude-sonnet-4-20250514", "p") + .build() + val engine = TramaiEngine(providerRegistry = registry) + val service = engine.create() + + val flow = service.stream("invoice-123") + engine.close() + + assertThatThrownBy { runBlocking { flow.toList() } } + .isInstanceOf(IllegalStateException::class.java) + .hasMessage("Tramai runtime is closed") + assertThat(provider.streamRequests).isEmpty() + } + + @Test + fun `mid-collection close terminates an in-flight stream`() = runBlocking { + val gate = CompletableDeferred() + val provider = NamedStreamingProvider("p") { + flow { + emit(StreamChunk.Token("first")) + gate.await() + emit(StreamChunk.Token("second")) + } + } + val registry = ProviderRegistry.builder() + .provider("p", provider) + .model("claude-sonnet-4-20250514", "p") + .build() + val engine = TramaiEngine(providerRegistry = registry) + val service = engine.create() + + val chunks = mutableListOf() + val firstChunkDelivered = CompletableDeferred() + val collection = async { + try { + service.stream("invoice-123").collect { + chunks += it + firstChunkDelivered.complete(Unit) + } + null + } catch (t: Throwable) { + t + } + } + // Wait until the first chunk is delivered (provider is streaming). + firstChunkDelivered.await() + engine.close() + gate.complete(Unit) + val terminal = collection.await() + + // The stream must not deliver the second chunk after close: either the + // collection failed with the fixed lifecycle error, or it stopped + // without the post-close chunk. The first chunk was provably delivered + // (firstChunkDelivered completes only after a chunk arrives). + assertThat(chunks.map { it }).contains(StreamChunk.Token("first")) + assertThat(chunks.map { it }).doesNotContain(StreamChunk.Token("second")) + if (terminal != null) { + // Either the engine-owned collection was cancelled (the stronger + // guarantee: close() cancels lifecycleJob, which terminates the + // provider stream's collection) or the per-chunk closed gate fired. + assertThat(terminal).satisfiesAnyOf( + { assertThat(it).isInstanceOf(IllegalStateException::class.java).hasMessage("Tramai runtime is closed") }, + { assertThat(it).isInstanceOf(kotlinx.coroutines.CancellationException::class.java) }, + ) + } + } + + @Test + fun `blocking invocation in long suspension is cancelled and joined by close`() = runBlocking { + val providerEntered = CompletableDeferred() + val providerReleased = CompletableDeferred() + val providerCleanedUp = CompletableDeferred() + val provider = RecordingProvider { + providerEntered.complete(Unit) + try { + awaitCancellation() + } finally { + withContext(kotlinx.coroutines.NonCancellable) { + providerCleanedUp.complete(Unit) + } + } + } + val engine = TramaiEngine(provider) + val service = engine.create() + + val caller = Thread { + runCatching { service.summarize("raw") } + } + caller.start() + providerEntered.await() + + engine.close() + + // close() must terminate the in-flight blocking call (it is a child of + // the engine-owned lifecycle job) and wait for its cleanup. + providerCleanedUp.await() // completed within close() + caller.join(5_000) + assertThat(caller.isAlive).isFalse() + providerReleased.complete(Unit) + } + + @Test + fun `streaming collection suspended indefinitely is cancelled and cleaned up by close`() = runBlocking { + val providerCleanedUp = CompletableDeferred() + val firstChunkDelivered = CompletableDeferred() + val provider = NamedStreamingProvider("p") { + flow { + try { + emit(StreamChunk.Token("first")) + firstChunkDelivered.complete(Unit) + awaitCancellation() + } finally { + withContext(kotlinx.coroutines.NonCancellable) { + providerCleanedUp.complete(Unit) + } + } + } + } + val registry = ProviderRegistry.builder() + .provider("p", provider) + .model("claude-sonnet-4-20250514", "p") + .build() + val engine = TramaiEngine(providerRegistry = registry) + val service = engine.create() + + val collectionDone = CompletableDeferred() + val collection = async { + runCatching { service.stream("invoice-123").collect {} } + collectionDone.complete(Unit) + } + // Deterministic admission gate: wait until the provider has emitted its + // first chunk and is suspended indefinitely, instead of a sleep. + firstChunkDelivered.await() + engine.close() + // close() cancels lifecycleJob -> the engine-owned collection job is + // cancelled -> the provider's finally runs -> the collection terminates. + collectionDone.await() + // Completed within close(): the engine-owned collection job was + // cancelled and the provider's NonCancellable cleanup ran. + providerCleanedUp.await() + } + + @Test + fun `close does not deadlock when caller supplied its own job and scope`() { + runBlocking { + val provider = RecordingProvider { ModelResponse(content = "ok") } + val engine = TramaiEngine( + provider = provider, + job = requireNotNull(coroutineContext[Job]), + scope = this, + ) + // No in-flight work: close() must cancel/join the engine-owned + // lifecycle job, never the caller's job — returning promptly. + engine.close() + } + } + + @Test + fun `self close from streaming owned coroutine does not deadlock`() = runBlocking { + lateinit var engine: TramaiEngine + val provider = NamedStreamingProvider("p") { + flow { + engine.close() + emit(StreamChunk.Token("unreachable")) + } + } + val registry = ProviderRegistry.builder() + .provider("p", provider) + .model("claude-sonnet-4-20250514", "p") + .build() + engine = TramaiEngine(providerRegistry = registry) + val service = engine.create() + + // The provider runs inside the engine-owned streaming collection + // coroutine (lifecycleScope child) and re-enters close(). The engine + // thread marker is carried by the SCOPE, so close() skips the join and + // returns; the collection is then cancelled and the flow terminates + // instead of self-joining forever. + withTimeout(2_000) { + runCatching { + service.stream("invoice").collect {} + } + } + } + + @Test + fun `stream start racing close never hangs the collector`() = runBlocking { + // Force the admission race deterministically: the flow's open check + // passes, then close() cancels lifecycleJob BEFORE the collection + // coroutine body begins executing. Channel termination must come from + // job completion (invokeOnCompletion), not from the body's finally. + repeat(200) { iteration -> + val provider = NamedStreamingProvider("p") { + flow { + emit(StreamChunk.Token("first")) + awaitCancellation() + } + } + val registry = ProviderRegistry.builder() + .provider("p", provider) + .model("claude-sonnet-4-20250514", "p") + .build() + val engine = TramaiEngine(providerRegistry = registry) + val service = engine.create() + + val collectionDone = CompletableDeferred() + val collection = async { + runCatching { service.stream("invoice-$iteration").collect {} } + collectionDone.complete(Unit) + } + // Close immediately: either the collection won (tracked + joined) + // or close() won (launch rejected / cancelled pre-start). Either + // way the collector must terminate — never hang. + engine.close() + withTimeout(5_000) { collectionDone.await() } + collection.await() + } + } + + @Test + fun `streaming bridge preserves backpressure when the collector is slow`() = runBlocking { + val attempted = java.util.concurrent.atomic.AtomicInteger(0) + val delivered = java.util.concurrent.atomic.AtomicInteger(0) + val collectorGate = CompletableDeferred() + val provider = NamedStreamingProvider("p") { + flow { + var i = 0 + while (true) { + attempted.incrementAndGet() + emit(StreamChunk.Token("chunk-${i++}")) + } + } + } + val registry = ProviderRegistry.builder() + .provider("p", provider) + .model("claude-sonnet-4-20250514", "p") + .build() + val engine = TramaiEngine(providerRegistry = registry) + val service = engine.create() + + val firstChunk = CompletableDeferred() + val collection = async { + service.stream("invoice-123").collect { + delivered.incrementAndGet() + if (!firstChunk.isCompleted) firstChunk.complete(Unit) + collectorGate.await() + } + } + firstChunk.await() + // The collector is deliberately blocked after chunk 1. Prove + // backpressure deterministically (no sleeps): the provider must + // ATTEMPT chunk 2 (it got past chunk 1) and then STALL at the send — + // a rendezvous bridge holds the second chunk until the collector + // consumes it; UNLIMITED would let the provider race ahead past 2. + withTimeout(2_000) { + while (attempted.get() < 2) { + delay(10) + } + } + assertThat(attempted.get()).isEqualTo(2) + assertThat(delivered.get()).isEqualTo(1) + collectorGate.complete(Unit) + engine.close() + // Collection is terminated by close() (rendezvous bridge cancelled + // through lifecycleJob) — the CE is the expected termination signal. + runCatching { collection.await() } + } + @Test fun `supports blocking interfaces`() { val provider = RecordingProvider { ModelResponse(content = "summary") } @@ -3570,7 +4013,7 @@ private class NamedStreamingProvider( } private class RecordingObserver : OperationObserver { - val records = mutableListOf() + val records = java.util.Collections.synchronizedList(mutableListOf()) override fun onCallStarted(context: OperationCallContext): OperationObservation { val record = Record(context = context) @@ -3613,7 +4056,7 @@ private class RecordingObserver : OperationObserver { var providerFailure: Throwable? = null, var parseSuccess: Boolean? = null, var completionCount: Int = 0, - val engineEvents: MutableList = mutableListOf(), + val engineEvents: MutableList = java.util.Collections.synchronizedList(mutableListOf()), var cancelled: Boolean = false, ) } @@ -3624,7 +4067,7 @@ private data class EngineEventRecord( ) private class RecordingEngineEventObserver : EngineEventObserver { - val events = mutableListOf() + val events = java.util.Collections.synchronizedList(mutableListOf()) override fun onEngineEvent( name: String, diff --git a/tramai-engine/src/test/kotlin/dev/tramai/engine/TrustedReplayEnvelopeRegistryTest.kt b/tramai-engine/src/test/kotlin/dev/tramai/engine/TrustedReplayEnvelopeRegistryTest.kt index aa421d90..33d03b54 100644 --- a/tramai-engine/src/test/kotlin/dev/tramai/engine/TrustedReplayEnvelopeRegistryTest.kt +++ b/tramai-engine/src/test/kotlin/dev/tramai/engine/TrustedReplayEnvelopeRegistryTest.kt @@ -402,6 +402,8 @@ class TrustedReplayEnvelopeRegistryTest { approvalGateCoordinator = null, approvalLifecycleAuditEmitter = NoOpApprovalLifecycleAuditEmitter, resumeOperationRegistry = registry, + lifecycleJob = kotlinx.coroutines.SupervisorJob(), + lifecycleScope = kotlinx.coroutines.CoroutineScope(kotlinx.coroutines.Dispatchers.Default), clock = Clock.systemUTC(), ) } diff --git a/tramai-orchestration/src/test/kotlin/dev/tramai/orchestration/TramaiWorkerTest.kt b/tramai-orchestration/src/test/kotlin/dev/tramai/orchestration/TramaiWorkerTest.kt index aa8065bd..ac21e41d 100644 --- a/tramai-orchestration/src/test/kotlin/dev/tramai/orchestration/TramaiWorkerTest.kt +++ b/tramai-orchestration/src/test/kotlin/dev/tramai/orchestration/TramaiWorkerTest.kt @@ -62,6 +62,44 @@ class TramaiWorkerTest { } } + @Test + fun `close deregisters the JVM shutdown hook and retains no reference`() = runBlocking { + val checkpointStore = InMemoryWorkflowCheckpointStore() + val leaseStore = InMemoryWorkflowLeaseStore() + val workflow = workerWorkflow("hook-leak") { + localStep(name = "noop", transform = { state, _ -> state }) + } + val worker = worker("hook-worker", leaseStore, checkpointStore, workflow) + + worker.start() + // The hook field is populated by start(): a JVM shutdown hook is registered. + val hookField = TramaiWorker::class.java.getDeclaredField("shutdownHook") + hookField.isAccessible = true + val hookAfterStart = hookField.get(worker) as? Thread + assertThat(hookAfterStart).isNotNull() + + try { + worker.close() + // close() -> shutdown() must drop the worker's reference... + val hookAfterClose = hookField.get(worker) as? Thread + assertThat(hookAfterClose).isNull() + // ...AND deregister the hook from the JVM. removeShutdownHook + // returns false when the hook was already deregistered; true would + // mean close() left a live hook registered (a real JVM-level leak + // that the field-null check alone cannot detect). If the + // implementation is broken, this call also removes the leaked hook + // from the test JVM before the assertion fails. + assertThat(Runtime.getRuntime().removeShutdownHook(requireNotNull(hookAfterStart))).isFalse() + } finally { + // Never leave a hook registered in the test JVM, even if an + // assertion above failed before the deregistration check ran. + val residual = hookField.get(worker) as? Thread + if (residual != null) { + worker.close() + } + } + } + @Test fun `worker crash leaves non replayable step in unknown state and takeover fails`() = runBlocking { val checkpointStore = InMemoryWorkflowCheckpointStore() diff --git a/tramai-sovereign/api/tramai-sovereign.api b/tramai-sovereign/api/tramai-sovereign.api index 6cb3a4d9..08d2a741 100644 --- a/tramai-sovereign/api/tramai-sovereign.api +++ b/tramai-sovereign/api/tramai-sovereign.api @@ -35,10 +35,11 @@ public final class dev/tramai/sovereign/SovereignProfileConfiguration { public final class dev/tramai/sovereign/SovereignProfileConfiguration$Companion { } -public final class dev/tramai/sovereign/SovereignTramai { +public final class dev/tramai/sovereign/SovereignTramai : java/lang/AutoCloseable { public static final field Companion Ldev/tramai/sovereign/SovereignTramai$Companion; public synthetic fun (Ldev/tramai/standalone/Tramai;Ljava/util/List;Ldev/tramai/sovereign/SovereignProfileConfiguration;Ldev/tramai/core/model/ModelArtifactVerificationSettings;Lkotlin/jvm/internal/DefaultConstructorMarker;)V public static final fun builder ()Ldev/tramai/sovereign/SovereignTramai$Builder; + public fun close ()V public final fun create (Lkotlin/reflect/KClass;)Ljava/lang/Object; public final fun evidencePack (Ldev/tramai/sovereign/evidence/ZeroEgressEvidenceV1;Ldev/tramai/sovereign/evidence/AuditChainEvidenceV1;Ldev/tramai/sovereign/evidence/SupplyChainEvidenceV1;Ldev/tramai/sovereign/evidence/ReleaseBundleEvidenceV1;Ldev/tramai/sovereign/evidence/AttestationEvidenceV1;)Ldev/tramai/sovereign/evidence/SovereignEvidencePackV1; public static synthetic fun evidencePack$default (Ldev/tramai/sovereign/SovereignTramai;Ldev/tramai/sovereign/evidence/ZeroEgressEvidenceV1;Ldev/tramai/sovereign/evidence/AuditChainEvidenceV1;Ldev/tramai/sovereign/evidence/SupplyChainEvidenceV1;Ldev/tramai/sovereign/evidence/ReleaseBundleEvidenceV1;Ldev/tramai/sovereign/evidence/AttestationEvidenceV1;ILjava/lang/Object;)Ldev/tramai/sovereign/evidence/SovereignEvidencePackV1; diff --git a/tramai-sovereign/src/main/kotlin/dev/tramai/sovereign/SovereignTramai.kt b/tramai-sovereign/src/main/kotlin/dev/tramai/sovereign/SovereignTramai.kt index 39353aac..8c2d72f9 100644 --- a/tramai-sovereign/src/main/kotlin/dev/tramai/sovereign/SovereignTramai.kt +++ b/tramai-sovereign/src/main/kotlin/dev/tramai/sovereign/SovereignTramai.kt @@ -79,19 +79,24 @@ class SovereignTramai private constructor( verificationReceipts: List, private val profile: SovereignProfileConfiguration, private val verificationSettings: ModelArtifactVerificationSettings, -) { +) : AutoCloseable { private val verificationReceipts: List = Collections.unmodifiableList(ArrayList(verificationReceipts)) + /** One cached wrapper around the delegate's single owned runtime. */ + private val ownedRuntime: SovereignTramaiRuntime by lazy { SovereignTramaiRuntime(delegate.runtime()) } /** * Creates a service proxy for the given service type. */ fun create(serviceType: KClass): T = delegate.create(serviceType) /** - * Creates a [SovereignTramaiRuntime] that owns exactly one engine and exposes - * both service creation and approval-resume operations. + * Returns the single [SovereignTramaiRuntime] wrapping the one owned + * engine; repeated calls return the same instance. */ - fun runtime(): SovereignTramaiRuntime = SovereignTramaiRuntime(delegate.runtime()) + fun runtime(): SovereignTramaiRuntime = ownedRuntime + + /** Closes the owned standalone runtime and its engine. */ + override fun close() = delegate.close() /** * Returns immutable verification receipts from build-time artifact verification. diff --git a/tramai-sovereign/src/test/kotlin/dev/tramai/sovereign/SovereignTramaiTest.kt b/tramai-sovereign/src/test/kotlin/dev/tramai/sovereign/SovereignTramaiTest.kt index 575af3f8..25aad92f 100644 --- a/tramai-sovereign/src/test/kotlin/dev/tramai/sovereign/SovereignTramaiTest.kt +++ b/tramai-sovereign/src/test/kotlin/dev/tramai/sovereign/SovereignTramaiTest.kt @@ -97,6 +97,39 @@ class SovereignTramaiTest { .provider(FakeProvider(), name = "local-provider", default = true) .model("test-model", "local-provider") + @Test + fun `sovereign create and runtime share one owned engine`() { + val tramai = validBuilder().build() + + tramai.create() + val runtime = tramai.runtime() + runtime.create() + + tramai.close() + assertThatThrownBy { tramai.create() } + .isInstanceOf(IllegalStateException::class.java) + .hasMessage("Tramai runtime is closed") + } + + @Test + fun `sovereign runtime returns the same wrapper instance`() { + val tramai = validBuilder().build() + + assertThat(tramai.runtime()).isSameAs(tramai.runtime()) + } + + @Test + fun `sovereign close propagates to the delegate runtime`() { + val tramai = validBuilder().build() + val runtime = tramai.runtime() + + tramai.close() + + assertThatThrownBy { runtime.create() } + .isInstanceOf(IllegalStateException::class.java) + .hasMessage("Tramai runtime is closed") + } + // ========================================================================= // Composition Validation // ========================================================================= diff --git a/tramai-spring/src/main/kotlin/dev/tramai/spring/TramaiAutoConfiguration.kt b/tramai-spring/src/main/kotlin/dev/tramai/spring/TramaiAutoConfiguration.kt index 2da56613..11b64b06 100644 --- a/tramai-spring/src/main/kotlin/dev/tramai/spring/TramaiAutoConfiguration.kt +++ b/tramai-spring/src/main/kotlin/dev/tramai/spring/TramaiAutoConfiguration.kt @@ -78,7 +78,7 @@ private data class TramaiBeanDependencies( @EnableConfigurationProperties(TramaiProperties::class) class TramaiAutoConfiguration { - @Bean + @Bean(destroyMethod = "close") @ConditionalOnMissingBean fun tramai( properties: TramaiProperties, diff --git a/tramai-spring/src/test/kotlin/dev/tramai/spring/TramaiAutoConfigurationTest.kt b/tramai-spring/src/test/kotlin/dev/tramai/spring/TramaiAutoConfigurationTest.kt index 2aeaaf05..4b6e877d 100644 --- a/tramai-spring/src/test/kotlin/dev/tramai/spring/TramaiAutoConfigurationTest.kt +++ b/tramai-spring/src/test/kotlin/dev/tramai/spring/TramaiAutoConfigurationTest.kt @@ -42,6 +42,43 @@ import kotlin.test.Test class TramaiAutoConfigurationTest { + @Test + fun `spring context destruction closes the runtime`() { + var tramai: Tramai? = null + ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(TramaiAutoConfiguration::class.java)) + .withUserConfiguration(TestApplication::class.java, ProviderConfiguration::class.java) + .withPropertyValues("tramai.default-provider=stub") + .run { context -> tramai = context.getBean(Tramai::class.java) } + + assertThatThrownBy { tramai!!.runtime() } + .isInstanceOf(IllegalStateException::class.java) + .hasMessage("Tramai runtime is closed") + } + + @Test + fun `multiple ai service beans share one runtime and all fail after context close`() { + lateinit var analyzer: TestInvoiceAnalyzer + lateinit var cached: CachedInvoiceAnalyzer + ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(TramaiAutoConfiguration::class.java)) + .withUserConfiguration(TestApplication::class.java, ProviderConfiguration::class.java) + .withPropertyValues("tramai.default-provider=stub") + .run { context -> + analyzer = context.getBean(TestInvoiceAnalyzer::class.java) + cached = context.getBean(CachedInvoiceAnalyzer::class.java) + } + // Context destroyed -> the shared Tramai bean runtime closed. If each + // factory bean had created a hidden independent engine, these proxies + // would still be usable; instead they must fail before provider work. + assertThatThrownBy { runBlocking { analyzer.analyze("invoice-1") } } + .isInstanceOf(IllegalStateException::class.java) + .hasMessage("Tramai runtime is closed") + assertThatThrownBy { runBlocking { cached.analyze("invoice-1") } } + .isInstanceOf(IllegalStateException::class.java) + .hasMessage("Tramai runtime is closed") + } + @Test fun `registers ai service beans and injects a custom provider bean`() { val contextRunner = ApplicationContextRunner() diff --git a/tramai-standalone/api/tramai-standalone.api b/tramai-standalone/api/tramai-standalone.api index 8c600c47..9ed3dcbb 100644 --- a/tramai-standalone/api/tramai-standalone.api +++ b/tramai-standalone/api/tramai-standalone.api @@ -1,7 +1,8 @@ -public final class dev/tramai/standalone/Tramai { +public final class dev/tramai/standalone/Tramai : java/lang/AutoCloseable { public static final field Companion Ldev/tramai/standalone/Tramai$Companion; public synthetic fun (Ldev/tramai/core/provider/ProviderRegistry;Ldev/tramai/engine/ToolRegistry;Ldev/tramai/core/observation/OperationObserver;Ldev/tramai/core/observation/OperationInterceptor;Ldev/tramai/engine/OperationResponseCache;Ldev/tramai/engine/CircuitBreakerSettings;Ldev/tramai/engine/RetryPolicySettings;Ldev/tramai/engine/TokenBudgetSettings;Ldev/tramai/core/security/DlpInterceptor;Ldev/tramai/core/security/DlpRedactionAuditEmitter;Ldev/tramai/engine/ToolResultFilteringSettings;Ldev/tramai/engine/EngineEventObserver;Ldev/tramai/core/observation/ToolFailureDiagnosticObserver;Ldev/tramai/core/security/PromptSanitizer;Ldev/tramai/core/memory/ChatMemory;Ldev/tramai/core/policy/PolicyDecisionAuditEmitter;Ldev/tramai/core/policy/PolicyEngine;Ldev/tramai/core/model/ModelRegistry;Ldev/tramai/core/model/ModelRegistrySettings;Ldev/tramai/engine/SuspendedInvocationStore;Ldev/tramai/core/approval/ApprovalContinuationStore;Ldev/tramai/core/approval/ToolArgumentsDigester;Ldev/tramai/core/approval/ApprovalGateCoordinator;Ldev/tramai/core/approval/ApprovalLifecycleAuditEmitter;Ljava/time/Clock;Lkotlin/jvm/internal/DefaultConstructorMarker;)V public static final fun builder ()Ldev/tramai/standalone/Tramai$Builder; + public fun close ()V public final fun create (Lkotlin/reflect/KClass;)Ljava/lang/Object; public final fun runtime ()Ldev/tramai/standalone/TramaiRuntime; } diff --git a/tramai-standalone/src/main/kotlin/dev/tramai/standalone/Tramai.kt b/tramai-standalone/src/main/kotlin/dev/tramai/standalone/Tramai.kt index 2b77b4c3..b389c71f 100644 --- a/tramai-standalone/src/main/kotlin/dev/tramai/standalone/Tramai.kt +++ b/tramai-standalone/src/main/kotlin/dev/tramai/standalone/Tramai.kt @@ -59,6 +59,10 @@ import kotlin.reflect.full.createType /** * Minimal composition module that wires core, engine, and structured output support. + * + * This instance owns the runtime and engine it creates. Closing it closes that runtime; + * providers, stores, clients, executors, and observers supplied to the builder remain + * caller-owned unless their own API explicitly transfers ownership. */ class Tramai private constructor( private val providerRegistry: ProviderRegistry, @@ -87,7 +91,7 @@ class Tramai private constructor( private val approvalGateCoordinator: ApprovalGateCoordinator? = null, private val approvalLifecycleAuditEmitter: ApprovalLifecycleAuditEmitter = NoOpApprovalLifecycleAuditEmitter, private val clock: Clock = Clock.systemUTC(), -) { +) : AutoCloseable { /** * Structured-output diagnostic observer, delivered additively (class-body * member, not a constructor parameter) so the published JVM constructor @@ -96,16 +100,35 @@ class Tramai private constructor( */ internal var structuredOutputFailureDiagnosticObserver: StructuredOutputFailureDiagnosticObserver = NoOpStructuredOutputFailureDiagnosticObserver + + // Class-body lifecycle state preserves published constructor descriptors. + private val lifecycleLock = Any() + private var ownedRuntime: TramaiRuntime? = null + private var closed = false + + private fun activeRuntime(): TramaiRuntime = synchronized(lifecycleLock) { + check(!closed) { "Tramai runtime is closed" } + ownedRuntime ?: TramaiRuntime(newEngine()).also { ownedRuntime = it } + } + /** * Creates a service proxy using the built-in Jackson structured output handler. */ - fun create(serviceType: KClass): T = newEngine().create(serviceType) + fun create(serviceType: KClass): T = activeRuntime().create(serviceType) /** * Creates a [TramaiRuntime] that owns exactly one engine and exposes * both service creation and approval-resume operations. */ - fun runtime(): TramaiRuntime = TramaiRuntime(newEngine()) + fun runtime(): TramaiRuntime = activeRuntime() + + override fun close() = synchronized(lifecycleLock) { + if (!closed) { + closed = true + ownedRuntime?.close() + ownedRuntime = null + } + } /** * Returns a configured [TramaiEngine] from the current builder state. diff --git a/tramai-standalone/src/test/kotlin/dev/tramai/standalone/TramaiTest.kt b/tramai-standalone/src/test/kotlin/dev/tramai/standalone/TramaiTest.kt index dcd615da..516a0f42 100644 --- a/tramai-standalone/src/test/kotlin/dev/tramai/standalone/TramaiTest.kt +++ b/tramai-standalone/src/test/kotlin/dev/tramai/standalone/TramaiTest.kt @@ -37,6 +37,125 @@ import kotlin.test.Test class TramaiTest { + @Test + fun `two create calls share one runtime lifecycle`() { + val tramai = configuredTramai() + + tramai.create() + val runtime = tramai.runtime() + tramai.create() + + assertThat(runtime).isSameAs(tramai.runtime()) + } + + @Test + fun `runtime returns the same runtime as create`() { + val tramai = configuredTramai() + val runtime = tramai.runtime() + + runtime.create() + + assertThat(runtime).isSameAs(tramai.runtime()) + } + + @Test + fun `close before first use rejects runtime creation`() { + val tramai = configuredTramai() + + tramai.close() + + assertThatThrownBy { tramai.runtime() } + .isInstanceOf(IllegalStateException::class.java) + .hasMessage("Tramai runtime is closed") + } + + @Test + fun `repeated close is harmless`() { + val tramai = configuredTramai() + tramai.runtime() + + tramai.close() + tramai.close() + + assertThatThrownBy { tramai.create() } + .isInstanceOf(IllegalStateException::class.java) + .hasMessage("Tramai runtime is closed") + } + + @Test + fun `externally supplied provider is not closed`() { + val provider = CloseTrackingProvider() + val tramai = Tramai { + provider(provider, default = true) + model("claude-sonnet-4-20250514", "external") + } + + // Force engine creation so the test proves the rule against a REAL + // runtime: closing an actual engine must leave an externally supplied + // provider alone (a bare close() without runtime() never created an + // engine, making the old assertion vacuous). + tramai.runtime() + tramai.close() + + assertThat(provider.closed).isFalse() + } + + @Test + fun `concurrent create calls create only one engine`() { + val tramai = configuredTramai() + val runtimes = java.util.Collections.synchronizedList(mutableListOf()) + val threads = (1..8).map { + Thread { + repeat(50) { + runtimes.add(tramai.runtime()) + } + }.also { it.start() } + } + threads.forEach { it.join() } + + assertThat(runtimes.map { it }.distinct()).hasSize(1) + } + + @Test + fun `create racing with close cannot resurrect the runtime`() { + repeat(50) { + val tramai = configuredTramai() + val outcomes = java.util.Collections.synchronizedList(mutableListOf()) + val creator = Thread { + try { + tramai.runtime() + outcomes.add(null) + } catch (e: Throwable) { + outcomes.add(e) + } + } + val closer = Thread { tramai.close() } + creator.start() + closer.start() + creator.join() + closer.join() + // Post-close, any further access must reject; the close may win the + // race before the creator obtains the runtime, in which case the + // creator sees the fixed rejection. + val terminal = try { + tramai.runtime() + null + } catch (e: Throwable) { + e + } + assertThat(terminal) + .describedAs("iteration $it: runtime after close must reject") + .isInstanceOf(IllegalStateException::class.java) + .hasMessage("Tramai runtime is closed") + outcomes.forEach { e -> + assertThat(e).satisfiesAnyOf( + { assertThat(e).isNull() }, + { assertThat(e).isInstanceOf(IllegalStateException::class.java).hasMessage("Tramai runtime is closed") }, + ) + } + } + } + @Test fun `structured failure diagnostic observer is frozen at build time`() { val provider = RecordingProvider("anthropic") { ModelResponse(content = "not json $SO_FIXTURE") } @@ -573,6 +692,11 @@ class TramaiTest { } } +private fun configuredTramai(): Tramai = Tramai { + provider(RecordingProvider("anthropic") { ModelResponse(content = "hello") }, default = true) + model("claude-sonnet-4-20250514", "anthropic") +} + private class RecordingStructuredDiagnostics : StructuredOutputFailureDiagnosticObserver { val events = mutableListOf() override suspend fun onFailure(event: StructuredOutputFailureDiagnosticEvent) { @@ -635,6 +759,13 @@ private class RecordingProvider( override fun providerId(): String = id } +private class CloseTrackingProvider : ModelProvider, AutoCloseable { + var closed = false + override suspend fun complete(request: ModelRequest): ModelResponse = ModelResponse(content = "unused") + override fun providerId(): String = "external" + override fun close() { closed = true } +} + private class LookupInput(val query: String) private class LookupResult(val resolved: String)