feat: production-grade upgrade — streaming, structured JSON, metrics, safe prompts

[rahulraonatarajan]

Summary

• Token-accurate MLX budgeting via MLXTokenAwareBackend wrapping a real Tokenizer
• Streaming output via StreamingLLMBackend protocol with ChainEvent async stream
• Structured JSON extraction via runJSON<T: Decodable>(...) with retry and code-fence stripping
• Rich results via ChainResult with source chunks and ChainMetrics (timing, call counts, token estimates)
• Safe prompt rendering via ChainPromptBuilder with opt-in <source> tag wrapping
• CodeBlockAwareChunker for Markdown/code-heavy notes
• 194 tests (87 new), all passing, no model downloads required
• Removed stale process docs (pr-body.md, pr-summary.md, repo-adoption-metadata.md, repo-improvements.md)

Why this matters

Turns mlx-swift-chain from a chain orchestration layer into a production-grade document processing package. Every new feature is opt-in with backward-compatible defaults (.raw prompt style, run() still returns String). Designed for real app workflows: meeting transcript → action items, crash report triage, log analysis, Markdown summarization.

New public types (11)

Type	Purpose
ChainResult	Rich result with .text, .sourceChunks, .metrics
ChainMetrics	Chunk count, map/reduce call counts, elapsed time, estimated tokens, throughput
ChainEvent	Streaming event: .chunk(String), .progress(Update), .result(ChainResult)
MLXTokenCounter	TokenCounter backed by a real MLXLMCommon Tokenizer
MLXTokenAwareBackend	MLX backend with TokenAwareBackend + StreamingLLMBackend conformance
StreamingLLMBackend	Protocol extending LLMBackend with token-by-token streaming
ChainPromptBuilder	Centralized prompt rendering with metadata-enriched <source> tags
PromptStyle	.raw (legacy) or .delimited (<source> tag wrapping)
CodeBlockAwareChunker	Splits at paragraph boundaries, never inside fenced code blocks
StructuredOutputError	.invalidJSON, .schemaMismatch, .retriesExhausted
JSONRepairHelper (internal)	Strips code fences, extracts outermost JSON

Modified public types

Type	Change
DocumentChain	Added runWithMetadata() and stream() protocol requirements with default implementations
ChainExecutionOptions	Added promptStyle field (default .raw)
ChainProgress.Update	Added optional partialMetrics field
ChainRunner	Added chainResult, partialText, runStreaming()
MLXBackend	Added StreamingLLMBackend conformance

Backward compatibility

All existing run() signatures, protocols, and types are unchanged. New protocol requirements on DocumentChain have default implementations, so existing conformances compile without changes. The existing 107 tests continue to pass alongside the 87 new tests.

Test plan

• swift build — 0 errors, 0 warnings
• swift test — 194 tests, 0 failures
• Public API canary tests confirm no symbol removals
• All tests use mock backends — no model downloads required

Made with Cursor

---

Note

Medium Risk
Adds new execution paths (streaming, prompt rendering, structured JSON retries) and expands the DocumentChain surface area; behavior remains backward-compatible by default but changes touch core chain execution and prompt construction.

Overview
Introduces streaming chain execution via StreamingLLMBackend and ChainEvent, with MLXBackend/MLXTokenAwareBackend streaming support and SwiftUI ChainRunner.runStreaming(...) + partialText.

Adds rich results + metrics: runWithMetadata(...) returns ChainResult (text + source chunks + optional ChainMetrics), chains now accumulate call counts/timing and (when TokenAwareBackend is available) token estimates.

Centralizes prompt construction in ChainPromptBuilder with a new PromptStyle (.raw default, .delimited <source> wrapping) wired through ChainExecutionOptions, and adds structured JSON extraction via DocumentChain.runJSON(...) with deterministic JSON repair + retry-on-decode-failure.

Improves long-document handling with token-accurate budgeting (MLXTokenCounter, MLXTokenAwareBackend) and more conservative token-to-word fallback rechunking, adds CodeBlockAwareChunker for fenced-code-safe Markdown splitting, expands tests significantly, and removes stale process docs under docs/.

^{Reviewed by Cursor Bugbot for commit 45dc7e1. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 45dc7e1. Configure here.}

[cursor]

MapReduceChain stream ignores caller's progress parameter

Medium Severity

MapReduceChain.stream() accepts a progress parameter but never uses it. Instead, it creates a local progressRelay and passes that to runWithMetadata. The caller's ChainProgress object never receives any updates. This is inconsistent with StuffChain.stream(), which correctly reports to the caller's progress. When AdaptiveChain routes to MapReduceChain.stream, the caller's progress observer is silently dropped.

 

^{Reviewed by Cursor Bugbot for commit 45dc7e1. Configure here.}

[cursor]

Concurrent map omits token count tracking in metrics

Medium Severity

concurrentMap only updates acc.mapCallCount after collecting results but never accumulates acc.inputTokens or acc.outputTokens. In contrast, sequentialMap tracks both input and output token counts per call when using a TokenAwareBackend. This means ChainMetrics.estimatedInputTokens and estimatedOutputTokens will be missing map-phase tokens whenever maxConcurrentMapTasks > 1.

 

^{Reviewed by Cursor Bugbot for commit 45dc7e1. Configure here.}