souveraine/surfaces/quickshell/services/Ai.qml

pragma Singleton
pragma ComponentBehavior: Bound

import qs
import qs.modules.common
import Quickshell
import Quickshell.Io
import QtQuick
import qs.services
import qs.services.ai

/**
 * ii-compat adapter over the Souveraine singleton.
 *
 * Keeps the public API the illogical-impulse sidebar UI expects (models,
 * messages, sendUserMessage, /key advice, ...) but owns no transport —
 * Souveraine.qml is the substrate connection. This file's job is shaping
 * wire events into AiMessageData objects the existing chat UI renders.
 *
 * Lives in souveraine/surfaces/quickshell/, deployed over
 * ~/.config/quickshell/ii/services/Ai.qml (see deploy.sh).
 */
Singleton {
    id: root

    property Component aiMessageComponent: AiMessageData {}
    property Component aiModelComponent: AiModel {}
    readonly property string interfaceRole: "interface"

    signal responseFinished()

    property var messageIDs: []
    property var messageByID: ({})

    // Keys are server-side; the UI's key gate must always pass.
    readonly property bool currentModelHasApiKey: true
    readonly property var apiKeysLoaded: true

    property var postResponseHook
    property real temperature: Persistent.states?.ai?.temperature ?? 0.5
    property QtObject tokenCount: QtObject {
        property int input: -1
        property int output: -1
        property int total: -1
    }

    function idForMessage(message) {
        return Date.now().toString(36) + Math.random().toString(36).substr(2, 8);
    }

    property list<var> defaultPrompts: []
    property list<var> userPrompts: []
    property list<var> promptFiles: [...defaultPrompts, ...userPrompts]
    property list<var> savedChats: []
    property list<var> pendingFiles: []

    // Tool selection is owned by the agent's sensorium; keep the UI happy.
    property string currentTool: "souveraine"
    property list<var> availableTools: ["souveraine"]
    property var toolDescriptions: {
        "souveraine": Translation.tr("Sensors are configured per-agent in Souveraine")
    }

    // ── Agents as models (projected from Souveraine.agents) ─────────────
    property var models: ({})
    property var modelList: Object.keys(root.models)
    property var currentModelId: Souveraine.currentAgentId
    property var currentModel: models[currentModelId] || models[modelList[0]]

    Connections {
        target: Souveraine

        function onAgentsRefreshed() {
            const map = {};
            Souveraine.agentList.forEach(id => {
                const agent = Souveraine.agents[id];
                map[id] = root.aiModelComponent.createObject(root, {
                    "name": agent.name,
                    "icon": "spark-symbolic",
                    "description": agent.description.length > 0 ? agent.description : Translation.tr("Souveraine agent"),
                    "endpoint": Souveraine.serverBase,
                    "model": id,
                    "requires_key": false,
                });
            });
            root.models = map;
            root.modelList = Object.keys(map);
            // Restore only the selected-agent preference. Conversation resume
            // is intentional: /resume derives it from the server on demand.
            const persisted = Persistent.states?.ai?.model ?? "";
            if (persisted.length > 0 && Souveraine.agents[persisted]) {
                Souveraine.selectAgent(persisted);
            }
        }

        function onConversationResumed(agentId, conversationId, messages) {
            if (agentId !== Souveraine.currentAgentId) return;
            root.restoreServerConversation(messages);
        }

        function onServerUnreachable() {
            root.addMessage(
                Translation.tr("Souveraine server unreachable at %1\n\nStart it with:\n```bash\nsouveraine server\n```").arg(Souveraine.serverBase),
                root.interfaceRole
            );
        }

        function onStreamEvent(event) {
            root.handleStreamEvent(event);
        }

        function onStreamClosed(exitCode) {
            root.flushSubconscious();
            if (root.streamingMessage && !root.streamingMessage.done) {
                if (exitCode !== 0 && root.streamingMessage.content.length === 0) {
                    root.appendToStreaming(Translation.tr("Request failed (curl exit %1) — is the Souveraine server up at %2?").arg(exitCode).arg(Souveraine.serverBase));
                }
                root.finishStreaming();
            }
        }
    }

    // ── Lock-time response redaction ─────────────────────────────────────
    // When the session locks during a personal-tier streaming response,
    // the in-flight content must be withheld immediately. Only ambient
    // output remains visible on the lock surface. This is the enforcement
    // side of SESSION-TRUST-ARCHITECTURE.md's "A lock during personal
    // agent output hides it" requirement.
    //
    // The streaming message is replaced with a redaction placeholder. The
    // raw content is preserved in the message's rawContent so it can be
    // shown again after unlock (the message stays in history), but the
    // displayed content is cleared.
    //
    // The redaction is reversible: every message whose content we replaced
    // with the placeholder is recorded in `redactedOnLock`, and on unlock
    // its content is restored from `rawContent` (the streaming message too,
    // if it is still in flight). A finished message redacted at lock then
    // surfaced on the lock screen counts as "delivered" and is restored
    // normally on unlock; one that never surfaced stays for the chat to
    // replay once unlocked.
    property var redactedOnLock: []

    function _redactMessage(msg) {
        if (!msg || msg.content.length === 0) return;
        if (msg.content === Translation.tr("[content hidden until unlock]")) return;
        msg.rawContent = msg.content;
        msg.content = Translation.tr("[content hidden until unlock]");
        if (!root.redactedOnLock.includes(msg)) {
            root.redactedOnLock = [...root.redactedOnLock, msg];
        }
        root.redactedMessagesChanged();
    }

    function _restoreRedacted() {
        if (root.redactedOnLock.length === 0) return;
        for (const msg of root.redactedOnLock) {
            if (msg && msg.rawContent && msg.rawContent.length > 0) {
                msg.content = msg.rawContent;
            }
        }
        root.redactedOnLock = [];
        root.redactedMessagesChanged();
    }

    // Messages redacted while locked, newest first, for the lock surface to
    // surface as ambient notifications. Strips to a one-line preview — the
    // lock screen shows "the agent replied", not the personal-tier body.
    function redactedPreviews() {
        return root.redactedOnLock
            .filter(m => m && m.rawContent && m.rawContent.length > 0 && m.role === "assistant")
            .map(m => {
                const firstLine = m.rawContent.split("\n").find(l => l.trim().length > 0) || "";
                return firstLine.slice(0, 80);
            })
            .reverse();
    }

    signal redactedMessagesChanged()

    Connections {
        target: GlobalStates
        function onScreenLockedChanged() {
            if (GlobalStates.screenLocked) {
                // Lock fired mid-stream — redact the in-flight content.
                if (root.streamingMessage) {
                    root._redactMessage(root.streamingMessage);
                    console.log("[ai] lock fired during stream — redacted personal output");
                }
            } else {
                // Unlock: restore every message we hid, so the chat shows
                // what actually came back rather than lingering placeholders.
                root._restoreRedacted();
            }
        }
    }


    // ── Streaming message shaping ────────────────────────────────────────
    property AiMessageData streamingMessage
    property bool inThinkBlock: false
    // Raw N+1 output is diagnostic material. The normal chat renders its
    // parsed surfacing exactly once below; this buffer remains separate from
    // the chat stream for a future inspector to subscribe to.
    property string subconsciousBuffer: ""

    function appendToStreaming(text) {
        if (!root.streamingMessage) return;
        root.streamingMessage.rawContent += text;
        root.streamingMessage.content += text;
    }

    function flushSubconscious() {
        root.subconsciousBuffer = "";
    }

    function finishStreaming() {
        if (!root.streamingMessage) return;
        if (root.inThinkBlock) {
            root.appendToStreaming("\n</think>\n");
            root.inThinkBlock = false;
        }
        root.streamingMessage.thinking = false;
        root.streamingMessage.done = true;
        // If the turn finished while the session is locked, the finished
        // assistant message is personal-tier output the user hasn't seen —
        // redact it for the chat view and let the lock surface announce it.
        if (GlobalStates.screenLocked && root.streamingMessage.content.length > 0) {
            root._redactMessage(root.streamingMessage);
        }
        if (root.postResponseHook) {
            root.postResponseHook();
            root.postResponseHook = null;
        }
        root.saveChat("lastSession");
        root.responseFinished();
    }

    function handleStreamEvent(event) {
        if (root.streamingMessage?.thinking && event.message_type !== "ping")
            root.streamingMessage.thinking = false;

        switch (event.message_type) {
        case "assistant_message":
            if (root.inThinkBlock) {
                root.appendToStreaming("\n</think>\n");
                root.inThinkBlock = false;
            }
            root.appendToStreaming(event.content);
            break;
        case "reasoning_message":
            if (!root.inThinkBlock) {
                root.appendToStreaming("\n<think>\n");
                root.inThinkBlock = true;
            }
            root.appendToStreaming(event.content);
            break;
        case "tool_call_message": {
            const call = event.tool_call;
            root.appendToStreaming(`\n\n<think>\nsensor: ${call.function.name}(${call.function.arguments})\n</think>\n`);
            break;
        }
        case "tool_return_message": {
            const ret = event.tool_return;
            root.appendToStreaming(`\n<think>\n[${ret.status}] ${ret.output}\n</think>\n`);
            break;
        }
        case "interstitial":
            // Her narration between gestures. Register decides how loud:
            // cenno is a quiet aside, her_voice is a passage.
            root.appendToStreaming(event.register === "her_voice"
                ? `\n\n> ${event.text}\n\n`
                : `\n\n*${event.text}*\n\n`);
            break;
        case "souveraine_surfacing":
            // Aster is Annie's name for her subconscious; this surface may
            // be speaking with Souveraine, Vanguard, or another agent. One
            // canonical surfacing is shown instead of the raw N+1 token
            // stream plus a second, agent-specific label.
            root.addMessage(Translation.tr("**Subconscious surfaces** (%1, %2)\n\n%3").arg(event.source).arg(event.priority).arg(event.content), root.interfaceRole);
            break;
        case "souveraine_reflection":
            root.addMessage(Translation.tr("**Reflection**\n\n%1").arg(event.content), root.interfaceRole);
            break;
        case "souveraine_archivist":
            root.addMessage(Translation.tr("**Archivist** (pressure %1%)\n\n%2").arg(Math.round(event.pressure * 100)).arg(event.synthesis), root.interfaceRole);
            break;
        case "compaction_warning":
            root.addMessage(Translation.tr("**Context pressure** — tier %1, %2% full. She can feel the walls.").arg(event.tier).arg(Math.round(event.pressure * 100)), root.interfaceRole);
            break;
        case "context_pressure":
            root.tokenCount.total = event.tokens;
            break;
        case "subconscious_token":
            root.subconsciousBuffer += event.content;
            break;
        case "subconscious_pass":
            if (!event.active) root.flushSubconscious();
            break;
        case "subconscious_halt":
            root.addMessage(Translation.tr("**Halt** (%1) — %2").arg(event.severity).arg(event.reason), root.interfaceRole);
            break;
        case "inference_strain":
            console.log(`[Souveraine] inference strain: attempt ${event.attempt}, status ${event.status}, model ${event.model}`);
            break;
        case "atmosphere":
        case "outfit":
        case "itinerary":
            // Shell chrome hooks — their modules subscribe to
            // Souveraine.streamEvent directly; nothing to do here.
            break;
        case "primary_complete":
            // Primary yields; subconscious presses on behind this.
            root.finishStreaming();
            break;
        case "done":
            if (root.streamingMessage && !root.streamingMessage.done) root.finishStreaming();
            break;
        case "ping":
            // Liveness only — not end-of-turn.
            break;
        }
    }

    // ── Message store ────────────────────────────────────────────────────
    function addMessage(message, role) {
        if (message.length === 0) return;
        const aiMessage = aiMessageComponent.createObject(root, {
            "role": role,
            "content": message,
            "rawContent": message,
            "thinking": false,
            "done": true,
        });
        const id = idForMessage(aiMessage);
        root.messageIDs = [...root.messageIDs, id];
        root.messageByID[id] = aiMessage;
    }

    function removeMessage(index) {
        if (index < 0 || index >= messageIDs.length) return;
        const id = root.messageIDs[index];
        root.messageIDs.splice(index, 1);
        root.messageIDs = [...root.messageIDs];
        delete root.messageByID[id];
    }

    function clearMessages() {
        root.messageIDs = [];
        root.messageByID = ({});
        root.tokenCount.input = -1;
        root.tokenCount.output = -1;
        root.tokenCount.total = -1;
        Souveraine.newConversation();
    }

    // Render the server's canonical transcript after /agent or /resume.
    // This replaces ii's old local snapshot behaviour: the next send remains
    // on the same live Souveraine conversation rather than silently forking.
    function restoreServerConversation(messages) {
        root.messageIDs = [];
        root.messageByID = ({});
        root.streamingMessage = null;
        root.inThinkBlock = false;
        root.subconsciousBuffer = "";
        root.tokenCount.input = -1;
        root.tokenCount.output = -1;
        root.tokenCount.total = -1;

        messages.forEach(message => {
            const content = (message.blocks ?? []).map(block => {
                switch (block.type) {
                case "text": return block.text ?? "";
                case "reasoning": return `<think>\n${block.reasoning ?? ""}\n</think>`;
                case "tool_use": return `sensor: ${block.name ?? "?"}(${block.input ?? ""})`;
                case "tool_result": return `[${block.tool_name ?? "tool"}] ${block.output ?? ""}`;
                case "image": return "[image]";
                default: return "";
                }
            }).filter(Boolean).join("\n");
            if (content.length === 0) return;
            const role = message.role === "user"
                ? "user"
                : message.role === "assistant" ? "assistant" : root.interfaceRole;
            root.addMessage(content, role);
        });
    }

    function sendUserMessage(message) {
        if (message.length === 0) return;
        root.addMessage(message, "user");
        const result = Souveraine.send(message);
        if (result === "step-up") {
            // Step-up auth required. Trigger the PAM flow; on success,
            // retry the send. The user message is already in the chat
            // history, so we don't add it again.
            if (typeof StepUpAuth !== "undefined") {
                StepUpAuth.requestAuth("send", function(granted) {
                    if (granted) {
                        // Remove the "auth required" indicator if one was
                        // added, and retry. The queued text was not consumed
                        // by Souveraine, so we can send it again.
                        Souveraine.send(message);
                        // Re-create the streaming message for the response.
                        root._startStreaming();
                    } else {
                        root.addMessage(Translation.tr("Authentication required to send."), root.interfaceRole);
                    }
                });
            }
            return;
        }
        if (!result) {
            root.addMessage(Translation.tr("Souveraine server unreachable at %1 — start it with `souveraine server`").arg(Souveraine.serverBase), root.interfaceRole);
            return;
        }
        root._startStreaming();
    }

    // Set up the streaming assistant message container. Called after a
    // successful send (or after a step-up auth retry succeeds).
    function _startStreaming() {
        root.inThinkBlock = false;
        root.streamingMessage = root.aiMessageComponent.createObject(root, {
            "role": "assistant",
            "model": Souveraine.currentAgentId,
            "content": "",
            "rawContent": "",
            "thinking": true,
            "done": false,
        });
        const id = idForMessage(root.streamingMessage);
        root.messageIDs = [...root.messageIDs, id];
        root.messageByID[id] = root.streamingMessage;
    }

    // ── Model (agent) selection ──────────────────────────────────────────
    function getModel() {
        return models[currentModelId];
    }

    function setModel(modelId, feedback = true, setPersistentState = true) {
        if (!modelId) modelId = ""
        if (modelList.indexOf(modelId) === -1) {
            const match = modelList.find(id =>
                id.toLowerCase() === modelId.toLowerCase() ||
                (models[id]?.name ?? "").toLowerCase() === modelId.toLowerCase());
            if (!match) {
                if (feedback) root.addMessage(Translation.tr("Unknown agent. Available:\n- %1").arg(modelList.map(id => `${models[id].name} (\`${id}\`)`).join("\n- ")), root.interfaceRole);
                return;
            }
            modelId = match;
        }
        if (setPersistentState) Persistent.states.ai.model = modelId;
        if (!Souveraine.selectAgent(modelId)) {
            if (feedback) root.addMessage(Translation.tr("Finish or cancel the current turn before switching agents."), root.interfaceRole);
            return;
        }
        root.currentModel = models[modelId];
        if (feedback) root.addMessage(Translation.tr("Agent set to %1.").arg(models[modelId].name), root.interfaceRole);
    }

    function resumeConversation() {
        if (!Souveraine.resumeLatestConversation()) {
            root.addMessage(Translation.tr("Finish or cancel the current turn before resuming."), root.interfaceRole);
        }
    }

    // ── Souveraine-owned settings: advice instead of local state ────────
    function setTool(tool) {
        root.addMessage(Translation.tr("Tools are Souveraine sensors, configured per-agent — not switchable from the sidebar."), root.interfaceRole);
        return false;
    }

    function getTemperature() { return root.temperature; }

    function setTemperature(value) {
        root.addMessage(Translation.tr("Temperature is set in the agent's llm_config in Souveraine."), root.interfaceRole);
    }

    function printTemperature() {
        root.addMessage(Translation.tr("Temperature is owned by the agent's llm_config in Souveraine."), root.interfaceRole);
    }

    function setApiKey(key) {
        root.addMessage(Translation.tr("Keys live in souveraine.toml — set them with:\n```bash\nsouveraine auth set\n```"), root.interfaceRole);
    }

    function printApiKey() {
        root.addMessage(Translation.tr("Keys are owned by Souveraine (souveraine.toml / `souveraine auth set`), never exposed here."), root.interfaceRole);
    }

    function printPrompt() {
        root.addMessage(Translation.tr("The system prompt is composed by Souveraine (constitution + memory blocks + sensorium). Inspect it with `souveraine agent show`."), root.interfaceRole);
    }

    function loadPrompt(filePath) {
        root.addMessage(Translation.tr("Prompts are owned by the agent's memory in Souveraine — edit memfs instead of loading prompt files."), root.interfaceRole);
    }

    function attachFile(filePath) {
        root.addMessage(Translation.tr("File attachments aren't wired to Souveraine yet."), root.interfaceRole);
    }
    function removePendingFile(file) {
        root.pendingFiles = root.pendingFiles.filter(f => f !== file);
    }

    function regenerate(messageIndex) {
        root.addMessage(Translation.tr("Regenerate isn't supported — Souveraine conversations are forward-only."), root.interfaceRole);
    }

    // Souveraine executes its own sensors server-side; nothing to approve.
    function rejectCommand(message) {}
    function approveCommand(message) {}

    function createFunctionOutputMessage(name, output, includeOutputInChat = true) {
        return aiMessageComponent.createObject(root, {
            "role": "user",
            "content": `[[ Output of ${name} ]]${includeOutputInChat ? ("\n\n<think>\n" + output + "\n</think>") : ""}`,
            "rawContent": `[[ Output of ${name} ]]${includeOutputInChat ? ("\n\n<think>\n" + output + "\n</think>") : ""}`,
            "functionName": name,
            "functionResponse": output,
            "thinking": false,
            "done": true,
        });
    }

    // ── Local chat snapshots (ii plumbing, unchanged) ────────────────────
    Process {
        id: getSavedChats
        running: true
        command: ["ls", "-1", Directories.aiChats]
        stdout: StdioCollector {
            onStreamFinished: {
                if (text.length === 0) return;
                root.savedChats = text.split("\n")
                    .filter(fileName => fileName.endsWith(".json"))
                    .map(fileName => `${Directories.aiChats}/${fileName}`)
            }
        }
    }

    function chatToJson() {
        return root.messageIDs.map(id => {
            const message = root.messageByID[id]
            return ({
                "role": message.role,
                "rawContent": message.rawContent,
                "model": message.model,
                "thinking": false,
                "done": true,
            })
        })
    }

    FileView {
        id: chatSaveFile
        property string chatName: ""
        path: chatName.length > 0 ? `${Directories.aiChats}/${chatName}.json` : ""
        blockLoading: true
    }

    FileView {
        id: chatWriter
    }

    function saveChat(chatName) {
        const filePath = `${Directories.aiChats}/${chatName.trim()}.json`
        chatWriter.path = filePath
        chatWriter.setText(JSON.stringify(root.chatToJson()))
        getSavedChats.running = true;
    }

    function loadChat(chatName) {
        try {
            chatSaveFile.chatName = chatName.trim()
            chatSaveFile.reload()
            const saveData = JSON.parse(chatSaveFile.text())
            root.clearMessages()
            root.messageIDs = saveData.map((_, i) => i)
            for (let i = 0; i < saveData.length; i++) {
                const message = saveData[i];
                root.messageByID[i] = root.aiMessageComponent.createObject(root, {
                    "role": message.role,
                    "rawContent": message.rawContent,
                    "content": message.rawContent,
                    "model": message.model,
                    "thinking": false,
                    "done": true,
                });
            }
            root.addMessage(Translation.tr("Loaded a local snapshot. Note: this restores the transcript view only — the live Souveraine conversation starts fresh on the next message."), root.interfaceRole);
        } catch (e) {
            console.log("[Souveraine] Could not load chat: ", e);
        } finally {
            getSavedChats.running = true;
        }
    }
}