feat: stream LLM tokens to the extension via /v1/ai/stream

- LocalEngineServer: add a streaming AI endpoint (newline-delimited JSON frames)
  over a no-Content-Length chunked response, plus an AIStreamHandler.
- AppState.handleServerAIStream: gate + iterate AIEngine.runStreaming, emitting
  each token — gives the runStreaming infrastructure its first real consumer.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
BigMarc 2026-06-20 20:26:22 +02:00
parent 201fe04d8a
commit 2f3dd02148
2 changed files with 68 additions and 1 deletions

View File

@ -180,6 +180,10 @@ final class AppState: ObservableObject {
aiHandler: { [weak self] (req: LocalEngineServer.AIRequest) async -> LocalEngineServer.AIOutcome in aiHandler: { [weak self] (req: LocalEngineServer.AIRequest) async -> LocalEngineServer.AIOutcome in
guard let self else { return .error("GrammarGem is shutting down.") } guard let self else { return .error("GrammarGem is shutting down.") }
return await self.handleServerAI(req) return await self.handleServerAI(req)
},
aiStreamHandler: { [weak self] (req, onChunk) -> LocalEngineServer.AIOutcome in
guard let self else { return .error("GrammarGem is shutting down.") }
return await self.handleServerAIStream(req, onChunk: onChunk)
}) })
server.start() server.start()
localServer = server localServer = server
@ -198,6 +202,24 @@ final class AppState: ObservableObject {
} }
} }
/// Streaming variant for the extension: yields tokens as they generate.
private func handleServerAIStream(_ req: LocalEngineServer.AIRequest,
onChunk: @escaping (String) -> Void) async -> LocalEngineServer.AIOutcome {
if case .denied(let reason) = gate.authorizeAI(req.action) { return .error(reason) }
guard ai.isReady else { return .error(AIEngineError.modelNotReady.localizedDescription) }
var accumulated = ""
do {
for try await chunk in ai.runStreaming(req.action, on: req.text, protectedTerms: dictionary.entries) {
accumulated += chunk
onChunk(chunk)
}
gate.recordAIActionUsed()
return .ok(accumulated)
} catch {
return .error(error.localizedDescription)
}
}
// MARK: - Hotkey handling // MARK: - Hotkey handling
/// (Re)register the global hotkeys from the current preferences. /// (Re)register the global hotkeys from the current preferences.

View File

@ -36,22 +36,28 @@ final class LocalEngineServer {
enum AIOutcome { case ok(String); case error(String) } enum AIOutcome { case ok(String); case error(String) }
/// AI action: runs the gate + model, returns the result or a user-facing error. /// AI action: runs the gate + model, returns the result or a user-facing error.
typealias AIHandler = (AIRequest) async -> AIOutcome typealias AIHandler = (AIRequest) async -> AIOutcome
/// Streaming AI action: runs the gate + model, invoking `onChunk` per token as it
/// generates, then returns the final outcome (full text, or an error message).
typealias AIStreamHandler = (AIRequest, @escaping (String) -> Void) async -> AIOutcome
private let port: UInt16 private let port: UInt16
private let token: String private let token: String
private let grammarHandler: GrammarHandler private let grammarHandler: GrammarHandler
private let aiHandler: AIHandler private let aiHandler: AIHandler
private let aiStreamHandler: AIStreamHandler
private var listener: NWListener? private var listener: NWListener?
private let queue = DispatchQueue(label: "com.grammargem.localserver") private let queue = DispatchQueue(label: "com.grammargem.localserver")
init(port: UInt16, token: String, init(port: UInt16, token: String,
grammarHandler: @escaping GrammarHandler, grammarHandler: @escaping GrammarHandler,
aiHandler: @escaping AIHandler) { aiHandler: @escaping AIHandler,
aiStreamHandler: @escaping AIStreamHandler) {
self.port = port self.port = port
self.token = token self.token = token
self.grammarHandler = grammarHandler self.grammarHandler = grammarHandler
self.aiHandler = aiHandler self.aiHandler = aiHandler
self.aiStreamHandler = aiStreamHandler
} }
/// A stable per-install token, persisted so the extension can be paired once. /// A stable per-install token, persisted so the extension can be paired once.
@ -145,6 +151,23 @@ final class LocalEngineServer {
case .error(let msg): sendJSON(conn, status: 422, ErrorResponse(error: msg)) case .error(let msg): sendJSON(conn, status: 422, ErrorResponse(error: msg))
} }
case ("POST", "/v1/ai/stream"):
guard let req = decodeAI(request.body) else {
sendJSON(conn, status: 400, ErrorResponse(error: "bad request")); return
}
// Newline-delimited JSON frames: {"chunk":"…"} per token, then a final
// {"done":true} or {"error":"…"}. The 200 header goes out before tokens,
// so an entitlement/readiness failure surfaces as a closing error frame.
beginStream(conn)
let outcome = await aiStreamHandler(req) { [weak self] chunk in
self?.streamFrame(conn, ["chunk": chunk])
}
switch outcome {
case .ok: streamFrame(conn, ["done": true])
case .error(let msg): streamFrame(conn, ["error": msg])
}
endStream(conn)
default: default:
sendJSON(conn, status: 404, ErrorResponse(error: "not found")) sendJSON(conn, status: 404, ErrorResponse(error: "not found"))
} }
@ -201,6 +224,28 @@ final class LocalEngineServer {
send(conn, status: status, body: body) send(conn, status: status, body: body)
} }
// MARK: - Streaming responses (newline-delimited JSON, no Content-Length)
private func beginStream(_ conn: NWConnection) {
var head = "HTTP/1.1 200 OK\r\n"
head += "Content-Type: application/x-ndjson\r\n"
head += "Cache-Control: no-cache\r\n"
head += "Access-Control-Allow-Origin: *\r\n"
head += "Access-Control-Allow-Headers: Content-Type, X-GrammarGem-Token\r\n"
head += "Connection: close\r\n\r\n"
conn.send(content: Data(head.utf8), completion: .contentProcessed { _ in })
}
private func streamFrame(_ conn: NWConnection, _ object: [String: Any]) {
guard var data = try? JSONSerialization.data(withJSONObject: object) else { return }
data.append(0x0A) // newline delimiter
conn.send(content: data, completion: .contentProcessed { _ in })
}
private func endStream(_ conn: NWConnection) {
conn.send(content: nil, completion: .contentProcessed { _ in conn.cancel() })
}
private func send(_ conn: NWConnection, status: Int, body: Data) { private func send(_ conn: NWConnection, status: Int, body: Data) {
var head = "HTTP/1.1 \(status) \(Self.reason(status))\r\n" var head = "HTTP/1.1 \(status) \(Self.reason(status))\r\n"
head += "Content-Type: application/json\r\n" head += "Content-Type: application/json\r\n"