From f4faa70098c76f7e2b50337217132cedcc607d39 Mon Sep 17 00:00:00 2001 From: Ayman Hamed Date: Wed, 5 Aug 2026 10:54:00 +0300 Subject: [PATCH] feat(context): keep the latest read of each file inline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stops the re-read loop on multi-file tasks: once the agent has read a file, its most-recent read stays in live context instead of being externalized to an artifact, so the model doesn't re-read the same file to "confirm." - ContextManager gains `keepLatestReadsInline` (default true) and `readToolNames` (default ["read_file"]). - modelMessages() precomputes the newest `rest` index per file path for keep-worthy reads (in readToolNames, not error, ≤ maxActiveResultChars), and the oldest-first eviction loop skips externalizing the whole exchange that holds a latest-per-path read (preserving tool_call/result pairing). - Bounded: only over inlineBudgetChars, size-capped per read, and older kept reads still fall out via Conversation.trim() — can't grow unbounded. Tests (5): latest read kept while older shell evicted; only newest read of a path kept; newest read of each distinct path kept; large read still externalized (size guard); flag off restores old behavior. Full suite 160 green. Co-Authored-By: Claude Opus 4.8 --- .../Context/ContextManager.swift | 88 ++++++++++++-- .../SwiftAgentKitTests.swift | 111 ++++++++++++++++++ 2 files changed, 189 insertions(+), 10 deletions(-) diff --git a/Sources/SwiftAgentKit/Context/ContextManager.swift b/Sources/SwiftAgentKit/Context/ContextManager.swift index b8c4d53..66f18df 100644 --- a/Sources/SwiftAgentKit/Context/ContextManager.swift +++ b/Sources/SwiftAgentKit/Context/ContextManager.swift @@ -49,18 +49,33 @@ public final class ContextManager: @unchecked Sendable { /// before. A single huge tool result still trips the budget (so it's offloaded). public var inlineBudgetChars: Int + /// Keep the most-recent read of each distinct file inline instead of + /// externalizing it, so the agent stops re-reading the same file in a loop. + /// Only the newest read per path is protected, and only if it's not an error + /// and is within `maxActiveResultChars` (large reads still externalize). + public var keepLatestReadsInline: Bool + + /// Tool names whose (non-error) result is a file read whose latest-per-path + /// output should be kept inline. A set so the generic manager isn't coupled to + /// any particular tools product. + public var readToolNames: Set + public init( store: any ArtifactStore = InMemoryArtifactStore(), maxActiveResultChars: Int = 8_000, ledgerEntries: Int = 20, summaryLength: Int = 320, - inlineBudgetChars: Int = 16_000 + inlineBudgetChars: Int = 16_000, + keepLatestReadsInline: Bool = true, + readToolNames: Set = ["read_file"] ) { self.store = store self.maxActiveResultChars = maxActiveResultChars self.ledgerEntries = ledgerEntries self.summaryLength = summaryLength self.inlineBudgetChars = inlineBudgetChars + self.keepLatestReadsInline = keepLatestReadsInline + self.readToolNames = readToolNames } /// The retrieval tools the model uses to pull full outputs back from the @@ -102,6 +117,22 @@ public final class ContextManager: @unchecked Sendable { let activeStart = activeExchangeStart(in: rest) ?? rest.count + // Map each tool-call id → its call, so a receipt can name the invocation + // (e.g. the shell command), not just the tool. Built here (before the + // eviction loop) so we can look up a read's file path while deciding what + // to keep. + var callsByID: [String: AgentToolCall] = [:] + for message in rest where message.role == .assistant { + for call in message.toolCalls ?? [] { callsByID[call.id] = call } + } + + // The `rest` index of the most-recent read of each distinct file path — + // these exchanges are kept inline so the model doesn't re-read the file. + let latestReadIndexByPath = keepLatestReadsInline + ? latestReadIndices(in: rest, upTo: activeStart, callsByID: callsByID) + : [:] + 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 → @@ -114,13 +145,21 @@ public final class ContextManager: @unchecked Sendable { while i < activeStart && remaining > inlineBudgetChars { if rest[i].role == .assistant, rest[i].toolCalls?.isEmpty == false { var j = i + 1 - var exchangeChars = rest[i].content.count var span = [i] while j < activeStart, rest[j].role == .tool { span.append(j) - exchangeChars += rest[j].toolResults?.reduce(0) { $0 + $1.result.count } ?? 0 j += 1 } + // 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 @@ -129,13 +168,6 @@ public final class ContextManager: @unchecked Sendable { } } - // Map each tool-call id → its call, so a receipt can name the invocation - // (e.g. the shell command), not just the tool. - var callsByID: [String: AgentToolCall] = [:] - for message in rest where message.role == .assistant { - for call in message.toolCalls ?? [] { callsByID[call.id] = call } - } - // Receipts for the externalized (older) tool results only. var receipts: [ToolReceipt] = [] for index in externalized.sorted() where rest[index].role == .tool { @@ -234,6 +266,42 @@ public final class ContextManager: @unchecked Sendable { return receipt } + /// For each distinct file path, the highest `rest` index (< `activeStart`) of a + /// keep-worthy read of that path: a `readToolNames` result that isn't an error + /// and is within `maxActiveResultChars`. Reads appear in order, so the last + /// seen per path wins. + private func latestReadIndices( + in rest: [AgentMessage], + upTo activeStart: Int, + callsByID: [String: AgentToolCall] + ) -> [String: Int] { + var byPath: [String: Int] = [:] + var index = 0 + while index < activeStart { + let message = rest[index] + if message.role == .tool { + for result in message.toolResults ?? [] { + guard readToolNames.contains(result.toolName ?? ""), + !result.isError, + result.result.count <= maxActiveResultChars, + let call = callsByID[result.toolCallId], + let path = Self.pathArgument(for: call) + else { continue } + byPath[path] = index // ascending index → newest per path + } + } + index += 1 + } + return byPath + } + + /// The file path a call targets (the `path` argument), trimmed; nil if absent. + private static func pathArgument(for call: AgentToolCall) -> String? { + guard let value = call.parameters["path"]?.value as? String else { return nil } + let trimmed = value.trimmingCharacters(in: .whitespaces) + return trimmed.isEmpty ? nil : trimmed + } + /// A short, single-line hint of a call's most salient argument (the command, /// path, query, etc.) for the ledger line. private static func argHint(for call: AgentToolCall) -> String? { diff --git a/Tests/SwiftAgentKitTests/SwiftAgentKitTests.swift b/Tests/SwiftAgentKitTests/SwiftAgentKitTests.swift index 58c655a..aaf497b 100644 --- a/Tests/SwiftAgentKitTests/SwiftAgentKitTests.swift +++ b/Tests/SwiftAgentKitTests/SwiftAgentKitTests.swift @@ -1882,6 +1882,117 @@ private func firstArtifactID(in text: String) -> String? { #expect(out.contains { $0.role == .assistant && ($0.toolCalls?.isEmpty == false) }) } +// MARK: - ContextManager: keep the latest read of each file inline + +private func readCall(_ id: String, path: String) -> AgentToolCall { + AgentToolCall(id: id, name: "read_file", parameters: ["path": AnyCodable(path)]) +} + +@Test func testKeepsLatestReadWhileEvictingOlderShell() async { + // Over budget: the older read of a file is KEPT inline (so the model doesn't + // re-read it) and the newer run_shell exchange is externalized instead. + let manager = ContextManager(inlineBudgetChars: 900) + let readResult = "READ_KEPT " + String(repeating: "r", count: 300) + let shellResult = "SHELL_MARKER " + String(repeating: "s", count: 600) + let messages: [AgentMessage] = [ + .user("start"), + .assistant(content: "", toolCalls: [readCall("r1", path: "/config.json")]), + .tool(results: [.success(toolCallId: "r1", toolName: "read_file", result: readResult)]), + .assistant(content: "", toolCalls: [AgentToolCall(id: "s1", name: "run_shell")]), + .tool(results: [.success(toolCallId: "s1", toolName: "run_shell", result: shellResult)]), + .assistant("thinking"), + .user("continue"), + ] + + let out = await manager.modelMessages(messages) { $0 } + + #expect(out.contains { $0.role == .tool && $0.content.contains("READ_KEPT") }) // read kept inline + #expect(!out.contains { $0.role == .tool && $0.content.contains("SHELL_MARKER") }) // shell externalized + #expect(out.contains { $0.role == .system && $0.content.contains("ledger") }) +} + +@Test func testKeepsOnlyNewestReadOfSamePath() async { + let manager = ContextManager(inlineBudgetChars: 700) + let messages: [AgentMessage] = [ + .user("start"), + .assistant(content: "", toolCalls: [readCall("r1", path: "/a.txt")]), + .tool(results: [.success(toolCallId: "r1", toolName: "read_file", + result: "OLD_A " + String(repeating: "o", count: 400))]), + .assistant(content: "", toolCalls: [readCall("r2", path: "/a.txt")]), + .tool(results: [.success(toolCallId: "r2", toolName: "read_file", + result: "NEW_A " + String(repeating: "n", count: 400))]), + .assistant("done"), + .user("continue"), + ] + + let out = await manager.modelMessages(messages) { $0 } + + #expect(out.contains { $0.role == .tool && $0.content.contains("NEW_A") }) // newest kept + #expect(!out.contains { $0.role == .tool && $0.content.contains("OLD_A") }) // older read externalized +} + +@Test func testKeepsNewestReadOfEachDistinctPath() async { + let manager = ContextManager(inlineBudgetChars: 900) + let messages: [AgentMessage] = [ + .user("start"), + .assistant(content: "", toolCalls: [readCall("r1", path: "/a.txt")]), + .tool(results: [.success(toolCallId: "r1", toolName: "read_file", result: "AAA_" + String(repeating: "a", count: 300))]), + .assistant(content: "", toolCalls: [readCall("r2", path: "/b.txt")]), + .tool(results: [.success(toolCallId: "r2", toolName: "read_file", result: "BBB_" + String(repeating: "b", count: 300))]), + .assistant(content: "", toolCalls: [AgentToolCall(id: "s1", name: "run_shell")]), + .tool(results: [.success(toolCallId: "s1", toolName: "run_shell", result: "SHELL_" + String(repeating: "s", count: 600))]), + .assistant("done"), + .user("continue"), + ] + + let out = await manager.modelMessages(messages) { $0 } + + #expect(out.contains { $0.role == .tool && $0.content.contains("AAA_") }) // /a.txt kept + #expect(out.contains { $0.role == .tool && $0.content.contains("BBB_") }) // /b.txt kept + #expect(!out.contains { $0.role == .tool && $0.content.contains("SHELL_") }) // shell externalized +} + +@Test func testLargeReadStillExternalized() async { + // A read bigger than maxActiveResultChars is NOT protected (already paged). + let manager = ContextManager(maxActiveResultChars: 100, inlineBudgetChars: 300) + let big = "BIGREAD_" + String(repeating: "x", count: 500) + let messages: [AgentMessage] = [ + .user("start"), + .assistant(content: "", toolCalls: [readCall("r1", path: "/huge.txt")]), + .tool(results: [.success(toolCallId: "r1", toolName: "read_file", result: big)]), + .assistant(content: "", toolCalls: [AgentToolCall(id: "s1", name: "run_shell")]), + .tool(results: [.success(toolCallId: "s1", toolName: "run_shell", result: "RECENT_kept")]), + .assistant("done"), + .user("continue"), + ] + + let out = await manager.modelMessages(messages) { $0 } + + #expect(!out.contains { $0.role == .tool && $0.content.contains("BIGREAD_") }) // size guard → externalized + #expect(out.contains { $0.role == .system && $0.content.contains("ledger") }) +} + +@Test func testKeepLatestReadsDisabledRestoresOldBehavior() async { + let manager = ContextManager(inlineBudgetChars: 900, keepLatestReadsInline: false) + let readResult = "READ_OLD " + String(repeating: "r", count: 300) + let messages: [AgentMessage] = [ + .user("start"), + .assistant(content: "", toolCalls: [readCall("r1", path: "/config.json")]), + .tool(results: [.success(toolCallId: "r1", toolName: "read_file", result: readResult)]), + .assistant(content: "", toolCalls: [AgentToolCall(id: "s1", name: "run_shell")]), + .tool(results: [.success(toolCallId: "s1", toolName: "run_shell", + result: "SHELL_recent " + String(repeating: "s", count: 600))]), + .assistant("done"), + .user("continue"), + ] + + let out = await manager.modelMessages(messages) { $0 } + + // Old behavior: oldest exchange (the read) is externalized, recent shell kept. + #expect(!out.contains { $0.role == .tool && $0.content.contains("READ_OLD") }) + #expect(out.contains { $0.role == .tool && $0.content.contains("SHELL_recent") }) +} + // MARK: - Skill store + learn_skill (self-improvement) private func tempSkillDir() -> URL {