From 201fe04d8a99f927026a95567462bfcf33069295 Mon Sep 17 00:00:00 2001 From: BigMarc Date: Sat, 20 Jun 2026 20:21:29 +0200 Subject: [PATCH] feat: backend-neutral AI seam, Ollama backend, extension bridge + trust/correctness fixes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wave 1 (correctness/trust): - Fix paid large-model dead path: engine reads the persisted selected repo via a thread-safe ModelManager.activeRepo(), so download/load/isReady now agree. - Lemon Squeezy store/product check fails closed in release builds — the placeholder-ID 'any key unlocks paid' bypass can no longer ship. - Replace force-unwrapped Hugging Face URLs with safe URLComponents construction. - Never read/replace/scan secure (password) fields via AX.isSecureTextField. Wave 2 (compounding seam): - Extract AIPromptKit (prompts/temperature/maxTokens/cleanup) out of MLXEngine so every backend produces identical prompts and output. - AIEngine gains protectedTerms (personal dictionary now actually injected into the LLM), warmup(), and runStreaming(); add makeAIEngine() factory. - Honor WritingMode.lengthCap for maxTokens; warm the model on launch/after download. - Guard custom modes against empty prompts and duplicate names. Wave 3: OllamaEngine — a URLSession-only AIEngine backend (cross-platform, opt-in via GRAMMARGEM_AI_BACKEND=ollama). Wave 4: LocalEngineServer — loopback-only, token-authed HTTP bridge (Network.framework, no new deps), off by default, exposing Harper + the LLM + the entitlement gate to a browser extension. Tests: 6 new unit tests for the AIPrompts seam and custom-mode guard (19 pass). Co-Authored-By: Claude Opus 4.8 (1M context) --- Sources/GrammarGem/AI/AIEngine.swift | 45 ++- Sources/GrammarGem/AI/AIPromptKit.swift | 111 ++++++++ Sources/GrammarGem/AI/CustomModesStore.swift | 14 +- Sources/GrammarGem/AI/MLXEngine.swift | 117 +++----- Sources/GrammarGem/AI/ModelManager.swift | 39 ++- Sources/GrammarGem/AI/OllamaEngine.swift | 176 ++++++++++++ Sources/GrammarGem/App/AppState.swift | 92 +++++- Sources/GrammarGem/App/GrammarGemApp.swift | 2 +- Sources/GrammarGem/Core/AppConfig.swift | 17 ++ .../Licensing/LemonSqueezyClient.swift | 16 +- Sources/GrammarGem/System/AXSupport.swift | 11 + .../GrammarGem/System/LocalEngineServer.swift | 266 ++++++++++++++++++ Sources/GrammarGem/System/TextCapture.swift | 17 +- .../GrammarGem/System/TextFieldScanner.swift | 4 +- .../System/TextReplacementCoordinator.swift | 11 +- Tests/GrammarGemTests/GrammarGemTests.swift | 60 ++++ 16 files changed, 900 insertions(+), 98 deletions(-) create mode 100644 Sources/GrammarGem/AI/AIPromptKit.swift create mode 100644 Sources/GrammarGem/AI/OllamaEngine.swift create mode 100644 Sources/GrammarGem/System/LocalEngineServer.swift diff --git a/Sources/GrammarGem/AI/AIEngine.swift b/Sources/GrammarGem/AI/AIEngine.swift index cb03c01..1911059 100644 --- a/Sources/GrammarGem/AI/AIEngine.swift +++ b/Sources/GrammarGem/AI/AIEngine.swift @@ -10,11 +10,52 @@ protocol AIEngine: AnyObject { var isReady: Bool { get } /// Run an action against `text` and return the rewritten result. - /// Throws on model-not-ready or generation failure. Never performs network I/O. - func run(_ action: AIAction, on text: String) async throws -> String + /// `protectedTerms` are vocabulary the model must reproduce verbatim (the + /// personal dictionary / style-guide terms). Throws on model-not-ready or + /// generation failure. Never performs network I/O on the on-device backend. + func run(_ action: AIAction, on text: String, protectedTerms: [String]) async throws -> String /// Best-effort local tone classification (used by tone-detection UI). func detectTone(_ text: String) async -> Tone + + /// Optionally preload the model so the first real action is instant. + func warmup() async + + /// Streaming variant: yields incremental output as it generates. The default + /// implementation wraps `run` as a single chunk, so a backend that doesn't + /// natively stream still conforms. + func runStreaming(_ action: AIAction, on text: String, protectedTerms: [String]) + -> AsyncThrowingStream +} + +extension AIEngine { + /// Convenience: run with no protected terms. + func run(_ action: AIAction, on text: String) async throws -> String { + try await run(action, on: text, protectedTerms: []) + } + + /// Default: no preload. + func warmup() async {} + + /// Default streaming: produce the full result, then emit it as one chunk. + func runStreaming(_ action: AIAction, on text: String, protectedTerms: [String]) + -> AsyncThrowingStream { + AsyncThrowingStream { continuation in + Task { + do { + continuation.yield(try await run(action, on: text, protectedTerms: protectedTerms)) + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + } + } + + /// Convenience streaming overload with no protected terms. + func runStreaming(_ action: AIAction, on text: String) -> AsyncThrowingStream { + runStreaming(action, on: text, protectedTerms: []) + } } enum AIEngineError: LocalizedError { diff --git a/Sources/GrammarGem/AI/AIPromptKit.swift b/Sources/GrammarGem/AI/AIPromptKit.swift new file mode 100644 index 0000000..0194c08 --- /dev/null +++ b/Sources/GrammarGem/AI/AIPromptKit.swift @@ -0,0 +1,111 @@ +import Foundation + +/// Backend-neutral prompt construction, generation parameters, and output +/// cleanup for Layer-2 actions. Extracted from `MLXEngine` so EVERY `AIEngine` +/// backend (MLX today, Ollama/llama.cpp next, the localhost extension server) is +/// guaranteed to produce identical prompts and identical output trimming — no +/// behavioral drift between runtimes — and so the personal dictionary / a future +/// style guide can inject protected vocabulary in exactly ONE place. +/// +/// Everything here is a pure function of its inputs (no MLX, no I/O), which also +/// makes per-action behavior unit-testable without a model present. +enum AIPrompts { + + /// The system prompt for an action. `protectedTerms` (personal-dictionary + /// entries + future style-guide vocabulary) are terms the model must never + /// alter, so AI rewrites stop "correcting" brand names and jargon. + static func systemPrompt(for action: AIAction, protectedTerms: [String] = []) -> String { + base(for: action) + protectedTermsClause(protectedTerms) + } + + private static func base(for action: AIAction) -> String { + switch action { + case .rewriteClarity: + return "Rewrite the user's text to be clearer and more concise. Preserve the " + + "original meaning and the author's voice. Output only the rewritten text." + case .rewrite: + return "Paraphrase the user's text. Keep the meaning but vary the wording and " + + "structure. Output only the rewritten text." + case .adjustTone(let tone): + return "Rewrite the user's text in a \(tone.rawValue) tone. Preserve the meaning. " + + "Output only the rewritten text." + case .ask(let instruction): + return "Apply this instruction to the user's text: \(instruction). " + + "Output only the resulting text." + case .translate(let language): + return "Translate the user's text into \(language). Output only the translation." + case .applyMode(let mode): + return mode.systemPrompt + } + } + + /// Appended to every system prompt: terms the model must reproduce verbatim + /// (brand names, jargon the user added to their personal dictionary). Capped + /// so a large dictionary can't blow the context window. + private static func protectedTermsClause(_ terms: [String]) -> String { + let cleaned = terms + .map { $0.trimmingCharacters(in: .whitespacesAndNewlines) } + .filter { !$0.isEmpty } + guard !cleaned.isEmpty else { return "" } + let list = cleaned.prefix(50).map { "\"\($0)\"" }.joined(separator: ", ") + return " Never change, translate, or \"correct\" these exact terms — reproduce them " + + "verbatim, including capitalization: \(list)." + } + + /// Low temperature for faithful rewrites; a touch higher for free paraphrase. + static func temperature(for action: AIAction) -> Float { + switch action { + case .rewrite: return 0.5 + case .ask: return 0.4 + default: return 0.2 + } + } + + /// Generation token budget. Honors a Mode's `lengthCap` (in characters) for + /// short-form modes (Post = 280, Team Chat = 60) so we don't spend the full + /// 512-token default on a one-line message; falls back to the default otherwise. + static func maxTokens(for action: AIAction) -> Int { + let fallback = 512 + guard case .applyMode(let mode) = action, let cap = mode.lengthCap else { return fallback } + // ~1 token per 3 characters of English, plus headroom; clamp to a floor. + return min(fallback, max(64, cap / 3 + 48)) + } + + /// Strip a leading "Sure, here's …:" preamble and surrounding quotes that small + /// instruct models sometimes add, so the result drops cleanly into the field. + static func clean(_ s: String) -> String { + var out = s.trimmingCharacters(in: .whitespacesAndNewlines) + let lower = out.lowercased() + for prefix in ["sure, here", "sure! here", "here is", "here's", "certainly"] { + if lower.hasPrefix(prefix), let colon = out.firstIndex(of: ":") { + let after = out[out.index(after: colon)...] + .trimmingCharacters(in: .whitespacesAndNewlines) + if !after.isEmpty { out = after } + break + } + } + if out.count >= 2, out.first == "\"", out.last == "\"" { + out = String(out.dropFirst().dropLast()) + } + return out.trimmingCharacters(in: .whitespacesAndNewlines) + } + + /// The tone-classifier system prompt (kept here so any backend reuses it). + static func toneClassifierPrompt() -> String { + let options = Tone.allCases.map(\.rawValue).joined(separator: ", ") + return """ + You are a tone classifier. Read the user's text and respond with exactly ONE \ + word — the single best-matching tone from this list: \(options). \ + Output only that word, lowercase, no punctuation or explanation. + """ + } + + /// Parse the one-word tone classification output into a `Tone`. + static func parseTone(_ raw: String) -> Tone { + let token = raw.lowercased().trimmingCharacters(in: CharacterSet.alphanumerics.inverted) + if let exact = Tone(rawValue: token) { return exact } + let lower = raw.lowercased() + for tone in Tone.allCases where lower.contains(tone.rawValue) { return tone } + return .professional + } +} diff --git a/Sources/GrammarGem/AI/CustomModesStore.swift b/Sources/GrammarGem/AI/CustomModesStore.swift index 76629fc..29dd876 100644 --- a/Sources/GrammarGem/AI/CustomModesStore.swift +++ b/Sources/GrammarGem/AI/CustomModesStore.swift @@ -25,19 +25,27 @@ final class CustomModesStore: ObservableObject { /// Built-in + custom modes (custom ones are always paid). func allModes() -> [WritingMode] { ModeRegistry.all + customModes } - func addCustom(name: String, prompt: String) { + /// Returns false (no-op) if the name or prompt is empty, or the name + /// collides with an existing built-in/custom mode — so an empty-prompt mode + /// can't slip through to the LLM as an empty system prompt. + @discardableResult + func addCustom(name: String, prompt: String) -> Bool { let trimmed = name.trimmingCharacters(in: .whitespacesAndNewlines) - guard !trimmed.isEmpty else { return } + let trimmedPrompt = prompt.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty, !trimmedPrompt.isEmpty else { return false } + guard !allModes().contains(where: { $0.name.caseInsensitiveCompare(trimmed) == .orderedSame }) + else { return false } let mode = WritingMode( id: "custom-\(UUID().uuidString.prefix(8))", name: trimmed, - systemPrompt: prompt.trimmingCharacters(in: .whitespacesAndNewlines), + systemPrompt: trimmedPrompt, isPaid: true, lengthCap: nil, autoFormat: "Custom" ) customModes.append(mode) persistModes() + return true } func removeCustom(id: String) { diff --git a/Sources/GrammarGem/AI/MLXEngine.swift b/Sources/GrammarGem/AI/MLXEngine.swift index e2e2c96..d46c12f 100644 --- a/Sources/GrammarGem/AI/MLXEngine.swift +++ b/Sources/GrammarGem/AI/MLXEngine.swift @@ -33,33 +33,63 @@ final class MLXEngine: AIEngine { // MARK: - AIEngine - func run(_ action: AIAction, on text: String) async throws -> String { + func run(_ action: AIAction, on text: String, protectedTerms: [String]) async throws -> String { let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) guard !trimmed.isEmpty else { return text } let container = try await container() - let params = GenerateParameters(temperature: temperature(for: action)) + let params = GenerateParameters(temperature: AIPrompts.temperature(for: action)) let raw = try await complete( - system: systemPrompt(for: action), user: trimmed, - maxTokens: 512, parameters: params, container: container) - return clean(raw) + system: AIPrompts.systemPrompt(for: action, protectedTerms: protectedTerms), + user: trimmed, + maxTokens: AIPrompts.maxTokens(for: action), + parameters: params, container: container) + return AIPrompts.clean(raw) + } + + /// Preload the model container so the first real action doesn't pay the + /// multi-second load. Safe to call repeatedly (the loader caches). + func warmup() async { + _ = try? await container() + } + + func runStreaming(_ action: AIAction, on text: String, protectedTerms: [String]) + -> AsyncThrowingStream { + AsyncThrowingStream { continuation in + Task { + let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { continuation.yield(text); continuation.finish(); return } + do { + let container = try await container() + var params = GenerateParameters(temperature: AIPrompts.temperature(for: action)) + params.maxTokens = AIPrompts.maxTokens(for: action) + let system = AIPrompts.systemPrompt(for: action, protectedTerms: protectedTerms) + try await container.perform { context in + let input = try await context.processor.prepare( + input: UserInput(chat: [.system(system), .user(trimmed)])) + let stream = try MLXLMCommon.generate( + input: input, parameters: params, context: context) + for await item in stream { + if let chunk = item.chunk { continuation.yield(chunk) } + } + } + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + } } func detectTone(_ text: String) async -> Tone { let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) guard !trimmed.isEmpty else { return .professional } - let options = Tone.allCases.map(\.rawValue).joined(separator: ", ") - let system = """ - You are a tone classifier. Read the user's text and respond with exactly ONE \ - word — the single best-matching tone from this list: \(options). \ - Output only that word, lowercase, no punctuation or explanation. - """ do { let container = try await container() let raw = try await complete( - system: system, user: trimmed, maxTokens: 8, + system: AIPrompts.toneClassifierPrompt(), user: trimmed, maxTokens: 8, parameters: GenerateParameters(temperature: 0.0), container: container) - return parseTone(raw) + return AIPrompts.parseTone(raw) } catch { return .professional } @@ -98,64 +128,9 @@ final class MLXEngine: AIEngine { } } - // MARK: - Prompting - - private func systemPrompt(for action: AIAction) -> String { - switch action { - case .rewriteClarity: - return "Rewrite the user's text to be clearer and more concise. Preserve the " - + "original meaning and the author's voice. Output only the rewritten text." - case .rewrite: - return "Paraphrase the user's text. Keep the meaning but vary the wording and " - + "structure. Output only the rewritten text." - case .adjustTone(let tone): - return "Rewrite the user's text in a \(tone.rawValue) tone. Preserve the meaning. " - + "Output only the rewritten text." - case .ask(let instruction): - return "Apply this instruction to the user's text: \(instruction). " - + "Output only the resulting text." - case .translate(let language): - return "Translate the user's text into \(language). Output only the translation." - case .applyMode(let mode): - return mode.systemPrompt - } - } - - /// Low temperature for faithful rewrites; a touch higher for free paraphrase. - private func temperature(for action: AIAction) -> Float { - switch action { - case .rewrite: return 0.5 - case .ask: return 0.4 - default: return 0.2 - } - } - - // MARK: - Output cleanup - - private func clean(_ s: String) -> String { - var out = s.trimmingCharacters(in: .whitespacesAndNewlines) - let lower = out.lowercased() - for prefix in ["sure, here", "sure! here", "here is", "here's", "certainly"] { - if lower.hasPrefix(prefix), let colon = out.firstIndex(of: ":") { - let after = out[out.index(after: colon)...] - .trimmingCharacters(in: .whitespacesAndNewlines) - if !after.isEmpty { out = after } - break - } - } - if out.count >= 2, out.first == "\"", out.last == "\"" { - out = String(out.dropFirst().dropLast()) - } - return out.trimmingCharacters(in: .whitespacesAndNewlines) - } - - private func parseTone(_ raw: String) -> Tone { - let token = raw.lowercased().trimmingCharacters(in: CharacterSet.alphanumerics.inverted) - if let exact = Tone(rawValue: token) { return exact } - let lower = raw.lowercased() - for tone in Tone.allCases where lower.contains(tone.rawValue) { return tone } - return .professional - } + // Prompt construction, temperature, token budget, and output cleanup now live + // in `AIPrompts` (backend-neutral; see AIPromptKit.swift) so every AIEngine + // backend produces identical prompts and identical output trimming. } /// Caches one loaded `ModelContainer` (keyed by directory) and coalesces diff --git a/Sources/GrammarGem/AI/ModelManager.swift b/Sources/GrammarGem/AI/ModelManager.swift index c30d67c..97f814e 100644 --- a/Sources/GrammarGem/AI/ModelManager.swift +++ b/Sources/GrammarGem/AI/ModelManager.swift @@ -14,7 +14,9 @@ final class ModelManager: ObservableObject { @Published private(set) var state: State = .notDownloaded @Published private(set) var statusText: String = "" - @Published var selectedRepo: String = AppConfig.Model.defaultRepo + @Published var selectedRepo: String = ModelManager.activeRepo() { + didSet { UserDefaults.standard.set(selectedRepo, forKey: Self.selectedRepoKey) } + } private var task: Task? private var generation = 0 @@ -55,6 +57,33 @@ final class ModelManager: ObservableObject { isModelPresent(repo: repo) ? modelDir(repo) : nil } + /// UserDefaults key backing the user's chosen model repo. + nonisolated static let selectedRepoKey = "GrammarGem.selectedModelRepo" + + /// The user's currently-selected repo, readable from ANY thread (it lives in + /// UserDefaults, which is thread-safe). The AI engine's model-directory + /// provider reads this so download, load, and `isReady` all agree on the SAME + /// repo — including the paid large model, which the old hardcoded-default + /// provider silently ignored. Falls back to the default repo when unset. + nonisolated static func activeRepo() -> String { + UserDefaults.standard.string(forKey: selectedRepoKey) ?? AppConfig.Model.defaultRepo + } + + /// Safe URL construction for the two Hugging Face endpoints. A malformed + /// repo/path yields nil → a thrown `ModelError`, never a force-unwrap crash. + nonisolated static func fileURL(repo: String, path: String) -> URL? { + var c = URLComponents(string: AppConfig.Model.huggingFaceBase) + c?.path = "/\(repo)/resolve/main/\(path)" + return c?.url + } + + nonisolated static func treeURL(repo: String) -> URL? { + var c = URLComponents(string: AppConfig.Model.huggingFaceBase) + c?.path = "/api/models/\(repo)/tree/main" + c?.queryItems = [URLQueryItem(name: "recursive", value: "1")] + return c?.url + } + /// Kick off (or resume to) a real download. Safe to call from the UI. func startDownload(repo: String? = nil) { let target = repo ?? selectedRepo @@ -92,7 +121,9 @@ final class ModelManager: ObservableObject { try Task.checkCancellation() set(gen, .downloading(progress: min(1, Double(done) / Double(total))), "Downloading \(file.path) (\(byteString(file.size)))…") - let url = URL(string: "\(AppConfig.Model.huggingFaceBase)/\(repo)/resolve/main/\(file.path)")! + guard let url = Self.fileURL(repo: repo, path: file.path) else { + throw ModelError.invalidURL(file.path) + } let (tmp, response) = try await URLSession.shared.download(from: url) if let http = response as? HTTPURLResponse, http.statusCode != 200 { throw ModelError.http(http.statusCode, file.path) @@ -148,7 +179,7 @@ final class ModelManager: ObservableObject { /// List the repo's files (+ sizes) from the Hugging Face tree API. private func fetchFileList(repo: String) async throws -> [(path: String, size: Int64)] { - let url = URL(string: "\(AppConfig.Model.huggingFaceBase)/api/models/\(repo)/tree/main?recursive=1")! + guard let url = Self.treeURL(repo: repo) else { throw ModelError.invalidURL("file list") } let (data, response) = try await URLSession.shared.data(from: url) if let http = response as? HTTPURLResponse, http.statusCode != 200 { throw ModelError.http(http.statusCode, "file list") @@ -166,11 +197,13 @@ final class ModelManager: ObservableObject { case empty case http(Int, String) case corrupt(String) + case invalidURL(String) var errorDescription: String? { switch self { case .empty: return "No model files were found for this repository." case .http(let code, let what): return "Download failed (\(code)) while fetching \(what)." case .corrupt(let what): return "Downloaded file looked wrong (\(what)); please retry." + case .invalidURL(let what): return "Couldn't build a valid download URL for \(what)." } } } diff --git a/Sources/GrammarGem/AI/OllamaEngine.swift b/Sources/GrammarGem/AI/OllamaEngine.swift new file mode 100644 index 0000000..842d870 --- /dev/null +++ b/Sources/GrammarGem/AI/OllamaEngine.swift @@ -0,0 +1,176 @@ +import Foundation + +/// A cross-platform Layer-2 backend that talks to a locally-running **Ollama** +/// server over HTTP. Conforms to the same `AIEngine` seam as `MLXEngine`, reuses +/// the backend-neutral `AIPrompts` (so prompts/output are identical across +/// runtimes), and depends only on `URLSession` — no MLX, no Metal — so it compiles +/// and runs anywhere, including Windows/Linux. This is what converts the AIEngine +/// abstraction from nominal to real and opens the non-Apple-Silicon market. +/// +/// Privacy is preserved: Ollama runs the model on the user's own machine on +/// loopback (127.0.0.1); no text leaves the device. +final class OllamaEngine: AIEngine { + private let base: URL + private let model: String + private let session: URLSession + + /// `isReady` is read synchronously by the coordinator but readiness for an HTTP + /// backend is an async probe, so we cache the last-known result behind a lock. + private let lock = NSLock() + private var cachedReady = false + + init(base: URL = AppConfig.Ollama.baseURL, + model: String = AppConfig.Ollama.defaultModel, + session: URLSession = .shared) { + self.base = base + self.model = model + self.session = session + } + + var isReady: Bool { + lock.lock(); defer { lock.unlock() } + return cachedReady + } + + private func setReady(_ value: Bool) { + lock.lock(); cachedReady = value; lock.unlock() + } + + // MARK: - AIEngine + + /// Probe the server and confirm the configured model is pulled. Readiness for + /// Ollama means "server up + model present" (vs. MLX's "weights on disk"). + func warmup() async { + setReady(await probeModelPresent()) + } + + func run(_ action: AIAction, on text: String, protectedTerms: [String]) async throws -> String { + let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return text } + let raw = try await chat( + system: AIPrompts.systemPrompt(for: action, protectedTerms: protectedTerms), + user: trimmed, + temperature: AIPrompts.temperature(for: action), + maxTokens: AIPrompts.maxTokens(for: action)) + setReady(true) + return AIPrompts.clean(raw) + } + + func detectTone(_ text: String) async -> Tone { + let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return .professional } + do { + let raw = try await chat( + system: AIPrompts.toneClassifierPrompt(), user: trimmed, + temperature: 0.0, maxTokens: 8) + return AIPrompts.parseTone(raw) + } catch { + return .professional + } + } + + func runStreaming(_ action: AIAction, on text: String, protectedTerms: [String]) + -> AsyncThrowingStream { + AsyncThrowingStream { continuation in + Task { + let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { continuation.yield(text); continuation.finish(); return } + do { + try await streamChat( + system: AIPrompts.systemPrompt(for: action, protectedTerms: protectedTerms), + user: trimmed, + temperature: AIPrompts.temperature(for: action), + maxTokens: AIPrompts.maxTokens(for: action), + onChunk: { continuation.yield($0) }) + setReady(true) + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + } + } + + // MARK: - HTTP + + private func probeModelPresent() async -> Bool { + do { + let (data, response) = try await session.data(from: base.appendingPathComponent("api/tags")) + guard (response as? HTTPURLResponse)?.statusCode == 200 else { return false } + let names = (try JSONDecoder().decode(TagsResponse.self, from: data)).models.map(\.name) + // Match exact or family (e.g. configured "qwen2.5:3b" vs a pulled "qwen2.5:3b-instruct"). + return names.contains(model) + || names.contains { $0.hasPrefix(model) || model.hasPrefix($0) } + } catch { + return false + } + } + + private func chat(system: String, user: String, + temperature: Float, maxTokens: Int) async throws -> String { + let request = try makeRequest(path: "api/chat", body: ChatRequest( + model: model, + messages: [.init(role: "system", content: system), .init(role: "user", content: user)], + stream: false, + options: .init(temperature: temperature, num_predict: maxTokens))) + let (data, response) = try await session.data(for: request) + guard let http = response as? HTTPURLResponse else { + throw AIEngineError.generationFailed("No HTTP response from Ollama at \(base.absoluteString)") + } + guard http.statusCode == 200 else { + throw AIEngineError.generationFailed("Ollama returned HTTP \(http.statusCode)") + } + return (try JSONDecoder().decode(ChatResponse.self, from: data)).message?.content ?? "" + } + + private func streamChat(system: String, user: String, + temperature: Float, maxTokens: Int, + onChunk: (String) -> Void) async throws { + let request = try makeRequest(path: "api/chat", body: ChatRequest( + model: model, + messages: [.init(role: "system", content: system), .init(role: "user", content: user)], + stream: true, + options: .init(temperature: temperature, num_predict: maxTokens))) + let (bytes, response) = try await session.bytes(for: request) + guard (response as? HTTPURLResponse)?.statusCode == 200 else { + throw AIEngineError.generationFailed("Ollama streaming request failed") + } + // Ollama streams newline-delimited JSON frames, one per token-ish chunk. + for try await line in bytes.lines { + guard let data = line.data(using: .utf8), + let frame = try? JSONDecoder().decode(ChatResponse.self, from: data), + let chunk = frame.message?.content, !chunk.isEmpty + else { continue } + onChunk(chunk) + } + } + + private func makeRequest(path: String, body: T) throws -> URLRequest { + var request = URLRequest(url: base.appendingPathComponent(path)) + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.httpBody = try JSONEncoder().encode(body) + return request + } + + // MARK: - Wire types + + private struct ChatRequest: Encodable { + let model: String + let messages: [Message] + let stream: Bool + let options: Options + struct Message: Encodable { let role: String; let content: String } + struct Options: Encodable { let temperature: Float; let num_predict: Int } + } + + private struct ChatResponse: Decodable { + let message: Message? + struct Message: Decodable { let content: String } + } + + private struct TagsResponse: Decodable { + let models: [Model] + struct Model: Decodable { let name: String } + } +} diff --git a/Sources/GrammarGem/App/AppState.swift b/Sources/GrammarGem/App/AppState.swift index 6d8b2d5..d69dc53 100644 --- a/Sources/GrammarGem/App/AppState.swift +++ b/Sources/GrammarGem/App/AppState.swift @@ -30,6 +30,8 @@ final class AppState: ObservableObject { let updater: UpdaterManager private let capture: TextCapture let coordinator: TextReplacementCoordinator + /// Loopback bridge to the browser extension (opt-in; nil unless enabled). + private var localServer: LocalEngineServer? // UI state @Published var lastStatus: String = "Ready" @@ -48,9 +50,9 @@ final class AppState: ObservableObject { grammar = engine // The on-device LLM loads weights from ModelManager's local snapshot (if a // complete download is present) — never the network. Resolved lazily on use. - ai = MLXEngine(modelDirectoryProvider: { - ModelManager.completedModelDirectory(repo: AppConfig.Model.defaultRepo) - }) + // Reads the user's *selected* repo (default or paid large model) so the + // model that gets loaded matches the one that gets downloaded. + ai = AppState.makeAIEngine() permissions = Permissions() model = ModelManager() license = LicenseManager() @@ -67,8 +69,29 @@ final class AppState: ObservableObject { exclusions = excl liveMonitor = LiveMonitor(grammar: engine, detector: detector, exclusions: excl, capture: capture) updater = UpdaterManager() + let dict = dictionary coordinator = TextReplacementCoordinator( - capture: capture, grammar: engine, ai: ai, gate: gate, detector: detector) + capture: capture, grammar: engine, ai: ai, gate: gate, detector: detector, + protectedTerms: { dict.entries }) + } + + // MARK: - Backend factory + + /// The single place the Layer-2 backend is chosen. Defaults to the MLX + /// on-device engine (Apple Silicon); a user/CI override selects the + /// cross-platform `OllamaEngine`. Adding a backend is a one-line branch here + /// instead of edits scattered across construction sites. The MLX provider + /// reads the user's *selected* repo so the loaded model matches the downloaded + /// one (incl. the paid large model). + static func makeAIEngine() -> AIEngine { + let override = ProcessInfo.processInfo.environment["GRAMMARGEM_AI_BACKEND"] + ?? UserDefaults.standard.string(forKey: "GrammarGem.aiBackend") + if override == "ollama" { + return OllamaEngine() + } + return MLXEngine(modelDirectoryProvider: { + ModelManager.completedModelDirectory(repo: ModelManager.activeRepo()) + }) } /// Called once at launch (from the AppDelegate). @@ -101,10 +124,19 @@ final class AppState: ObservableObject { // Auto-install the on-device AI model in the background on first use, so // the AI features are ready without the user hunting for a button. The - // download is resumable/cancelable from the AI Model screen. - if !model.isModelPresent(repo: model.selectedRepo) { + // download is resumable/cancelable from the AI Model screen. If a model is + // already present, warm it up now so the first AI action is instant. + if model.isModelPresent(repo: model.selectedRepo) { + Task { [weak self] in await self?.ai.warmup() } + } else { model.startDownload() } + // Warm the model up as soon as a (first-run) download completes. + model.$state.dropFirst() + .sink { [weak self] state in + if case .ready = state { Task { [weak self] in await self?.ai.warmup() } } + } + .store(in: &cancellables) permissions.refresh() appAwareEnabled = gate.appAwareEnabled @@ -115,9 +147,57 @@ final class AppState: ObservableObject { } Task { await license.validateIfNeeded() } + startLocalServerIfEnabled() Log.app.info("GrammarGem started. Licensed: \(self.license.isLicensed, privacy: .public)") } + // MARK: - Browser-extension bridge (opt-in) + + /// Start the loopback engine server only when the user has opted in, so the app + /// ships no listening socket by default. The server reuses the SAME engines, + /// prompts, and entitlement gate as the menu-bar app. + private func startLocalServerIfEnabled() { + let enabled = ProcessInfo.processInfo.environment["GRAMMARGEM_LOCAL_SERVER"] == "1" + || UserDefaults.standard.bool(forKey: AppConfig.LocalServer.enableKey) + guard enabled, localServer == nil else { return } + + let server = LocalEngineServer( + port: AppConfig.LocalServer.port, + token: LocalEngineServer.loadOrCreateToken(), + grammarHandler: { [weak self] text in + await MainActor.run { + guard let self else { return (text, []) } + let corrected = self.grammar.correct(text) + let suggestions = self.grammar.check(text).map { + LocalEngineServer.WireSuggestion( + location: $0.location, length: $0.length, + original: $0.original, replacement: $0.replacement, + kind: $0.kind.rawValue, message: $0.message) + } + return (corrected, suggestions) + } + }, + aiHandler: { [weak self] (req: LocalEngineServer.AIRequest) async -> LocalEngineServer.AIOutcome in + guard let self else { return .error("GrammarGem is shutting down.") } + return await self.handleServerAI(req) + }) + server.start() + localServer = server + } + + /// Run an AI action requested by the extension — same gate + model as the app. + private func handleServerAI(_ req: LocalEngineServer.AIRequest) async -> LocalEngineServer.AIOutcome { + if case .denied(let reason) = gate.authorizeAI(req.action) { return .error(reason) } + guard ai.isReady else { return .error(AIEngineError.modelNotReady.localizedDescription) } + do { + let out = try await ai.run(req.action, on: req.text, protectedTerms: dictionary.entries) + gate.recordAIActionUsed() + return .ok(out) + } catch { + return .error(error.localizedDescription) + } + } + // MARK: - Hotkey handling /// (Re)register the global hotkeys from the current preferences. diff --git a/Sources/GrammarGem/App/GrammarGemApp.swift b/Sources/GrammarGem/App/GrammarGemApp.swift index 5a2bcf4..e3726ee 100644 --- a/Sources/GrammarGem/App/GrammarGemApp.swift +++ b/Sources/GrammarGem/App/GrammarGemApp.swift @@ -25,7 +25,7 @@ enum MLXSelfTest { static func run() { let sample = "me and him was going to the store for to buy some milks" let engine = MLXEngine(modelDirectoryProvider: { - ModelManager.completedModelDirectory(repo: AppConfig.Model.defaultRepo) + ModelManager.completedModelDirectory(repo: ModelManager.activeRepo()) }) guard engine.isReady else { FileHandle.standardError.write(Data("SELFTEST_FAIL: model not present\n".utf8)) diff --git a/Sources/GrammarGem/Core/AppConfig.swift b/Sources/GrammarGem/Core/AppConfig.swift index 7df8bab..9f0f998 100644 --- a/Sources/GrammarGem/Core/AppConfig.swift +++ b/Sources/GrammarGem/Core/AppConfig.swift @@ -37,6 +37,23 @@ enum AppConfig { static let huggingFaceBase = "https://huggingface.co" } + /// The cross-platform Layer-2 backend: a locally-running Ollama server. Used + /// when MLX isn't available (Windows/Linux) or when the user opts in. Inference + /// stays on-device — Ollama runs the model on the user's own machine. + enum Ollama { + static let baseURL = URL(string: "http://127.0.0.1:11434")! + static let defaultModel = "qwen2.5:3b" // matches the MLX default model family + } + + /// Loopback engine server for the browser extension (off by default; see + /// `LocalEngineServer`). Exposes Harper + the local LLM to the extension over + /// 127.0.0.1 with a per-install token — no cloud round-trip. + enum LocalServer { + static let port: UInt16 = 47800 + static let tokenKey = "GrammarGem.localServerToken" + static let enableKey = "GrammarGem.enableLocalServer" + } + /// Free-tier limits + license validation cadence (see spec §5/§6). enum Limits { static let freeDailyAIActions = 10 diff --git a/Sources/GrammarGem/Licensing/LemonSqueezyClient.swift b/Sources/GrammarGem/Licensing/LemonSqueezyClient.swift index 15f2f69..58e5399 100644 --- a/Sources/GrammarGem/Licensing/LemonSqueezyClient.swift +++ b/Sources/GrammarGem/Licensing/LemonSqueezyClient.swift @@ -58,16 +58,24 @@ final class LemonSqueezyClient { private func verifyAndResolveTier(meta: LSMeta?, licenseKey: LSLicenseKey?) throws -> Tier { guard let meta else { throw LicenseError.wrongProduct } - // Store + product check. If AppConfig still holds placeholder (non-numeric) - // IDs, we log and skip the check so the stub flow runs — TODO(real-integration): - // once real numeric IDs are set, this becomes a hard gate. + // Store + product check. With real numeric IDs configured this is a hard + // gate that prevents a key from another Lemon Squeezy product unlocking + // GrammarGem. If AppConfig still holds placeholder (non-numeric) IDs the + // gate can't run — allow the bypass ONLY in DEBUG so local stub flows work; + // in a release build FAIL CLOSED so the "any key unlocks paid" bypass can + // never ship. TODO(real-integration): set real IDs in AppConfig.LemonSqueezy. if let expectedStore = Int(AppConfig.LemonSqueezy.expectedStoreID), let expectedProduct = Int(AppConfig.LemonSqueezy.expectedProductID) { guard meta.store_id == expectedStore, meta.product_id == expectedProduct else { throw LicenseError.wrongProduct } } else { - Log.licensing.warning("Skipping store/product hard-verify — placeholder IDs in AppConfig.") + #if DEBUG + Log.licensing.warning("Skipping store/product hard-verify — placeholder IDs in AppConfig (DEBUG only).") + #else + Log.licensing.error("Refusing to verify license: placeholder Lemon Squeezy IDs in a release build.") + throw LicenseError.wrongProduct + #endif } // Resolve tier from the variant id (reverse-lookup the configured map). diff --git a/Sources/GrammarGem/System/AXSupport.swift b/Sources/GrammarGem/System/AXSupport.swift index 48f155c..7858204 100644 --- a/Sources/GrammarGem/System/AXSupport.swift +++ b/Sources/GrammarGem/System/AXSupport.swift @@ -38,6 +38,17 @@ enum AX { AXUIElementSetMessagingTimeout(element, seconds) } + /// Whether an element is a secure text field (a password input). GrammarGem + /// must NEVER read or rewrite these — so a user who forgot to exclude a + /// password manager or login form can't have a decrypted password captured, + /// sent through the grammar/AI engines, and written back. This is enforced in + /// code rather than relying on the user's exclusion blocklist. + static func isSecureTextField(_ element: AXUIElement) -> Bool { + // AXSecureTextField is the subrole AppKit/most Electron inputs report for + // password fields. String-compared (the constant isn't always imported). + copyString(element, kAXSubroleAttribute) == "AXSecureTextField" + } + /// The currently-focused UI element across the system — robust to the /// system-wide element failing. /// diff --git a/Sources/GrammarGem/System/LocalEngineServer.swift b/Sources/GrammarGem/System/LocalEngineServer.swift new file mode 100644 index 0000000..f6da35a --- /dev/null +++ b/Sources/GrammarGem/System/LocalEngineServer.swift @@ -0,0 +1,266 @@ +import Foundation +import Network + +/// A minimal, loopback-only HTTP server that exposes GrammarGem's local engines +/// (Harper grammar + the on-device LLM) to a **browser extension** — the +/// highest-volume surface for a grammar tool — with NO cloud round-trip. Built on +/// Network.framework so it adds no third-party HTTP dependency, keeping the +/// privacy/trust surface minimal. +/// +/// Security posture (deliberate): +/// - Binds to 127.0.0.1 only (never a routable interface), via `requiredLocalEndpoint`. +/// - Every engine endpoint requires a per-install bearer token (`X-GrammarGem-Token`); +/// a web page that doesn't know the token cannot drive the engines. +/// - OFF by default — only starts when the user opts in, so the app ships no +/// listening socket unless explicitly enabled. +/// +/// It reuses the SAME engines as the menu-bar app via injected async handlers, so +/// there is one engine, one set of prompts (`AIPrompts`), and one entitlement gate +/// across the app, the extension, and (later) other backends. +final class LocalEngineServer { + + struct AIRequest { let action: AIAction; let text: String } + + struct WireSuggestion: Encodable { + let location: Int + let length: Int + let original: String + let replacement: String + let kind: String + let message: String + } + + /// Grammar pass: corrected text + suggestion spans for the extension to underline. + typealias GrammarHandler = (String) async -> (corrected: String, suggestions: [WireSuggestion]) + /// AI action outcome (a plain enum so a user-facing message needn't be an `Error`). + enum AIOutcome { case ok(String); case error(String) } + /// AI action: runs the gate + model, returns the result or a user-facing error. + typealias AIHandler = (AIRequest) async -> AIOutcome + + private let port: UInt16 + private let token: String + private let grammarHandler: GrammarHandler + private let aiHandler: AIHandler + + private var listener: NWListener? + private let queue = DispatchQueue(label: "com.grammargem.localserver") + + init(port: UInt16, token: String, + grammarHandler: @escaping GrammarHandler, + aiHandler: @escaping AIHandler) { + self.port = port + self.token = token + self.grammarHandler = grammarHandler + self.aiHandler = aiHandler + } + + /// A stable per-install token, persisted so the extension can be paired once. + static func loadOrCreateToken() -> String { + let d = UserDefaults.standard + if let existing = d.string(forKey: AppConfig.LocalServer.tokenKey), !existing.isEmpty { + return existing + } + let token = UUID().uuidString + UUID().uuidString.replacingOccurrences(of: "-", with: "") + d.set(token, forKey: AppConfig.LocalServer.tokenKey) + return token + } + + func start() { + guard listener == nil else { return } + do { + let params = NWParameters.tcp + guard let nwPort = NWEndpoint.Port(rawValue: port) else { return } + // Loopback-only: refuse to bind to any routable interface. + params.requiredLocalEndpoint = .hostPort(host: "127.0.0.1", port: nwPort) + let listener = try NWListener(using: params) + listener.newConnectionHandler = { [weak self] conn in self?.accept(conn) } + listener.start(queue: queue) + self.listener = listener + Log.app.info("Local engine server on 127.0.0.1:\(self.port, privacy: .public) (extension bridge)") + } catch { + Log.app.error("Local engine server failed to start: \(error.localizedDescription, privacy: .public)") + } + } + + func stop() { + listener?.cancel() + listener = nil + } + + // MARK: - Connection lifecycle + + private func accept(_ conn: NWConnection) { + conn.start(queue: queue) + receive(conn, buffer: Data()) + } + + private func receive(_ conn: NWConnection, buffer: Data) { + conn.receive(minimumIncompleteLength: 1, maximumLength: 1 << 16) { [weak self] data, _, isComplete, error in + guard let self else { conn.cancel(); return } + var buffer = buffer + if let data { buffer.append(data) } + if let request = HTTPRequest(buffer), request.isComplete { + Task { await self.route(request, on: conn) } + return + } + if isComplete || error != nil { conn.cancel(); return } + self.receive(conn, buffer: buffer) // headers/body not fully arrived yet + } + } + + // MARK: - Routing + + private func route(_ request: HTTPRequest, on conn: NWConnection) async { + // CORS preflight — extensions issue this before a cross-origin POST. + if request.method == "OPTIONS" { + send(conn, status: 204, body: Data()) + return + } + + if request.path == "/health" || request.path == "/v1/health" { + sendJSON(conn, status: 200, HealthResponse(ok: true, app: AppConfig.appName, version: AppConfig.appVersion)) + return + } + + // Token gate for every engine endpoint. + guard request.headers["x-grammargem-token"] == token else { + sendJSON(conn, status: 401, ErrorResponse(error: "unauthorized")) + return + } + + switch (request.method, request.path) { + case ("POST", "/v1/grammar"): + guard let text = decodeText(request.body) else { + sendJSON(conn, status: 400, ErrorResponse(error: "missing 'text'")); return + } + let result = await grammarHandler(text) + sendJSON(conn, status: 200, GrammarResponse(corrected: result.corrected, suggestions: result.suggestions)) + + case ("POST", "/v1/ai"): + guard let req = decodeAI(request.body) else { + sendJSON(conn, status: 400, ErrorResponse(error: "bad request")); return + } + switch await aiHandler(req) { + case .ok(let out): sendJSON(conn, status: 200, AIResponse(result: out)) + case .error(let msg): sendJSON(conn, status: 422, ErrorResponse(error: msg)) + } + + default: + sendJSON(conn, status: 404, ErrorResponse(error: "not found")) + } + } + + // MARK: - Body decoding + + private func decodeText(_ body: Data) -> String? { + struct P: Decodable { let text: String } + guard let p = try? JSONDecoder().decode(P.self, from: body), !p.text.isEmpty else { return nil } + return p.text + } + + private func decodeAI(_ body: Data) -> AIRequest? { + struct P: Decodable { + let action: String + let text: String + let instruction: String? + let language: String? + let mode: String? + let tone: String? + } + guard let p = try? JSONDecoder().decode(P.self, from: body), !p.text.isEmpty else { return nil } + let action: AIAction + switch p.action { + case "rewriteClarity": action = .rewriteClarity + case "rewrite": action = .rewrite + case "ask": + guard let i = p.instruction, !i.trimmingCharacters(in: .whitespaces).isEmpty else { return nil } + action = .ask(i) + case "translate": + guard let l = p.language, !l.isEmpty else { return nil } + action = .translate(language: l) + case "tone": + action = .adjustTone(Tone(rawValue: p.tone ?? "") ?? .professional) + case "mode": + guard let id = p.mode, let m = ModeRegistry.mode(id: id) else { return nil } + action = .applyMode(m) + default: + return nil + } + return AIRequest(action: action, text: p.text) + } + + // MARK: - Responses + + private struct HealthResponse: Encodable { let ok: Bool; let app: String; let version: String } + private struct GrammarResponse: Encodable { let corrected: String; let suggestions: [WireSuggestion] } + private struct AIResponse: Encodable { let result: String } + private struct ErrorResponse: Encodable { let error: String } + + private func sendJSON(_ conn: NWConnection, status: Int, _ value: T) { + let body = (try? JSONEncoder().encode(value)) ?? Data("{}".utf8) + send(conn, status: status, body: body) + } + + private func send(_ conn: NWConnection, status: Int, body: Data) { + var head = "HTTP/1.1 \(status) \(Self.reason(status))\r\n" + head += "Content-Type: application/json\r\n" + head += "Content-Length: \(body.count)\r\n" + // The extension calls cross-origin; a secret token (not the origin) is the gate. + head += "Access-Control-Allow-Origin: *\r\n" + head += "Access-Control-Allow-Headers: Content-Type, X-GrammarGem-Token\r\n" + head += "Access-Control-Allow-Methods: POST, GET, OPTIONS\r\n" + head += "Connection: close\r\n\r\n" + var out = Data(head.utf8) + out.append(body) + conn.send(content: out, completion: .contentProcessed { _ in conn.cancel() }) + } + + private static func reason(_ status: Int) -> String { + switch status { + case 200: return "OK" + case 204: return "No Content" + case 400: return "Bad Request" + case 401: return "Unauthorized" + case 404: return "Not Found" + case 422: return "Unprocessable Entity" + default: return "Error" + } + } +} + +/// A just-enough HTTP/1.1 request parser for small JSON POSTs. Returns nil until +/// the headers terminator (`\r\n\r\n`) has arrived; `isComplete` is false until the +/// full `Content-Length` body is buffered, so the receive loop knows to read more. +private struct HTTPRequest { + let method: String + let path: String + let headers: [String: String] + let body: Data + let isComplete: Bool + + init?(_ raw: Data) { + guard let headerEnd = raw.range(of: Data("\r\n\r\n".utf8)) else { return nil } + guard let headerString = String(data: raw[..= 2 else { return nil } + method = String(parts[0]) + path = String(parts[1]) + lines.removeFirst() + + var hdrs: [String: String] = [:] + for line in lines { + guard let colon = line.firstIndex(of: ":") else { continue } + let key = line[..= contentLength + } +} diff --git a/Sources/GrammarGem/System/TextCapture.swift b/Sources/GrammarGem/System/TextCapture.swift index 80e4f48..1ff2241 100644 --- a/Sources/GrammarGem/System/TextCapture.swift +++ b/Sources/GrammarGem/System/TextCapture.swift @@ -24,10 +24,12 @@ final class TextCapture { enum CaptureError: LocalizedError { case noSelection case notPermitted + case secureField var errorDescription: String? { switch self { case .noSelection: return "Select some text first, then press the hotkey." case .notPermitted: return "GrammarGem needs Accessibility permission to read and replace text." + case .secureField: return "GrammarGem never reads password fields." } } } @@ -41,8 +43,15 @@ final class TextCapture { // MARK: - Capture /// Try AX first; on empty/unavailable, fall back to the clipboard path. + /// Refuses outright if the focused element is a password field. func capture() throws -> Capture { - if let ax = captureViaAccessibility() { + let focused = AXIsProcessTrusted() ? AX.focusedElement() : nil + if let focused, AX.isSecureTextField(focused) { + // Bail before the clipboard fallback too — never synthesize ⌘C into a + // password field. + throw CaptureError.secureField + } + if let ax = captureViaAccessibility(focused) { return ax } if let clip = try captureViaClipboard() { @@ -51,9 +60,8 @@ final class TextCapture { throw CaptureError.noSelection } - private func captureViaAccessibility() -> Capture? { - guard AXIsProcessTrusted() else { return nil } - guard let focused = AX.focusedElement() else { return nil } + private func captureViaAccessibility(_ focused: AXUIElement?) -> Capture? { + guard let focused else { return nil } var valueRef: CFTypeRef? guard AXUIElementCopyAttributeValue( @@ -92,6 +100,7 @@ final class TextCapture { func focusedFieldText() -> (text: String, element: AXUIElement)? { guard AXIsProcessTrusted() else { return nil } guard let focused = AX.focusedElement() else { return nil } + guard !AX.isSecureTextField(focused) else { return nil } // never fix a password field // Only treat genuine editable text roles as fixable, so we don't report // "nothing selected" failures on sliders/steppers/custom controls. diff --git a/Sources/GrammarGem/System/TextFieldScanner.swift b/Sources/GrammarGem/System/TextFieldScanner.swift index ff5ac1d..4e320ea 100644 --- a/Sources/GrammarGem/System/TextFieldScanner.swift +++ b/Sources/GrammarGem/System/TextFieldScanner.swift @@ -47,7 +47,9 @@ enum TextFieldScanner { if role == kAXTextFieldRole as String || role == kAXTextAreaRole as String || role == kAXComboBoxRole as String { - if let value = AX.copyString(el, kAXValueAttribute), + // Never surface password fields in the live monitor. + if !AX.isSecureTextField(el), + let value = AX.copyString(el, kAXValueAttribute), value.trimmingCharacters(in: .whitespacesAndNewlines).count >= 3 { let label = AX.copyString(el, kAXTitleAttribute) ?? AX.copyString(el, kAXDescriptionAttribute) diff --git a/Sources/GrammarGem/System/TextReplacementCoordinator.swift b/Sources/GrammarGem/System/TextReplacementCoordinator.swift index 3746686..2cc0771 100644 --- a/Sources/GrammarGem/System/TextReplacementCoordinator.swift +++ b/Sources/GrammarGem/System/TextReplacementCoordinator.swift @@ -12,14 +12,19 @@ final class TextReplacementCoordinator { private let ai: AIEngine private let gate: FeatureGate private let detector: AppDetector + /// Vocabulary the LLM must reproduce verbatim (personal dictionary / style + /// guide). Read fresh per action so dictionary edits take effect immediately. + private let protectedTerms: () -> [String] init(capture: TextCapture, grammar: GrammarEngine, ai: AIEngine, - gate: FeatureGate, detector: AppDetector) { + gate: FeatureGate, detector: AppDetector, + protectedTerms: @escaping () -> [String] = { [] }) { self.capture = capture self.grammar = grammar self.ai = ai self.gate = gate self.detector = detector + self.protectedTerms = protectedTerms } /// The "fix" hotkey — instant grammar correction (Layer 1, always free). @@ -52,7 +57,7 @@ final class TextReplacementCoordinator { do { cap = try await captureOffMain() } catch { return .failed(error.localizedDescription) } do { - let result = try await ai.run(action, on: cap.text) + let result = try await ai.run(action, on: cap.text, protectedTerms: protectedTerms()) gate.recordAIActionUsed() let ok = await replaceOffMain(result, cap) return ok ? .replaced(result) : .failed("Couldn't write the result back.") @@ -73,7 +78,7 @@ final class TextReplacementCoordinator { do { cap = try await captureOffMain() } catch { return .failed(error.localizedDescription) } do { - let result = try await ai.run(action, on: cap.text) + let result = try await ai.run(action, on: cap.text, protectedTerms: protectedTerms()) gate.recordAIActionUsed() let ok = await replaceOffMain(result, cap) return ok ? .replaced(result) : .failed("Couldn't write the result back.") diff --git a/Tests/GrammarGemTests/GrammarGemTests.swift b/Tests/GrammarGemTests/GrammarGemTests.swift index 41bc914..0341edd 100644 --- a/Tests/GrammarGemTests/GrammarGemTests.swift +++ b/Tests/GrammarGemTests/GrammarGemTests.swift @@ -149,6 +149,66 @@ final class GrammarGemTests: XCTestCase { XCTAssertNil(ModeRegistry.mode(forBundleID: "com.unknown.app")) } + // MARK: - AIPrompts (backend-neutral seam) + + func testProtectedTermsInjectedIntoSystemPrompt() { + let base = AIPrompts.systemPrompt(for: .rewriteClarity) + let withTerms = AIPrompts.systemPrompt(for: .rewriteClarity, protectedTerms: ["GrammarGem", "Harper"]) + XCTAssertFalse(base.contains("GrammarGem")) + XCTAssertTrue(withTerms.contains("\"GrammarGem\"")) + XCTAssertTrue(withTerms.contains("\"Harper\"")) + XCTAssertTrue(withTerms.contains("verbatim")) + } + + func testProtectedTermsClauseSkipsWhenEmpty() { + let blank = AIPrompts.systemPrompt(for: .rewrite, protectedTerms: ["", " "]) + XCTAssertEqual(blank, AIPrompts.systemPrompt(for: .rewrite)) + } + + func testMaxTokensHonorsModeLengthCap() { + // Default actions get the full budget… + XCTAssertEqual(AIPrompts.maxTokens(for: .rewriteClarity), 512) + // …but a short-form Mode (Team Chat, lengthCap 60) gets far fewer. + let slack = ModeRegistry.slack + XCTAssertNotNil(slack.lengthCap) + let capped = AIPrompts.maxTokens(for: .applyMode(slack)) + XCTAssertLessThan(capped, 512) + XCTAssertGreaterThanOrEqual(capped, 64) + // A mode with no cap (Polish) falls back to the default budget. + XCTAssertEqual(AIPrompts.maxTokens(for: .applyMode(ModeRegistry.polish)), 512) + } + + func testCleanStripsPreambleAndQuotes() { + XCTAssertEqual(AIPrompts.clean("Sure, here's the rewrite: Hello world"), "Hello world") + XCTAssertEqual(AIPrompts.clean("\"Quoted result\""), "Quoted result") + XCTAssertEqual(AIPrompts.clean(" already clean "), "already clean") + } + + func testParseToneIsRobust() { + XCTAssertEqual(AIPrompts.parseTone("professional"), .professional) + XCTAssertEqual(AIPrompts.parseTone(" Friendly.\n"), .friendly) + XCTAssertEqual(AIPrompts.parseTone("the tone is academic here"), .academic) + XCTAssertEqual(AIPrompts.parseTone("???"), .professional) // safe fallback + } + + // MARK: - Custom modes + + @MainActor + func testCustomModeRejectsEmptyPromptAndDuplicates() { + // Deterministic regardless of prior runs (the store persists to UserDefaults). + UserDefaults.standard.removeObject(forKey: "GrammarGem.customModes") + defer { UserDefaults.standard.removeObject(forKey: "GrammarGem.customModes") } + + let store = CustomModesStore() + XCTAssertFalse(store.addCustom(name: "Brand Voice", prompt: " "), + "empty prompt must be rejected") + XCTAssertTrue(store.addCustom(name: "Brand Voice", prompt: "Write in Acme's warm voice.")) + XCTAssertFalse(store.addCustom(name: "brand voice", prompt: "different prompt"), + "duplicate name (case-insensitive) must be rejected") + XCTAssertFalse(store.addCustom(name: "Polish", prompt: "collides with a built-in"), + "name colliding with a built-in mode must be rejected") + } + // MARK: - Device fingerprint func testDeviceFingerprintIsStableAndHashed() {