diff --git a/Sources/SwiftAgentKit/Context/ContextManager.swift b/Sources/SwiftAgentKit/Context/ContextManager.swift index 5f7cd03..9589354 100644 --- a/Sources/SwiftAgentKit/Context/ContextManager.swift +++ b/Sources/SwiftAgentKit/Context/ContextManager.swift @@ -104,6 +104,55 @@ public final class ContextManager: @unchecked Sendable { /// aren't worth a disk write. public var eagerPersistMinChars = 1_000 + /// Low watermark for eviction hysteresis: a budget breach evicts down to + /// this fraction of `inlineBudgetChars`, then the evicted set FREEZES + /// until the budget is breached again. Keeps the model-facing prefix + /// byte-stable between eviction events so provider prompt caches hit. + public var evictionTargetFraction: Double = 0.5 + + /// Head-message ids of exchanges evicted by previous calls. Sticky — + /// never un-evicted — so the prefix cannot flap. + private var stickyEvicted: Set = [] + + private func currentStickyEvicted() -> Set { + lock.lock(); defer { lock.unlock() } + return stickyEvicted + } + + private func rememberEvicted(_ ids: [UUID]) { + lock.lock(); defer { lock.unlock() } + stickyEvicted.formUnion(ids) + } + + /// Completed tool exchanges before `activeStart`: the assistant tool-call + /// turn (`head`), the indices of the whole exchange, and its char size. + private struct ExchangeSpan { + let head: Int + let indices: [Int] + let chars: Int + } + + private func exchangeSpans(in rest: [AgentMessage], upTo activeStart: Int) -> [ExchangeSpan] { + var spans: [ExchangeSpan] = [] + var i = 0 + while i < activeStart { + guard rest[i].role == .assistant, rest[i].toolCalls?.isEmpty == false else { i += 1; continue } + var j = i + 1 + var indices = [i] + while j < activeStart, rest[j].role == .tool { + indices.append(j) + j += 1 + } + let chars = indices.reduce(0) { sum, idx in + sum + rest[idx].content.count + + (rest[idx].toolResults?.reduce(0) { $0 + $1.result.count } ?? 0) + } + spans.append(ExchangeSpan(head: i, indices: indices, chars: chars)) + i = j + } + return spans + } + /// Save completed tool results to the store AS THEY FINISH, not only when /// sifting later spills them. Without this, a short run that never exceeds /// the inline budget stores nothing — and a restart-surviving store @@ -173,39 +222,41 @@ public final class ContextManager: @unchecked Sendable { : [:] let protectedIndices = Set(latestReadIndexByPath.values) - // Over budget: externalize whole tool exchanges OLDEST-FIRST until we're - // back under budget, keeping the most RECENT tool results inline. This - // preserves the working set an iterative task needs (run → read error → - // fix → rerun) while still capping growth. An exchange is an + // Over budget: externalize whole tool exchanges OLDEST-FIRST, keeping + // the most RECENT tool results inline. An exchange is an // assistant-with-toolCalls turn plus its following tool-result messages; // evicting whole exchanges keeps tool_call/result pairing valid. + // + // CACHE-STABILITY HYSTERESIS: continuous eviction (evict just enough, + // every turn) changes the model-facing prefix on every call, so + // provider prompt caches never hit. Instead, evictions are STICKY + // (remembered by message id, never undone) and a budget breach evicts + // down to `evictionTargetFraction` of the budget — then the evicted + // set, the ledger, and the whole prefix stay byte-stable until roughly + // half a budget of new content accumulates. var externalized = Set() // indices in `rest` to move to the ledger var remaining = totalChars - var i = 0 - while i < activeStart && remaining > inlineBudgetChars { - if rest[i].role == .assistant, rest[i].toolCalls?.isEmpty == false { - var j = i + 1 - var span = [i] - while j < activeStart, rest[j].role == .tool { - span.append(j) - j += 1 - } + let spans = exchangeSpans(in: rest, upTo: activeStart) + let sticky = currentStickyEvicted() + for span in spans where sticky.contains(rest[span.head].id) { + span.indices.forEach { externalized.insert($0) } + remaining -= span.chars + } + if remaining > inlineBudgetChars { + let fraction = min(1, max(0, evictionTargetFraction)) + let target = Int(Double(inlineBudgetChars) * fraction) + var newlyEvicted: [UUID] = [] + for span in spans where !externalized.contains(span.head) { + guard remaining > target else { break } // Keep the whole exchange inline if it holds a latest-per-path // read (preserving tool_call/result pairing); evict the rest. - if span.contains(where: { protectedIndices.contains($0) }) { - i = j - continue - } - let exchangeChars = span.reduce(0) { sum, idx in - sum + rest[idx].content.count - + (rest[idx].toolResults?.reduce(0) { $0 + $1.result.count } ?? 0) - } - span.forEach { externalized.insert($0) } - remaining -= exchangeChars - i = j - } else { - i += 1 + // Sticky evictions above are exempt — never un-evict. + if span.indices.contains(where: { protectedIndices.contains($0) }) { continue } + span.indices.forEach { externalized.insert($0) } + remaining -= span.chars + newlyEvicted.append(rest[span.head].id) } + if !newlyEvicted.isEmpty { rememberEvicted(newlyEvicted) } } // Receipts for the externalized (older) tool results only. diff --git a/Tests/SwiftAgentKitTests/SwiftAgentKitTests.swift b/Tests/SwiftAgentKitTests/SwiftAgentKitTests.swift index e2f2b63..11929b9 100644 --- a/Tests/SwiftAgentKitTests/SwiftAgentKitTests.swift +++ b/Tests/SwiftAgentKitTests/SwiftAgentKitTests.swift @@ -2704,3 +2704,56 @@ private struct NoopTool: AgentTool { let listed = await store.list(limit: 10) #expect(listed.count == 1) // no duplicate for the same call } + +// MARK: - Cache-stable eviction (hysteresis) + +private func completedShellExchange(id i: Int, size: Int) -> [AgentMessage] { + [ + .assistant(content: "", toolCalls: [AgentToolCall(id: "hx\(i)", name: "run_shell")]), + .tool(results: [.success(toolCallId: "hx\(i)", toolName: "run_shell", + result: String(repeating: "o", count: size))]), + ] +} + +@Test func testEvictionHysteresisKeepsPrefixByteStable() async { + // The cache contract: between eviction events, everything the model sees + // before the newest content is BYTE-IDENTICAL to the previous call — + // system block (incl. ledger) and all prior messages. Continuous eviction + // (the old behavior) broke this on every call once over budget. + let manager = ContextManager(inlineBudgetChars: 1_000) + var messages: [AgentMessage] = [.system("You are helpful."), .user("do the long task")] + for i in 0..<6 { messages.append(contentsOf: completedShellExchange(id: i, size: 400)) } + messages.append(.assistant("progress so far")) + + let first = await manager.modelMessages(messages) { $0 } + + // One SMALL new completed exchange (within the hysteresis slack). + messages.append(contentsOf: completedShellExchange(id: 99, size: 100)) + messages.append(.assistant("more progress")) + let second = await manager.modelMessages(messages) { $0 } + + // The entire first output is a byte-identical prefix of the second. + #expect(second.count > first.count) + for (index, message) in first.enumerated() { + #expect(second[index].role == message.role) + #expect(second[index].content == message.content) + } +} + +@Test func testEvictionBreachEvictsToLowWatermark() async { + // A breach must evict PAST the budget down to the target fraction, buying + // slack so subsequent turns don't each trigger a fresh eviction. + let manager = ContextManager(inlineBudgetChars: 1_000) // target = 500 + var messages: [AgentMessage] = [.user("go")] + for i in 0..<6 { messages.append(contentsOf: completedShellExchange(id: i, size: 400)) } + messages.append(.assistant("done step")) + + let out = await manager.modelMessages(messages) { $0 } + + // 6 exchanges ≈ 2400 chars; target 500 → at least 5 evicted (ledger'd). + let system = out.first { $0.role == .system }?.content ?? "" + let ledgerLines = system.split(separator: "\n").filter { $0.hasPrefix("- ") }.count + #expect(ledgerLines >= 5) + // Inline tool results that remain: at most one exchange. + #expect(out.filter { $0.role == .tool }.count <= 1) +}