[kyo-ai] make streaming and completion deadline tests deterministic

LLMStreamTest ("a stalled stream fails typed at the configured timeout", "a
slow consumer does not spend the streaming deadline") and LLMTest ("a
client-side timeout halts fast", "the configured timeout bounds a completion
call and its retries") set a short config.timeout and raced it against a real
mock-server round-trip. On the slow single-threaded arm64 Native runner the
local production alone can exceed the budget, so the deadline fires at the
wrong point and the leaves fail intermittently.

Drive the deadline on a controlled Clock (Clock.withTimeControl) and order the
real I/O with latches: the deadline now fires on an explicit time advance
rather than by racing wall-clock against production, so the leaves are
deterministic on every platform. Test-only; no production changes.

Author: fwbrasil

kyo-ai/shared/src/test/scala/kyo/LLMStreamTest.scala

@@ -40,53 +40,80 @@ class LLMStreamTest extends kyo.test.Test[Any]:
 
     "a stalled stream fails typed at the configured timeout after delivering its fragments" in {
         TestCompletionServer.runStreaming { server =>
-            // The provider emits part of the envelope and then stops producing without a terminator. The
-            // fragments must REACH the consumer and the call must still end at its deadline, so the
-            // assertion is arrived-then-failed: a bound covering only the response headers would fail
-            // before anything was delivered, and a stream with no bound would never end at all.
-            val config = serverConfig(server.baseUrl).timeout(300.millis)
-            for
-                seen <- AtomicRef.init(Chunk.empty[String])
-                _    <- server.enqueueStreamStall(Chunk(argDelta("""{"resultValue":""""), argDelta("partial")))
-                result <- Abort.run[AIException](
-                    LLM.run(config)(Scope.run(AI.stream[String].map(_.foreach(s => seen.updateAndGet(_.append(s)).unit))))
-                )
-                delivered <- seen.get
-            yield
-                assert(
-                    delivered.mkString == "partial",
-                    s"the emitted fragments must reach the consumer before the deadline: $delivered"
-                )
-                result match
-                    case Result.Failure(_: AICompletionTimeoutException) => succeed
-                    case other => fail(s"expected the streaming deadline to fire on a stalled stream, got: $other")
-                end match
-            end for
+            // The provider emits part of the envelope then stops without a terminator. The fragments must REACH
+            // the consumer and the call must still end at its deadline. Timing is made deterministic with a
+            // controlled Clock: the deadline sleeps on that Clock, so it fires exactly when the test advances
+            // virtual time, never by racing a real short timeout against the real SSE round-trip. The `arrived`
+            // latch orders "fragments delivered" (real I/O) BEFORE "deadline fired" (a clock advance) causally,
+            // so the leaf cannot flake on a slow runner where delivery alone exceeds a fixed millisecond budget.
+            val config = serverConfig(server.baseUrl).timeout(10.seconds)
+            Clock.withTimeControl { control =>
+                for
+                    arrived <- Latch.init(1)
+                    seen    <- AtomicRef.init(Chunk.empty[String])
+                    _       <- server.enqueueStreamStall(Chunk(argDelta("""{"resultValue":""""), argDelta("partial")))
+                    fiber <- Fiber.init {
+                        Abort.run[AIException] {
+                            LLM.run(config)(Scope.run(AI.stream[String].map(_.foreach { s =>
+                                seen.updateAndGet(_.append(s)).flatMap(cur => Kyo.when(cur.mkString == "partial")(arrived.release))
+                            })))
+                        }
+                    }
+                    _         <- arrived.await
+                    _         <- control.advance(config.timeout + 1.second, 500.millis)
+                    result    <- fiber.get
+                    delivered <- seen.get
+                yield
+                    assert(
+                        delivered.mkString == "partial",
+                        s"the emitted fragments must reach the consumer before the deadline: $delivered"
+                    )
+                    result match
+                        case Result.Failure(_: AICompletionTimeoutException) => succeed
+                        case other => fail(s"expected the streaming deadline to fire on a stalled stream, got: $other")
+                    end match
+                end for
+            }
         }
     }
 
     "a slow consumer does not spend the streaming deadline" in {
         TestCompletionServer.runStreaming { server =>
-            // The deadline is a budget on the provider's production, not on wall-clock: a consumer that
-            // pauses between elements for longer than the timeout must still receive the WHOLE stream, so
-            // the assertion is on the delivered value. Asserting only that the call completed would also
-            // be satisfied by an empty stream, which proves nothing about the bound.
-            val config   = serverConfig(server.baseUrl).timeout(300.millis)
+            // The deadline is a budget on the provider's PRODUCTION, not on the consumer's wall-clock. Production
+            // buffers the whole (tiny) stream and completes independently of the consumer, so once it has, virtual
+            // time can be pushed FAR past the deadline while the consumer is still blocked and the deadline must not
+            // fire. The consumer's slowness is a latch, not a sleep; the `firstSeen` latch plus a short real settle
+            // order "production done" before the advance, so the leaf is deterministic rather than racing a fixed
+            // millisecond consumer delay against the deadline.
+            val config   = serverConfig(server.baseUrl).timeout(10.seconds)
             val expected = List(Answer("ok"), Answer("two"))
             val args     = elementArgs(expected)
             val split    = Chunk(args.substring(0, 12), args.substring(12, 24), args.substring(24))
-            for
-                seen <- AtomicRef.init(Chunk.empty[Answer])
-                _    <- server.enqueueStream(split.map(argDelta))
-                _ <- LLM.run(config)(
-                    Scope.run(AI.stream[Answer].map(_.foreach(a => Async.sleep(700.millis).andThen(seen.updateAndGet(_.append(a)).unit))))
+            Clock.withTimeControl { control =>
+                for
+                    firstSeen <- Latch.init(1)
+                    proceed   <- Latch.init(1)
+                    seen      <- AtomicRef.init(Chunk.empty[Answer])
+                    _         <- server.enqueueStream(split.map(argDelta))
+                    fiber <- Fiber.init {
+                        LLM.run(config)(Scope.run(AI.stream[Answer].map(_.foreach { a =>
+                            seen.updateAndGet(_.append(a)).flatMap(cur =>
+                                Kyo.when(cur.size == 1)(firstSeen.release).andThen(proceed.await)
+                            )
+                        })))
+                    }
+                    _         <- firstSeen.await
+                    _         <- control.advance(Duration.Zero, 300.millis)
+                    _         <- control.advance(config.timeout * 2)
+                    _         <- proceed.release
+                    _         <- fiber.get
+                    delivered <- seen.get
+                yield assert(
+                    delivered == Chunk.from(expected),
+                    s"a consumer slower than the timeout must still receive the whole stream, got: $delivered"
                 )
-                delivered <- seen.get
-            yield assert(
-                delivered == Chunk.from(expected),
-                s"a consumer slower than the timeout must still receive the whole stream, got: $delivered"
-            )
-            end for
+                end for
+            }
         }
     }
 

kyo-ai/shared/src/test/scala/kyo/LLMTest.scala

@@ -563,40 +563,83 @@ class LLMTest extends kyo.test.Test[Any]:
         }
     }
 
+    /** Waits until the server has captured at least `n` requests, advancing only real wall-clock time (never virtual
+      * time) under the controlled Clock, so a leaf can defer advancing past the deadline until the request is genuinely
+      * in flight rather than racing a fixed millisecond budget for it to reach the server.
+      */
+    private def awaitCaptured(server: TestCompletionServer, control: Clock.TimeControl, n: Int)(using Frame): Unit < Async =
+        Loop(0) { i =>
+            server.captured.flatMap { caps =>
+                if caps.size >= n then Loop.done
+                else if i >= 200 then Loop.done
+                else control.advance(Duration.Zero, 20.millis).andThen(Loop.continue(i + 1))
+            }
+        }
+
     "a client-side timeout halts fast as AICompletionTimeoutException" in {
         TestCompletionServer.run { server =>
-            val config = serverConfig(server.baseUrl).timeout(100.millis).retrySchedule(Schedule.repeat(1))
-            server.enqueueNeverRespond.andThen {
-                Abort.run[AIException](LLM.run(config)(AI.gen[String])).map { result =>
-                    server.captured.map { caps =>
+            // Deterministic timing: the deadline runs on a controlled Clock, so it fires exactly when the test advances
+            // virtual time, never by racing a real short timeout against the real request reaching the server. The
+            // request is sent over real I/O and confirmed captured BEFORE the advance, so "request in flight" and
+            // "deadline fired" are causally ordered; the leaf cannot flake on a slow runner where the request alone
+            // takes longer than a fixed millisecond budget to arrive.
+            val config = serverConfig(server.baseUrl).timeout(10.seconds).retrySchedule(Schedule.repeat(1))
+            Clock.withTimeControl { control =>
+                server.enqueueNeverRespond.andThen {
+                    for
+                        fiber  <- Fiber.init(Abort.run[AIException](LLM.run(config)(AI.gen[String])))
+                        _      <- awaitCaptured(server, control, 1)
+                        _      <- control.advance(config.timeout + 1.second, 500.millis)
+                        result <- fiber.get
+                        caps   <- server.captured
+                    yield
                         result match
                             case Result.Failure(_: AICompletionTimeoutException) => ()
                             case other => fail(s"expected AICompletionTimeoutException, got: $other")
                         assert(caps.size == 1, s"a per-call timeout must halt without retry, expected 1 request, got: ${caps.size}")
-                    }
+                    end for
                 }
             }
         }
     }
 
     "the configured timeout bounds a completion call and its retries, not one attempt" in {
         TestCompletionServer.run { server =>
-            // Every attempt fails transiently and is retried on a schedule whose backoff outlasts the
-            // configured timeout. The deadline covers the retry clause, so the call surfaces the timeout;
-            // a deadline that covered only one attempt would let the schedule run to exhaustion and
-            // surface the throttle instead.
+            // The deadline covers the whole call INCLUDING its retry backoffs, not one attempt. Made deterministic
+            // with a controlled Clock rather than by racing a real deadline against a real schedule: the first retry
+            // backoff is set LONGER than the deadline, so after one throttled attempt the call is waiting to retry
+            // when virtual time is advanced past the deadline. A call-scoped deadline (armed once at call start)
+            // fires during that backoff, so exactly one attempt was made; an attempt-scoped deadline would have been
+            // cancelled with the completed attempt and never fire, leaving the call to retry to exhaustion. The one
+            // attempt is confirmed over real I/O, and a short real settle lets its throttle response schedule the
+            // backoff, before the advance; the advance itself carries the causal firing, so nothing races real time.
+            val callTimeout = 5.seconds
             val config = serverConfig(server.baseUrl)
-                .timeout(300.millis)
-                .retrySchedule(Schedule.exponentialBackoff(initial = 200.millis, factor = 2, maxBackoff = 2.seconds).take(10))
+                .timeout(callTimeout)
+                .retrySchedule(Schedule.exponentialBackoff(initial = callTimeout * 2, factor = 2, maxBackoff = 1.minute).take(10))
             def throttle(remaining: Int): Unit < Async =
                 if remaining == 0 then ()
                 else server.enqueueStatus(429, """{"error":{"message":"rate limited"}}""").andThen(throttle(remaining - 1))
-            throttle(12).andThen {
-                Abort.run[AIException](LLM.run(config)(AI.gen[String])).map { result =>
-                    result match
-                        case Result.Failure(_: AICompletionTimeoutException) => succeed
-                        case other =>
-                            fail(s"expected the call deadline to fire while retries were still pending, got: $other")
+            Clock.withTimeControl { control =>
+                throttle(12).andThen {
+                    for
+                        fiber  <- Fiber.init(Abort.run[AIException](LLM.run(config)(AI.gen[String])))
+                        _      <- awaitCaptured(server, control, 1)
+                        _      <- control.advance(Duration.Zero, 300.millis)
+                        _      <- control.advance(callTimeout + 1.second, 500.millis)
+                        result <- fiber.get
+                        caps   <- server.captured
+                    yield
+                        result match
+                            case Result.Failure(_: AICompletionTimeoutException) => ()
+                            case other =>
+                                fail(s"expected the call deadline to fire while retries were still pending, got: $other")
+                        end match
+                        assert(
+                            caps.size == 1,
+                            s"the deadline must fire during the first retry's backoff, so exactly one attempt is made, got ${caps.size}"
+                        )
+                    end for
                 }
             }
         }