diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index 54146615..26b01df1 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -1498,7 +1498,8 @@ struct MLXServer: AsyncParsableCommand { } // ── Build Hummingbird router ── - let router = Router() + // The non-streaming handlers need the connection to notice a client leaving. + let router = Router(context: ConnectionRequestContext.self) // ── CORS middleware ── if let origin = corsOrigin { @@ -1587,7 +1588,7 @@ struct MLXServer: AsyncParsableCommand { } } memoryPressure.resume() - router.post("/v1/chat/completions") { request, _ -> Response in + router.post("/v1/chat/completions") { request, context -> Response in do { let bodyData = try await collectBody(request) return try await handleChatCompletion( @@ -1595,7 +1596,8 @@ struct MLXServer: AsyncParsableCommand { draftModelRef: draftModelRef, numDraftTokens: numDraftTokensConfig, dflashModel: dflashModel, dflashBlockSize: dflashBlockSizeConfig, dflashTargetModel: dflashTargetModel, - mtpAssistant: loadedMTPAssistant + mtpAssistant: loadedMTPAssistant, + channel: context.channel ) } catch { return Response( @@ -1607,11 +1609,12 @@ struct MLXServer: AsyncParsableCommand { } // Text completions — handler extracted to avoid type-checker timeout - router.post("/v1/completions") { request, _ -> Response in + router.post("/v1/completions") { request, context -> Response in do { let bodyData = try await collectBody(request) return try await handleTextCompletion( - request: request, bodyData: bodyData, config: config, container: container, semaphore: semaphore, stats: stats + request: request, bodyData: bodyData, config: config, container: container, semaphore: semaphore, stats: stats, + channel: context.channel ) } catch { return Response( @@ -2475,7 +2478,8 @@ func handleChatCompletion( dflashModel: DFlashDraftModel? = nil, dflashBlockSize: Int? = nil, dflashTargetModel: (any DFlashTargetModel)? = nil, - mtpAssistant: (any DualModelMTP)? = nil + mtpAssistant: (any DualModelMTP)? = nil, + channel: (any Channel)? = nil ) async throws -> Response { let chatReq: ChatCompletionRequest do { @@ -2941,14 +2945,21 @@ func handleChatCompletion( emitPrefillProgress: emitPrefillProgress ) } else { - let (stream, onPrefillDone) = try await startGeneration() - return try await handleChatNonStreaming( - stream: stream, modelId: modelId, stopSequences: stopSequences, - promptTokenCount: promptTokenCount, enableThinking: enableThinking, - thinkingPreOpened: thinkingPreOpened, jsonMode: jsonMode, tokenEcho: config.tokenEcho, - slot: slot, stats: stats, genStart: genStart, prefillStart: prefillStart, - onPrefillDone: onPrefillDone - ) + // A client may give up during a long answer, and nothing on this path would + // notice: the SSE path learns of a disconnect from a failed write, but a + // non-streaming response writes nothing until the end. The generation then ran + // to max_tokens for nobody and held the slot the whole time (and with the + // default --parallel 1, every later request behind it). + return try await cancellingOnClientDisconnect(channel) { [enableThinking, thinkingPreOpened] in + let (stream, onPrefillDone) = try await startGeneration() + return try await handleChatNonStreaming( + stream: stream, modelId: modelId, stopSequences: stopSequences, + promptTokenCount: promptTokenCount, enableThinking: enableThinking, + thinkingPreOpened: thinkingPreOpened, jsonMode: jsonMode, tokenEcho: config.tokenEcho, + slot: slot, stats: stats, genStart: genStart, prefillStart: prefillStart, + onPrefillDone: onPrefillDone + ) + } } } @@ -3416,7 +3427,9 @@ func handleChatStreaming( // it arrives. The prompt is saved here, not on the first chunk: that raced the // decode step the producer had already started on the same KVCacheSimple. An // early exit (stop sequence, disconnect) never reaches `.info` and skips the save. - if !firstToken, let onPrefillDone { await onPrefillDone() } + // The save does not wait for a chunk: a turn that is only a tool call never + // yields one, and its prompt is exactly what the next turn of the loop needs. + if let onPrefillDone { await onPrefillDone() } // The chunk count undercounts: tokens the decoder buffers (tool-call // bodies, partial text) emit no chunk. `info` carries the real count. completionTokenCount = info.generationTokenCount @@ -3592,7 +3605,7 @@ func handleChatNonStreaming( case .info(let info): // The producer's last yield; see the matching comment in `handleChatStreaming`. // The stop-sequence `break` above exits before `.info`, so it skips the save. - if !firstToken, let onPrefillDone { await onPrefillDone() } + if let onPrefillDone { await onPrefillDone() } generationStopReason = info.stopReason completionTokenCount = info.generationTokenCount // real count; see handleChatStreaming } @@ -3725,7 +3738,8 @@ func handleTextCompletion( config: ServerConfig, container: ModelContainer, semaphore: AsyncSemaphore, - stats: ServerStats + stats: ServerStats, + channel: (any Channel)? = nil ) async throws -> Response { let compReq: TextCompletionRequest do { @@ -3790,11 +3804,14 @@ func handleTextCompletion( genStart: genStart, emitPrefillProgress: emitPrefillProgress ) } else { - let stream = try await container.generate(input: lmInput, parameters: params) - return try await handleTextNonStreaming( - stream: stream, modelId: modelId, stopSequences: stopSequences, - promptTokenCount: promptTokenCount, slot: slot, stats: stats, genStart: genStart - ) + // Same as the chat path: stop generating if the client goes away. + return try await cancellingOnClientDisconnect(channel) { + let stream = try await container.generate(input: lmInput, parameters: params) + return try await handleTextNonStreaming( + stream: stream, modelId: modelId, stopSequences: stopSequences, + promptTokenCount: promptTokenCount, slot: slot, stats: stats, genStart: genStart + ) + } } } @@ -4105,6 +4122,106 @@ final class GenerationSlot: @unchecked Sendable { } } +// ── Client disconnect (non-streaming) ───────────────────────────────────────── + +/// Request context that keeps the connection's channel. `BasicRequestContext` drops it, +/// and the non-streaming handlers need it to find out when the client goes away. +struct ConnectionRequestContext: RequestContext { + var coreContext: CoreRequestContextStorage + let channel: any Channel + + init(source: ApplicationRequestContextSource) { + coreContext = .init(source: source) + channel = source.channel + } +} + +/// Reports, once, that the client of a connection has stopped listening. +/// +/// Hummingbird enables remote half-closure on its connections, so a client that closes +/// its socket mid-request shows up as `ChannelEvent.inputClosed`, not as a closed +/// channel; a reset shows up as `channelInactive`. The streaming handlers need neither, +/// because their next write fails and tears down the body stream, but a non-streaming +/// handler writes nothing until the end, so this is its only signal. +/// +/// Installed at the head of the pipeline for the span of one request and removed after +/// it, so a keep-alive connection closing between requests is not taken for a client +/// abandoning one. Both callbacks run on the channel's event loop, and `fired` is only +/// touched there. +final class ClientDisconnectWatch: ChannelInboundHandler, RemovableChannelHandler, @unchecked Sendable { + typealias InboundIn = ByteBuffer + + private let onDisconnect: @Sendable () -> Void + private var fired = false + + init(onDisconnect: @escaping @Sendable () -> Void) { + self.onDisconnect = onDisconnect + } + + func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) { + if (event as? ChannelEvent) == .inputClosed { fire() } + context.fireUserInboundEventTriggered(event) + } + + func channelInactive(context: ChannelHandlerContext) { + fire() + context.fireChannelInactive() + } + + private func fire() { + guard !fired else { return } + fired = true + onDisconnect() + } +} + +/// A cancel request that may arrive before the task it is meant for exists. +final class DeferredCancellation: @unchecked Sendable { + private let lock = NSLock() + private var cancelled = false + private var cancelTask: (@Sendable () -> Void)? + + func cancel() { + lock.lock() + cancelled = true + let cancelTask = self.cancelTask + lock.unlock() + cancelTask?() + } + + func attach(_ task: Task) { + lock.lock() + cancelTask = { task.cancel() } + let cancelled = self.cancelled + lock.unlock() + if cancelled { task.cancel() } + } +} + +/// Runs `body`, a non-streaming generation, and cancels it if the client disconnects +/// meanwhile. Cancellation reaches the producer through the generation stream, the +/// same way it does when an SSE consumer is cancelled: the model stops after the token +/// in flight, the handler falls out of its loop and releases the slot, and the response +/// it builds goes to a closed socket. With no channel (tests), `body` just runs. +func cancellingOnClientDisconnect( + _ channel: (any Channel)?, _ body: @escaping @Sendable () async throws -> T +) async throws -> T { + guard let channel else { return try await body() } + let cancellation = DeferredCancellation() + let watch = ClientDisconnectWatch { cancellation.cancel() } + // Installed before the work starts, so a disconnect in between is not missed. If + // the channel is already closed the install fails, and the client has left. + do { + try await channel.pipeline.addHandler(watch, position: .first).get() + } catch { + cancellation.cancel() + } + defer { channel.pipeline.removeHandler(watch, promise: nil) } + let work = Task { try await body() } + cancellation.attach(work) + return try await work.value +} + // ── CORS Middleware ─────────────────────────────────────────────────────────── struct CORSMiddleware: RouterMiddleware { diff --git a/tests/SwiftLMTests/ClientDisconnectTests.swift b/tests/SwiftLMTests/ClientDisconnectTests.swift new file mode 100644 index 00000000..9612b19e --- /dev/null +++ b/tests/SwiftLMTests/ClientDisconnectTests.swift @@ -0,0 +1,181 @@ +import XCTest +import Foundation +import Hummingbird +import MLXLMCommon +import NIOCore +import NIOEmbedded +@testable import SwiftLM + +// MARK: - Non-streaming requests stop when the client disconnects +// +// The SSE handlers cancel generation when the client goes away, because their next write +// fails and tears down the body stream. A non-streaming handler writes nothing until the +// end, so the server never noticed the client had left: the answer ran to max_tokens for +// nobody and held the generation slot the whole time. These drive the pieces with an +// EmbeddedChannel and a synthetic generation stream, no model needed. +final class ClientDisconnectTests: XCTestCase { + + // MARK: The watch + + /// Counts calls from any thread; the watch fires on the event loop. + private final class Counter: @unchecked Sendable { + private let lock = NSLock() + private var count = 0 + func increment() { lock.lock(); count += 1; lock.unlock() } + var value: Int { lock.lock(); defer { lock.unlock() }; return count } + } + + /// Records the user events that reach the tail of the pipeline. + private final class EventRecorder: ChannelInboundHandler { + typealias InboundIn = ByteBuffer + var events: [Any] = [] + func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) { + events.append(event) + context.fireUserInboundEventTriggered(event) + } + } + + private func installedWatch(on channel: EmbeddedChannel, fired: Counter) throws -> EventRecorder { + let watch = ClientDisconnectWatch { fired.increment() } + let recorder = EventRecorder() + try channel.pipeline.syncOperations.addHandler(watch, position: .first) + try channel.pipeline.syncOperations.addHandler(recorder, position: .last) + return recorder + } + + func testWatchFiresOnceOnInputClosedAndForwardsTheEvent() throws { + let channel = EmbeddedChannel() + let fired = Counter() + let recorder = try installedWatch(on: channel, fired: fired) + + channel.pipeline.fireUserInboundEventTriggered(ChannelEvent.inputClosed) + XCTAssertEqual(fired.value, 1) + XCTAssertEqual(recorder.events.count, 1, "the event must keep flowing to the handlers behind the watch") + + // Neither a repeat nor the eventual close fires it again. + channel.pipeline.fireUserInboundEventTriggered(ChannelEvent.inputClosed) + _ = try channel.finish() + XCTAssertEqual(fired.value, 1) + } + + func testWatchFiresOnChannelInactive() throws { + let channel = EmbeddedChannel() + let fired = Counter() + _ = try installedWatch(on: channel, fired: fired) + + try channel.close().wait() + XCTAssertEqual(fired.value, 1) + } + + func testWatchIgnoresOtherEvents() throws { + let channel = EmbeddedChannel() + let fired = Counter() + let recorder = try installedWatch(on: channel, fired: fired) + + channel.pipeline.fireUserInboundEventTriggered(ChannelEvent.outputClosed) + XCTAssertEqual(fired.value, 0) + XCTAssertEqual(recorder.events.count, 1) + } + + // MARK: The wrapper around a non-streaming generation + + /// An open-ended generation: a chunk every few milliseconds until the consumer goes. + /// `terminated` is fulfilled when the stream is torn down, which is what cancels + /// the real producer. + private func endlessGeneration(terminated: XCTestExpectation) -> (AsyncStream, Task) { + let (stream, continuation) = AsyncStream.makeStream() + continuation.onTermination = { _ in terminated.fulfill() } + let producer = Task { + var token = 0 + while !Task.isCancelled { + if case .terminated = continuation.yield(.chunk("x", tokenId: token)) { return } + token += 1 + try? await Task.sleep(for: .milliseconds(2)) + } + } + return (stream, producer) + } + + private func acquiredSlot(_ semaphore: AsyncSemaphore) async -> GenerationSlot { + await semaphore.wait() + return GenerationSlot(semaphore: semaphore) + } + + func testNonStreamingGenerationStopsWhenTheClientDisconnects() async throws { + let channel = EmbeddedChannel() + let terminated = expectation(description: "generation stream torn down") + let (stream, producer) = endlessGeneration(terminated: terminated) + defer { producer.cancel() } + let semaphore = AsyncSemaphore(limit: 1) + let slot = await acquiredSlot(semaphore) + let stats = ServerStats() + await stats.requestStarted() + + let response = Task { + try await cancellingOnClientDisconnect(channel) { + try await handleChatNonStreaming( + stream: stream, modelId: "test-model", stopSequences: [], promptTokenCount: 3, + tokenEcho: false, slot: slot, stats: stats, genStart: Date(), prefillStart: Date()) + } + } + // Let it consume a few chunks, then the client goes away. + try await Task.sleep(for: .milliseconds(40)) + channel.pipeline.fireUserInboundEventTriggered(ChannelEvent.inputClosed) + + _ = try await response.value + await fulfillment(of: [terminated], timeout: 2) + + // The slot is free again and the request is no longer counted as active. + let released = Task { await semaphore.wait(); return true } + let reacquired = await withTaskGroup(of: Bool.self) { group -> Bool in + group.addTask { await released.value } + group.addTask { try? await Task.sleep(for: .seconds(2)); return false } + let first = await group.next() ?? false + group.cancelAll() + return first + } + XCTAssertTrue(reacquired, "the generation slot was not released after the client disconnected") + let snapshot = await stats.snapshot() + XCTAssertEqual(snapshot.requestsActive, 0) + } + + func testNonStreamingGenerationCompletesNormallyAndRemovesTheWatch() async throws { + let channel = EmbeddedChannel() + let (stream, continuation) = AsyncStream.makeStream() + for token in 1...3 { continuation.yield(.chunk("word ", tokenId: token)) } + continuation.yield(.info(GenerateCompletionInfo( + promptTokenCount: 3, generationTokenCount: 3, promptTime: 0, generationTime: 0, + stopReason: .stop))) + continuation.finish() + let semaphore = AsyncSemaphore(limit: 1) + let slot = await acquiredSlot(semaphore) + let stats = ServerStats() + await stats.requestStarted() + + let response = try await cancellingOnClientDisconnect(channel) { + try await handleChatNonStreaming( + stream: stream, modelId: "test-model", stopSequences: [], promptTokenCount: 3, + tokenEcho: false, slot: slot, stats: stats, genStart: Date(), prefillStart: Date()) + } + XCTAssertEqual(response.status, .ok) + + // The watch only covers the request: once the answer is built it is gone, so + // the keep-alive connection closing later touches nothing. + XCTAssertThrowsError(try channel.pipeline.syncOperations.handler(type: ClientDisconnectWatch.self)) + channel.pipeline.fireUserInboundEventTriggered(ChannelEvent.inputClosed) + _ = try channel.finish() + } + + func testWithoutAChannelTheBodyJustRuns() async throws { + let value = try await cancellingOnClientDisconnect(nil) { 42 } + XCTAssertEqual(value, 42) + } + + func testDeferredCancellationCancelsATaskAttachedAfterTheRequest() async { + let cancellation = DeferredCancellation() + cancellation.cancel() + let task = Task { try await Task.sleep(for: .seconds(10)) } + cancellation.attach(task) + if case .success = await task.result { XCTFail("the task should have been cancelled") } + } +} diff --git a/tests/SwiftLMTests/ToolCallPromptCacheTests.swift b/tests/SwiftLMTests/ToolCallPromptCacheTests.swift new file mode 100644 index 00000000..df990a20 --- /dev/null +++ b/tests/SwiftLMTests/ToolCallPromptCacheTests.swift @@ -0,0 +1,140 @@ +import XCTest +import Foundation +import Hummingbird +import HTTPTypes +import MLXLMCommon +import NIOCore +@testable import SwiftLM + +// MARK: - A tool-call-only turn saves the prompt cache (#245) +// +// Both chat handlers save the prompt cache when the producer's `.info` arrives, and used +// to do so only once a text chunk had been seen. A turn the model answers with nothing +// but a tool call never yields a chunk: the parser holds the whole `` body and +// emits one `.toolCall`. So the save was skipped and the next turn of the agent loop, +// tool result and follow-up, re-prefilled the whole conversation. These drive the +// handlers with a synthetic generation stream, no model needed. +final class ToolCallPromptCacheTests: XCTestCase { + + private final class Counter: @unchecked Sendable { + private let lock = NSLock() + private var count = 0 + func increment() { lock.lock(); count += 1; lock.unlock() } + var value: Int { lock.lock(); defer { lock.unlock() }; return count } + } + + private final class BodySink: @unchecked Sendable { + var data = Data() + } + + private struct CollectingWriter: ResponseBodyWriter { + let sink: BodySink + mutating func write(_ buffer: ByteBuffer) async throws { + sink.data.append(contentsOf: buffer.readableBytesView) + } + consuming func finish(_ trailingHeaders: HTTPFields?) async throws {} + } + + private func toolCall() -> Generation { + .toolCall(ToolCall(function: .init(name: "get_weather", arguments: ["city": "Paris"]))) + } + + private func info(tokens: Int) -> Generation { + .info(GenerateCompletionInfo( + promptTokenCount: 3, generationTokenCount: tokens, + promptTime: 0, generationTime: 0, stopReason: .stop)) + } + + private func finishedStream(_ events: [Generation]) -> AsyncStream { + let (stream, continuation) = AsyncStream.makeStream() + for event in events { continuation.yield(event) } + continuation.finish() + return stream + } + + private func acquiredSlot() async -> GenerationSlot { + let semaphore = AsyncSemaphore(limit: 1) + await semaphore.wait() + return GenerationSlot(semaphore: semaphore) + } + + /// Runs `events` through the non-streaming handler and returns how many times the + /// prompt-cache save ran, plus the tool calls the response carries. + private func nonStreamingSaves(for events: [Generation]) async throws -> (saves: Int, toolCalls: Int) { + let saves = Counter() + let slot = await acquiredSlot() + let response = try await handleChatNonStreaming( + stream: finishedStream(events), modelId: "test-model", stopSequences: [], + promptTokenCount: 3, tokenEcho: false, slot: slot, stats: ServerStats(), + genStart: Date(), prefillStart: Date(), onPrefillDone: { saves.increment() }) + let sink = BodySink() + try await response.body.write(CollectingWriter(sink: sink)) + let json = try XCTUnwrap(JSONSerialization.jsonObject(with: sink.data) as? [String: Any]) + let choices = try XCTUnwrap(json["choices"] as? [[String: Any]]) + let message = try XCTUnwrap(choices.first?["message"] as? [String: Any]) + let toolCalls = (message["tool_calls"] as? [[String: Any]]) ?? [] + return (saves.value, toolCalls.count) + } + + /// Same for the streaming handler; the body is drained so the consumer runs to the end. + private func streamingSaves(for events: [Generation]) async throws -> (saves: Int, toolCalls: Int) { + let saves = Counter() + let stream = finishedStream(events) + let slot = await acquiredSlot() + let response = handleChatStreaming( + startGeneration: { (stream, { saves.increment() }) }, modelId: "test-model", + stopSequences: [], includeUsage: false, promptTokenCount: 3, tokenEcho: false, + slot: slot, stats: ServerStats(), genStart: Date(), prefillStart: Date(), + emitPrefillProgress: false) + let sink = BodySink() + try await response.body.write(CollectingWriter(sink: sink)) + let body = String(decoding: sink.data, as: UTF8.self) + var toolCalls = 0 + for raw in body.components(separatedBy: "\n") { + let line = raw.trimmingCharacters(in: .whitespacesAndNewlines) + guard line.hasPrefix("data: {") else { continue } + let event = try XCTUnwrap( + JSONSerialization.jsonObject(with: Data(line.dropFirst(6).utf8)) as? [String: Any]) + let choices = (event["choices"] as? [[String: Any]]) ?? [] + if let delta = choices.first?["delta"] as? [String: Any], + let calls = delta["tool_calls"] as? [[String: Any]] { + toolCalls += calls.count + } + } + return (saves.value, toolCalls) + } + + // MARK: The bug itself + + func testNonStreamingToolCallOnlyTurnSavesThePrompt() async throws { + let result = try await nonStreamingSaves(for: [toolCall(), info(tokens: 12)]) + XCTAssertEqual(result.toolCalls, 1) + XCTAssertEqual(result.saves, 1, "a turn that is only a tool call must still save its prompt") + } + + func testStreamingToolCallOnlyTurnSavesThePrompt() async throws { + let result = try await streamingSaves(for: [toolCall(), info(tokens: 12)]) + XCTAssertEqual(result.toolCalls, 1) + XCTAssertEqual(result.saves, 1, "a turn that is only a tool call must still save its prompt") + } + + // MARK: What stays the same + + func testTextTurnsStillSaveOnce() async throws { + let text: [Generation] = [.chunk("Sunny", tokenId: 1), .chunk(".", tokenId: 2), info(tokens: 2)] + let nonStreaming = try await nonStreamingSaves(for: text) + XCTAssertEqual(nonStreaming.saves, 1) + let streaming = try await streamingSaves(for: text) + XCTAssertEqual(streaming.saves, 1) + } + + /// An early exit (stop sequence, disconnect, error) never reaches `.info` and must + /// keep skipping the save: the cache may still be mid-write. + func testAStreamThatEndsWithoutInfoDoesNotSave() async throws { + let cut: [Generation] = [toolCall()] + let nonStreaming = try await nonStreamingSaves(for: cut) + XCTAssertEqual(nonStreaming.saves, 0) + let streaming = try await streamingSaves(for: cut) + XCTAssertEqual(streaming.saves, 0) + } +} diff --git a/tests/test-contract.sh b/tests/test-contract.sh index 25fbb74c..35c17556 100755 --- a/tests/test-contract.sh +++ b/tests/test-contract.sh @@ -407,6 +407,96 @@ else fail "json_object stream differs from non-streaming: '$(echo "$ST_JSON" | head -c 60)' vs '$(echo "$NS_JSON" | head -c 60)'" fi +# ── 14. An abandoned non-streaming request stops generating ─────────────────── +# The SSE handlers cancel generation when the client goes away (their next write fails). +# A non-streaming request writes nothing until the end, so the server never noticed the +# client had left: the answer ran to max_tokens for nobody and held the slot, and with +# the default --parallel 1 every later request queued behind it. Clients that time out +# and retry (the OpenAI SDKs' default) stacked up several such generations. +log "Test 14: a non-streaming request the client abandons stops generating" +metric() { curl -sf "$URL/metrics" | awk -v m="$1" '$1 == m { print $2 }'; } +TOKENS_BEFORE=$(metric swiftlm_tokens_generated_total) +# curl closes the socket after half a second; left alone, the request would generate 1500 tokens. +curl -s --max-time 0.5 "$URL/v1/chat/completions" -H 'Content-Type: application/json' \ + -d "{\"model\":\"x\",\"messages\":[{\"role\":\"user\",\"content\":\"$LONG_ASK\"}],\"max_tokens\":1500,\"temperature\":0.9}" >/dev/null 2>&1 || true +# Once the server has noticed, the request is no longer active and its token count is the +# few chunks it managed before the cancel, not a full answer. +ACTIVE=1 +for _ in $(seq 1 20); do + ACTIVE=$(metric swiftlm_requests_active) + [ "$ACTIVE" = "0" ] && break + sleep 0.25 +done +TOKENS_AFTER=$(metric swiftlm_tokens_generated_total) +ABANDONED_TOKENS=$((TOKENS_AFTER - TOKENS_BEFORE)) +if [ "$ACTIVE" = "0" ] && [ "$ABANDONED_TOKENS" -lt 400 ]; then + pass "the abandoned request stopped after $ABANDONED_TOKENS tokens and released its slot" +else + fail "the abandoned request generated $ABANDONED_TOKENS tokens; requests_active=$ACTIVE 5 s after the client left" +fi + +# ── 15. A tool-call-only turn saves the prompt cache ───────────────────────── +# #245: the save at `.info` waited for a text chunk, and a turn the model answers with +# nothing but a tool call never produces one. The next turn of an agent loop (tool +# result, follow-up) then re-prefilled the whole conversation. The server log shows the +# hit: the follow-up must reuse at least the whole prompt of the tool-call turn. +log "Test 15: a tool-call-only turn saves the prompt cache" +SERVER_LOG=/tmp/SwiftLM-test-contract.log +hit_count() { grep -c 'Prompt cache HIT' "$SERVER_LOG" || true; } +last_hit_reused() { grep -o 'Prompt cache HIT[^:]*: [0-9]*' "$SERVER_LOG" | tail -1 | grep -o '[0-9]*$'; } +# Reads a non-streaming response or an SSE stream (with include_usage) and prints +# " ". +tool_turn_summary() { + python3 -c ' +import json, sys +raw = sys.stdin.read() +if raw.lstrip().startswith("{"): + r = json.loads(raw); m = r["choices"][0]["message"] + name = ((m.get("tool_calls") or [{}])[0].get("function") or {}).get("name", "") + print(name, r.get("usage", {}).get("prompt_tokens", 0)); sys.exit() +name, prompt = "", 0 +for line in raw.splitlines(): + if not line.startswith("data: {"): continue + ev = json.loads(line[6:]) + for c in ev.get("choices") or []: + for tc in (c.get("delta") or {}).get("tool_calls") or []: + name = name or (tc.get("function") or {}).get("name", "") + prompt = (ev.get("usage") or {}).get("prompt_tokens") or prompt +print(name, prompt)' +} +# $1: label, $2: a system prompt no other test uses, $3: extra fields for the tool-call turn. +tool_turn_then_follow_up() { + local LABEL="$1" SYS="$2" EXTRA="$3" + local HISTORY="{\"role\":\"system\",\"content\":\"$SYS\"},{\"role\":\"user\",\"content\":\"What is the weather in Paris? Use the tool.\"}" + local SUMMARY NAME PROMPT_TOKENS HITS_BEFORE REUSED + SUMMARY=$(chat "{\"model\":\"x\",\"messages\":[$HISTORY],\"tools\":$TOOLS,\"max_tokens\":200,\"temperature\":0$EXTRA}" | tool_turn_summary) + NAME=${SUMMARY%% *}; PROMPT_TOKENS=${SUMMARY##* } + if [ "$NAME" != "get_weather" ]; then + skip "$LABEL: model did not answer with a tool call (capability, not contract)" + return + fi + HITS_BEFORE=$(hit_count) + # The follow-up replays the tool call and adds its result; its prompt starts with the + # whole prompt of the tool-call turn, which must now be in the cache. + chat "{\"model\":\"x\",\"messages\":[$HISTORY,{\"role\":\"assistant\",\"content\":\"\",\"tool_calls\":[{\"id\":\"call_15\",\"type\":\"function\",\"function\":{\"name\":\"get_weather\",\"arguments\":\"{\\\"city\\\": \\\"Paris\\\"}\"}}]},{\"role\":\"tool\",\"tool_call_id\":\"call_15\",\"content\":\"18 C, cloudy\"}],\"tools\":$TOOLS,\"max_tokens\":40,\"temperature\":0}" >/dev/null + for _ in $(seq 1 20); do + [ "$(hit_count)" -gt "$HITS_BEFORE" ] && break + sleep 0.25 + done + if [ "$(hit_count)" -le "$HITS_BEFORE" ]; then + fail "$LABEL: the follow-up after a tool-call turn got no prompt cache hit at all" + return + fi + REUSED=$(last_hit_reused) + if [ "${REUSED:-0}" -ge "$PROMPT_TOKENS" ]; then + pass "$LABEL: the follow-up reused $REUSED cached tokens, the tool-call turn's prompt was $PROMPT_TOKENS" + else + fail "$LABEL: the follow-up reused only $REUSED cached tokens; the tool-call turn's prompt was $PROMPT_TOKENS" + fi +} +tool_turn_then_follow_up "non-streaming" "You are a weather desk. Call get_weather for any weather question and say nothing else." "" +tool_turn_then_follow_up "streaming" "You are a weather kiosk. Call get_weather for any weather question and say nothing else." ',"stream":true,"stream_options":{"include_usage":true}' + log "═══════════════════════════════════════" log "Results: $PASS passed, $FAIL failed, $SKIP skipped, $TOTAL total" log "═══════════════════════════════════════"