Skip to content

fix: resolve AI chat hang caused by main-actor-bound streaming (#735) #738

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
datlechin merged 3 commits into from
Apr 14, 2026
Merged

fix: resolve AI chat hang caused by main-actor-bound streaming (#735) #738

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c96cb76
fix: resolve AI chat hang caused by main-actor-bound streaming (#735)
datlechin Apr 14, 2026
475782c
merge: resolve conflict with main
datlechin Apr 14, 2026
a9c3502
fix: address review feedback — throttled scroll, DRY schema fetch, cl…
datlechin Apr 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

## [Unreleased]

### Fixed

- Fix AI chat hanging the app during streaming, schema fetch, and conversation loading (#735)

## [0.31.4] - 2026-04-14

### Added
Expand Down
57 changes: 31 additions & 26 deletions TablePro/Core/AI/AIProviderFactory.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
//

import Foundation
import os

/// Factory for creating AI provider instances
enum AIProviderFactory {
Expand All @@ -16,46 +17,50 @@ enum AIProviderFactory {
let config: AIProviderConfig
}

private static var cachedProviders: [UUID: (apiKey: String?, provider: AIProvider)] = [:]
private static let cacheLock = OSAllocatedUnfairLock(
initialState: [UUID: (apiKey: String?, provider: AIProvider)]()
)

/// Create or return a cached AI provider for the given configuration
static func createProvider(
for config: AIProviderConfig,
apiKey: String?
) -> AIProvider {
if let cached = cachedProviders[config.id], cached.apiKey == apiKey {
return cached.provider
}
cacheLock.withLock { cache in
if let cached = cache[config.id], cached.apiKey == apiKey {
return cached.provider
}

let provider: AIProvider
switch config.type {
case .claude:
provider = AnthropicProvider(
endpoint: config.endpoint,
apiKey: apiKey ?? ""
)
case .gemini:
provider = GeminiProvider(
endpoint: config.endpoint,
apiKey: apiKey ?? ""
)
case .openAI, .openRouter, .ollama, .custom:
provider = OpenAICompatibleProvider(
endpoint: config.endpoint,
apiKey: apiKey,
providerType: config.type
)
let provider: AIProvider
switch config.type {
case .claude:
provider = AnthropicProvider(
endpoint: config.endpoint,
apiKey: apiKey ?? ""
)
case .gemini:
provider = GeminiProvider(
endpoint: config.endpoint,
apiKey: apiKey ?? ""
)
case .openAI, .openRouter, .ollama, .custom:
provider = OpenAICompatibleProvider(
endpoint: config.endpoint,
apiKey: apiKey,
providerType: config.type
)
}
cache[config.id] = (apiKey, provider)
return provider
}
cachedProviders[config.id] = (apiKey, provider)
return provider
}

static func invalidateCache() {
cachedProviders.removeAll()
cacheLock.withLock { $0.removeAll() }
}

static func invalidateCache(for configID: UUID) {
cachedProviders.removeValue(forKey: configID)
cacheLock.withLock { $0.removeValue(forKey: configID) }
}

static func resolveProvider(
Expand Down
24 changes: 9 additions & 15 deletions TablePro/Core/AI/AISchemaContext.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,20 +6,14 @@
//

import Foundation
import os
import TableProPluginKit

/// Builds schema context for AI system prompts
struct AISchemaContext {
private static let logger = Logger(
subsystem: "com.TablePro",
category: "AISchemaContext"
)

// MARK: - Public

/// Build a system prompt including database context
@MainActor static func buildSystemPrompt(
static func buildSystemPrompt(
databaseType: DatabaseType,
databaseName: String,
tables: [TableInfo],
Expand All @@ -28,7 +22,9 @@ struct AISchemaContext {
currentQuery: String?,
queryResults: String?,
settings: AISettings,
identifierQuote: String = "\""
identifierQuote: String = "\"",
editorLanguage: EditorLanguage,
queryLanguageName: String
) -> String {
var parts: [String] = []

Expand Down Expand Up @@ -56,7 +52,7 @@ struct AISchemaContext {
if settings.includeCurrentQuery,
let query = currentQuery,
!query.isEmpty {
let lang = PluginManager.shared.editorLanguage(for: databaseType).codeBlockTag
let lang = editorLanguage.codeBlockTag
parts.append("\n## Current Query\n```\(lang)\n\(query)\n```")
}

Expand All @@ -66,11 +62,9 @@ struct AISchemaContext {
parts.append("\n## Recent Query Results\n\(results)")
}

let editorLang = PluginManager.shared.editorLanguage(for: databaseType)
let langName = PluginManager.shared.queryLanguageName(for: databaseType)
let langTag = editorLang.codeBlockTag
let langTag = editorLanguage.codeBlockTag

switch editorLang {
switch editorLanguage {
case .sql:
parts.append(
"\nProvide SQL queries appropriate for"
Expand All @@ -82,10 +76,10 @@ struct AISchemaContext {
)
default:
parts.append(
"\nProvide \(langName) queries using `\(langTag)` fenced code blocks."
"\nProvide \(queryLanguageName) queries using `\(langTag)` fenced code blocks."
)
parts.append(
"Use \(langName) syntax, not SQL."
"Use \(queryLanguageName) syntax, not SQL."
)
}

Expand Down
80 changes: 30 additions & 50 deletions TablePro/Core/AI/AnthropicProvider.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,16 +60,36 @@ final class AnthropicProvider: AIProvider {

guard line.hasPrefix("data: ") else { continue }
let jsonString = String(line.dropFirst(6))
guard jsonString != "[DONE]" else { break }

if let text = parseContentBlockDelta(jsonString) {
continuation.yield(.text(text))
}
if let tokens = parseInputTokens(jsonString) {
inputTokens = tokens
}
if let tokens = parseOutputTokens(jsonString) {
outputTokens = tokens
guard jsonString != "[DONE]",
let data = jsonString.data(using: .utf8),
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
let type = json["type"] as? String
else { continue }

switch type {
case "content_block_delta":
if let delta = json["delta"] as? [String: Any],
let text = delta["text"] as? String {
continuation.yield(.text(text))
}
case "message_start":
if let message = json["message"] as? [String: Any],
let usage = message["usage"] as? [String: Any],
let tokens = usage["input_tokens"] as? Int {
inputTokens = tokens
}
case "message_delta":
if let usage = json["usage"] as? [String: Any],
let tokens = usage["output_tokens"] as? Int {
outputTokens = tokens
}
case "error":
if let errorObj = json["error"] as? [String: Any],
let message = errorObj["message"] as? String {
throw AIProviderError.streamingFailed(message)
}
default:
break
}
}

Expand Down Expand Up @@ -196,44 +216,4 @@ final class AnthropicProvider: AIProvider {
request.httpBody = try JSONSerialization.data(withJSONObject: body)
return request
}

private func parseContentBlockDelta(_ jsonString: String) -> String? {
guard let data = jsonString.data(using: .utf8),
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
let type = json["type"] as? String,
type == "content_block_delta",
let delta = json["delta"] as? [String: Any],
let text = delta["text"] as? String
else {
return nil
}
return text
}

private func parseInputTokens(_ jsonString: String) -> Int? {
guard let data = jsonString.data(using: .utf8),
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
let type = json["type"] as? String,
type == "message_start",
let message = json["message"] as? [String: Any],
let usage = message["usage"] as? [String: Any],
let inputTokens = usage["input_tokens"] as? Int
else {
return nil
}
return inputTokens
}

private func parseOutputTokens(_ jsonString: String) -> Int? {
guard let data = jsonString.data(using: .utf8),
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
let type = json["type"] as? String,
type == "message_delta",
let usage = json["usage"] as? [String: Any],
let outputTokens = usage["output_tokens"] as? Int
else {
return nil
}
return outputTokens
}
}
115 changes: 39 additions & 76 deletions TablePro/Core/AI/OpenAICompatibleProvider.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,38 +63,52 @@ final class OpenAICompatibleProvider: AIProvider {
for try await line in bytes.lines {
if Task.isCancelled { break }

let jsonString: String
if self.providerType == .ollama {
// Ollama: raw newline-delimited JSON (no SSE "data: " prefix)
guard !line.isEmpty else { continue }
Self.logger.debug("Ollama stream line: \(line.prefix(200), privacy: .public)")

if let text = self.parseChatCompletionDelta(line) {
continuation.yield(.text(text))
}
if let usage = self.parseUsageFromChunk(line) {
inputTokens = usage.inputTokens
outputTokens = usage.outputTokens
}
// Ollama signals completion with "done":true
if let data = line.data(using: .utf8),
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
json["done"] as? Bool == true
{
break
}
jsonString = line
} else {
// OpenAI/OpenRouter/Custom: SSE with "data: " prefix
guard line.hasPrefix("data: ") else { continue }
let jsonString = String(line.dropFirst(6))
guard jsonString != "[DONE]" else { break }

if let text = self.parseChatCompletionDelta(jsonString) {
continuation.yield(.text(text))
}
if let usage = self.parseUsageFromChunk(jsonString) {
inputTokens = usage.inputTokens
outputTokens = usage.outputTokens
}
let payload = String(line.dropFirst(6))
guard payload != "[DONE]" else { break }
jsonString = payload
}

// Single JSON parse per SSE line
guard let data = jsonString.data(using: .utf8),
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any]
else { continue }

// Text extraction
if let choices = json["choices"] as? [[String: Any]],
let delta = choices.first?["delta"] as? [String: Any],
let content = delta["content"] as? String {
continuation.yield(.text(content))
} else if let message = json["message"] as? [String: Any],
let content = message["content"] as? String,
!content.isEmpty {
continuation.yield(.text(content))
}

// Usage extraction
if let usage = json["usage"] as? [String: Any],
let promptTokens = usage["prompt_tokens"] as? Int,
let completionTokens = usage["completion_tokens"] as? Int {
inputTokens = promptTokens
outputTokens = completionTokens
} else if let done = json["done"] as? Bool, done,
let promptEval = json["prompt_eval_count"] as? Int,
let evalCount = json["eval_count"] as? Int {
inputTokens = promptEval
outputTokens = evalCount
}

// Ollama signals completion with "done":true
if json["done"] as? Bool == true {
break
}
}

Expand Down Expand Up @@ -250,57 +264,6 @@ final class OpenAICompatibleProvider: AIProvider {
return request
}

// MARK: - Response Parsing

private func parseChatCompletionDelta(_ jsonString: String) -> String? {
guard let data = jsonString.data(using: .utf8),
let json = try? JSONSerialization.jsonObject(with: data)
as? [String: Any]
else {
return nil
}

// OpenAI/OpenRouter format
if let choices = json["choices"] as? [[String: Any]],
let delta = choices.first?["delta"] as? [String: Any],
let content = delta["content"] as? String {
return content
}

// Ollama format
if let message = json["message"] as? [String: Any],
let content = message["content"] as? String,
!content.isEmpty {
return content
}

return nil
}

private func parseUsageFromChunk(_ jsonString: String) -> AITokenUsage? {
guard let data = jsonString.data(using: .utf8),
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any]
else {
return nil
}

// OpenAI/OpenRouter format: usage object in the chunk
if let usage = json["usage"] as? [String: Any],
let promptTokens = usage["prompt_tokens"] as? Int,
let completionTokens = usage["completion_tokens"] as? Int {
return AITokenUsage(inputTokens: promptTokens, outputTokens: completionTokens)
}

// Ollama format: done=true with eval counts
if let done = json["done"] as? Bool, done,
let promptEval = json["prompt_eval_count"] as? Int,
let evalCount = json["eval_count"] as? Int {
return AITokenUsage(inputTokens: promptEval, outputTokens: evalCount)
}

return nil
}

// MARK: - Model Fetching

private func fetchOpenAIModels() async throws -> [String] {
Expand Down
Loading
Loading