`CoreAIPipelinedEngine` producer remains active after EOS, causing the next response to crash in `drain()`

[timokoethe]

Model name

gemma3

Command run

// as simple as possible:
let model = try await CoreAILanguageModel(resourcesAt: modelUrl!)
let session = LanguageModelSession(model: model)
        
let response = try await session.respond(to: "What is the capital city of America?")
print(response)

// The next one fails:
let response2 = try await session.respond(to: "What is the capital city of Canada?")
print(response2)

macOS / iOS target

macOS 27.0 beta 1

Xcode version

Xcode 27.0 beta 1

Python / uv version

uv 0.11.19, python 3.14.5

Full error output

== RUNNABLE ERROR:

    CrashReportError: Fatal Error in CoreAIPipelinedEngine.swift
    
    Application crashed due to fatalError in CoreAIPipelinedEngine.swift at line 151.
    
    Engine not returned after drain() — tokenSequence Task stuck?
    
    Process:             CoreAIChat [38979]
    Path:                <none>
    
    Date/Time:           2026-06-14 17:20:32 +0000

Anything else?

Steps to reproduce:

1. In a SwiftUI App or Xcode-project: Load a dynamically shaped Core AI language model (eg. Gemma-3-4b-it-4bit-dynamic)
2. Create one LanguageModelSession with the model..
3. Call session.respond() and wait for the first response to complete.
4. Immediately call session.respond() again using the same session.
5. The first response completes normally after emitting EOS.
6. The second call waits in CoreAIPipelinedEngine.reset() / drain().
7. After approximately five seconds, the process crashes with:Fatal error: Engine not returned after drain() — tokenSequence Task stuck?

Why does this fail?

I started some research in the repository and I found the following:
CoreAIPipelinedEngine.generate() starts an independent producer task that keeps generating tokens up to maxTokens while holding exclusive ownership of the engine.
When respondVanilla() detects EOS, it records .eos and stops consuming the stream. However, this does not cancel or await the producer task. The producer therefore continues running in the background while engineInUse remains true.

The second respond() call invokes reset(), which waits in drain() for the previous producer to release the engine. If it does not finish within approximately five seconds, drain() terminates the process with fatalError.
A single Task.yield() at the end of respondVanilla() does not guarantee that the producer has completed or released the engine.

Intended Behavior

I am not 100% sure whether consumers are intentionally expected to drain the remaining token stream before reusing the engine, or whether early termination should automatically cancel and await the producer task.

If applications are expected to handle this themselves, maybe that lifecycle requirement should be documented. However, draining the stream would still allow the engine to generate and discard all remaining tokens up to maxTokens, wasting GPU time and all belonging to that.

Thanks in advance👍

[timokoethe]

I saw, this may be related to #32. However, the current diff still calls drain() before cancelling the pipelined producer, and CoreAILanguageModel does not appear to call cancel() before reset(). Therefore, the EOS reproduction described here may still trigger the same crash.

[stikves]

Yes, #32 is specifically to solve this issue.

The clients no longer need to call cancel() before reset() anymore, as reset() implicitly calls cancel():
https://github.com/apple/coreai-models/pull/32/changes#diff-4c5b00e84f56b905811ae4ddcd68aaa3b06cbe18e2f66011019821ce247b0a89R381

[stikves]

The PR is now merged.

Can you verify on your end that it fixed the behavior?

[timokoethe]

Thank you very much! I checked it again using the current state of the main branch (commit dd12488) but I still get the exact same error again: Application crashed due to fatalError in CoreAIPipelinedEngine.swift at line 169. Engine not returned after drain() — tokenSequence Task stuck?

Here the code snippet again:

let modelUrl: URL? = Bundle.main.url(forResource: "gemma_3_4b_it_4bit_dynamic", withExtension: nil)

let model = try await CoreAILanguageModel(resourcesAt: modelUrl!)
let session = LanguageModelSession(model: model)
      
let response = try await session.respond(to: "What is the capital city of America?")
print(response.content) // this is still printed perfectly

let response2 = try await session.respond(to: "What is the capital city of Canada?") // here the error is thrown
print(response2.content)

So I dived a bit deeper into the code:
In reset() at CoreAIPipelinedEngine.swift line 188, drain() is called before the active generation is cancelled. When the consumer stops early after receiving an EOS token, the producer task can remain active.
drain() then waits for that task to release the engine, but cancellation only happens afterward, so it can never unblock and eventually triggers the fatal error.
reset() should maybe cancel and await the active generation before draining or resetting the engine state.

Furthermore, the current tests miss this because MockEngine.reset() cancels first, while CoreAIPipelinedEngine.reset() drains before cancelling, causing the EOS timeout.

Sorry for the inconvenience. Please let me know if you need any additional details, another minimal reproducible example, or if you’d like me to investigate and propose a code-level fix..
Thank you for looking into this 😄

[stikves]

Looks like we are doing drain and cancel in the wrong order. Should be a simple follow up (but will be included in a larger PR for now)

[timokoethe]

Thank you! If it helps, I can also test the larger PR against this issue once it’s up and leave a brief comment confirming whether it resolves this specific problem on my side 😃

[stikves]

It might be better to isolate this fix
#47

Yes, it would be helpful if you can verify on your side

[RichNasz]

Confirmed fix works for multi-turn agentic sessions with tools (cherry-picked PR #47)

We hit this same crash building a multi-turn agentic loop where AgentSessionLoop creates a single LanguageModelSession with tools and calls respond() repeatedly — evaluating each response against validation criteria and re-prompting with feedback until the output converges or a max iteration count is reached.

Environment: macOS 27.0 beta (26A5353q), gemma3-12b 4-bit dynamic, Swift 6.2

Reproduction path: LanguageModelSession(model:, tools:, instructions:) → session.respond(to: goal) succeeds → evaluator returns retry feedback → session.respond(to: feedback) crashes in drain().

Cherry-picked e53b8b3 + c482711 onto main (a43371e). The fix resolves the crash — our loop now successfully iterates 3+ times on the same session with tool dispatch active in each turn.

One additional observation: before discovering this issue, we tried working around the crash by creating a new LanguageModelSession per iteration with the same model. This also crashes, since the CoreAIPipelinedEngine is shared across sessions on the same CoreAILanguageModel instance — confirming the fix must be at the engine level, not the session level. creates a single LanguageModelSession with tools and calls respond() repeatedly — evaluating each response against validation criteria and re-prompting with feedback until the output converges or a max iteration count is reached.