byre/mac/Sources/Byre/TranscriptBuilder.swift

import AppKit
import Foundation
import Observation
import OpenAPIRuntime

let quietReply = "NO_REPLY"

struct Usage {
    var used: Int
    var window: Int
    var model: String
}

struct Call {
    let name: String
    let target: String
}

enum NotePart {
    case text(String)
    case task(String)
    case url(String)
}

enum NoteStyle {
    case plain, error, ask
}

struct Note {
    var parts: [NotePart]
    var when: Date?
    var style: NoteStyle = .plain
    var cancelMergeAfterReview: String?
}

enum SubagentState {
    case running, done, error
}

enum RowKind {
    case user(actor: String, when: Date, images: [String], md: [MDBlock]?)
    case byrePrompt(brief: String, when: Date, md: [MDBlock])
    case agent(md: [MDBlock])
    case subagentMessage(when: Date, md: [MDBlock])
    case stream
    case thinking(md: [MDBlock])
    case tool(name: String, brief: String, input: String)
    case subagent(title: String, description: String, prompt: [MDBlock])
    case group
    case orphan(String)
    case turnEnd(error: String?)
    case note(Note)
    case handoff
    case snapshot(Snapshot)
}

@MainActor
private var rowIds = 0

@MainActor
@Observable
final class RowList {
    var rows: [Row] = []
}

@MainActor
@Observable
final class Row: Identifiable {
    let id: Int
    let kind: RowKind
    let taskId: String?
    var hidden = false
    var tag = false
    var result: String?
    var unfinished = false
    var status = "starting…"
    var statusState = SubagentState.running
    var active = false
    weak var owner: Row?
    var calls: [Call] = []
    var summary: Note?
    var review: Review?
    var text = ""
    var md: [MDBlock] = []
    var summarySeq: Int64?
    let body: RowList?

    var pending: Bool {
        switch kind {
        case .tool, .subagent:
            return result == nil && !unfinished
        case .group:
            return body!.rows.contains { $0.pending }
        default:
            return false
        }
    }

    var inActive: Bool {
        var o = owner
        while let r = o {
            if r.active {
                return true
            }
            o = r.owner
        }
        return false
    }

    func within(_ ancestor: Row) -> Bool {
        var o = owner
        while let r = o {
            if r === ancestor {
                return true
            }
            o = r.owner
        }
        return false
    }

    init(_ kind: RowKind, taskId: String?, body: Bool = false) {
        rowIds += 1
        id = rowIds
        self.kind = kind
        self.taskId = taskId
        self.body = body ? RowList() : nil
    }
}

struct PendingSend: Identifiable {
    let id: Int
    let text: String
    let images: [PendingImage]
    let queued: Bool
    let md: [MDBlock]
}

@MainActor
private final class Run {
    let first: Row
    let container: RowList
    var group: Row?
    var calls: [Call]
    let parent: String

    init(first: Row, container: RowList, calls: [Call], parent: String) {
        self.first = first
        self.container = container
        self.calls = calls
        self.parent = parent
    }
}

@MainActor
private final class Subagent {
    let row: Row
    var run: Run?

    init(row: Row) {
        self.row = row
    }
}

@MainActor
private final class Ctx {
    var tools: [String: Row] = [:]
    var subagents: [String: Subagent] = [:]
    var reviewers: Set<String> = []
    var reviewButtons: [Int: Row] = [:]
    var run: Run?
    var held: RowList?
    var open = false
    var drawn: [Row]?
    var quiet = false
    var handoff: Row?
}

@MainActor
@Observable
final class TranscriptBuilder {
    let root = RowList()
    private(set) var pending: [PendingSend] = []
    private(set) var context: [String: Usage] = [:]
    private(set) var open = false
    private(set) var lastMergeAfterReview: Int?
    @ObservationIgnored private var ctx = Ctx()
    @ObservationIgnored private var sink: RowList?
    @ObservationIgnored private var stream: Row?
    @ObservationIgnored private var streamTimer: Task<Void, Never>?
    @ObservationIgnored private let taskId: String?
    @ObservationIgnored private let worktree: () -> String
    @ObservationIgnored private let running: () -> Bool
    @ObservationIgnored private var pendingIds = 0

    init(taskId: String?, worktree: @escaping () -> String, running: @escaping () -> Bool) {
        self.taskId = taskId
        self.worktree = worktree
        self.running = running
    }

    func stop() {
        dropStream()
    }

    func append(_ events: [ByreEvent]) -> Bool {
        let changed = render(events)
        open = ctx.open
        return changed
    }

    func prepend(_ events: [ByreEvent]) {
        let live = ctx
        let box = RowList()
        ctx = Ctx()
        sink = box
        _ = render(events)
        sink = nil
        ctx = live
        root.rows.insert(contentsOf: box.rows, at: 0)
    }

    func showPending(_ text: String, _ images: [PendingImage]) -> () -> Void {
        pendingIds += 1
        let entry = PendingSend(id: pendingIds, text: text, images: images, queued: running(), md: trimmed(text).isEmpty ? [] : parseMarkdown(text))
        pending.append(entry)
        ctx.run = nil
        return { [weak self] in
            self?.pending.removeAll { $0.id == entry.id }
        }
    }

    private func confirmPending(_ text: String) {
        if let i = pending.firstIndex(where: { $0.text == text }) {
            pending.remove(at: i)
        }
    }

    private func release() {
        guard let buffered = ctx.held else {
            return
        }
        ctx.held = nil
        ctx.open = true
        (sink ?? root).rows.append(contentsOf: buffered.rows)
    }

    @discardableResult
    private func put(_ row: Row, keepHeld: Bool = false) -> RowList {
        if let handoff = ctx.handoff {
            handoff.body!.rows.append(row)
            return handoff.body!
        }
        if !keepHeld {
            release()
        }
        ctx.drawn?.append(row)
        if let held = ctx.held {
            held.rows.append(row)
            return held
        }
        let list = sink ?? root
        list.rows.append(row)
        return list
    }

    private func fold() {
        let drawn = ctx.drawn ?? []
        (sink ?? root).rows.removeAll { r in drawn.contains { $0 === r } }
        ctx.drawn = nil
        ctx.held = nil
        ctx.open = false
        ctx.quiet = false
    }

    private func subagentTag(_ parent: String) -> Bool {
        !parent.isEmpty && ctx.subagents[parent] == nil
    }

    private func addToRun(_ row: Row, _ call: Call, _ parent: String) {
        if let sub = ctx.subagents[parent] {
            sub.run = extendRun(sub.run, row, call, parent) { r in
                sub.row.body!.rows.append(r)
                return sub.row.body!
            }
            return
        }
        ctx.run = extendRun(ctx.run, row, call, parent) { r in
            self.put(r, keepHeld: true)
        }
    }

    private func extendRun(_ run: Run?, _ row: Row, _ call: Call, _ parent: String, _ place: (Row) -> RowList) -> Run {
        guard let run, run.parent == parent else {
            let list = place(row)
            return Run(first: row, container: list, calls: [call], parent: parent)
        }
        run.calls.append(call)
        if run.group == nil {
            let group = Row(.group, taskId: taskId, body: true)
            group.owner = ctx.subagents[parent]?.row
            for list in [run.container, sink ?? root] {
                if let i = list.rows.firstIndex(where: { $0 === run.first }) {
                    list.rows[i] = group
                    break
                }
            }
            ctx.drawn?.append(group)
            group.body!.rows.append(run.first)
            run.group = group
        }
        run.group!.body!.rows.append(row)
        run.group!.calls = run.calls
        run.group!.tag = subagentTag(run.parent)
        return run
    }

    private func dropStream() {
        streamTimer?.cancel()
        streamTimer = nil
        guard let stream else {
            return
        }
        root.rows.removeAll { $0 === stream }
        self.stream = nil
    }

    func appendDelta(_ ev: ByreEvent) {
        if ctx.handoff != nil {
            return
        }
        if let sid = ev.agentSessionId, ctx.reviewers.contains(sid) {
            return
        }
        let row: Row
        if let stream {
            row = stream
        } else {
            row = Row(.stream, taskId: taskId)
            row.hidden = true
            stream = row
            ctx.run = nil
            put(row, keepHeld: true)
        }
        row.text += str(ev.payload.value, "text") ?? ""
        if streamTimer != nil {
            return
        }
        streamTimer = Task { [weak self] in
            try? await Task.sleep(for: .milliseconds(100))
            guard let self, !Task.isCancelled, let stream = self.stream else {
                return
            }
            self.streamTimer = nil
            stream.hidden = self.ctx.held != nil || quietReply.hasPrefix(trimmed(stream.text))
            stream.md = parseMarkdown(stream.text)
        }
    }

    private func openHandoff() {
        let row = Row(.handoff, taskId: taskId, body: true)
        row.summary = Note(parts: [.text("writing a handoff…")])
        put(row)
        ctx.handoff = row
    }

    private func closeHandoff(_ end: Row?) {
        guard let handoff = ctx.handoff else {
            return
        }
        ctx.handoff = nil
        if let end, case .note(let note) = end.kind {
            handoff.summary = note
        } else {
            handoff.summary = Note(parts: [.text("handoff not written")])
        }
    }

    private func render(_ events: [ByreEvent]) -> Bool {
        var contextChanged = false
        for ev in events {
            let p = ev.payload.value
            let parent = parentOf(p)
            if ev.kind == "session.created" && str(p, "role") == "reviewer", let sid = ev.agentSessionId, !sid.isEmpty {
                ctx.reviewers.insert(sid)
                continue
            }
            if let sid = ev.agentSessionId, ctx.reviewers.contains(sid) {
                continue
            }
            if ev.kind == "agent.vendor" {
                claudeTaskStatus(p)
                continue
            }
            if ev.kind == "session.subagents" {
                if int(p, "subagents") == 0 {
                    for sub in ctx.subagents.values where sub.row.active {
                        endSubagent(sub.row, .error, "ended with its process")
                    }
                }
                continue
            }
            if sink == nil && ((ev.kind == "agent.message" && parent.isEmpty) || ev.kind == "turn.ended" || ev.kind == "user.message") {
                dropStream()
            }
            if sink == nil && ev.kind == "user.message" && ev.actor != "byre" {
                confirmPending(str(p, "text") ?? "")
            }
            if ev.kind == "session.context" && sink != nil {
                continue
            }
            if ev.kind == "session.context", let sid = ev.agentSessionId, !sid.isEmpty {
                let prev = context[sid]
                let window = int(p, "window").flatMap { $0 == 0 ? nil : $0 } ?? prev?.window ?? 0
                context[sid] = Usage(used: int(p, "used") ?? 0, window: window, model: str(p, "model") ?? "")
                contextChanged = true
                continue
            }
            if ev.kind == "snapshot.created" {
                ctx.run = nil
                let snap = try! decodePayload(ev.payload, as: Snapshot.self)
                let row = Row(.snapshot(snap), taskId: taskId)
                ctx.reviewButtons[snap.n] = row
                put(row)
                continue
            }
            if ev.kind == "review.started" || ev.kind == "review.finished" || ev.kind == "review.finding" {
                let r = try! decodePayload(ev.payload, as: Review.self)
                ctx.reviewButtons[r.snapshotN]?.review = r
            }
            let sub = ctx.subagents[parent]
            if ev.kind == "agent.tool_call" && isSubagentCall(p) {
                let row = renderSubagent(p)
                if let sub {
                    sub.run = nil
                    sub.row.body!.rows.append(row)
                    continue
                }
                ctx.run = nil
                put(row, keepHeld: true)
                continue
            }
            if ev.kind == "agent.tool_call" {
                let name = toolLabel(jsString(p["name"]))
                addToRun(renderToolCall(p), Call(name: name, target: toolTarget(p["input"] ?? nil, worktree())), parent)
                continue
            }
            guard let row = renderEvent(ev) else {
                continue
            }
            if let sub {
                sub.run = nil
                sub.row.body!.rows.append(row)
                continue
            }
            ctx.run = nil
            if ev.kind == "agent.message" && parent.isEmpty {
                ctx.drawn = nil
            }
            if ev.kind == "turn.ended" {
                ctx.open = false
            }
            if ev.kind == "turn.ended" && bool(p, "handoff") && ctx.handoff != nil {
                closeHandoff(row)
                continue
            }
            if ev.kind == "turn.ended" || (ev.kind == "user.message" && !bool(p, "mid_turn")) {
                closeHandoff(nil)
            }
            if ev.kind == "user.message" && bool(p, "handoff") {
                openHandoff()
            }
            if ev.kind == "turn.ended" && ctx.quiet && ctx.drawn != nil {
                fold()
                continue
            }
            if ev.kind == "turn.ended" || ev.kind == "user.message" {
                ctx.drawn = nil
            }
            if ev.kind == "user.message" && ev.actor == "byre" && !(ev.owner ?? "").isEmpty && !bool(p, "mid_turn") && !bool(p, "handoff") {
                ctx.held = RowList()
                ctx.open = false
                ctx.drawn = []
                continue
            }
            if ev.kind == "user.message" && !bool(p, "mid_turn") {
                ctx.open = true
            }
            if ev.kind == "turn.ended" && bool(p, "summary") && taskId != nil {
                row.summarySeq = ev.seq
            }
            if ev.kind == "task.merge_after_review" && sink == nil {
                lastMergeAfterReview = row.id
            }
            put(row, keepHeld: ev.kind == "agent.thinking")
        }
        return contextChanged
    }

    private func claudeTaskStatus(_ p: JSONDict) {
        guard str(p, "type") == "system", let sub = ctx.subagents[jsString(p["tool_use_id"])] else {
            return
        }
        switch str(p, "subtype") {
        case "task_started":
            sub.row.active = true
            sub.row.status = bool(p, "is_backgrounded") ? "running in the background" : "running"
        case "task_progress":
            sub.row.status = str(p, "description") ?? "running"
        case "task_notification":
            endSubagent(sub.row, str(p, "status") == "completed" ? .done : .error, jsString(p["status"]))
        default:
            break
        }
    }

    private func endSubagent(_ row: Row, _ state: SubagentState, _ status: String) {
        row.active = false
        row.statusState = state
        row.status = status
        for tool in ctx.tools.values where tool.result == nil && tool.within(row) {
            tool.unfinished = true
        }
    }

    private func renderToolCall(_ p: JSONDict) -> Row {
        let row = Row(.tool(name: toolLabel(jsString(p["name"])), brief: toolBrief(p["input"] ?? nil, worktree()), input: prettyJSON(p["input"] ?? nil)), taskId: taskId)
        row.tag = subagentTag(parentOf(p))
        row.owner = ctx.subagents[parentOf(p)]?.row
        ctx.tools[jsString(p["id"])] = row
        return row
    }

    private func renderSubagent(_ p: JSONDict) -> Row {
        let input = dict(p, "input") ?? [:]
        let id = jsString(p["id"])
        let type = str(input, "subagent_type") ?? "subagent"
        let row = Row(.subagent(title: "\(jsString(p["name"])) \(type)", description: str(input, "description") ?? "", prompt: parseMarkdown(str(input, "prompt") ?? "")), taskId: taskId, body: true)
        row.tag = subagentTag(parentOf(p))
        row.owner = ctx.subagents[parentOf(p)]?.row
        ctx.tools[id] = row
        ctx.subagents[id] = Subagent(row: row)
        return row
    }

    private func renderEvent(_ ev: ByreEvent) -> Row? {
        let p = ev.payload.value
        let when = ev.createdAt
        func note(_ parts: [NotePart], _ style: NoteStyle = .plain, time: Bool = true) -> Row {
            Row(.note(Note(parts: parts, when: time ? when : nil, style: style)), taskId: taskId)
        }

        switch ev.kind {
        case "user.message":
            ctx.quiet = false
            let text = str(p, "text") ?? ""
            if ev.actor == "byre" {
                let first = text.components(separatedBy: "\n").first ?? ""
                let brief = first.count > 120 ? "\(first.prefix(120))…" : first
                let row = Row(.byrePrompt(brief: brief, when: when, md: parseMarkdown(text)), taskId: taskId)
                row.hidden = bool(p, "continuation")
                return row
            }
            let images = (p["images"] as? [(any Sendable)?] ?? []).map { jsString($0) }
            return Row(.user(actor: ev.actor, when: when, images: images, md: trimmed(text).isEmpty ? nil : parseMarkdown(text)), taskId: taskId)

        case "agent.message":
            let text = jsString(p["text"])
            if !parentOf(p).isEmpty {
                let row = Row(.subagentMessage(when: when, md: parseMarkdown(text)), taskId: taskId)
                row.tag = subagentTag(parentOf(p))
                return row
            }
            if trimmed(trimmed(text).components(separatedBy: "\n").last ?? "") == quietReply {
                ctx.quiet = true
                return nil
            }
            ctx.quiet = false
            return Row(.agent(md: parseMarkdown(text)), taskId: taskId)

        case "agent.thinking":
            let row = Row(.thinking(md: parseMarkdown(jsString(p["text"]))), taskId: taskId)
            row.tag = subagentTag(parentOf(p))
            return row

        case "agent.tool_result":
            let body = resultText(p["content"] ?? nil)
            guard let target = ctx.tools[jsString(p["tool_use_id"])] else {
                return Row(.orphan(body), taskId: taskId)
            }
            target.result = body
            return nil

        case "turn.ended":
            for row in ctx.tools.values where row.result == nil && !(row.active || row.inActive) {
                row.unfinished = true
            }
            if bool(p, "handoff") {
                return note([.text("handoff written; continuing in a new session with a fresh context")])
            }
            return Row(.turnEnd(error: bool(p, "is_error") ? "turn ended: \(jsString(p["subtype"]).replacingOccurrences(of: "_", with: " "))" : nil), taskId: taskId)

        case "session.created":
            let model = str(p, "model") ?? ""
            return note([.text("\(ev.actor) started a \(jsString(p["vendor"])) session\(model.isEmpty ? "" : " (\(model))")")])

        case "session.model":
            return note([.text("\(ev.actor) switched the model to \(jsString(p["model"]))")])

        case "session.state":
            guard let error = str(p, "error"), !error.isEmpty else {
                return nil
            }
            return note([.text("session \(jsString(p["state"])): \(error)")], .error, time: false)

        case "inbox.created":
            let q = dict(p, "payload") ?? [:]
            switch str(p, "kind") {
            case "question":
                return note([.text("asked: \(jsString(q["question"]))")], .ask, time: false)
            case "session_error", "rate_limited":
                return note([.text(jsString(q["error"]))], .error, time: false)
            default:
                return nil
            }

        case "task.created":
            if let from = str(p, "forked_from") {
                return note([.text("\(ev.actor) forked this task from "), .task(from), .text(" at snapshot #\(jsString(p["forked_at"])); its agent continues that task's conversation as it stood there")])
            }
            if let parentTask = str(p, "parent_task_id") {
                return note([.text("subtask of "), .task(parentTask), .text(", created by its agent")])
            }
            return nil

        case "task.parent_merge_started":
            return note([.text("\(ev.actor) asked "), .task(jsString(p["parent_task_id"])), .text("'s agent to merge this subtask into its branch")])

        case "task.parent_merge_stopped":
            return note([.text("the merge into the parent task didn't start")], .error)

        case "task.merge_title":
            return note([.text("\(ev.actor) set the merge commit title to \"\(jsString(p["merge_title"]))\"")])

        case "review.finished":
            let r = try! decodePayload(ev.payload, as: Review.self)
            let outcome = r.state == .failed ? "failed: \(r.error)" : r.state == .skipped ? "skipped: no changes" : reviewCounts(r)
            return note([.text("review of snapshot #\(r.snapshotN): \(outcome)")], r.state == .failed ? .error : .plain)

        case "task.merge_queued":
            return note([.text("\(ev.actor) asked to merge; held until the merge in progress in this project is done")])

        case "task.merge_stopped":
            return note([.text("\(ev.actor) stopped the merge for review: \(jsString(p["reason"]))")])

        case "task.merge_after_review":
            var n = Note(parts: [.text("\(ev.actor) asked to merge once the review of snapshot #\(jsString(p["snapshot_n"])) finishes with no findings")], when: when)
            n.cancelMergeAfterReview = jsString(p["id"])
            return Row(.note(n), taskId: taskId)

        case "task.merge_after_review_cancelled":
            return note([.text("\(ev.actor) cancelled the merge after review")])

        case "task.merge_after_review_dropped":
            return note([.text("didn't merge after the review: \(jsString(p["reason"]))")])

        case "preview.exposed":
            return note([.text("\(ev.actor) exposed \"\(jsString(p["name"]))\" at "), .url(String(int(p, "port") ?? 0))])

        case "task.merged":
            let commit = str(p, "commit") ?? ""
            if let parentTask = str(p, "parent_task_id") {
                return note([.text("\(ev.actor) merged the task into "), .task(parentTask), .text(", now at \(commit.prefix(7))")])
            }
            let pushed = str(p, "pushed") ?? ""
            return note([.text("\(ev.actor) merged the task\(commit.isEmpty ? "" : " as \(commit.prefix(7))")\(pushed.isEmpty ? "" : " and pushed it to \(pushed)")")])

        case "task.summary":
            return note([.text("\(ev.actor) set the task summary")])

        case "task.discard_started":
            return note([.text("\(ev.actor) asked to discard the task; the agent writes the task summary first")])

        case "task.discarded":
            return note([.text("\(ev.actor) discarded the task")])

        default:
            return nil
        }
    }
}

private let severities = ["critical", "high", "medium", "low", "info"]

func isReviewEvent(_ ev: ByreEvent) -> Bool {
    ev.kind.hasPrefix("review.")
}

func findings(_ r: Review) -> [JSONDict] {
    (r.result?.value["findings"] as? [(any Sendable)?] ?? []).map { $0 as? JSONDict ?? [:] }
}

func findingCount(_ r: Review) -> Int {
    r.state == .done ? findings(r).indices.filter { r.findingStates.additionalProperties[String($0)] == nil }.count : 0
}

func latestReview(_ events: [ByreEvent]) -> Review? {
    var best: Review?
    for ev in events where isReviewEvent(ev) {
        let r = try! decodePayload(ev.payload, as: Review.self)
        if best == nil || r.snapshotN >= best!.snapshotN {
            best = r
        }
    }
    return best
}

func reviewCounts(_ r: Review) -> String {
    severityCounts(findings(r))
}

func severityCounts(_ fs: [JSONDict]) -> String {
    let parts = severities.map { s in (s, fs.filter { str($0, "severity") == s }.count) }.filter { $0.1 > 0 }
    return parts.isEmpty ? "no findings" : parts.map { "\($0.1) \($0.0)" }.joined(separator: ", ")
}

private func parentOf(_ p: JSONDict) -> String {
    str(p, "parent_tool_use_id") ?? ""
}

private func isSubagentCall(_ p: JSONDict) -> Bool {
    let name = str(p, "name")
    return name == "Agent" || name == "Task"
}

func toolLabel(_ name: String) -> String {
    let prefix = "mcp__byre__"
    if !name.hasPrefix(prefix) {
        return name
    }
    return name.dropFirst(prefix.count).split(separator: "_", omittingEmptySubsequences: false).map { $0.prefix(1).uppercased() + $0.dropFirst() }.joined(separator: " ")
}

private func toolBrief(_ input: (any Sendable)?, _ root: String) -> String {
    guard let i = input as? JSONDict else {
        return ""
    }
    let brief = ["command", "file_path", "pattern", "url", "question", "description"].lazy.compactMap { i[$0] ?? nil }.first
    let text = relative(brief.map { jsString($0) } ?? "", root)
    return text.count > 120 ? "\(text.prefix(120))…" : text
}

private func toolTarget(_ input: (any Sendable)?, _ root: String) -> String {
    guard let path = (input as? JSONDict)?["file_path"] as? String else {
        return ""
    }
    return relative(path, root)
}

private func relative(_ path: String, _ root: String) -> String {
    !root.isEmpty && path.hasPrefix("\(root)/") ? String(path.dropFirst(root.count + 1)) : path
}

private func resultText(_ content: (any Sendable)?) -> String {
    if let s = content as? String {
        return s
    }
    if let list = content as? [(any Sendable)?] {
        return list.map { c in
            if let d = c as? JSONDict, let text = d["text"] {
                return jsString(text)
            }
            return compactJSON(c)
        }.joined(separator: "\n")
    }
    return prettyJSON(content)
}