diff --git a/Package.swift b/Package.swift index 0e796c2..3a36f13 100644 --- a/Package.swift +++ b/Package.swift @@ -20,7 +20,7 @@ let package = Package( .library(name: "SwiftAgentKitMac", targets: ["SwiftAgentKitMac"]), ], dependencies: [ - .package(url: "https://github.com/ayman3000/LLMProviderKit.git", from: "0.1.0-alpha.41"), + .package(url: "https://github.com/ayman3000/LLMProviderKit.git", from: "0.1.0-alpha.42"), .package(url: "https://github.com/swiftlang/swift-syntax.git", from: "603.0.0"), .package(url: "https://github.com/modelcontextprotocol/swift-sdk.git", from: "0.12.1"), .package(url: "https://github.com/swiftlang/swift-docc-plugin.git", from: "1.4.0"), diff --git a/Sources/SwiftAgentKit/Context/ArtifactStore.swift b/Sources/SwiftAgentKit/Context/ArtifactStore.swift index d63f57b..5dc58f3 100644 --- a/Sources/SwiftAgentKit/Context/ArtifactStore.swift +++ b/Sources/SwiftAgentKit/Context/ArtifactStore.swift @@ -85,7 +85,9 @@ public struct ArtifactMatch: Sendable, Equatable { /// A store is shared by reference between agents (parent and sub-agents), so /// conformers must be reference types (classes or actors). public protocol ArtifactStore: AnyObject, Sendable { - /// Store a full output and return its artifact record (with a fresh id). + /// Store a full output and return its artifact record. Built-in stores + /// derive the id from the tool-call id and content (`Artifact.makeID`) + /// and return the existing artifact instead of storing it twice. /// `toolName` is the producing tool (nil when unknown) — persistence /// filters key off it (e.g. keep run_shell logs, skip re-derivable reads). func save(_ content: String, description: String, toolCallID: String?, toolName: String?) async -> Artifact @@ -112,6 +114,24 @@ public protocol ListableArtifactStore: ArtifactStore { func list(limit: Int) async -> [ArtifactSummary] } +public extension Artifact { + /// The id a store gives an artifact: `artifact-` + 12 lowercase hex. + /// + /// With a tool-call id it is derived from that id AND the content, so the + /// same evicted output gets the same id after a relaunch (its receipt is + /// the same bytes, which keeps the prompt cache warm) and a store can + /// reuse the artifact instead of writing a duplicate. The content is part + /// of the key because some providers reuse call ids (`call_0`) across + /// runs: a different output under the same call id gets its own artifact. + /// Without a call id the id is random, as before. + static func makeID(toolCallID: String?, content: String) -> String { + guard let toolCallID, !toolCallID.isEmpty else { + return "artifact-" + UUID().uuidString.lowercased().replacingOccurrences(of: "-", with: "").prefix(12) + } + return "artifact-" + PromptDigest.hex(toolCallID + "\u{0}" + content).prefix(12) + } +} + /// In-memory artifact store (process lifetime). Good for apps and tests; swap in /// a file-backed store for durability. public actor InMemoryArtifactStore: ArtifactStore, ListableArtifactStore { @@ -120,8 +140,10 @@ public actor InMemoryArtifactStore: ArtifactStore, ListableArtifactStore { public init() {} public func save(_ content: String, description: String, toolCallID: String?, toolName: String?) -> Artifact { - let id = "artifact-" + UUID().uuidString.lowercased().replacingOccurrences(of: "-", with: "").prefix(12) - let artifact = Artifact(id: String(id), toolCallID: toolCallID, toolName: toolName, + let id = Artifact.makeID(toolCallID: toolCallID, content: content) + // Same call, same output: already stored — reuse it. + if let existing = artifacts[id] { return existing } + let artifact = Artifact(id: id, toolCallID: toolCallID, toolName: toolName, description: description, content: content) artifacts[artifact.id] = artifact return artifact diff --git a/Sources/SwiftAgentKit/Context/FileArtifactStore.swift b/Sources/SwiftAgentKit/Context/FileArtifactStore.swift index b9ae154..4c62705 100644 --- a/Sources/SwiftAgentKit/Context/FileArtifactStore.swift +++ b/Sources/SwiftAgentKit/Context/FileArtifactStore.swift @@ -61,8 +61,11 @@ public actor FileArtifactStore: ArtifactStore, ListableArtifactStore { // MARK: - ArtifactStore public func save(_ content: String, description: String, toolCallID: String?, toolName: String?) -> Artifact { - let id = "artifact-" + UUID().uuidString.lowercased().replacingOccurrences(of: "-", with: "").prefix(12) - let artifact = Artifact(id: String(id), toolCallID: toolCallID, toolName: toolName, + let id = Artifact.makeID(toolCallID: toolCallID, content: content) + // Same call, same output — saved earlier this session or before a + // relaunch: reuse it instead of writing a duplicate to disk. + if let existing = get(id) { return existing } + let artifact = Artifact(id: id, toolCallID: toolCallID, toolName: toolName, description: description, content: content) memory[artifact.id] = artifact if persistFilter(toolName) { diff --git a/Sources/SwiftAgentKit/Core/Agent.swift b/Sources/SwiftAgentKit/Core/Agent.swift index 98761d1..bcfbe55 100644 --- a/Sources/SwiftAgentKit/Core/Agent.swift +++ b/Sources/SwiftAgentKit/Core/Agent.swift @@ -161,6 +161,15 @@ public struct AgentConfig: Sendable { /// Compact before the next call once the last prompt passes this share of /// the window. nil = only on a context-length error. public var compactAtFraction: Double? + /// Build the system prompt's dynamic parts — persistent memory, the app's + /// per-run context, the skill index — once, and reuse them on every + /// later run (`Agent.frozenRunContext`) until `invalidateRunContext()` or + /// a compaction. Off (default): rebuilt at the start of every run, as + /// before. On, a conversation's system prompt stays the same bytes from + /// run to run, so provider prompt caches cover the whole conversation; + /// what the agent files during it reaches the model through the tool + /// result that filed it, not through the prompt. + public var freezeRunContext: Bool public init( provider: any LLMProvider, @@ -194,10 +203,12 @@ public struct AgentConfig: Sendable { toolGroups: [DeferredToolGroup] = [], maxReasoningSeconds: TimeInterval? = nil, contextCompactor: (any ContextCompactor)? = nil, - compactAtFraction: Double? = nil + compactAtFraction: Double? = nil, + freezeRunContext: Bool = false ) { self.contextCompactor = contextCompactor self.compactAtFraction = compactAtFraction + self.freezeRunContext = freezeRunContext self.provider = provider self.model = model self.temperature = temperature @@ -395,6 +406,15 @@ public actor Agent { /// Set at the start of a run; the next request built records its digest. private var firstRequestPending = false + /// The frozen parts of the system prompt (`AgentConfig.freezeRunContext`): + /// nil until a run builds them, and after `invalidateRunContext()` or a + /// compaction. + public private(set) var frozenRunContext: FrozenRunContext? + /// Bumped by `invalidateRunContext()`, `restoreRunContext` and + /// compaction. `runContextParts` captures this before its first await; + /// if it changed by the time the build finishes, the build was + /// overtaken by one of those and must not freeze its (now stale) result. + private var runContextGeneration = 0 /// The overflow safety net's cut: the id of the last stored message calls /// leave out (with every older evictable step). Sticky until the request /// passes the bound again; cleared by compaction, or when that message is @@ -600,6 +620,60 @@ public actor Agent { + lines.joined(separator: "\n") } + /// The tool line of the system prompt (empty without tools). Deferred + /// tools are left out, loaded or not. + private func toolInstructionLine(_ registered: [any AgentTool]) -> String { + guard !registered.isEmpty else { return "" } + let hiddenNames = Set(deferredGroups.flatMap(\.toolNames)) + let toolNames = registered.map { $0.name }.filter { !hiddenNames.contains($0) }.joined(separator: ", ") + return """ + + You have access to the following tools: \(toolNames). + IMPORTANT: When the user's request requires action (reading files, running commands, searching, creating, etc.), you MUST use the available tools instead of answering from memory. Call the appropriate tool to get real information, then use the tool results to formulate your answer. Do not guess or hallucinate results — always call the tool to get the actual data. + """ + } + + /// Memory, the per-run context and the skill index for this run: the + /// frozen ones when frozen for this base, otherwise read now (and frozen + /// when `freezeRunContext` is on). + private func runContextParts(base: String, toolLine: String, groupIndex: String, + toolNames: [String]) async -> FrozenRunContext { + let baseDigest = PromptDigest.hex([base, toolLine, groupIndex].joined(separator: "\u{1F}")) + if config.freezeRunContext, let frozen = frozenRunContext, frozen.baseDigest == baseDigest { + return frozen + } + let generation = runContextGeneration + var memory = "" + if let memoryStore { + memory = await memoryStore.loadContextBlock(project: memoryProject) + } + var runContext = "" + if let runContextProvider { + runContext = await runContextProvider(toolNames) + } + let parts = FrozenRunContext(baseDigest: baseDigest, memory: memory, runContext: runContext, + skillIndex: await skillRegistry.skillIndex()) + // An invalidate/restore/compaction that landed during these awaits + // already bumped the generation; the stale parts above are still + // used for this run (the running reply keeps its prompt) but are + // not frozen, so the next run rebuilds instead of reusing them. + if config.freezeRunContext && generation == runContextGeneration { frozenRunContext = parts } + return parts + } + + /// The system prompt, assembled exactly as before freezing existed: the + /// base, memory and the per-run context separated by a blank line, then + /// the tool line, the skill index and the tool-group index. + static func renderSystemPrompt(base: String, parts: FrozenRunContext, + toolLine: String, groupIndex: String) -> String { + var prompt = base + for block in [parts.memory, parts.runContext] where !block.isEmpty { + if !prompt.isEmpty { prompt += "\n\n" } + prompt += block + } + return prompt + toolLine + parts.skillIndex + groupIndex + } + /// The deferred group a tool belongs to, when it isn't loaded yet. private func unloadedGroup(ofTool name: String) -> DeferredToolGroup? { deferredGroups.first { $0.toolNames.contains(name) && !loadedToolGroupIDs.contains($0.id) } @@ -644,6 +718,28 @@ public actor Agent { try requireIdle() runContextProvider = provider } + + /// Build memory, the per-run context and the skill index again at the + /// start of the next run. Safe during a run: the running reply keeps the + /// prompt it started with. + public func invalidateRunContext() { + frozenRunContext = nil + runContextGeneration += 1 + } + + /// Use a frozen context saved with the conversation (after a relaunch). + /// A no-op when `freezeRunContext` is off: nothing ever reads this + /// context without it, so storing one would just be a value no run + /// reads back and an app that persists `frozenRunContext` after each + /// run would keep saving a copy that is not in use. Idle-only, like the + /// other setters. When freeze is on, it is used only if it was made for + /// this engine's base; otherwise the next run builds fresh. + public func restoreRunContext(_ context: FrozenRunContext) throws { + try requireIdle() + guard config.freezeRunContext else { return } + frozenRunContext = context + runContextGeneration += 1 + } public func setSkillStore(_ store: (any AgentSkillStore)?) throws { try requireIdle() skillStore = store @@ -997,51 +1093,20 @@ public actor Agent { runTaskID = task.id conversation.append(task) - // Get registered tools and strengthen system prompt (must happen before skill injection) + // Get registered tools, then the system prompt (before skill injection). let registeredToolsEarly = await tools.allTools() - var effectiveSystemPrompt = config.systemPrompt ?? "" - - // Persistent memory must be part of every model call. Loading it at run - // time ensures facts saved by earlier runs are immediately available. - if let memoryStore { - let memoryContext = await memoryStore.loadContextBlock(project: memoryProject) - if !memoryContext.isEmpty { - if !effectiveSystemPrompt.isEmpty { - effectiveSystemPrompt += "\n\n" - } - effectiveSystemPrompt += memoryContext - } - } - - // The app's per-run context (built now, so what changed since the last - // run — a lesson filed after it — is used from this one). - if let runContextProvider { - let extra = await runContextProvider(registeredToolsEarly.map(\.name)) - if !extra.isEmpty { - if !effectiveSystemPrompt.isEmpty { effectiveSystemPrompt += "\n\n" } - effectiveSystemPrompt += extra - } - } - - if !registeredToolsEarly.isEmpty { - let hiddenNames = Set(deferredGroups.flatMap(\.toolNames)) - let toolNames = registeredToolsEarly.map { $0.name }.filter { !hiddenNames.contains($0) }.joined(separator: ", ") - let toolInstruction = """ - - You have access to the following tools: \(toolNames). - IMPORTANT: When the user's request requires action (reading files, running commands, searching, creating, etc.), you MUST use the available tools instead of answering from memory. Call the appropriate tool to get real information, then use the tool results to formulate your answer. Do not guess or hallucinate results — always call the tool to get the actual data. - """ - effectiveSystemPrompt += toolInstruction - } - - // Model-driven skills: a compact, query-independent index of every - // skill goes into the system prompt; the model loads full - // instructions on demand with `use_skill`. The index is byte-stable - // (alphabetical, changes only when skills change) so the prompt - // prefix stays cache-friendly across steps and turns. - let skillIndex = await skillRegistry.skillIndex() + toolGroupIndex() - if !skillIndex.isEmpty || !effectiveSystemPrompt.isEmpty { - conversation.setSystemMessage(.system(effectiveSystemPrompt + skillIndex)) + let basePrompt = config.systemPrompt ?? "" + let toolLine = toolInstructionLine(registeredToolsEarly) + let groupIndex = toolGroupIndex() + // Persistent memory, the app's per-run context and the skill index: + // read now, so facts saved by earlier runs are available — or, with + // `freezeRunContext`, kept from the conversation's first run. + let parts = await runContextParts(base: basePrompt, toolLine: toolLine, groupIndex: groupIndex, + toolNames: registeredToolsEarly.map(\.name)) + let systemPrompt = Self.renderSystemPrompt(base: basePrompt, parts: parts, + toolLine: toolLine, groupIndex: groupIndex) + if !systemPrompt.isEmpty { + conversation.setSystemMessage(.system(systemPrompt)) } var totalTurns = 0 @@ -1790,6 +1855,10 @@ public actor Agent { estimate: conversation.estimateTokens) guard !middle.isEmpty, let checkpoint = await compactor.summarize(middle: middle, reason: reason) else { return nil } conversation.replaceNonSystemMessages(ContextCompaction.assemble(checkpoint: checkpoint, tail: tail)) + // Compaction is the sanctioned break: memory, the per-run context and + // the skill index are read again at the start of the next run. + frozenRunContext = nil + runContextGeneration += 1 overflowCutAfter = nil let after = conversation.estimateTotalTokens(conversation.allMessages()) emit(.contextCompacted(tokensBefore: before, tokensAfter: after, reason: reason)) diff --git a/Sources/SwiftAgentKit/Core/FrozenRunContext.swift b/Sources/SwiftAgentKit/Core/FrozenRunContext.swift new file mode 100644 index 0000000..0d8c601 --- /dev/null +++ b/Sources/SwiftAgentKit/Core/FrozenRunContext.swift @@ -0,0 +1,29 @@ +// +// FrozenRunContext.swift +// SwiftAgentKit +// +// The parts of the system prompt an agent re-reads at the start of a run — +// persistent memory, the app's per-run context, the skill index — kept as +// text, so a conversation can send the same system prompt run after run +// (provider prompt caching) and an app can save them with the conversation +// and restore them after a relaunch. +// + +import Foundation + +public struct FrozenRunContext: Codable, Sendable, Equatable { + /// `PromptDigest.hex` of what the engine builds itself: the configured + /// prompt, the tool line and the tool-group index. A frozen context is + /// used only with the same base, so a changed engine never sends a stale mix. + public let baseDigest: String + public let memory: String + public let runContext: String + public let skillIndex: String + + public init(baseDigest: String, memory: String, runContext: String, skillIndex: String) { + self.baseDigest = baseDigest + self.memory = memory + self.runContext = runContext + self.skillIndex = skillIndex + } +} diff --git a/Sources/SwiftAgentKit/SubAgents/SubAgentSpawner.swift b/Sources/SwiftAgentKit/SubAgents/SubAgentSpawner.swift index 3b64ea1..964323a 100644 --- a/Sources/SwiftAgentKit/SubAgents/SubAgentSpawner.swift +++ b/Sources/SwiftAgentKit/SubAgents/SubAgentSpawner.swift @@ -106,6 +106,9 @@ public final class SubAgentSpawner: @unchecked Sendable { config.subAgentModel = nil config.subAgentReasoningEffort = nil config.enableSubAgents = false // defense in depth vs. recursion + // A child's prompt is built fresh, from current memory, on every + // spawn; it runs one bounded task, so there is nothing to keep. + config.freezeRunContext = false config.maxTurns = min(config.maxTurns, max(1, config.maxSubAgentTurns)) config.tools = [] // registered explicitly below diff --git a/Tests/SwiftAgentKitReplayTests/Fixtures/wire/anthropic-tool-image.json b/Tests/SwiftAgentKitReplayTests/Fixtures/wire/anthropic-tool-image.json index f72d1f4..09f350e 100644 --- a/Tests/SwiftAgentKitReplayTests/Fixtures/wire/anthropic-tool-image.json +++ b/Tests/SwiftAgentKitReplayTests/Fixtures/wire/anthropic-tool-image.json @@ -2,7 +2,15 @@ "max_tokens" : 32000, "messages" : [ { - "content" : "screenshot the page", + "content" : [ + { + "cache_control" : { + "type" : "ephemeral" + }, + "text" : "screenshot the page", + "type" : "text" + } + ], "role" : "user" }, { diff --git a/Tests/SwiftAgentKitTests/CacheStableRequestTests.swift b/Tests/SwiftAgentKitTests/CacheStableRequestTests.swift index 5547502..aa15b9f 100644 --- a/Tests/SwiftAgentKitTests/CacheStableRequestTests.swift +++ b/Tests/SwiftAgentKitTests/CacheStableRequestTests.swift @@ -354,3 +354,117 @@ extension CacheStableRequestTests { } } } + +extension CacheStableRequestTests { + /// Steps 1–3 together, on the OpenAI wire: two runs of one chat with a + /// fact filed between them and a ContextSift batch in the second, then a + /// relaunch. One system text throughout; append-only except the one + /// batch; after the relaunch the system message and the tool definitions + /// are byte-identical. The filed fact reaches the prompt only after + /// `invalidateRunContext()` (what compaction or a user edit does). + @Test func aChatStaysCacheStableAcrossRunsABatchAndARelaunch() async throws { + let memory = TextMemoryStore("FACTS v1") + // Artifacts on disk, as the app keeps them: the relaunch opens a new + // store on the same folder. + let artifactDir = FileManager.default.temporaryDirectory + .appendingPathComponent("sak-cache-stable-\(UUID().uuidString)") + defer { try? FileManager.default.removeItem(at: artifactDir) } + func makeAgent(_ provider: ScriptedProvider) async throws -> Agent { + // Progress notes off: with the defaults, turn 3 of each run would + // carry one, each a tail-only break, and this test counts batches. + let agent = Agent(config: AgentConfig(provider: provider, model: "mock", systemPrompt: "BASE", + maxTurns: 6, tools: [StepTool()], + contextManager: ContextManager(store: FileArtifactStore(directory: artifactDir), + inlineBudgetChars: 2_600), + loopDetection: nil, progressNudgeFractions: [], + freezeRunContext: true)) + try await agent.setMemoryStore(memory) + return agent + } + let provider = ScriptedProvider(turns: Self.stepCalls(1...2) + [[]] + Self.stepCalls(3...5) + [[]]) + let agent = try await makeAgent(provider) + _ = try await agent.run("first task") + memory.text = "FACTS v2" // the keeper files a fact between the runs + _ = try await agent.run("second task") + + let requests = provider.captured + #expect(requests.count == 7) + let systems = Set(requests.map { $0.messages.first { $0.role == .system }?.content ?? "" }) + #expect(systems.count == 1, "one system text for the whole chat") + #expect(systems.first?.contains("FACTS v1") == true) + let breaks = try Self.breakIndices(requests) + #expect(breaks.compactMap { $0 }.count == 1, "exactly one ContextSift batch") + let batch = try #require(breaks.firstIndex { $0 != nil }) + let at = try #require(breaks[batch]) + #expect(requests[batch + 1].messages[at].content.hasPrefix(ContextManager.receiptHeader)) + let tools = try requests.map { Self.toolsSection(try Self.wireBody($0)) } + #expect(Set(tools.compactMap { $0 }).count == 1) + // The streaming body breaks at exactly the same pair, and nowhere else. + let streamed = try requests.map(Self.streamWireBody) + for (pair, index) in breaks.enumerated() { + let prefix = try #require(Self.messagesPrefix(streamed[pair])) + #expect(streamed[pair + 1].starts(with: prefix) == (index == nil), "streamed pair \(pair)") + } + + // Relaunch: the frozen context goes through JSON (as the app saves it); + // a new agent gets the same stored history and restores it. + let saved = try JSONEncoder().encode(try #require(await agent.frozenRunContext)) + let relaunched = ScriptedProvider(turns: []) + let again = try await makeAgent(relaunched) + again.conversation.append(agent.conversation.allMessages().filter { $0.role != .system }) + try await again.restoreRunContext(try JSONDecoder().decode(FrozenRunContext.self, from: saved)) + _ = try await again.run("third task") + let before = try #require(requests.last) + let after = try #require(relaunched.captured.first) + #expect(after.messages.first { $0.role == .system }?.content == systems.first) + #expect(Self.toolsSection(try Self.wireBody(after)) == Self.toolsSection(try Self.wireBody(before))) + // Across the relaunch the prefix holds up to the first evicted step: + // the system prompt and the first task are the same bytes. + let kept = try #require(Self.messagesPrefix(try Self.wireBody( + LLMRequest(model: "mock", messages: Array(after.messages.prefix(2)), tools: after.tools)))) + #expect(try Self.wireBody(before).starts(with: kept) && Self.wireBody(after).starts(with: kept)) + #expect(try Self.streamWireBody(after).starts(with: try #require(Self.messagesPrefix( + try Self.streamWireBody(LLMRequest(model: "mock", messages: Array(after.messages.prefix(2)), + tools: after.tools)))))) + // Receipts are the same bytes after the relaunch: each artifact id + // comes from its tool-call id and output, and the new store reuses + // the artifact already on disk instead of saving it again. + let receipts = before.messages.filter { $0.content.hasPrefix(ContextManager.receiptHeader) } + #expect(!receipts.isEmpty) + for receipt in receipts { + #expect(after.messages.contains(receipt), "a receipt changed across the relaunch") + } + let artifactFiles = try FileManager.default.contentsOfDirectory(atPath: artifactDir.path) + #expect(Set(artifactFiles.map { ($0 as NSString).deletingPathExtension }).count == artifactFiles.count / 2) + #expect(artifactFiles.count == 2 * 4, "steps 1–4 saved once each, none again on the relaunch") + // So the prefix now holds through every receipt, up to the first + // step the OLD engine kept inline (step 4). + let firstDiff = try #require(zip(before.messages, after.messages).enumerated() + .first { $0.element.0 != $0.element.1 }?.offset) + #expect(before.messages[firstDiff].toolCalls?.first?.id == "s4") + #expect(after.messages[firstDiff].content.hasPrefix(ContextManager.receiptHeader)) + let held = try #require(Self.messagesPrefix(try Self.wireBody( + LLMRequest(model: "mock", messages: Array(before.messages.prefix(firstDiff)), tools: before.tools)))) + #expect(try Self.wireBody(after).starts(with: held)) + // Remaining gap (accepted follow-up, not fixed here): the evicted set + // lives in memory, so the new engine evicts afresh and also evicts + // step 4, which the old engine still sent inline. When the evicted + // boundary is persisted this block stops recording an issue and fails. + withKnownIssue("the evicted set is rebuilt after a relaunch and evicts one more step") { + let prefix = try #require(Self.messagesPrefix(try Self.wireBody(before))) + #expect(try Self.wireBody(after).starts(with: prefix), "the relaunch rewrote an earlier byte") + } + + // The filed fact appears only once the context is invalidated + // (compaction or a user edit), and from then on. + let all = requests + relaunched.captured + #expect(!all.contains { String(decoding: (try? Self.wireBody($0)) ?? Data(), as: UTF8.self).contains("FACTS v2") }) + await again.invalidateRunContext() + _ = try await again.run("fourth task") + let fresh = try #require(relaunched.captured.last) + #expect(relaunched.captured.count == 2) + #expect(fresh.messages.first { $0.role == .system }?.content.contains("FACTS v2") == true) + #expect(Self.toolsSection(try Self.wireBody(fresh)) == Self.toolsSection(try Self.wireBody(after))) + } +} + diff --git a/Tests/SwiftAgentKitTests/FrozenRunContextTests.swift b/Tests/SwiftAgentKitTests/FrozenRunContextTests.swift new file mode 100644 index 0000000..595da4d --- /dev/null +++ b/Tests/SwiftAgentKitTests/FrozenRunContextTests.swift @@ -0,0 +1,266 @@ +import Testing +import Foundation +import LLMProviderKit +@testable import SwiftAgentKit + +/// A memory store whose block is whatever the test says it is now. +final class TextMemoryStore: AgentMemoryStore, @unchecked Sendable { + private let lock = NSLock() + private var current: String + init(_ text: String) { current = text } + var text: String { + get { lock.withLock { current } } + set { lock.withLock { current = newValue } } + } + func save(_ entry: AgentMemoryEntry) async throws {} + func delete(id: String) async throws {} + func loadAll() async throws -> [AgentMemoryEntry] { [] } + func load(kind: AgentMemoryKind) async throws -> [AgentMemoryEntry] { [] } + func loadContextBlock() async -> String { text } +} + +/// Counts how often the per-run block is built. +final class RunContextCounter: @unchecked Sendable { + private let lock = NSLock() + private var n = 0 + func next() -> Int { lock.withLock { n += 1; return n } } +} + +/// A memory store that invalidates the agent's frozen run context from +/// inside `loadContextBlock`, simulating an edit landing while a build is +/// suspended on this very await (actors are reentrant). +final class InvalidatingMemoryStore: AgentMemoryStore, @unchecked Sendable { + private let text: String + /// Set after the agent exists; `loadContextBlock` calls back into it. + var agent: Agent? + init(_ text: String) { self.text = text } + func save(_ entry: AgentMemoryEntry) async throws {} + func delete(id: String) async throws {} + func loadAll() async throws -> [AgentMemoryEntry] { [] } + func load(kind: AgentMemoryKind) async throws -> [AgentMemoryEntry] { [] } + func loadContextBlock() async -> String { + if let agent { await agent.invalidateRunContext() } + return text + } +} + +private struct FixedCompactor: ContextCompactor { + func summarize(middle: [AgentMessage], reason: CompactionReason) async -> String? { "SUMMARY" } +} + +/// With `freezeRunContext`, a chat's system prompt is the same bytes from run +/// to run: memory, the app's per-run block and the skill index are built once +/// and kept until invalidated, compacted, or replaced by a restored copy. +struct FrozenRunContextTests { + /// An agent whose system prompt has every dynamic part. + static func agent(freeze: Bool, provider: ScriptedProvider, memory: TextMemoryStore, + systemPrompt: String = "BASE") async throws -> Agent { + let counter = RunContextCounter() + let agent = Agent(config: AgentConfig(provider: provider, model: "mock", systemPrompt: systemPrompt, + maxTurns: 1, contextCompactor: FixedCompactor(), + compactAtFraction: 0.9, freezeRunContext: freeze)) + try await agent.setMemoryStore(memory) + try await agent.setRunContextProvider { _ in "RUN \(counter.next())" } + await agent.registerSkill(AgentSkill(name: "alpha", description: "First skill.", instructions: "Do alpha.")) + return agent + } + + static func systems(_ provider: ScriptedProvider) -> [String] { + provider.captured.map { $0.messages.first { $0.role == .system }?.content ?? "" } + } + + @Test func byDefaultEveryRunReadsMemoryAgain() async throws { + let provider = ScriptedProvider(turns: []) + let memory = TextMemoryStore("MEMORY v1") + let agent = try await Self.agent(freeze: false, provider: provider, memory: memory) + _ = try await agent.run("one") + memory.text = "MEMORY v2" + _ = try await agent.run("two") + let systems = Self.systems(provider) + #expect(systems[0].contains("MEMORY v1") && systems[0].contains("RUN 1")) + #expect(systems[1].contains("MEMORY v2") && systems[1].contains("RUN 2")) + #expect(await agent.frozenRunContext == nil) + } + + @Test func aFrozenRunContextKeepsTheFirstRunsParts() async throws { + let provider = ScriptedProvider(turns: []) + let memory = TextMemoryStore("MEMORY v1") + let agent = try await Self.agent(freeze: true, provider: provider, memory: memory) + _ = try await agent.run("one") + memory.text = "MEMORY v2" // the keeper filed something after the first reply + await agent.registerSkill(AgentSkill(name: "beta", description: "Second skill.", instructions: "Do beta.")) + _ = try await agent.run("two") + let systems = Self.systems(provider) + #expect(systems[0] == systems[1]) + #expect(systems[1].contains("MEMORY v1") && systems[1].contains("RUN 1")) + #expect(!systems[1].contains("beta")) + let frozen = try #require(await agent.frozenRunContext) + #expect(frozen.memory == "MEMORY v1") + #expect(frozen.runContext == "RUN 1") + #expect(frozen.skillIndex.contains("alpha")) + } + + @Test func invalidatingRebuildsOnTheNextRun() async throws { + let provider = ScriptedProvider(turns: []) + let memory = TextMemoryStore("MEMORY v1") + let agent = try await Self.agent(freeze: true, provider: provider, memory: memory) + _ = try await agent.run("one") + memory.text = "MEMORY v2" + await agent.invalidateRunContext() + _ = try await agent.run("two") + let systems = Self.systems(provider) + #expect(systems[1].contains("MEMORY v2") && systems[1].contains("RUN 2")) + } + + @Test func aCompactionClearsIt() async throws { + let provider = ScriptedProvider(turns: []) + let agent = try await Self.agent(freeze: true, provider: provider, memory: TextMemoryStore("MEMORY v1")) + _ = try await agent.run("one") + #expect(await agent.frozenRunContext != nil) + agent.conversation.append([.user(String(repeating: "a", count: 20_000)), + .assistant(String(repeating: "b", count: 20_000)), + .user("now"), .assistant("ok")]) + #expect(await agent.compactHistory(reason: .nearLimit) != nil) + #expect(await agent.frozenRunContext == nil) + } + + /// The relaunch case: the frozen context goes through JSON (as an app + /// saves it) into a new agent, which sends the same system prompt even + /// though memory changed in between. + @Test func aRestoredContextIsUsedOnTheNextRun() async throws { + let first = ScriptedProvider(turns: []) + let memory = TextMemoryStore("MEMORY v1") + let a = try await Self.agent(freeze: true, provider: first, memory: memory) + _ = try await a.run("one") + let saved = try JSONEncoder().encode(try #require(await a.frozenRunContext)) + memory.text = "MEMORY v2" + let second = ScriptedProvider(turns: []) + let b = try await Self.agent(freeze: true, provider: second, memory: memory) + try await b.restoreRunContext(try JSONDecoder().decode(FrozenRunContext.self, from: saved)) + _ = try await b.run("two") + #expect(Self.systems(second)[0] == Self.systems(first)[0]) + } + + /// A changed engine (another configured prompt, other tools) never sends + /// a stale mix: the restored context is ignored and built fresh. + @Test func aChangedBaseIgnoresARestoredContext() async throws { + let first = ScriptedProvider(turns: []) + let memory = TextMemoryStore("MEMORY v1") + let a = try await Self.agent(freeze: true, provider: first, memory: memory) + _ = try await a.run("one") + let saved = try #require(await a.frozenRunContext) + memory.text = "MEMORY v2" + let second = ScriptedProvider(turns: []) + let b = try await Self.agent(freeze: true, provider: second, memory: memory, systemPrompt: "OTHER BASE") + try await b.restoreRunContext(saved) + _ = try await b.run("two") + #expect(Self.systems(second)[0].contains("MEMORY v2")) + #expect(await b.frozenRunContext?.memory == "MEMORY v2") + } + + /// An invalidate landing while the build is suspended on the memory + /// read (actors are reentrant across that await) must not be clobbered + /// by the stale parts the build started with. Run 1 still gets the + /// stale parts it was already mid-build on (the running reply keeps its + /// prompt), but nothing is frozen, so run 2 rebuilds and sees the edit. + @Test func anInvalidateDuringTheBuildIsNotOverwritten() async throws { + let provider = ScriptedProvider(turns: []) + let memory = InvalidatingMemoryStore("MEMORY v1") + let counter = RunContextCounter() + let agent = Agent(config: AgentConfig(provider: provider, model: "mock", systemPrompt: "BASE", + maxTurns: 1, freezeRunContext: true)) + try await agent.setMemoryStore(memory) + try await agent.setRunContextProvider { _ in "RUN \(counter.next())" } + memory.agent = agent + _ = try await agent.run("one") + // Run 1 saw the memory it started reading, even though the agent + // was invalidated mid-build. + #expect(Self.systems(provider)[0].contains("MEMORY v1")) + // The stale build must not have frozen itself back in after the + // invalidate: nothing to reuse, so the next run reads memory again. + #expect(await agent.frozenRunContext == nil) + memory.agent = nil // run 2 must not invalidate itself again + _ = try await agent.run("two") + #expect(Self.systems(provider)[1].contains("RUN 2")) + #expect(await agent.frozenRunContext != nil) + } + + /// `restoreRunContext` with freeze off is a no-op: nothing ever reads + /// `frozenRunContext` when `config.freezeRunContext` is false (it is an + /// immutable `let`, so it can't turn on later), so storing a value there + /// would be a context no run ever uses, silently kept forever by an app + /// that persists `agent.frozenRunContext` after each run. + @Test func restoreRunContextIsANoOpWithFreezeOff() async throws { + let provider = ScriptedProvider(turns: []) + let agent = try await Self.agent(freeze: false, provider: provider, memory: TextMemoryStore("MEMORY")) + let saved = FrozenRunContext(baseDigest: "whatever", memory: "OTHER", runContext: "OTHER", skillIndex: "") + try await agent.restoreRunContext(saved) + #expect(await agent.frozenRunContext == nil) + _ = try await agent.run("one") + #expect(await agent.frozenRunContext == nil) + #expect(!Self.systems(provider)[0].contains("OTHER")) + } + + /// The app persists `FrozenRunContext` as JSON (`chat-prompts/*.json`). + /// Renaming a property would silently fail every decode on a relaunch + /// (the app would just rebuild fresh, never raising an error) — so the + /// wire shape is pinned here, not left to the synthesized Codable + /// conformance drifting unnoticed. + @Test func codableKeysArePinned() throws { + let json = #"{"baseDigest":"D","memory":"M","runContext":"R","skillIndex":"S"}"# + let decoded = try JSONDecoder().decode(FrozenRunContext.self, from: Data(json.utf8)) + #expect(decoded == FrozenRunContext(baseDigest: "D", memory: "M", runContext: "R", skillIndex: "S")) + let reencoded = try JSONSerialization.jsonObject(with: JSONEncoder().encode(decoded)) as? [String: String] + #expect(reencoded == ["baseDigest": "D", "memory": "M", "runContext": "R", "skillIndex": "S"]) + } + + /// With freeze on, calling `load_tools` mid-chat must not change the + /// frozen system prompt: the loaded tools reach the model only through + /// the tools array, never through the prompt text. + @Test func loadToolsMidChatLeavesTheFrozenSystemPromptUnchanged() async throws { + let browser = DeferredToolGroup(id: "browser", description: "control a browser", + toolNames: ["browser_open"]) + let provider = ScriptedProvider(turns: [ + [LLMToolCall(id: "l1", name: "load_tools", arguments: #"{"groups":["browser"]}"#)], [] + ]) + let agent = Agent(config: AgentConfig(provider: provider, model: "mock", systemPrompt: "BASE", + maxTurns: 6, toolGroups: [browser], freezeRunContext: true)) + try await agent.setMemoryStore(TextMemoryStore("MEMORY")) + _ = try await agent.run("open it") + let systems = Self.systems(provider) + #expect(systems.count == 2) + #expect(systems[0] == systems[1], "load_tools must not change the frozen prompt mid-run") + let frozenAfterLoad = try #require(await agent.frozenRunContext) + _ = try await agent.run("again") + #expect(Self.systems(provider)[2] == systems[0], "run 2 still sends the same frozen prompt") + #expect(await agent.frozenRunContext == frozenAfterLoad, "the frozen copy was not rebuilt") + } + + /// Exactly today's assembly: base, memory, per-run block (blank line + /// between each), then the tool line, the skill index, the group index. + @Test func theSystemPromptKeepsItsOrderAndSeparators() { + let all = FrozenRunContext(baseDigest: "d", memory: "M", runContext: "R", skillIndex: "\n\nS") + #expect(Agent.renderSystemPrompt(base: "B", parts: all, toolLine: "\nT", groupIndex: "\n\nG") + == "B\n\nM\n\nR\nT\n\nS\n\nG") + let memoryOnly = FrozenRunContext(baseDigest: "d", memory: "M", runContext: "", skillIndex: "") + #expect(Agent.renderSystemPrompt(base: "", parts: memoryOnly, toolLine: "", groupIndex: "") == "M") + let none = FrozenRunContext(baseDigest: "d", memory: "", runContext: "", skillIndex: "") + #expect(Agent.renderSystemPrompt(base: "B", parts: none, toolLine: "\nT", groupIndex: "") == "B\nT") + } +} + +extension FrozenRunContextTests { + /// A child's prompt is new on every spawn and carries memory as it is + /// then — a parent's frozen context does not hold it back. + @Test func aSubAgentReadsCurrentMemoryAndNeverFreezes() async throws { + let memory = TextMemoryStore("MEMORY v1") + let parent = Agent(config: AgentConfig(provider: ScriptedProvider(turns: []), model: "mock", + systemPrompt: "BASE", maxTurns: 1, freezeRunContext: true)) + try await parent.setMemoryStore(memory) + _ = try await parent.run("one") + memory.text = "MEMORY v2" + let child = await SubAgentSpawner(parent: parent).makeChild() + #expect(child.config.systemPrompt?.contains("MEMORY v2") == true) + #expect(child.config.freezeRunContext == false) + } +} diff --git a/Tests/SwiftAgentKitTests/SwiftAgentKitTests.swift b/Tests/SwiftAgentKitTests/SwiftAgentKitTests.swift index 2df1a21..d679b48 100644 --- a/Tests/SwiftAgentKitTests/SwiftAgentKitTests.swift +++ b/Tests/SwiftAgentKitTests/SwiftAgentKitTests.swift @@ -2630,6 +2630,57 @@ private func tempArtifactDir() -> URL { #expect(listed.first?.toolName == "run_shell") } +/// Saves one evicted step through a FRESH ContextManager (what a relaunch +/// builds) and returns the model-facing messages. +private func evictOneStep(_ store: any ArtifactStore) async -> [LLMMessage] { + let manager = ContextManager(store: store, summaryLength: 40, inlineBudgetChars: 0) + let messages: [AgentMessage] = [ + .system("sys"), + .user("build it"), + .assistant(content: "", toolCalls: [AgentToolCall(id: "c1", name: "run_shell")]), + .tool(results: [.success(toolCallId: "c1", toolName: "run_shell", + result: "BUILD LOG " + String(repeating: "x", count: 500))]), + .assistant("built"), + .user("next"), + ] + return await manager.modelMessages(messages) { $0 } +} + +@Test func testArtifactForAnEvictedResultIsSavedOnceAcrossFreshManagers() async throws { + let dir = tempArtifactDir() + defer { try? FileManager.default.removeItem(at: dir) } + let store = FileArtifactStore(directory: dir) + let first = await evictOneStep(store) + let second = await evictOneStep(store) + // A relaunch: a new store on the same folder, and a new manager. + let third = await evictOneStep(FileArtifactStore(directory: dir)) + + #expect(first.map(\.content) == second.map(\.content), "receipts are byte-identical") + #expect(first.map(\.content) == third.map(\.content), "receipts survive a relaunch byte-identical") + let id = try #require(first.lazy.compactMap { firstArtifactID(in: $0.content) }.first) + #expect(await store.list(limit: 10).map(\.id) == [id]) + let files = try FileManager.default.contentsOfDirectory(atPath: dir.path).sorted() + #expect(files == [id + ".json", id + ".txt"], "one artifact on disk, not one per relaunch") +} + +@Test func testInMemoryArtifactStoreReusesTheArtifactForACall() async { + let store = InMemoryArtifactStore() + let a = await store.save("full output", description: "t", toolCallID: "c1", toolName: "run_shell") + let b = await store.save("full output", description: "t", toolCallID: "c1", toolName: "run_shell") + #expect(a.id == b.id) + #expect(await store.list(limit: 10).count == 1) + #expect(a.id.range(of: #"^artifact-[0-9a-f]{12}$"#, options: .regularExpression) != nil) + // Same id, different output (a provider that reuses call ids across + // runs): a different artifact, never the stale one. + let c = await store.save("other output", description: "t", toolCallID: "c1", toolName: "run_shell") + #expect(c.id != a.id) + #expect(await store.get(c.id)?.content == "other output") + // No call id: a fresh artifact each time, as before. + let d = await store.save("full output", description: "t", toolCallID: nil, toolName: "run_shell") + let e = await store.save("full output", description: "t", toolCallID: nil, toolName: "run_shell") + #expect(d.id != e.id) +} + // MARK: - Progress nudge (turn-budget checkpoint) @Test func testNudgeTurnSchedule() {