You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
DFlash bridge keeps emitting tokens after a stop token in the same block #265
Found while reviewing #262 (merged as 7858bf5, fixes #261). Related to #224, which made the stop path cancel the DFlash stream.
Behaviour. The stop-token branch in dflashGenerationStream (Sources/SwiftLM/Server.swift, the .token case) calls onStopToken. The default handler returns .more, so the generation is not ended and the bridge keeps forwarding the committed tokens of the same DFlash block to onToken. The standard path (Libraries/MLXLMCommon/Evaluate.swift) ends on any stop id, whatever the handler's disposition. So under --dflash, text after the end-of-turn token can reach the client, and a <tool_call> written after the stop token can be parsed as a tool call.
Reproduction. On main at 7858bf5, a scratch test in DFlashStreamBridgeTests (not committed) feeds .token 5, then stop 99, then 7 and 8, with stopTokenIDs: [99] and the default SpellingHandler:
continuation.yield(.token(tokenID:5,generatedTokens:1,acceptanceRatio:0,cyclesCompleted:1))continuation.yield(.token(tokenID:99,generatedTokens:2,acceptanceRatio:0,cyclesCompleted:1))continuation.yield(.token(tokenID:7,generatedTokens:3,acceptanceRatio:0,cyclesCompleted:1))continuation.yield(.token(tokenID:8,generatedTokens:4,acceptanceRatio:0,cyclesCompleted:1))
// ... summary, finish ...
XCTAssertEqual(chunks,["t5 "],"tokens after the stop token were emitted")
Result: chunks == ["t5 ", "t7 ", "t8 "], so the test fails. The existing testStopTokenGoesToOnStopTokenNotOnToken ends on the stop token, so it cannot catch this.
Suggested fix. Treat a stop id as the end of the generation, unless the handler already cancelled:
letisStop=stopTokenIDs.contains(tokenID)letdisposition=isStop?handler.onStopToken(tokenID,logProbabilities:nil,emit:emit):handler.onToken(tokenID,logProbabilities:nil,emit:emit)if case .more=disposition,!isStop{continue}
// stop ids fall through to the existing finish block: onGenerationEnd, .info with .stop, cancel
Add the regression test above to DFlashStreamBridgeTests. It should also check that the bridge yields exactly one .info with stopReason: .stop.
Severity: medium. Only tokens from the same DFlash block (up to blockTokens - 1) can leak, but the leak includes tool-call text.
Activity