byre/mac/Sources/Byre/Store.swift

import Foundation
import Observation

let taskTabs: [TaskTab] = [.transcript, .changes, .files, .terminal, .inbox, .reviews, .workflow, .project]

private let historyPage = 1000
private let leadPage = 50
private let visitsKey = "byre.visits"

struct TaskState {
    enum Kind {
        case answer, merging, held, error, working, idle
    }

    let kind: Kind
    let label: String
    let title: String
}

private let workingTitles: [ByreTask.StatePayload: String] = [
    .running: "The agent is mid-turn or its subagents are running",
    .reviewPending: "A review of the turn's work is about to start",
    .reviewing: "A review is running",
    .discarding: "Discarding",
]

@MainActor
@Observable
final class Store {
    let api: API
    var me: User?
    var projects: [Project] = []
    var tasks: [String: [ByreTask]] = [:]
    var sessions: [String: [AgentSession]] = [:]
    var inbox: [InboxItem] = []
    private(set) var selectedTaskId: String?
    var pinned: Int?
    var reviewed: Int?
    var connected = false
    private(set) var focus = Focus(tabs: [.transcript], tab: .transcript)
    var leadSessions: [AgentSession] = []
    let leadFeed: Feed
    var jump: (taskId: String, seq: Int64)?
    var summaryReads: [String: Int64] = [:]
    private(set) var leadMore = false
    var todoStarting: Set<String> = []
    private(set) var drafts: Drafts?
    let images = ImageDrafts()
    let events = Listeners<ByreEvent>()
    let paneFocus: PaneFocus

    private var taskFeeds: [String: Feed] = [:]
    @ObservationIgnored private var historyLoads: Set<ObjectIdentifier> = []
    @ObservationIgnored private var requestSeqs: [String: Int] = [:]
    @ObservationIgnored private var globalStream: EventStream?
    @ObservationIgnored private var leadStream: EventStream?
    @ObservationIgnored private var stopped = false

    init(api: API) {
        self.api = api
        leadFeed = Feed(shown: { true })
        paneFocus = PaneFocus(api: api)
    }

    func stop() async {
        stopped = true
        globalStream?.stop()
        leadStream?.stop()
        drafts?.flush()
        await paneFocus.leave()
    }

    func loadDrafts() {
        drafts = Drafts(origin: api.origin, login: me!.login)
        drafts!.prune(self)
    }

    func startTodo(_ projectId: String, _ title: String, discuss: Bool, _ overlays: Overlays) {
        let key = "\(projectId)\n\(title)"
        if todoStarting.contains(key) {
            return
        }
        todoStarting.insert(key)
        Task {
            do {
                _ = try await api.client.startTodoItem(path: .init(projectId: projectId), body: .json(.init(title: title, model: models[0], discuss: discuss))).created
            } catch {
                overlays.showAlert(errorText(error))
            }
            todoStarting.remove(key)
        }
    }

    func startSuggestion(_ projectId: String, _ title: String, brief: String, _ overlays: Overlays) {
        let key = "\(projectId)\n\(title)"
        if todoStarting.contains(key) {
            return
        }
        todoStarting.insert(key)
        Task {
            do {
                _ = try await api.client.createTask(path: .init(projectId: projectId), body: .json(.init(title: title, prompt: brief, model: models[0]))).created
            } catch {
                overlays.showAlert(errorText(error))
            }
            todoStarting.remove(key)
        }
    }

    func discussSuggestion(_ title: String, brief: String, _ overlays: Overlays) {
        let lead = currentLead()
        let live = lead != nil && lead!.state != .completed
        Task {
            do {
                _ = try await api.client.messageLead(body: .json(.init(text: "Let's discuss your suggestion \"\(title)\" before starting it:\n\n\(brief)", model: live ? nil : models[0]))).accepted
            } catch {
                overlays.showAlert(errorText(error))
            }
        }
    }

    var transcript: [ByreEvent] {
        selectedFeed?.events ?? []
    }

    var selectedFeed: Feed? {
        selectedTaskId.flatMap { taskFeeds[$0] }
    }

    func activeTasks() -> [(project: Project, task: ByreTask)] {
        projects.flatMap { project in
            (tasks[project.id] ?? []).filter { $0.status == .active }.map { (project, $0) }
        }
    }

    func findTask(_ id: String) -> ByreTask? {
        for list in tasks.values {
            if let task = list.first(where: { $0.id == id }) {
                return task
            }
        }
        return nil
    }

    func currentSession(_ taskId: String) -> AgentSession? {
        sessions[taskId]?.last(where: { $0.state != .completed }) ?? sessions[taskId]?.last
    }

    func taskState(_ task: ByreTask) -> TaskState {
        let open = inbox.filter { $0.taskId == task.id }
        if open.contains(where: { $0.kind == .question }) {
            return TaskState(kind: .answer, label: "Needs your answer", title: "The agent asked a question in the inbox")
        }
        if task.merging {
            return TaskState(kind: .merging, label: "Merging", title: "The merge turn is running")
        }
        if task.mergingIntoParent {
            return TaskState(kind: .merging, label: "Merging", title: "Waiting for the parent task's agent to merge it")
        }
        if task.mergeQueuedAt != nil {
            return TaskState(kind: .held, label: "Merge held", title: "Waiting for the merge in progress in this project to finish")
        }
        if task.mergeAfterReview {
            return TaskState(kind: .held, label: "Merge after review", title: "Merges when the running review finishes with no findings")
        }
        if !open.isEmpty {
            return TaskState(kind: .error, label: "Error", title: "The agent's session failed or hit a rate limit; see the inbox")
        }
        if let working = workingTitles[task.state] {
            return TaskState(kind: .working, label: "Working", title: working)
        }
        return TaskState(kind: .idle, label: "Idle", title: "Waiting for a message")
    }

    func currentLead() -> AgentSession? {
        leadSessions.last(where: { $0.state != .completed }) ?? leadSessions.last
    }

    func snapshots() -> [Snapshot] {
        transcript.filter { $0.kind == "snapshot.created" }.map { try! decodePayload($0.payload, as: Snapshot.self) }
    }

    func pinSnapshot(_ n: Int?) {
        pinned = n
    }

    func pickReview(_ n: Int) {
        reviewed = n
    }

    func showReview(_ n: Int) {
        pickReview(n)
        focusTab(.reviews)
    }

    func loadMe() async throws {
        me = try await api.client.getMe().ok.body.json
    }

    func loadProjects() async throws {
        projects = try await api.client.listProjects().ok.body.json
        try await each(projects.map(\.id)) { id in
            try await self.loadTasks(id)
        }
    }

    private func claimRequest(_ key: String) -> () -> Bool {
        let seq = (requestSeqs[key] ?? 0) + 1
        requestSeqs[key] = seq
        return { [weak self] in
            self?.requestSeqs[key] == seq
        }
    }

    func loadTasks(_ projectId: String) async throws {
        let latest = claimRequest("tasks:\(projectId)")
        let list = try await api.client.listTasks(path: .init(projectId: projectId)).ok.body.json
        if !latest() {
            return
        }
        tasks[projectId] = list
        for t in list where t.status == .active {
            _ = taskFeed(t.id)
        }
        dropClosedFeeds()
        drafts?.prune(self)
        try await each(list.filter { $0.status == .active }.map(\.id)) { id in
            try await self.loadSessions(id)
        }
    }

    func loadSessions(_ taskId: String) async throws {
        let latest = claimRequest("sessions:\(taskId)")
        let list = try await api.client.listSessions(path: .init(taskId: taskId)).ok.body.json
        if !latest() {
            return
        }
        sessions[taskId] = list
    }

    func loadInbox() async throws {
        let latest = claimRequest("inbox")
        let items = try await api.client.listInbox().ok.body.json
        if !latest() {
            return
        }
        inbox = items
        drafts?.prune(self)
    }

    func loadLeadSessions() async throws {
        leadSessions = try await api.client.listLeadSessions().ok.body.json
    }

    func loadFocus() async throws {
        let latest = claimRequest("focus")
        let f = try await api.client.getFocus().ok.body.json
        if !latest() {
            return
        }
        applyFocus(f)
    }

    func applyFocus(_ f: Focus) {
        focus = f
        if f.taskId != selectedTaskId {
            selectTask(f.taskId)
        }
    }

    func focusedProject() -> Project? {
        projects.first { $0.id == focus.projectId }
    }

    func setFocus(taskId: String?? = .none, projectId: String?? = .none, tab: TaskTab? = nil) async throws {
        let task = taskId ?? focus.taskId
        let project = task != nil ? findTask(task!)?.projectId : (projectId ?? focus.projectId)
        var next = Focus(tabs: focus.tabs, tab: tab ?? focus.tab)
        if let task, !task.isEmpty {
            next.taskId = task
        }
        if let project, !project.isEmpty {
            next.projectId = project
        }
        let latest = claimRequest("focus")
        applyFocus(next)
        let f = try await api.client.setFocus(body: .json(next)).ok.body.json
        if !latest() {
            return
        }
        applyFocus(f)
    }

    func focusTask(_ taskId: String, _ tab: TaskTab = .transcript) {
        spawn("set focus") {
            try await self.setFocus(taskId: taskId, tab: tab)
        }
    }

    func focusProject(_ projectId: String) {
        spawn("set focus") {
            try await self.setFocus(taskId: .some(nil), projectId: projectId, tab: .project)
        }
    }

    func focusTab(_ tab: TaskTab) {
        spawn("set focus") {
            try await self.setFocus(tab: tab)
        }
    }

    private func loadHistory(_ taskId: String, _ feed: Feed) async throws {
        var got: [ByreEvent] = []
        while true {
            let page = try await api.client.listTaskEvents(path: .init(taskId: taskId), query: .init(after: got.last?.seq ?? 0, limit: historyPage)).ok.body.json
            got += page
            if page.count < historyPage {
                break
            }
        }
        feed.load(got)
    }

    private func taskFeed(_ taskId: String) -> Feed {
        let feed = taskFeeds[taskId] ?? Feed(shown: { [weak self] in
            self?.selectedTaskId == taskId
        })
        taskFeeds[taskId] = feed
        let key = ObjectIdentifier(feed)
        if feed.ready || historyLoads.contains(key) {
            return feed
        }
        historyLoads.insert(key)
        Task {
            do {
                try await loadHistory(taskId, feed)
            } catch {
                log.error("load task history \(taskId, privacy: .public): \(errorText(error), privacy: .public)")
            }
            historyLoads.remove(key)
        }
        return feed
    }

    private func dropClosedFeeds() {
        for taskId in taskFeeds.keys {
            if taskId != selectedTaskId, let task = findTask(taskId), task.status != .active {
                taskFeeds[taskId] = nil
            }
        }
    }

    func transcriptLoading() -> Bool {
        selectedTaskId != nil && selectedFeed?.ready != true
    }

    func leadLoading() -> Bool {
        !leadFeed.ready
    }

    func lastVisits() -> [String: Double] {
        UserDefaults.standard.dictionary(forKey: visitsKey) as? [String: Double] ?? [:]
    }

    private func markVisited(_ taskId: String?) {
        guard let taskId else {
            return
        }
        var visits = lastVisits()
        visits[taskId] = Date().timeIntervalSince1970 * 1000
        UserDefaults.standard.set(visits, forKey: visitsKey)
    }

    func loadSummaryReads() async throws {
        summaryReads = try await api.client.getSummaryReads().ok.body.json.additionalProperties
    }

    func unreadSummary(_ task: ByreTask) -> Bool {
        task.workflow == nil && task.summarySeq > (summaryReads[task.id] ?? 0)
    }

    private func noteSummaryRead(_ taskId: String, _ seq: Int64) -> Bool {
        if (summaryReads[taskId] ?? 0) >= seq {
            return false
        }
        summaryReads[taskId] = seq
        return true
    }

    func markSummaryRead(_ taskId: String, _ seq: Int64) {
        if !noteSummaryRead(taskId, seq) {
            return
        }
        spawn("mark summary read") {
            _ = try await self.api.client.markSummaryRead(path: .init(taskId: taskId), body: .json(.init(seq: seq))).noContent
        }
    }

    func openSummary(_ task: ByreTask) {
        jump = (task.id, task.summarySeq)
        focusTask(task.id, .transcript)
    }

    func selectTask(_ taskId: String?) {
        markVisited(selectedTaskId)
        markVisited(taskId)
        selectedTaskId = taskId
        if let taskId {
            _ = taskFeed(taskId)
        }
        pinned = nil
        reviewed = nil
        dropClosedFeeds()
    }

    private func fetchLeadTail() async -> EventPage? {
        while !stopped {
            do {
                return try await api.client.listLeadEvents(query: .init(turns: leadPage)).ok.body.json
            } catch {
                log.error("load lead history: \(errorText(error), privacy: .public)")
                try? await Task.sleep(for: .seconds(2))
            }
        }
        return nil
    }

    func startLeadStream() async {
        guard let page = await fetchLeadTail(), !stopped else {
            return
        }
        leadMore = page.more
        leadStream = EventStream(api: api, after: page.events.last?.seq ?? 0, lead: true, onEvent: { [weak self] ev in
            self?.leadFeed.push(ev)
        }, onState: { _ in }, onReady: { [weak self] in
            self?.leadFeed.load(page.events)
        })
    }

    func startGlobalStream(after: Int64) {
        globalStream = EventStream(api: api, after: after, lead: false, onEvent: { [weak self] ev in
            self?.handleEvent(ev)
        }, onState: { [weak self] up in
            self?.connected = up
            if up {
                self?.paneFocus.reconnected()
            }
        }, onReady: {})
    }

    func leadHasOlder() -> Bool {
        leadFeed.ready && leadMore
    }

    func loadOlderLead() async throws -> [ByreEvent] {
        guard let first = leadFeed.events.first else {
            return []
        }
        let page = try await api.client.listLeadEvents(query: .init(before: first.seq, turns: leadPage)).ok.body.json
        leadMore = page.more
        leadFeed.prepend(page.events)
        return page.events
    }

    func handleEvent(_ ev: ByreEvent) {
        if let taskId = ev.taskId {
            taskFeeds[taskId]?.push(ev)
        }
        if ev.seq == 0 {
            events.send(ev)
            return
        }

        switch ev.kind {
        case "user.updated":
            if ev.payload.value["login"] as? String == me?.login {
                spawn("load me") {
                    try await self.loadMe()
                }
            }
        case "project.created":
            spawn("load projects") {
                try await self.loadProjects()
            }
        case "task.created", "task.merge_title", "task.merge_queued", "task.merge_unqueued", "task.merge_after_review", "task.merge_after_review_ended", "task.merge_after_review_cancelled", "task.parent_merge_started", "task.parent_merge_stopped", "task.merge_started", "task.merge_ended", "task.discard_started", "task.merged", "task.discarded", "task.state":
            if let projectId = ev.projectId {
                spawn("load tasks") {
                    try await self.loadTasks(projectId)
                }
            }
        case "session.created", "session.state", "session.subagents", "session.model":
            if ev.owner != nil {
                spawn("load lead sessions") {
                    try await self.loadLeadSessions()
                }
            } else if let taskId = ev.taskId {
                spawn("load sessions") {
                    try await self.loadSessions(taskId)
                }
            }
        case "inbox.created":
            notifyInbox(self, try! decodePayload(ev.payload, as: InboxItem.self))
            spawn("load inbox") {
                try await self.loadInbox()
            }
        case "inbox.updated":
            spawn("load inbox") {
                try await self.loadInbox()
            }
        case "turn.ended":
            notifyTurn(self, ev)
            if ev.payload.value["summary"] as? Bool == true, let taskId = ev.taskId {
                updateTask(taskId) { task in
                    if ev.seq > task.summarySeq {
                        task.summarySeq = ev.seq
                    }
                }
            }
        case "summary.read":
            if let taskId = ev.payload.value["task_id"] as? String, let seq = ev.payload.value["seq"] as? Int {
                _ = noteSummaryRead(taskId, Int64(seq))
            }
        case "focus.updated":
            if ev.actor != me?.login {
                _ = claimRequest("focus")
                applyFocus(try! decodePayload(ev.payload, as: Focus.self))
            }
        default:
            break
        }

        events.send(ev)
    }

    private func updateTask(_ id: String, _ change: (inout ByreTask) -> Void) {
        for (projectId, list) in tasks {
            if let i = list.firstIndex(where: { $0.id == id }) {
                change(&tasks[projectId]![i])
            }
        }
    }

    func spawn(_ what: String, _ op: @escaping @MainActor () async throws -> Void) {
        Task {
            do {
                try await op()
            } catch {
                log.error("\(what, privacy: .public): \(errorText(error), privacy: .public)")
            }
        }
    }

    private func each(_ ids: [String], _ op: @escaping @Sendable @MainActor (String) async throws -> Void) async throws {
        try await withThrowingTaskGroup(of: Void.self) { group in
            for id in ids {
                group.addTask {
                    try await op(id)
                }
            }
            try await group.waitForAll()
        }
    }
}