spb/zyquo-agent Public MIT

1//2//  MemoryManager.swift3//  Zyquo Agent4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//8//  Context compression for long autonomous runs (docs/AGENT-RESEARCH.md §3):9//10//  • Token accounting — chars/4 heuristic, calibrated by the provider's11//    reported input tokens when available; the model's contextWindow bounds it.12//  • Compaction at 85% — older completed steps' messages are replaced by ONE13//    synthetic summary message produced by a summarization call to the SAME14//    provider/model (with a mechanical fallback when that call fails). Always15//    kept verbatim: the system prompt (never in history), the original task16//    message, the most recent N steps, and — re-injected fresh inside the17//    summary — the current plan and MEMORY.md. A thrash guard enforces a18//    minimum step distance between compactions.19//  • Output offloading — tool results above a size threshold are written to20//    `.zyquo/outputs/step<N>-<tool>-<id>.txt` and replaced in-context by a21//    head-of-output stub pointing at the file (readable via read_file /22//    search_files).23//  • MEMORY.md — created in the workspace root at task start; the system24//    prompt gives the model ownership of it via the file tools; its content25//    is re-read and re-injected at each compaction.26//2728import Foundation2930/// Compaction/offloading tunables (Settings › Agent in Phase 6).31struct MemoryConfiguration: Sendable {32    /// Fraction of the model's context window that triggers compaction.33    var compactionThreshold = 0.8534    /// Steps whose messages are always kept verbatim (most recent).35    var keepRecentSteps = 636    /// Thrash guard: minimum steps between two compactions.37    var minStepsBetweenCompactions = 1038    /// Tool-result content above this many bytes is offloaded to a file.39    var offloadThresholdBytes = 8_19240    /// Lines of the original output kept in the in-context stub.41    var offloadStubLines = 4042    /// Max tokens requested from the summarization call.43    var summaryMaxTokens = 1_2004445    static let `default` = MemoryConfiguration()46}4748/// One message of the loop's conversation history, tagged with the step that49/// produced it so compaction can operate on whole steps (preserving50/// tool_use/tool_result pairing, which is per-step by construction).51struct HistoryEntry: Sendable {52    /// 0 = the initial task message; N = messages of step N.53    var stepIndex: Int54    var message: Message55    /// Never compacted (the original task message).56    var pinned: Bool = false57    /// A synthetic summary produced by an earlier compaction (eligible for58    /// re-summarization by the next one).59    var isSummary: Bool = false60}6162/// One compaction, for the transcript and the UI.63struct CompactionRecord: Codable, Sendable {64    var id: UUID = UUID()65    var timestamp: Date = Date()66    /// Estimated context tokens immediately before / after.67    var beforeTokens: Int68    var afterTokens: Int69    /// How many step's message groups were folded into the summary, through70    /// which step index.71    var summarizedSteps: Int72    var throughStep: Int73    /// The synthetic summary text injected into context.74    var summary: String75}7677/// Owned by the AgentLoop actor (class: mutable state survives async78/// summarization calls without inout-across-await restrictions).79final class MemoryManager {80    let configuration: MemoryConfiguration81    private let workspaceRoot: URL82    private let outputsDirectory: URL83    private(set) var lastCompactionStep = 084    /// Latest input-token figure reported by the provider (calibrates the85    /// heuristic estimate).86    var lastReportedInputTokens: Int?8788    init(workspaceRoot: URL, outputsDirectory: URL, configuration: MemoryConfiguration = .default) {89        self.workspaceRoot = workspaceRoot90        self.outputsDirectory = outputsDirectory91        self.configuration = configuration92    }9394    // MARK: - Token accounting9596    /// chars/4 heuristic — deliberately conservative and provider-neutral.97    static func estimateTokens(_ text: String) -> Int {98        text.isEmpty ? 0 : max(1, text.count / 4)99    }100101    /// Estimated tokens of one message (text + reasoning + tool calls +102    /// tool results + fixed per-message overhead).103    static func estimateTokens(of message: Message) -> Int {104        var total = 8 // role/framing overhead105        total += estimateTokens(message.text)106        if let reasoning = message.reasoning { total += estimateTokens(reasoning) }107        for call in message.toolCalls ?? [] {108            total += estimateTokens(call.name) + estimateTokens(call.argumentsJSON) + 8109        }110        for result in message.toolResults ?? [] {111            total += estimateTokens(result.content) + 8112        }113        return total114    }115116    func estimateTokens(entries: [HistoryEntry]) -> Int {117        entries.reduce(0) { $0 + Self.estimateTokens(of: $1.message) }118    }119120    /// Fixed per-request overhead: system prompt + tool schemas.121    static func estimateFixedTokens(systemPrompt: String, toolSpecs: [ToolSpec]) -> Int {122        var total = estimateTokens(systemPrompt)123        for spec in toolSpecs {124            total += estimateTokens(spec.name) + estimateTokens(spec.description) + estimateTokens(spec.parametersJSONSchema)125        }126        return total127    }128129    /// Best current estimate of the next request's input size.130    func currentContextTokens(entries: [HistoryEntry], fixedTokens: Int) -> Int {131        max(fixedTokens + estimateTokens(entries: entries), lastReportedInputTokens ?? 0)132    }133134    /// True when usage crossed the threshold, the thrash guard allows it, and135    /// there is at least one whole old step to fold away.136    func shouldCompact(entries: [HistoryEntry], fixedTokens: Int, contextWindow: Int, currentStep: Int) -> Bool {137        guard contextWindow > 0 else { return false }138        let usage = currentContextTokens(entries: entries, fixedTokens: fixedTokens)139        guard Double(usage) >= configuration.compactionThreshold * Double(contextWindow) else { return false }140        guard currentStep - lastCompactionStep >= configuration.minStepsBetweenCompactions else { return false }141        let cutoff = currentStep - configuration.keepRecentSteps142        return entries.contains { !$0.pinned && $0.stepIndex <= cutoff && $0.stepIndex >= 1 }143    }144145    // MARK: - Compaction146147    struct CompactionOutcome {148        var entries: [HistoryEntry]149        var record: CompactionRecord150    }151152    /// Folds every non-pinned message of steps ≤ (currentStep − keepRecent)153    /// — including previous summaries — into one synthetic user message:154    /// summary + fresh plan + fresh MEMORY.md. Returns nil when there is155    /// nothing to compact.156    func compact(157        entries: [HistoryEntry],158        planContext: String?,159        steps: [AgentStep],160        client: any ProviderClient,161        apiKey: String,162        model: AIModel,163        currentStep: Int,164        fixedTokens: Int165    ) async -> CompactionOutcome? {166        let cutoff = currentStep - configuration.keepRecentSteps167        let toSummarize = entries.filter { !$0.pinned && $0.stepIndex <= cutoff }168        guard !toSummarize.isEmpty else { return nil }169        let kept = entries.filter { $0.pinned || $0.stepIndex > cutoff }170171        let beforeTokens = fixedTokens + estimateTokens(entries: entries)172173        // Summarize via the same provider/model; mechanical fallback on error.174        let material = Self.renderMaterial(toSummarize)175        var summary: String176        do {177            summary = try await Self.requestSummary(178                material: material,179                client: client,180                apiKey: apiKey,181                model: model,182                maxTokens: configuration.summaryMaxTokens183            )184        } catch {185            summary = Self.mechanicalSummary(steps: steps, throughStep: cutoff)186            summary += "\n(Note: the summarization call failed — this is a mechanical digest. Full history: .zyquo/transcript.json.)"187        }188189        var checkpoint = "[CONTEXT CHECKPOINT — steps 1–\(max(cutoff, 1)) were compacted into this summary]\n"190        checkpoint += summary.trimmingCharacters(in: .whitespacesAndNewlines)191        checkpoint += "\n\nCURRENT PLAN:\n\(planContext ?? "(no plan recorded yet)")"192        let memoryText = readMemoryFile()?.trimmingCharacters(in: .whitespacesAndNewlines)193        checkpoint += "\n\nMEMORY.md:\n\((memoryText?.isEmpty == false ? memoryText! : "(empty)"))"194195        let summaryEntry = HistoryEntry(196            stepIndex: max(cutoff, 1),197            message: Message(role: .user, text: checkpoint),198            isSummary: true199        )200201        // Pinned early entries (the task message) stay first, then the202        // summary, then the recent steps in their original order.203        let earlyPinned = kept.filter { $0.stepIndex <= cutoff }204        let recent = kept.filter { $0.stepIndex > cutoff }205        let newEntries = earlyPinned + [summaryEntry] + recent206207        lastCompactionStep = currentStep208        let afterTokens = fixedTokens + estimateTokens(entries: newEntries)209        let record = CompactionRecord(210            beforeTokens: beforeTokens,211            afterTokens: afterTokens,212            summarizedSteps: Set(toSummarize.map(\.stepIndex)).count,213            throughStep: max(cutoff, 1),214            summary: summary215        )216        return CompactionOutcome(entries: newEntries, record: record)217    }218219    /// Renders the messages being folded away into summarization material,220    /// capped so the summarization call itself cannot blow the window.221    static func renderMaterial(_ entries: [HistoryEntry]) -> String {222        var lines: [String] = []223        for entry in entries {224            let message = entry.message225            var line = "[step \(entry.stepIndex)] \(message.role.rawValue):"226            if !message.text.isEmpty { line += " \(message.text)" }227            for call in message.toolCalls ?? [] {228                line += "\n  → \(call.name) \(String(call.argumentsJSON.prefix(400)))"229            }230            for result in message.toolResults ?? [] {231                let status = result.isError ? "ERROR" : "ok"232                line += "\n  ← [\(status)] \(String(result.content.prefix(600)))"233            }234            lines.append(line)235        }236        var material = lines.joined(separator: "\n")237        let cap = 60_000238        if material.count > cap {239            // Keep the head (task framing) and the tail (recent detail).240            let head = String(material.prefix(cap / 3))241            let tail = String(material.suffix(cap * 2 / 3))242            material = head + "\n… [middle elided] …\n" + tail243        }244        return material245    }246247    /// The summarization call: same provider/model, non-streaming, no tools.248    private static func requestSummary(249        material: String,250        client: any ProviderClient,251        apiKey: String,252        model: AIModel,253        maxTokens: Int254    ) async throws -> String {255        let instruction = """256            You compress an autonomous agent's working history into a compact \257            structured record. From the transcript below, produce EXACTLY \258            these three sections, tersely, preserving concrete values \259            (paths, filenames, commands, numbers, error messages):260261            Sub-tasks completed:262            Key facts, paths & decisions:263            Open issues:264265            TRANSCRIPT:266            \(material)267            """268        var parameters = ChatParameters(maxTokens: maxTokens)269        parameters.temperature = nil270        let request = ChatRequest(271            model: model,272            systemPrompt: "You are a precise summarizer inside the Zyquo Agent memory system. Output only the requested record — no preamble.",273            messages: [Message(role: .user, text: instruction)],274            parameters: parameters,275            stream: false276        )277        let reply = try await client.complete(request, apiKey: apiKey)278        let text = reply.text.trimmingCharacters(in: .whitespacesAndNewlines)279        guard !text.isEmpty else {280            throw ProviderError.invalidResponse(client.providerID, detail: "empty summarization reply")281        }282        return text283    }284285    /// Fallback digest built from the transcript's AgentSteps when the286    /// summarization call fails (offline, rate-limited, mock runs).287    static func mechanicalSummary(steps: [AgentStep], throughStep: Int) -> String {288        var lines = ["Sub-tasks completed:"]289        for step in steps where step.index >= 1 && step.index <= throughStep {290            var line = "- Step \(step.index):"291            let thought = step.text.trimmingCharacters(in: .whitespacesAndNewlines)292            if !thought.isEmpty { line += " \(String(thought.prefix(120)))" }293            if !step.toolInvocations.isEmpty {294                let calls = step.toolInvocations.map { invocation -> String in295                    let ok = invocation.result?.isError == true ? "error" : "ok"296                    return "\(invocation.call.name)(\(String(invocation.call.argumentsJSON.prefix(60))))→\(ok)"297                }298                line += " [\(calls.joined(separator: ", "))]"299            }300            lines.append(line)301        }302        lines.append("Key facts, paths & decisions:\n- (not extracted — mechanical digest)")303        lines.append("Open issues:\n- (unknown)")304        return lines.joined(separator: "\n")305    }306307    // MARK: - Output offloading308309    /// Writes oversized tool-result content to `.zyquo/outputs/…` and returns310    /// a stub result (head of the output + pointer). Small results pass311    /// through untouched, as do failures to write the file.312    func offloadIfNeeded(_ result: ToolResult, stepIndex: Int, toolName: String) -> ToolResult {313        guard result.content.utf8.count > configuration.offloadThresholdBytes else { return result }314315        let safeTool = toolName.map { $0.isLetter || $0.isNumber ? $0 : "-" }316        let safeID = result.toolCallID.suffix(8).map { $0.isLetter || $0.isNumber ? $0 : "-" }317        let fileName = "step\(stepIndex)-\(String(safeTool))-\(String(safeID)).txt"318        let fileURL = outputsDirectory.appendingPathComponent(fileName)319320        do {321            try FileManager.default.createDirectory(at: outputsDirectory, withIntermediateDirectories: true)322            try result.content.write(to: fileURL, atomically: true, encoding: .utf8)323        } catch {324            // Offloading is an optimization — never fail the tool call over it.325            return result326        }327328        let allLines = result.content.components(separatedBy: "\n")329        var head = allLines.prefix(configuration.offloadStubLines).joined(separator: "\n")330        if head.count > 4_000 { head = String(head.prefix(4_000)) }331        let sizeKB = Double(result.content.utf8.count) / 1000332        let relativePath = ".zyquo/outputs/\(fileName)"333        let stub = head + """334            \n… [output truncated: \(String(format: "%.1f", sizeKB)) KB, \(allLines.count) lines total — \335            full output: \(relativePath) — use read_file or search_files on that path]336            """337        return ToolResult(toolCallID: result.toolCallID, content: stub, isError: result.isError)338    }339340    // MARK: - MEMORY.md341342    var memoryFileURL: URL { workspaceRoot.appendingPathComponent("MEMORY.md") }343344    /// Creates MEMORY.md with its contract header if absent (task start).345    func ensureMemoryFile() {346        guard !FileManager.default.fileExists(atPath: memoryFileURL.path) else { return }347        let template = """348            # MEMORY.md — Zyquo Agent task memory349350            Durable facts, decisions, and open questions for this task.351            Maintained by the agent (see system prompt): read at task start,352            append important knowledge as it is learned. Content here survives353            context compaction.354355            """356        try? template.write(to: memoryFileURL, atomically: true, encoding: .utf8)357    }358359    /// Current MEMORY.md content (capped — memory is notes, not a dump).360    func readMemoryFile() -> String? {361        guard let text = try? String(contentsOf: memoryFileURL, encoding: .utf8) else { return nil }362        return text.count > 16_000 ? String(text.prefix(16_000)) + "\n… [MEMORY.md truncated]" : text363    }364}365