aish/repl.lua

-- repl.lua — readline loop, input dispatch, prompt rendering.
-- Wires ffi/readline + router + executor + broker + context + renderer.
-- See docs/PHASE0.md §5 (dispatch), §9 (prompt + readline).

local rl       = require("ffi.readline")
local router   = require("router")
local executor = require("executor")
local broker   = require("broker")
local renderer = require("renderer")
local Context  = require("context")
local history  = require("history")
local mcp      = require("mcp")
local safety   = require("safety")
local json     = require("dkjson")

local M = {}

local HELP = [[
Meta commands:
  :quit / :q              exit aish (session flushed and closed)
  :clear                  clear screen (history kept)
  :reset                  clear in-memory conversation history
  :model <name>           switch active model
  :models                 list configured models (* = active)
  :history                show conversation turns
  :exec <cmd>             force shell execution
  :ask <text>             force AI query
  :sessions               list session log files
  :save <name>            rename current session log to <name>.jsonl
  :resume <name>          load <name>.jsonl turns into the in-memory context
  :mcp list               show connected MCP servers
  :mcp tools              list tools across all sessions
  :mcp tool <alias.name>  show one tool's inputSchema
  :mcp connect <url> [a]  open an MCP session at runtime
  :mcp disconnect <alias> drop an MCP session
  :help                   this message
]]

function M.run(config)
    assert(config and config.models, "repl.run: config.models required")
    local active_name = config.default_model or next(config.models)
    local active_cfg  = config.models[active_name]
    if not active_cfg then
        error("aish: default_model '" .. tostring(active_name)
              .. "' not found in config.models")
    end

    local ctx = Context.new(config.context or {})

    -- Phase 2: MCP sessions. Populated from config.mcp.servers at startup
    -- (best-effort — failures are status-logged once, session absent from
    -- mcp_sessions until manual :mcp connect; no auto-retry per PHASE2.md
    -- §4 Lifecycle). Tools cached per-session for the session lifetime
    -- (lmcp announces capabilities.tools.listChanged = false).
    local mcp_sessions = {}  -- { [alias] = session }
    local function connect_mcp(alias, server_cfg)
        local sess = mcp.connect(server_cfg.url, {
            alias      = alias,
            auth_token = server_cfg.auth_token,
            auth_env   = server_cfg.auth_env,
        })
        local ok, kind, err = sess:initialize()
        if not ok then
            renderer.status(("mcp %s: %s (%s)")
                :format(alias, tostring(err), kind))
            return false
        end
        mcp_sessions[alias] = sess
        if sess.version_warning then
            renderer.status("mcp " .. alias .. ": " .. sess.version_warning)
        end
        return true, #sess:list_tools()
    end

    if config.mcp and config.mcp.servers then
        for alias, server_cfg in pairs(config.mcp.servers) do
            local ok, n = connect_mcp(alias, server_cfg)
            if ok then
                renderer.status(("mcp %s: %d tools"):format(alias, n))
            end
        end
    end

    -- Assemble OpenAI-shape `tools` array across all live sessions, with
    -- alias.name namespacing per PHASE2.md §5. Empty array → broker omits
    -- the field entirely (§12 risk row 1).
    local function tools_schema()
        local out = {}
        for alias, sess in pairs(mcp_sessions) do
            for _, t in ipairs(sess:list_tools()) do
                out[#out + 1] = {
                    type = "function",
                    ["function"] = {
                        name        = alias .. "." .. t.name,
                        description = t.description or "",
                        parameters  = t.inputSchema
                                        or { type = "object", properties = {} },
                    },
                }
            end
        end
        return out
    end

    -- §4 "Content flattening": tool results may carry multiple blocks; v1
    -- concatenates text and ignores non-text with a one-shot status.
    local non_text_warned = false
    local function flatten_content(content)
        local parts = {}
        local saw_non_text = false
        for _, b in ipairs(content or {}) do
            if b.type == "text" then
                parts[#parts + 1] = b.text or ""
            else
                saw_non_text = true
            end
        end
        if saw_non_text and not non_text_warned then
            non_text_warned = true
            renderer.status("tool returned non-text content blocks "
                            .. "(image/resource ignored in v1)")
        end
        return table.concat(parts, "\n")
    end

    -- Split <alias>.<tool>, look up session, call. Returns (content_string,
    -- is_error). Errors of all flavors (rpc, transport, missing alias)
    -- yield a synthesized "[aish] tool ... failed: ..." string so the
    -- caller always has a body for the role:"tool" turn — the strict-
    -- template alternation rationale per PHASE0.md §6 and the C5/C7 fold
    -- in PHASE2.md §4.
    local function dispatch_tool_call(name, args)
        local alias, tool_name = name:match("^([^.]+)%.(.+)$")
        if not alias then
            return ("[aish] tool name has no alias prefix: %s"):format(name), true
        end
        local sess = mcp_sessions[alias]
        if not sess then
            return ("[aish] no MCP server connected for alias '%s'")
                    :format(alias), true
        end
        local result, kind, err = sess:call_tool(tool_name, args)
        if not result then
            if kind == "rpc_error" then
                local msg = (type(err) == "table" and err.message)
                              or tostring(err)
                return ("[aish] tool dispatch failed: %s"):format(msg), true
            else
                return ("[aish] tool transport error: %s")
                        :format(tostring(err)), true
            end
        end
        -- result has content[] and possibly isError=true. flatten_content
        -- handles the text-blocks-only flattening. We pass through the
        -- content body regardless of isError (per PHASE2-baseline.md §3:
        -- some tools set isError=false on actual failures, content text
        -- is authoritative).
        return flatten_content(result.content), (kind == "handler_error")
    end

    -- Session log (PHASE1.md §6). Always open one on startup; auto-write
    -- every user/assistant turn; close on :quit. If history.dir is set but
    -- unwritable, log a status and continue without persistence.
    local history_dir   = (config.history and config.history.dir) or nil
    local sessions_dir  = history_dir and (history_dir .. "/sessions") or nil
    local session_path  = sessions_dir
        and (sessions_dir .. "/" .. os.date("!%Y-%m-%dT%H-%M-%SZ") .. ".jsonl")
    local session
    if session_path then
        local sess, serr = history.open(session_path, {
            started      = os.date("!%Y-%m-%dT%H:%M:%SZ"),
            model        = active_name,
            aish_version = "phase1",
        })
        if sess then
            session = sess
        else
            renderer.status("session log disabled: " .. tostring(serr))
        end
    end
    local function log_turn(turn)
        if session then session:append(turn) end
    end

    local function prompt()
        return ("[aish:%s]> "):format(active_name)
    end

    -- Phase 1 reserved-key wiring (PHASE1.md §7). The mechanism is real; the
    -- handlers are placeholders that emit a status. Phase 3 (Norris) is the
    -- first consumer that replaces the body with real work.
    rl.bind("\\C-n", function()
        renderer.status("Norris mode not yet implemented (Phase 3)")
    end)

    local function status_evictions(n)
        if n and n > 0 then
            renderer.status(("oldest %d turns evicted"):format(n))
        end
    end

    -- Run a shell command, framing output and (per config.shell.capture_output)
    -- buffering it for the NEXT user turn — context.append_exec_output keeps
    -- a [exec output] block pending until ask_ai flushes it via append_user.
    -- Direct user-role injection violated chat-template alternation (mistral-
    -- nemo's Jinja rejects user/user back-to-back); see PHASE0.md §6.
    local function run_shell(cmd)
        local chd, err = executor.maybe_chdir(cmd)
        if chd ~= nil then
            if chd then
                local pwd = io.popen("pwd"):read("*l") or "?"
                renderer.status("cwd -> " .. pwd)
            else
                renderer.status("cd: " .. tostring(err))
            end
            return
        end
        renderer.exec_begin()
        local out, code = executor.exec(cmd)
        renderer.exec_end(code)
        if config.shell and config.shell.capture_output then
            ctx:append_exec_output(out)
        end
    end

    -- Send user text to the active model and process the response. If MCP
    -- tools are connected and the model emits tool_calls, dispatch each
    -- call (with safety confirm gate), append role:"tool" turns, and
    -- re-call the broker — looping until the model returns pure text or
    -- max_tool_depth is hit. CMD: extraction runs ONCE on the final
    -- pure-text response (the §6 substrate invariant is unchanged).
    local max_tool_depth = (config.mcp and config.mcp.max_tool_depth) or 8

    local function ask_ai(text)
        local prev_pending = ctx.pending_exec_output
        ctx:append_user(text)
        log_turn(ctx.turns[#ctx.turns])

        local depth = 0
        local final_resp = ""
        local first_iteration = true

        while true do
            local text_parts      = {}
            local tool_calls_seen = {}
            local ok, err = broker.chat_stream(active_cfg, ctx:to_messages(),
                function(kind, payload)
                    if kind == "text" then
                        text_parts[#text_parts + 1] = payload
                        renderer.assistant_delta(payload)
                    elseif kind == "tool_call" then
                        tool_calls_seen[#tool_calls_seen + 1] = payload
                    end
                end,
                { tools = tools_schema() })
            renderer.assistant_flush()

            if not ok then
                renderer.status("broker error: " .. tostring(err))
                if first_iteration then
                    -- Back out the user turn so :resume / retry is clean.
                    table.remove(ctx.turns)
                    ctx.pending_exec_output = prev_pending
                end
                return
            end
            first_iteration = false

            local resp_text = table.concat(text_parts)

            if #tool_calls_seen == 0 then
                -- Pure text response — end of this AI turn.
                ctx:append({ role = "assistant", content = resp_text })
                log_turn(ctx.turns[#ctx.turns])
                final_resp = resp_text
                break
            end

            -- Record the assistant turn with text AND tool_calls. Content
            -- may be "" (C3: model often emits no prose before a call).
            ctx:append({
                role       = "assistant",
                content    = resp_text,
                tool_calls = tool_calls_seen,
            })
            log_turn(ctx.turns[#ctx.turns])

            -- Process each tool_call. Every iteration appends EXACTLY one
            -- role:"tool" turn per call (keeps alternation legal even on
            -- decline/error per C5/C7).
            for _, call in ipairs(tool_calls_seen) do
                local args_table, args_err
                if call.arguments and call.arguments ~= "" then
                    args_table, _, args_err = json.decode(call.arguments)
                else
                    args_table = {}
                end

                local tool_content, is_error
                if args_err then
                    tool_content = ("[aish] tool arguments not parseable as "
                                    .. "JSON: %s"):format(tostring(args_err))
                    is_error = true
                    renderer.tool_call_begin(call.name, call.arguments)
                    renderer.tool_call_end(tool_content, true)
                elseif not safety.confirm_tool_call(call.name, args_table,
                                                   config) then
                    tool_content = "[aish] tool call declined by user"
                    is_error = true
                    renderer.status(tool_content)
                else
                    renderer.tool_call_begin(call.name, call.arguments)
                    local content, errflag = dispatch_tool_call(call.name,
                                                                args_table)
                    tool_content = content
                    is_error = errflag
                    renderer.tool_call_end(content, errflag)
                end

                ctx:append({
                    role         = "tool",
                    tool_call_id = call.id,
                    content      = tool_content,
                })
                log_turn(ctx.turns[#ctx.turns])
            end

            depth = depth + 1
            if depth >= max_tool_depth then
                renderer.status(("tool-call depth limit reached (%d); "
                                 .. "stopping sub-loop"):format(max_tool_depth))
                final_resp = resp_text
                break
            end
            -- loop body re-runs broker.chat_stream with the now-extended ctx
        end

        status_evictions(ctx:enforce_budget())

        -- CMD: extraction on the final pure-text response only.
        for _, cmd in ipairs(executor.extract_cmd_lines(final_resp)) do
            local doit
            if config.shell and config.shell.confirm_cmd then
                local ans = rl.readline(("execute '%s'? [y/N] "):format(cmd)) or ""
                doit = (ans:lower():sub(1, 1) == "y")
            else
                doit = true
            end
            if doit then run_shell(cmd) end
        end
    end

    local function shutdown_session()
        if session then session:close(); session = nil end
    end

    -- Meta dispatch table.
    local meta = {
        quit  = function() shutdown_session(); os.exit(0) end,
        q     = function() shutdown_session(); os.exit(0) end,
        clear = function() io.write("\27[H\27[2J"); io.flush() end,
        reset = function()
            ctx:reset(); renderer.status("context reset")
        end,
        model = function(args)
            local name = args:match("^%s*(%S+)")
            if not name or not config.models[name] then
                renderer.status("usage: :model <name>; not found: " .. tostring(name))
                return
            end
            active_name, active_cfg = name, config.models[name]
            renderer.status("model -> " .. name)
        end,
        models = function()
            renderer.status(("models (active: %s):"):format(active_name))
            for name, cfg in pairs(config.models) do
                local mark = (name == active_name) and "*" or " "
                io.write(("  %s %-8s  %s @ %s\n"):format(
                    mark, name, cfg.model or "?", cfg.endpoint or "?"))
            end
        end,
        history = function()
            if #ctx.turns == 0 then
                renderer.status("(empty)"); return
            end
            for i, t in ipairs(ctx.turns) do
                io.write(("[%d] %s: %s\n"):format(
                    i, t.role, t.content:gsub("\n", " ")))
            end
        end,
        exec = function(args)
            args = (args or ""):match("^%s*(.-)%s*$")
            if args == "" then renderer.status("usage: :exec <cmd>"); return end
            run_shell(args)
        end,
        ask  = function(args)
            args = (args or ""):match("^%s*(.-)%s*$")
            if args == "" then renderer.status("usage: :ask <text>"); return end
            ask_ai(args)
        end,
        sessions = function()
            if not sessions_dir then renderer.status("(no history.dir configured)"); return end
            local names = history.list_sessions(sessions_dir)
            if #names == 0 then renderer.status("(no sessions in " .. sessions_dir .. ")"); return end
            for _, n in ipairs(names) do
                local mark = (session_path and session_path:match("[^/]+$") == n)
                             and "*" or " "
                io.write(("  %s %s\n"):format(mark, n))
            end
        end,
        save = function(args)
            local name = args:match("^%s*(%S+)")
            if not name then renderer.status("usage: :save <name>"); return end
            if not (session and session_path and sessions_dir) then
                renderer.status("no active session to save")
                return
            end
            name = name:gsub("%.jsonl$", "")
            local new_path = sessions_dir .. "/" .. name .. ".jsonl"
            if new_path == session_path then
                renderer.status("already named " .. name)
                return
            end
            session:close()
            local ok, rerr = os.rename(session_path, new_path)
            if not ok then
                renderer.status("rename failed: " .. tostring(rerr))
                -- best-effort reopen of original path so logging continues
                session = history.open(session_path)
                return
            end
            session_path = new_path
            session = history.open(session_path)  -- reopen for continued append
            renderer.status("saved as " .. name .. ".jsonl")
        end,
        resume = function(args)
            local name = args:match("^%s*(%S+)")
            if not name then renderer.status("usage: :resume <name>"); return end
            if not sessions_dir then renderer.status("(no history.dir configured)"); return end
            -- Refuse to silently clobber an active conversation; the user has
            -- to :reset first to express intent. The current session log on
            -- disk is unaffected by either choice.
            if #ctx.turns > 0 then
                renderer.status(("resume into non-empty ctx refused (%d turns); :reset first")
                                :format(#ctx.turns))
                return
            end
            name = name:gsub("%.jsonl$", "")
            local path = sessions_dir .. "/" .. name .. ".jsonl"
            local turns, _meta_hdr = history.load(path)
            if not turns then
                renderer.status("resume failed: cannot load " .. path)
                return
            end
            ctx:reset()
            for _, t in ipairs(turns) do ctx:append(t) end
            renderer.status(("resumed %d turns from %s"):format(#turns, name))
        end,
        mcp = function(args)
            local sub, sub_args = args:match("^%s*(%S*)%s*(.*)$")
            if sub == "list" or sub == "" then
                if next(mcp_sessions) == nil then
                    renderer.status("(no MCP sessions)"); return
                end
                for alias, sess in pairs(mcp_sessions) do
                    io.write(("  %s  %s  (%d tools)\n"):format(
                        alias, sess.url, #sess:list_tools()))
                end
            elseif sub == "tools" then
                local any = false
                for alias, sess in pairs(mcp_sessions) do
                    for _, t in ipairs(sess:list_tools()) do
                        any = true
                        local desc = (t.description or ""):gsub("\n", " ")
                        io.write(("  %s.%-18s %s\n"):format(
                            alias, t.name, desc:sub(1, 60)))
                    end
                end
                if not any then renderer.status("(no tools)") end
            elseif sub == "tool" then
                local name = sub_args:match("^%s*(%S+)")
                if not name then
                    renderer.status("usage: :mcp tool <alias.name>"); return
                end
                local alias, tname = name:match("^([^.]+)%.(.+)$")
                if not alias then
                    renderer.status("tool name missing alias prefix: " .. name)
                    return
                end
                local sess = mcp_sessions[alias]
                if not sess then
                    renderer.status("unknown alias: " .. alias)
                    return
                end
                local found
                for _, t in ipairs(sess:list_tools()) do
                    if t.name == tname then found = t; break end
                end
                if not found then
                    renderer.status("unknown tool: " .. name); return
                end
                io.write(("  %s.%s\n"):format(alias, found.name))
                io.write(("  description: %s\n"):format(found.description or "(none)"))
                io.write("  inputSchema:\n    ")
                io.write((json.encode(found.inputSchema or {}, {indent = true})
                          :gsub("\n", "\n    ")))
                io.write("\n")
            elseif sub == "connect" then
                local url, alias = sub_args:match("^%s*(%S+)%s*(%S*)")
                if not url or url == "" then
                    renderer.status("usage: :mcp connect <url> [alias]"); return
                end
                if alias == "" then
                    alias = url:match("https?://([^:/]+)") or url
                end
                if mcp_sessions[alias] then
                    renderer.status("already connected: " .. alias); return
                end
                local ok, n = connect_mcp(alias, { url = url })
                if ok then
                    renderer.status(("mcp %s: connected (%d tools)")
                                    :format(alias, n))
                end
            elseif sub == "disconnect" then
                local alias = sub_args:match("^%s*(%S+)")
                if not alias then
                    renderer.status("usage: :mcp disconnect <alias>"); return
                end
                local sess = mcp_sessions[alias]
                if not sess then
                    renderer.status("not connected: " .. alias); return
                end
                sess:close()
                mcp_sessions[alias] = nil
                renderer.status("disconnected " .. alias)
            else
                renderer.status("usage: :mcp {list|tools|tool|connect|disconnect}")
            end
        end,
        help = function() io.write(HELP) end,
    }

    -- Main loop.
    while true do
        local line = rl.readline(prompt())
        if line == nil then          -- EOF (Ctrl-D on empty line)
            io.write("\n")
            shutdown_session()
            break
        end
        if line:gsub("%s", "") == "" then
            -- empty / whitespace-only: skip silently
        else
            rl.add_history(line)
            local kind, payload = router.classify(line, config)
            if kind == "meta" then
                local name, rest = payload:match("^(%S+)%s*(.*)$")
                local handler = name and meta[name]
                if handler then
                    handler(rest or "")
                else
                    renderer.status("unknown meta command: :" .. tostring(name))
                end
            elseif kind == "shell" then
                run_shell(payload)
            else  -- "ai"
                ask_ai(payload)
            end
        end
    end
end

-- Phase 0 module export. Meta-command list shown above lives in HELP and
-- is implemented inline in run().
return M