From 6cc9fdacb280bfeb64bf02e90a66424eda256fd0 Mon Sep 17 00:00:00 2001 From: Giona Granchelli Date: Mon, 10 Aug 2026 15:17:19 +0200 Subject: [PATCH 01/12] fix(0.6.0): establish explicit runtime lifecycle ownership Epic 1.3 (Runtime Lifecycle Ownership): every runtime-created engine has exactly one reachable lifecycle owner; closing it deterministically prevents further work and terminates TramAI-owned work. - Tramai now owns ONE lazily-created TramaiRuntime (one engine) shared by all create()/runtime() calls; lifecycle state lives in class-body fields so the published JVM constructor descriptor is unchanged. Tramai is now AutoCloseable; close() is idempotent and synchronized; after close, create()/runtime() fail fast with a fixed IllegalStateException. - Engine proxies fail after close BEFORE provider execution (closed flag checked at the invocation handler seam). - TramaiEngine.close() cancels once and joins (except from its own coroutines, avoiding self-close deadlock), and explicitly cancels tracked suspend-invocation jobs: suspend bridges launch as children of the CALLER job (preserving parent-cancellation propagation) while the engine tracks them so close() owns in-flight work. - SovereignTramai propagates the same ownership: create()/runtime() share the delegate's owned runtime; SovereignTramai is AutoCloseable closing the delegate. - Spring: the Tramai bean uses destroyMethod = close so context destruction closes the shared runtime; multiple @AiService beans share one engine. - Resource ownership rule documented: TramAI closes only resources it creates; externally supplied providers/stores/clients/observers remain caller-owned. - Tests: shared lifecycle, single engine under concurrency, no resurrection after close, idempotent close, proxy-after-close fails before provider, in-flight suspend terminates on close, self-close no deadlock, Spring destruction + shared-engine, sovereign equivalence, external deps not closed. api dumps updated additively (AutoCloseable only). --- docs/modules/tramai-engine.md | 2 +- .../kotlin/dev/tramai/engine/TramaiEngine.kt | 44 +++++- .../dev/tramai/engine/TramaiEngineTest.kt | 48 +++++++ tramai-sovereign/api/tramai-sovereign.api | 3 +- .../dev/tramai/sovereign/SovereignTramai.kt | 5 +- .../tramai/sovereign/SovereignTramaiTest.kt | 26 ++++ .../tramai/spring/TramaiAutoConfiguration.kt | 2 +- .../spring/TramaiAutoConfigurationTest.kt | 37 +++++ tramai-standalone/api/tramai-standalone.api | 3 +- .../kotlin/dev/tramai/standalone/Tramai.kt | 29 +++- .../dev/tramai/standalone/TramaiTest.kt | 126 ++++++++++++++++++ 11 files changed, 314 insertions(+), 11 deletions(-) diff --git a/docs/modules/tramai-engine.md b/docs/modules/tramai-engine.md index 6077fc0c..130688e8 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 the engine's coroutine job hierarchy and waits for externally initiated shutdown | | `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..2e594fb9 100644 --- a/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt +++ b/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt @@ -92,6 +92,7 @@ import kotlinx.coroutines.CancellationException import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.Job +import kotlinx.coroutines.asContextElement import kotlinx.coroutines.SupervisorJob import kotlinx.coroutines.TimeoutCancellationException import kotlinx.coroutines.currentCoroutineContext @@ -171,6 +172,14 @@ 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() + /** + * 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 +393,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 +415,9 @@ class TramaiEngine( chatMemory = chatMemory, conversationIdProvider = conversationIdProvider, scope = scope, + isClosed = closed, + engineThreadMarker = engineThreadMarker, + activeInvocationJobs = activeInvocationJobs, serviceDefinition = definition, policyEngine = resolvedPolicyEngine, migrationWarningGuard = migrationWarningGuard, @@ -468,6 +481,9 @@ class TramaiEngine( chatMemory = chatMemory, conversationIdProvider = conversationIdProvider, scope = scope, + isClosed = closed, + engineThreadMarker = engineThreadMarker, + activeInvocationJobs = activeInvocationJobs, serviceDefinition = definition, policyEngine = resolvedPolicyEngine, migrationWarningGuard = migrationWarningGuard, @@ -531,10 +547,19 @@ class TramaiEngine( resumeApproval(command) as R /** - * Cancels the engine-owned coroutine job hierarchy. + * Cancels and, except from one of its own coroutines, waits for the engine-owned + * coroutine hierarchy. Dependencies supplied by callers are not closed. */ override fun close() { - job.cancel() + if (closed.compareAndSet(false, true)) { + job.cancel() + // Suspend invocations are children of their CALLER's job, not the + // engine scope job; cancel them explicitly so close() owns them. + activeInvocationJobs.forEach { it.cancel() } + if (engineThreadMarker.get() != true) { + runBlocking { job.join() } + } + } } } @@ -576,6 +601,9 @@ internal class TramaiInvocationHandler( private val chatMemory: ChatMemory?, private val conversationIdProvider: ConversationIdProvider, private val scope: 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 +667,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}") @@ -663,11 +693,19 @@ internal class TramaiInvocationHandler( ?: throw ConfigurationException("Suspend invocation for ${operation.method.name} is missing its continuation") val callArguments = args.dropLast(1) - scope.launch(continuation.context) { + // Launch as a child of the CALLER's job (continuation.context) so parent + // cancellation propagates synchronously into the in-flight invocation + // (validated by the ToolSafeFailureContract / StructuredOutputFailureBoundary + // parent-cancellation tests), while tracking the launched job so engine + // close() also owns it (cancel below). Engine close cancels the engine + // scope job AND every tracked invocation job. + val launched = scope.launch(continuation.context + engineThreadMarker.asContextElement(true)) { runCatching { execute(operation, callArguments, conversationId) } .onSuccess { continuation.resumeWith(Result.success(it)) } .onFailure { continuation.resumeWith(Result.failure(it)) } } + activeInvocationJobs += launched + launched.invokeOnCompletion { activeInvocationJobs -= launched } return COROUTINE_SUSPENDED } 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..2ee164e1 100644 --- a/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt +++ b/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt @@ -59,12 +59,14 @@ import dev.tramai.structured.JacksonStructuredOutputHandler import kotlinx.coroutines.async import kotlinx.coroutines.awaitAll import kotlinx.coroutines.awaitCancellation +import kotlinx.coroutines.CompletableDeferred 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 +122,52 @@ 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 `renders classified document payloads into prompts without wrapper metadata`() { val provider = RecordingProvider { ModelResponse(content = "hardcoded response") } 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..75e21a27 100644 --- a/tramai-sovereign/src/main/kotlin/dev/tramai/sovereign/SovereignTramai.kt +++ b/tramai-sovereign/src/main/kotlin/dev/tramai/sovereign/SovereignTramai.kt @@ -79,7 +79,7 @@ class SovereignTramai private constructor( verificationReceipts: List, private val profile: SovereignProfileConfiguration, private val verificationSettings: ModelArtifactVerificationSettings, -) { +) : AutoCloseable { private val verificationReceipts: List = Collections.unmodifiableList(ArrayList(verificationReceipts)) /** @@ -93,6 +93,9 @@ class SovereignTramai private constructor( */ fun runtime(): SovereignTramaiRuntime = SovereignTramaiRuntime(delegate.runtime()) + /** 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..4c270510 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,32 @@ 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 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..d18c5345 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,120 @@ 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") + } + + 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 +687,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 +754,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) From 48d1bed89573c71f5f966d3a3b9cb740534bc3d9 Mon Sep 17 00:00:00 2001 From: Giona Granchelli Date: Mon, 10 Aug 2026 16:11:11 +0200 Subject: [PATCH 02/12] fix(0.6.0): exactly-once suspend resume on pre-start cancel, sovereign runtime cache, race coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit agy round-1 review fixes (PR #226): - Suspend invokeSuspend now resumes the caller continuation exactly once even when close() cancels the tracked job BEFORE the dispatcher starts it: the block records its outcome before resuming, and invokeOnCompletion resumes with a cancellation when the block never ran — otherwise the caller's suspension would freeze forever. - close() cancels tracked caller-parented invocation jobs but never joins them: their completion is dispatched on the CALLER's dispatcher, which may be blocked waiting on this very close() (joining would deadlock). The engine scope job is still cancelled-and-joined. - SovereignTramai.runtime() caches the wrapper around the delegate's single owned runtime (repeated calls return the same instance; identity test). - New tests: close racing a fast suspend invocation never leaves work against a closed engine (100 iterations; provider-start vs close-complete ordering asserted); multiple Spring AI-service beans share one runtime and all fail after context close; sovereign runtime identity. --- .../kotlin/dev/tramai/engine/TramaiEngine.kt | 53 +++++++++++++++---- .../dev/tramai/engine/TramaiEngineTest.kt | 52 ++++++++++++++++++ .../dev/tramai/sovereign/SovereignTramai.kt | 8 +-- .../tramai/sovereign/SovereignTramaiTest.kt | 7 +++ 4 files changed, 106 insertions(+), 14 deletions(-) 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 2e594fb9..1de9c5ee 100644 --- a/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt +++ b/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt @@ -555,7 +555,17 @@ class TramaiEngine( job.cancel() // Suspend invocations are children of their CALLER's job, not the // engine scope job; cancel them explicitly so close() owns them. - activeInvocationJobs.forEach { it.cancel() } + // 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. + // Never JOIN caller-parented jobs here: their completion is + // dispatched on the caller's dispatcher, which may be blocked + // waiting on this very close() — joining would deadlock. Cancel + // propagates to the in-flight work; invokeSuspend's exactly-once + // resume guarantees the caller's suspension completes. + synchronized(activeInvocationJobs) { + activeInvocationJobs.toList().forEach { it.cancel() } + } if (engineThreadMarker.get() != true) { runBlocking { job.join() } } @@ -696,16 +706,37 @@ internal class TramaiInvocationHandler( // Launch as a child of the CALLER's job (continuation.context) so parent // cancellation propagates synchronously into the in-flight invocation // (validated by the ToolSafeFailureContract / StructuredOutputFailureBoundary - // parent-cancellation tests), while tracking the launched job so engine - // close() also owns it (cancel below). Engine close cancels the engine - // scope job AND every tracked invocation job. - val launched = scope.launch(continuation.context + engineThreadMarker.asContextElement(true)) { - runCatching { execute(operation, callArguments, conversationId) } - .onSuccess { continuation.resumeWith(Result.success(it)) } - .onFailure { continuation.resumeWith(Result.failure(it)) } - } - activeInvocationJobs += launched - launched.invokeOnCompletion { activeInvocationJobs -= launched } + // parent-cancellation tests). Engine close() owns the work too: the + // launch+add is synchronized with close()'s cancel snapshot, and the + // closed flag is re-checked INSIDE the lock — so a close that won the + // race rejects this launch instead of leaving an untracked in-flight job. + // 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 = scope.launch(continuation.context + engineThreadMarker.asContextElement(true)) { + val result = runCatching { execute(operation, callArguments, conversationId) } + 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 } 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 2ee164e1..e71829f4 100644 --- a/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt +++ b/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt @@ -168,6 +168,58 @@ class TramaiEngineTest { 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 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()) + } + closer.start() + try { + service.analyze("invoice-1") + outcome.complete(null) + } catch (t: Throwable) { + outcome.complete(t) + } + 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") } 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 75e21a27..8c2d72f9 100644 --- a/tramai-sovereign/src/main/kotlin/dev/tramai/sovereign/SovereignTramai.kt +++ b/tramai-sovereign/src/main/kotlin/dev/tramai/sovereign/SovereignTramai.kt @@ -82,16 +82,18 @@ class SovereignTramai private constructor( ) : 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() 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 4c270510..25aad92f 100644 --- a/tramai-sovereign/src/test/kotlin/dev/tramai/sovereign/SovereignTramaiTest.kt +++ b/tramai-sovereign/src/test/kotlin/dev/tramai/sovereign/SovereignTramaiTest.kt @@ -111,6 +111,13 @@ class SovereignTramaiTest { .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() From 6809f5060fe3e9f73b70677d528f3721a9863886 Mon Sep 17 00:00:00 2001 From: Giona Granchelli Date: Mon, 10 Aug 2026 16:22:18 +0200 Subject: [PATCH 03/12] fix(0.6.0): never deliver blocking results from a closed engine; changelog agy round-1 P2-3 + P3-3 (PR #226): - Blocking proxy invocations re-check the closed flag after the caller-owned runBlocking completes, so a call that raced close() surfaces the fixed 'Tramai runtime is closed' IllegalStateException instead of delivering a result computed against an already-closed engine. Test added. - CHANGELOG entry for PR #226 including the AutoCloseable supertype note (source-compatible; affects compiled negative instanceof checks). --- CHANGELOG.md | 2 ++ .../kotlin/dev/tramai/engine/TramaiEngine.kt | 8 +++++- .../dev/tramai/engine/TramaiEngineTest.kt | 27 +++++++++++++++++++ 3 files changed, 36 insertions(+), 1 deletion(-) 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/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt b/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt index 1de9c5ee..75949583 100644 --- a/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt +++ b/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt @@ -686,9 +686,15 @@ internal class TramaiInvocationHandler( return if (operation.isSuspend) { invokeSuspend(operation, args.orEmpty(), conversationId) } else { - runBlocking { + val result = runBlocking { execute(operation, args.orEmpty().toList(), conversationId) } + // The engine may have closed while this blocking call was in + // flight (caller-owned runBlocking, not cancellable from here). + // 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 } } 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 e71829f4..9b1c916c 100644 --- a/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt +++ b/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt @@ -243,6 +243,33 @@ 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 + // failed fast / after close (fixed lifecycle error). 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") }, + ) + } + @Test fun `supports blocking interfaces`() { val provider = RecordingProvider { ModelResponse(content = "summary") } From d2e28f41bda05c09b5cb81ad9aa24c13a3cef329 Mon Sep 17 00:00:00 2001 From: Giona Granchelli Date: Mon, 10 Aug 2026 16:38:31 +0200 Subject: [PATCH 04/12] fix(0.6.0): close guard on resumeApproval, registerService, and streaming flows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Independent review findings (PR #226): - P1: resumeApproval and registerService ran provider work deterministically after close() — the closed guard existed only on create(), the proxy invoke seam, and the suspend launch. Both entry points now fail fast with the fixed 'Tramai runtime is closed' IllegalStateException. - P2-1: the suspend invocation block could deliver a success computed against a closed engine (caller-parented job, not joined by close). The launched block now re-checks the closed flag after execute() and converts a success into the fixed lifecycle error, mirroring the blocking path. - P2-2: streaming flows escaped close() entirely — a flow obtained before close() and collected after ran the full provider pipeline. The flow body now fails fast on collection against a closed engine. - Tests: registerService/resumeApproval fail fast on a closed engine; streaming flow collected after close fails before provider executes (provider untouched). --- .../kotlin/dev/tramai/engine/TramaiEngine.kt | 12 ++++- .../dev/tramai/engine/TramaiEngineTest.kt | 49 +++++++++++++++++++ 2 files changed, 60 insertions(+), 1 deletion(-) 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 75949583..8be417e8 100644 --- a/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt +++ b/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt @@ -460,6 +460,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, @@ -522,6 +523,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 @@ -724,7 +726,14 @@ internal class TramaiInvocationHandler( val launched = synchronized(activeInvocationJobs) { check(!isClosed.get()) { "Tramai runtime is closed" } val job = scope.launch(continuation.context + engineThreadMarker.asContextElement(true)) { - val result = runCatching { execute(operation, callArguments, conversationId) } + 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 + // (caller-parented job, not joined by close()). 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) } @@ -784,6 +793,7 @@ internal class TramaiInvocationHandler( ?: (emptyList() to initialMessages) return flow { + check(!isClosed.get()) { "Tramai runtime is closed" } val correlationId = java.util.UUID.randomUUID().toString() enforceBeforeProviderResolution(operation, correlationId, securityContext) val candidates = providerRegistry.resolveCandidates(operation.operation) 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 9b1c916c..52d97dd6 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 @@ -270,6 +271,54 @@ class TramaiEngineTest { ) } + @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 `supports blocking interfaces`() { val provider = RecordingProvider { ModelResponse(content = "summary") } From 85f3f5c487df96cf79535098805144bb2b52066f Mon Sep 17 00:00:00 2001 From: Giona Granchelli Date: Mon, 10 Aug 2026 16:53:12 +0200 Subject: [PATCH 05/12] fix(0.6.0): terminate in-flight streams on close MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-3 review P2 (PR #226): mid-collection close left a live provider stream delivering chunks after close() — the flow-body start guard only covered collection-after-close, and the collector's job is not cancelled by close(), so cooperative cancellation never fired. Every emitted chunk is now gated on the engine being open (emitWhileOpen), so a cold flow being collected at close() time terminates deterministically within one chunk latency with the fixed 'Tramai runtime is closed' error. Test: mid-collection close terminates an in-flight stream (first chunk delivered, close, gate release -> second chunk never delivered). --- .../kotlin/dev/tramai/engine/TramaiEngine.kt | 15 ++++-- .../dev/tramai/engine/TramaiEngineTest.kt | 47 +++++++++++++++++++ 2 files changed, 59 insertions(+), 3 deletions(-) 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 8be417e8..63dd99bc 100644 --- a/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt +++ b/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt @@ -794,6 +794,15 @@ internal class TramaiInvocationHandler( return flow { check(!isClosed.get()) { "Tramai runtime is closed" } + // Every chunk is gated on the engine still being open: a cold flow + // being collected at close() time must not keep delivering chunks + // (the collector's job is not cancelled by close(), so cooperative + // cancellation alone is insufficient). Deterministic termination + // within one chunk latency. + suspend fun emitWhileOpen(chunk: StreamChunk) { + check(!isClosed.get()) { "Tramai runtime is closed" } + emit(chunk) + } val correlationId = java.util.UUID.randomUUID().toString() enforceBeforeProviderResolution(operation, correlationId, securityContext) val candidates = providerRegistry.resolveCandidates(operation.operation) @@ -824,7 +833,7 @@ internal class TramaiInvocationHandler( memoryMessages = effectiveMessages, historySize = history.size, conversationId = conversationId, - emitChunk = { emit(it) }, + emitChunk = { emitWhileOpen(it) }, ), correlationId = correlationId, securityContext = securityContext, @@ -855,13 +864,13 @@ internal class TramaiInvocationHandler( lastFailure = result.error } is StreamingRouteResult.TerminalError -> { - emit(result.errorChunk) + emitWhileOpen(result.errorChunk) return@flow } } } - emit(noAvailableStreamingRouteChunk(operation, lastFailure, lastCircuitOpen)) + emitWhileOpen(noAvailableStreamingRouteChunk(operation, lastFailure, lastCircuitOpen)) } } 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 52d97dd6..415de4fe 100644 --- a/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt +++ b/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt @@ -319,6 +319,53 @@ class TramaiEngineTest { 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. + assertThat(chunks.map { it }).doesNotContain(StreamChunk.Token("second")) + if (terminal != null) { + assertThat(terminal) + .isInstanceOf(IllegalStateException::class.java) + .hasMessage("Tramai runtime is closed") + } + } + @Test fun `supports blocking interfaces`() { val provider = RecordingProvider { ModelResponse(content = "summary") } From 5feef47f942c744ba12a94abb2bcbe363a036b3c Mon Sep 17 00:00:00 2001 From: Giona Granchelli Date: Mon, 10 Aug 2026 16:59:14 +0200 Subject: [PATCH 06/12] test(0.6.0): assert first chunk explicitly in mid-stream close test --- .../src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 415de4fe..2998d530 100644 --- a/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt +++ b/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt @@ -357,7 +357,9 @@ class TramaiEngineTest { // 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. + // 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) { assertThat(terminal) From e4cdcae53d975105e8d6ef83983312322ed6d5f6 Mon Sep 17 00:00:00 2001 From: Giona Granchelli Date: Mon, 10 Aug 2026 20:40:08 +0200 Subject: [PATCH 07/12] =?UTF-8?q?fix(0.6.0):=20engine-owned=20lifecycle=20?= =?UTF-8?q?job=20=E2=80=94=20close()=20cancels=20AND=20joins=20all=20engin?= =?UTF-8?q?e-initiated=20work?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses Giona's round-5 review (PR #226): output suppression is not work termination; close() must not return while engine-created invocations are still active. - The engine now owns an internal lifecycleJob/lifecycleScope. The caller- supplied job/scope constructor parameters are NEVER cancelled or joined (fixes the caller-supplied-job close() deadlock): close() cancels and joins lifecycleJob plus every tracked invocation. - Blocking calls run as children of lifecycleJob (runBlocking(lifecycleJob)), so close() terminates a blocking provider still executing and waits for its NonCancellable cleanup before returning. - Suspend invocations run on the engine's own dispatcher (caller's Job element retained for parent-cancellation propagation; the interceptor is stripped so close() joining cannot deadlock a single-threaded caller loop). Parent-cancellation contract tests still pass. - Streaming collections run in lifecycleScope and bridge chunks to the collector's emit through a channel; close() cancels the collection job and waits for provider cleanup. Per-chunk closed gate retained. - close() now joins tracked invocation jobs (cancellation request is not termination; NonCancellable cleanup must complete before close returns). - lifecycleScope carries a CoroutineExceptionHandler so orphaned background work failures log instead of leaking onto a shared global handler. - Tests: blocking long-suspension cancelled+joined by close; streaming collection suspended indefinitely cancelled+cleaned up; close with caller-supplied job/scope does not deadlock; external-provider test now forces engine creation (was vacuous); observer fixtures thread-safe for Default-dispatcher invocations; close-race test suspends instead of blocking Thread.join. Roadmap Epic 1.3 marked complete. --- docs/ROADMAP-0.6.0.md | 2 + .../kotlin/dev/tramai/engine/TramaiEngine.kt | 278 +++++++++++------- .../dev/tramai/engine/TramaiEngineTest.kt | 120 +++++++- .../TrustedReplayEnvelopeRegistryTest.kt | 2 + .../dev/tramai/standalone/TramaiTest.kt | 5 + 5 files changed, 300 insertions(+), 107 deletions(-) diff --git a/docs/ROADMAP-0.6.0.md b/docs/ROADMAP-0.6.0.md index 1050f4af..7db0b01a 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. 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 63dd99bc..7583971f 100644 --- a/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt +++ b/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt @@ -90,6 +90,7 @@ 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 @@ -103,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 @@ -174,6 +176,29 @@ class TramaiEngine( 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() + private val lifecycleScope: CoroutineScope = CoroutineScope( + lifecycleJob + Dispatchers.Default + 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 them instead. + System.getLogger("dev.tramai.engine.TramaiEngine").log( + System.Logger.Level.WARNING, + "Engine-owned coroutine failed after close or abandonment", + error, + ) + }, + ) /** * Suspend-invocation jobs launched for caller continuations. They are * children of the CALLER's job (so parent cancellation propagates), but the @@ -415,6 +440,8 @@ class TramaiEngine( chatMemory = chatMemory, conversationIdProvider = conversationIdProvider, scope = scope, + lifecycleJob = lifecycleJob, + lifecycleScope = lifecycleScope, isClosed = closed, engineThreadMarker = engineThreadMarker, activeInvocationJobs = activeInvocationJobs, @@ -482,6 +509,8 @@ class TramaiEngine( chatMemory = chatMemory, conversationIdProvider = conversationIdProvider, scope = scope, + lifecycleJob = lifecycleJob, + lifecycleScope = lifecycleScope, isClosed = closed, engineThreadMarker = engineThreadMarker, activeInvocationJobs = activeInvocationJobs, @@ -549,27 +578,38 @@ class TramaiEngine( resumeApproval(command) as R /** - * Cancels and, except from one of its own coroutines, waits for the engine-owned - * coroutine hierarchy. Dependencies supplied by callers are not closed. + * 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() { if (closed.compareAndSet(false, true)) { - job.cancel() + 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. - // Never JOIN caller-parented jobs here: their completion is - // dispatched on the caller's dispatcher, which may be blocked - // waiting on this very close() — joining would deadlock. Cancel - // propagates to the in-flight work; invokeSuspend's exactly-once - // resume guarantees the caller's suspension completes. - synchronized(activeInvocationJobs) { - activeInvocationJobs.toList().forEach { it.cancel() } + val tracked = synchronized(activeInvocationJobs) { + activeInvocationJobs.toList() } + tracked.forEach { it.cancel() } if (engineThreadMarker.get() != true) { - runBlocking { job.join() } + // 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. Caller-parented invocation jobs run + // on the caller's dispatcher; joining is safe as long as + // close() is not called from a coroutine dispatched on that + // same single-threaded dispatcher (documented caller + // constraint, matching the self-close marker guard below). + runBlocking { + lifecycleJob.join() + tracked.forEach { it.join() } + } } } } @@ -613,6 +653,8 @@ 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(), @@ -688,13 +730,19 @@ internal class TramaiInvocationHandler( return if (operation.isSuspend) { invokeSuspend(operation, args.orEmpty(), conversationId) } else { - val result = 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 (caller-owned runBlocking, not cancellable from here). - // Never deliver a result computed against a closed engine: the - // caller sees the fixed lifecycle error instead. + // 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 } @@ -711,26 +759,33 @@ internal class TramaiInvocationHandler( ?: throw ConfigurationException("Suspend invocation for ${operation.method.name} is missing its continuation") val callArguments = args.dropLast(1) - // Launch as a child of the CALLER's job (continuation.context) so parent - // cancellation propagates synchronously into the in-flight invocation - // (validated by the ToolSafeFailureContract / StructuredOutputFailureBoundary - // parent-cancellation tests). Engine close() owns the work too: the - // launch+add is synchronized with close()'s cancel snapshot, and the - // closed flag is re-checked INSIDE the lock — so a close that won the - // race rejects this launch instead of leaving an untracked in-flight job. - // 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. + // 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 = scope.launch(continuation.context + engineThreadMarker.asContextElement(true)) { + 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 - // (caller-parented job, not joined by close()). The caller sees - // the fixed lifecycle error instead (mirrors the blocking path). + // 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")) } @@ -794,83 +849,108 @@ internal class TramaiInvocationHandler( return flow { check(!isClosed.get()) { "Tramai runtime is closed" } - // Every chunk is gated on the engine still being open: a cold flow - // being collected at close() time must not keep delivering chunks - // (the collector's job is not cancelled by close(), so cooperative - // cancellation alone is insufficient). Deterministic termination - // within one chunk latency. - suspend fun emitWhileOpen(chunk: StreamChunk) { - check(!isClosed.get()) { "Tramai runtime is closed" } - emit(chunk) - } - 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 = { emitWhileOpen(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, + // 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 + // channel (emit itself must stay in the collector's coroutine). + val chunks = kotlinx.coroutines.channels.Channel(kotlinx.coroutines.channels.Channel.UNLIMITED) + 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 -> { - emitWhileOpen(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 (failure: Throwable) { + collectFailure.set(failure) + throw failure + } finally { + chunks.close() } } - - emitWhileOpen(noAvailableStreamingRouteChunk(operation, lastFailure, lastCircuitOpen)) + 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 2998d530..ee9e2f22 100644 --- a/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt +++ b/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt @@ -59,8 +59,11 @@ 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 @@ -182,6 +185,7 @@ class TramaiEngineTest { 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 @@ -190,6 +194,7 @@ class TramaiEngineTest { val closer = Thread { engine.close() closeCompletedAt.set(System.nanoTime()) + closeDone.complete(Unit) } closer.start() try { @@ -198,6 +203,10 @@ class TramaiEngineTest { } 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() @@ -261,13 +270,17 @@ class TramaiEngineTest { closer.join() // Either the call completed before close won the race (result), or it - // failed fast / after close (fixed lifecycle error). It must NEVER be - // a summary delivered from an engine that was already closed. + // 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") }, ) } @@ -362,9 +375,100 @@ class TramaiEngineTest { assertThat(chunks.map { it }).contains(StreamChunk.Token("first")) assertThat(chunks.map { it }).doesNotContain(StreamChunk.Token("second")) if (terminal != null) { - assertThat(terminal) - .isInstanceOf(IllegalStateException::class.java) - .hasMessage("Tramai runtime is closed") + // 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 provider = NamedStreamingProvider("p") { + flow { + try { + emit(StreamChunk.Token("first")) + 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) + } + // Let the first chunk arrive (provider is suspended indefinitely). + delay(200) + 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() } } @@ -3795,7 +3899,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) @@ -3838,7 +3942,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, ) } @@ -3849,7 +3953,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-standalone/src/test/kotlin/dev/tramai/standalone/TramaiTest.kt b/tramai-standalone/src/test/kotlin/dev/tramai/standalone/TramaiTest.kt index d18c5345..516a0f42 100644 --- a/tramai-standalone/src/test/kotlin/dev/tramai/standalone/TramaiTest.kt +++ b/tramai-standalone/src/test/kotlin/dev/tramai/standalone/TramaiTest.kt @@ -90,6 +90,11 @@ class TramaiTest { 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() From 86de9eb4e0f94374581522d44c79c30a7531bc5b Mon Sep 17 00:00:00 2001 From: Giona Granchelli Date: Mon, 10 Aug 2026 22:12:55 +0200 Subject: [PATCH 08/12] docs(0.6.0): precise close() contract in tramai-engine module docs Copilot thread r3749862916: 'waits for externally initiated shutdown' was ambiguous. State exactly what close() does: cancels and joins engine-owned work (blocking, suspend, streaming) and never touches the caller-supplied job/scope constructor parameters. --- docs/modules/tramai-engine.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/modules/tramai-engine.md b/docs/modules/tramai-engine.md index 130688e8..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 and waits for externally initiated shutdown | +| `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 | From d1d4a4ef4458fc835352641e2d936d7e10700f2b Mon Sep 17 00:00:00 2001 From: Giona Granchelli Date: Mon, 10 Aug 2026 23:22:16 +0200 Subject: [PATCH 09/12] =?UTF-8?q?fix(0.6.0):=20rendezvous=20streaming=20br?= =?UTF-8?q?idge=20=E2=80=94=20backpressure,=20pre-start=20close=20race,=20?= =?UTF-8?q?safe=20logging?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses Giona's round-6 review of the streaming lifecycle bridge. - Channel.UNLIMITED -> Channel.RENDEZVOUS: a slow collector now blocks the provider instead of letting it race ahead into unbounded buffering (backpressure semantics preserved; take(1)/slow-collector behavior unchanged). Regression: 'streaming bridge preserves backpressure when the collector is slow' proves the provider cannot emit chunk 2..N while the collector is blocked on chunk 1. - Channel close now depends on JOB completion, not on the collection body having started: collectJob.invokeOnCompletion { chunks.close(cause) }. If close() cancels lifecycleJob after the flow's open check but before the launched body runs, the collector terminates instead of hanging forever on receive. Regression: 'stream start racing close never hangs the collector' (200 iterations of the admission race, each bounded). - Streaming failures are captured into collectFailure and surfaced to the collector via the channel drain instead of being rethrown: an arbitrary (possibly sensitive, externally supplied) throwable no longer reaches the lifecycle CoroutineExceptionHandler and the normal logger. The handler now logs fixed safe metadata (exception type name only), never the raw throwable — consistent with Epic 1.2 safe-error-boundary work. - Fixed stale close() comment: invocation jobs run on the engine's own dispatcher (caller ContinuationInterceptor stripped), not the caller's. --- .../kotlin/dev/tramai/engine/TramaiEngine.kt | 53 +++++++++---- .../dev/tramai/engine/TramaiEngineTest.kt | 74 +++++++++++++++++++ 2 files changed, 112 insertions(+), 15 deletions(-) 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 7583971f..1e91948e 100644 --- a/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt +++ b/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt @@ -189,13 +189,15 @@ class TramaiEngine( lifecycleJob + Dispatchers.Default + 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 them instead. + // 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", - error, + "Engine-owned coroutine failed after close or abandonment (type: ${error::class.qualifiedName})", ) }, ) @@ -601,11 +603,14 @@ class TramaiEngine( // 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. Caller-parented invocation jobs run - // on the caller's dispatcher; joining is safe as long as - // close() is not called from a coroutine dispatched on that - // same single-threaded dispatcher (documented caller - // constraint, matching the self-close marker guard below). + // 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() } @@ -855,8 +860,11 @@ internal class TramaiInvocationHandler( // (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 - // channel (emit itself must stay in the collector's coroutine). - val chunks = kotlinx.coroutines.channels.Channel(kotlinx.coroutines.channels.Channel.UNLIMITED) + // 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 { @@ -928,13 +936,28 @@ internal class TramaiInvocationHandler( } 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) - throw failure - } finally { - chunks.close() } } + // 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" } 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 ee9e2f22..7e4f6430 100644 --- a/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt +++ b/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt @@ -472,6 +472,80 @@ class TramaiEngineTest { } } + @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 emitted = java.util.concurrent.atomic.AtomicInteger(0) + val collectorGate = CompletableDeferred() + val provider = NamedStreamingProvider("p") { + flow { + var i = 0 + while (true) { + emit(StreamChunk.Token("chunk-${i++}")) + emitted.incrementAndGet() + } + } + } + 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 { + if (!firstChunk.isCompleted) firstChunk.complete(Unit) + collectorGate.await() + } + } + firstChunk.await() + // The collector is deliberately blocked after chunk 1. With a + // rendezvous bridge the provider cannot race ahead producing chunk + // 2..N (UNLIMITED would let it buffer unboundedly). + delay(300) + assertThat(emitted.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") } From 9c0a736cbf4ee089dd20a93b147e52984dad9175 Mon Sep 17 00:00:00 2001 From: Giona Granchelli Date: Tue, 11 Aug 2026 20:07:33 +0200 Subject: [PATCH 10/12] =?UTF-8?q?fix(0.6.0):=20engine-thread=20marker=20on?= =?UTF-8?q?=20lifecycle=20scope=20=E2=80=94=20streaming=20self-close=20saf?= =?UTF-8?q?e?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses Giona's round-7 review. - The engine-thread marker now lives on lifecycleScope itself (engineThreadMarker.asContextElement(true) in the scope context), so EVERY engine-owned child — including the streaming collection job — carries the self-close protection automatically. Previously only blocking and suspend paths had it explicitly, so engine.close() called from inside a streaming provider/interceptor/observer would self-join forever (lifecycleJob.cancel -> join of a job blocked inside close()). - Regression: 'self close from streaming owned coroutine does not deadlock' (provider flow calls engine.close() mid-collection; withTimeout(2s) proves termination). - Roadmap: Epic 1.3 gains the leak-test evidence matrix (task 6): engine jobs, worker jobs, subprocesses, HTTP response streams, shutdown hooks each mapped to their concrete tests. --- docs/ROADMAP-0.6.0.md | 10 +++++++ .../kotlin/dev/tramai/engine/TramaiEngine.kt | 7 ++++- .../dev/tramai/engine/TramaiEngineTest.kt | 28 +++++++++++++++++++ 3 files changed, 44 insertions(+), 1 deletion(-) diff --git a/docs/ROADMAP-0.6.0.md b/docs/ROADMAP-0.6.0.md index 7db0b01a..3dd351b1 100644 --- a/docs/ROADMAP-0.6.0.md +++ b/docs/ROADMAP-0.6.0.md @@ -375,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 | PR #226 streaming lifecycle tests: `streaming collection suspended indefinitely is cancelled and cleaned up by close`, `mid-collection close terminates an in-flight stream`; example `kotlin-springboot-example` E2E smoke test | +| Shutdown hooks | Spring `destroyMethod` close + `tramai-spring` context-shutdown tests; `repeated close is harmless` idempotency test in tramai-standalone | + --- ## Epic 1.4: HTTP network-boundary correctness 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 1e91948e..f685d8d6 100644 --- a/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt +++ b/tramai-engine/src/main/kotlin/dev/tramai/engine/TramaiEngine.kt @@ -185,8 +185,13 @@ class TramaiEngine( * 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 + CoroutineExceptionHandler { _, error -> + 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. 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 7e4f6430..1aab3fc6 100644 --- a/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt +++ b/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt @@ -472,6 +472,34 @@ class TramaiEngineTest { } } + @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 From 25b998689dc3a396e9162b8c80da130ad47c57a0 Mon Sep 17 00:00:00 2001 From: Giona Granchelli Date: Tue, 11 Aug 2026 23:15:41 +0200 Subject: [PATCH 11/12] test(0.6.0): shutdown-hook leak proof + deterministic timing + correct stream evidence Addresses Giona's round-8 review (P2/P3; no P1 remains). - New regression 'close deregisters the JVM shutdown hook and retains no reference' (TramaiWorkerTest): proves start() registers a hook and close() -> shutdown() removes it (private field reflection: non-null -> null). - Roadmap HTTP-stream evidence row now cites the provider-level InputStream cleanup tests in OpenAiProviderTest (close after DONE, malformed chunk, collector stop after first token, mid-stream I/O failure) instead of only the engine Flow bridge tests. - Timing determinism: the suspended-stream cleanup test's delay(200) is now a firstChunkDelivered CompletableDeferred gate; the backpressure test's delay(300) is replaced with a structural proof (provider ATTEMPTS chunk 2 then stalls at the rendezvous send: attempted==2, delivered==1). --- docs/ROADMAP-0.6.0.md | 2 +- .../dev/tramai/engine/TramaiEngineTest.kt | 30 +++++++++++++------ .../tramai/orchestration/TramaiWorkerTest.kt | 23 ++++++++++++++ 3 files changed, 45 insertions(+), 10 deletions(-) diff --git a/docs/ROADMAP-0.6.0.md b/docs/ROADMAP-0.6.0.md index 3dd351b1..e53354a1 100644 --- a/docs/ROADMAP-0.6.0.md +++ b/docs/ROADMAP-0.6.0.md @@ -382,7 +382,7 @@ This phase is intentionally completed before large decomposition work. | 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 | PR #226 streaming lifecycle tests: `streaming collection suspended indefinitely is cancelled and cleaned up by close`, `mid-collection close terminates an in-flight stream`; example `kotlin-springboot-example` E2E smoke test | +| 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 | Spring `destroyMethod` close + `tramai-spring` context-shutdown tests; `repeated close is harmless` idempotency test in tramai-standalone | --- 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 1aab3fc6..357ea649 100644 --- a/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt +++ b/tramai-engine/src/test/kotlin/dev/tramai/engine/TramaiEngineTest.kt @@ -422,10 +422,12 @@ class TramaiEngineTest { @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) { @@ -446,8 +448,9 @@ class TramaiEngineTest { runCatching { service.stream("invoice-123").collect {} } collectionDone.complete(Unit) } - // Let the first chunk arrive (provider is suspended indefinitely). - delay(200) + // 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. @@ -536,14 +539,15 @@ class TramaiEngineTest { @Test fun `streaming bridge preserves backpressure when the collector is slow`() = runBlocking { - val emitted = java.util.concurrent.atomic.AtomicInteger(0) + 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++}")) - emitted.incrementAndGet() } } } @@ -557,16 +561,24 @@ class TramaiEngineTest { 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. With a - // rendezvous bridge the provider cannot race ahead producing chunk - // 2..N (UNLIMITED would let it buffer unboundedly). - delay(300) - assertThat(emitted.get()).isEqualTo(1) + // 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 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..c725e255 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,29 @@ 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() + + worker.close() + // close() -> shutdown() must remove the hook and drop the reference: + // no retained shutdown hook survives a graceful worker close. + val hookAfterClose = hookField.get(worker) as? Thread + assertThat(hookAfterClose).isNull() + } + @Test fun `worker crash leaves non replayable step in unknown state and takeover fails`() = runBlocking { val checkpointStore = InMemoryWorkflowCheckpointStore() From 40deb73535e52a356c1113ad55bdd5e1617adbf0 Mon Sep 17 00:00:00 2001 From: Giona Granchelli Date: Tue, 11 Aug 2026 23:38:15 +0200 Subject: [PATCH 12/12] test(0.6.0): shutdown-hook regression proves JVM deregistration, not just reference clearing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses Giona's round-8 follow-up (P2): the previous regression proved close() nulls the worker's shutdownHook field but NOT that the JVM registry was actually deregistered — a mutation removing Runtime.removeShutdownHook would still pass. Strengthened: after close(), Runtime.removeShutdownHook(hookAfterStart) must return FALSE (hook already deregistered). Verified mutation-sensitive: with the production deregistration commented out, the test FAILS; restored, it passes. close() moved into try/finally so an assertion failure can never leave a hook registered in the test JVM. Roadmap shutdown-hook row now cites this test (JVM-level deregistration + no retained Thread reference) instead of only the indirect Spring/standalone evidence. --- docs/ROADMAP-0.6.0.md | 2 +- .../tramai/orchestration/TramaiWorkerTest.kt | 25 +++++++++++++++---- 2 files changed, 21 insertions(+), 6 deletions(-) diff --git a/docs/ROADMAP-0.6.0.md b/docs/ROADMAP-0.6.0.md index e53354a1..46a083b5 100644 --- a/docs/ROADMAP-0.6.0.md +++ b/docs/ROADMAP-0.6.0.md @@ -383,7 +383,7 @@ This phase is intentionally completed before large decomposition work. | 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 | Spring `destroyMethod` close + `tramai-spring` context-shutdown tests; `repeated close is harmless` idempotency test in tramai-standalone | +| 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 | --- 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 c725e255..ac21e41d 100644 --- a/tramai-orchestration/src/test/kotlin/dev/tramai/orchestration/TramaiWorkerTest.kt +++ b/tramai-orchestration/src/test/kotlin/dev/tramai/orchestration/TramaiWorkerTest.kt @@ -78,11 +78,26 @@ class TramaiWorkerTest { val hookAfterStart = hookField.get(worker) as? Thread assertThat(hookAfterStart).isNotNull() - worker.close() - // close() -> shutdown() must remove the hook and drop the reference: - // no retained shutdown hook survives a graceful worker close. - val hookAfterClose = hookField.get(worker) as? Thread - assertThat(hookAfterClose).isNull() + 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