path: root/src/extension_loader.zig

//! Extension discovery: walk well-known directories, locate Lua extensions,
//! and register their tools with a `ToolRegistry`.
//!
//! Search order (later entries shadow earlier ones by extension *name*):
//!   1. `${XDG_CONFIG_HOME:-$HOME/.config}/panto/extensions/`  ("user")
//!   2. `./.panto/extensions/`                                  ("project")
//!
//! Layout per directory: each entry is either
//!   - `<name>.lua`        -- single-file extension; the extension name is
//!                            the basename without the `.lua` suffix.
//!   - `<name>/init.lua`   -- directory extension; the extension name is the
//!                            directory name. The directory is added to the
//!                            extension's `package.path` so it can `require`
//!                            sibling Lua files.
//!
//! Conflict rules:
//!   - Within a single directory, two entries with the same extension name
//!     are an error.
//!   - Project shadows user by extension name (debug-logged, not an error).
//!   - Tool-name collisions *between* loaded extensions are an error: a tool
//!     name is a contract the LLM relies on, and surprising overrides at
//!     load time are worse than failing fast.
//!
//! Symlinks: followed normally (no special handling).
//! Hidden files: dotfiles (`.foo.lua`) are skipped to leave room for editor
//! swap files and the like.

const std = @import("std");
const panto = @import("panto");
const lua_tool = @import("lua_tool.zig");
const lua_bridge = @import("lua_bridge.zig");

const c = lua_bridge.c;
const Allocator = std.mem.Allocator;
const Io = std.Io;

/// A discovered extension before loading. Owns its strings.
const Found = struct {
    /// Logical name (basename without `.lua`, or directory name).
    name: []u8,
    /// Absolute path to the Lua script to execute.
    script_path: []u8,
    /// For directory-style extensions, the directory containing the script
    /// (used to extend `package.path`). Null for single-file extensions.
    package_root: ?[]u8,
    /// Which search-path source this came from. Informational, used for
    /// shadowing log messages and tool-conflict error context.
    source: Source,

    pub fn deinit(self: *Found, allocator: Allocator) void {
        allocator.free(self.name);
        allocator.free(self.script_path);
        if (self.package_root) |p| allocator.free(p);
    }
};

pub const Source = enum {
    user,
    project,

    pub fn label(self: Source) []const u8 {
        return switch (self) {
            .user => "user",
            .project => "project",
        };
    }
};

/// Discover and register every extension found in the standard paths
/// derived from the environment + current working directory. Returns the
/// number of *tools* (not extensions) registered.
///
/// `environ_map` is consulted for `HOME` and `XDG_CONFIG_HOME`. The
/// project directory is always taken as `cwd()/.panto/extensions`.
pub fn discoverAndLoad(
    allocator: Allocator,
    io: Io,
    environ_map: *const std.process.Environ.Map,
    registry: *panto.ToolRegistry,
) !usize {
    const user_dir = try userExtensionsDir(allocator, environ_map);
    defer if (user_dir) |d| allocator.free(d);

    const project_dir = try projectExtensionsDir(allocator, io);
    defer allocator.free(project_dir);

    return loadFromDirs(allocator, io, registry, user_dir, project_dir);
}

/// Lower-level entry point: load from explicit user/project paths. Useful
/// for tests that want deterministic behavior without touching
/// process-global state (HOME, cwd). Either path may be null; missing
/// directories on disk are silently skipped.
pub fn loadFromDirs(
    allocator: Allocator,
    io: Io,
    registry: *panto.ToolRegistry,
    user_dir: ?[]const u8,
    project_dir: ?[]const u8,
) !usize {
    // 1. Collect candidates from both sources, project last so it wins.
    var found: std.array_list.Managed(Found) = .init(allocator);
    defer {
        for (found.items) |*f| f.deinit(allocator);
        found.deinit();
    }

    if (user_dir) |d| try scanDir(allocator, io, d, .user, &found);
    if (project_dir) |d| try scanDir(allocator, io, d, .project, &found);

    // 2. Apply project-shadows-user by name. The latest occurrence wins.
    try applyShadowing(allocator, &found);

    // 3. Load each surviving extension. Tool-name conflicts surface as
    //    `ToolRegistry.register` errors and abort startup.
    var total_tools: usize = 0;
    for (found.items) |f| {
        const n = loadOne(allocator, registry, f) catch |err| {
            // In test builds, log at warn level so a deliberate failure
            // test doesn't trip the test runner's err-count check.
            if (@import("builtin").is_test) {
                std.log.warn(
                    "extension '{s}' ({s}: {s}) failed to load: {t}",
                    .{ f.name, f.source.label(), f.script_path, err },
                );
            } else {
                std.log.err(
                    "extension '{s}' ({s}: {s}) failed to load: {t}",
                    .{ f.name, f.source.label(), f.script_path, err },
                );
            }
            return err;
        };
        std.log.debug(
            "extension: loaded {d} tool(s) from '{s}' ({s})",
            .{ n, f.name, f.source.label() },
        );
        total_tools += n;
    }
    return total_tools;
}

// ---------------------------------------------------------------------------
// Path resolution
// ---------------------------------------------------------------------------

/// Returns the absolute path of the user extensions directory, or null if
/// `HOME` is unset and `XDG_CONFIG_HOME` is not provided either. Caller
/// owns the returned slice.
fn userExtensionsDir(
    allocator: Allocator,
    environ_map: *const std.process.Environ.Map,
) !?[]u8 {
    if (environ_map.get("XDG_CONFIG_HOME")) |xdg| {
        return try std.fs.path.join(allocator, &.{ xdg, "panto", "extensions" });
    }
    if (environ_map.get("HOME")) |home| {
        return try std.fs.path.join(allocator, &.{ home, ".config", "panto", "extensions" });
    }
    return null;
}

/// Returns the absolute path of the project extensions directory.
/// Caller owns the returned slice.
fn projectExtensionsDir(allocator: Allocator, io: Io) ![]u8 {
    const cwd = try std.process.currentPathAlloc(io, allocator);
    defer allocator.free(cwd);
    return try std.fs.path.join(allocator, &.{ cwd, ".panto", "extensions" });
}

// ---------------------------------------------------------------------------
// Directory scanning
// ---------------------------------------------------------------------------

/// Scan `dir_path` and append any candidate extensions to `out`. Missing
/// directories are not an error: extensions are an optional feature.
fn scanDir(
    allocator: Allocator,
    io: Io,
    dir_path: []const u8,
    source: Source,
    out: *std.array_list.Managed(Found),
) !void {
    var dir = Io.Dir.cwd().openDir(io, dir_path, .{ .iterate = true }) catch |err| switch (err) {
        error.FileNotFound, error.NotDir => return,
        else => |e| return e,
    };
    defer dir.close(io);

    var local_names: std.StringHashMap(void) = .init(allocator);
    defer {
        var it = local_names.keyIterator();
        while (it.next()) |k| allocator.free(k.*);
        local_names.deinit();
    }

    var iter = dir.iterate();
    while (try iter.next(io)) |entry| {
        // Skip dotfiles (editor swap files, .DS_Store, hidden dirs, ...).
        if (entry.name.len == 0 or entry.name[0] == '.') continue;

        const maybe_found: ?Found = switch (entry.kind) {
            .file, .sym_link => try classifyFile(allocator, dir_path, entry.name, source),
            .directory => try classifyDirectory(allocator, io, dir, dir_path, entry.name, source),
            else => null,
        };
        const f = maybe_found orelse continue;

        // Within one directory, duplicate names are an error. (Can happen
        // if both `foo.lua` and `foo/init.lua` exist.)
        const gop = try local_names.getOrPut(f.name);
        if (gop.found_existing) {
            if (@import("builtin").is_test) {
                std.log.warn(
                    "extension name '{s}' is provided by multiple entries in {s}",
                    .{ f.name, dir_path },
                );
            } else {
                std.log.err(
                    "extension name '{s}' is provided by multiple entries in {s}",
                    .{ f.name, dir_path },
                );
            }
            // Free the duplicate's resources before bailing.
            var dup = f;
            dup.deinit(allocator);
            return error.DuplicateExtensionInDirectory;
        }
        // Hash map key is borrowed from f.name; we need an independent copy
        // since `out` owns f and we'd otherwise double-free.
        gop.key_ptr.* = try allocator.dupe(u8, f.name);

        try out.append(f);
    }
}

fn classifyFile(
    allocator: Allocator,
    dir_path: []const u8,
    entry_name: []const u8,
    source: Source,
) !?Found {
    if (!std.mem.endsWith(u8, entry_name, ".lua")) return null;
    const base = entry_name[0 .. entry_name.len - ".lua".len];
    if (base.len == 0) return null;

    const script_path = try std.fs.path.join(allocator, &.{ dir_path, entry_name });
    errdefer allocator.free(script_path);
    const name = try allocator.dupe(u8, base);
    errdefer allocator.free(name);

    return Found{
        .name = name,
        .script_path = script_path,
        .package_root = null,
        .source = source,
    };
}

/// Treat `entry_name/init.lua` as a directory-style extension if present.
/// Returns null if no `init.lua` exists inside.
fn classifyDirectory(
    allocator: Allocator,
    io: Io,
    parent: Io.Dir,
    dir_path: []const u8,
    entry_name: []const u8,
    source: Source,
) !?Found {
    // Probe for init.lua before doing any allocation.
    var sub = parent.openDir(io, entry_name, .{}) catch return null;
    defer sub.close(io);

    sub.access(io, "init.lua", .{}) catch |err| switch (err) {
        error.FileNotFound => return null,
        else => return null, // permission, etc. — quietly skip
    };

    const package_root = try std.fs.path.join(allocator, &.{ dir_path, entry_name });
    errdefer allocator.free(package_root);
    const script_path = try std.fs.path.join(allocator, &.{ package_root, "init.lua" });
    errdefer allocator.free(script_path);
    const name = try allocator.dupe(u8, entry_name);
    errdefer allocator.free(name);

    return Found{
        .name = name,
        .script_path = script_path,
        .package_root = package_root,
        .source = source,
    };
}

// ---------------------------------------------------------------------------
// Shadowing
// ---------------------------------------------------------------------------

/// In-place: for every name that appears more than once, keep only the
/// *last* occurrence and drop earlier ones with a debug log.
///
/// We do this in two passes. The first pass populates a `name → winning
/// index` map (last-write-wins by construction). The second pass walks
/// the list once, partitioning into kept and dropped sets *without*
/// freeing strings yet — the map still references them via its keys, so
/// freeing mid-walk would create dangling keys in the hash table.
fn applyShadowing(allocator: Allocator, list: *std.array_list.Managed(Found)) !void {
    var latest: std.StringHashMap(usize) = .init(allocator);

    for (list.items, 0..) |f, i| {
        try latest.put(f.name, i);
    }

    var keep: std.array_list.Managed(Found) = .init(allocator);
    var drop: std.array_list.Managed(Found) = .init(allocator);
    errdefer {
        latest.deinit();
        for (keep.items) |*f| f.deinit(allocator);
        keep.deinit();
        for (drop.items) |*f| f.deinit(allocator);
        drop.deinit();
    }
    try keep.ensureTotalCapacity(list.items.len);
    try drop.ensureTotalCapacity(list.items.len);

    for (list.items, 0..) |f, i| {
        const winner = latest.get(f.name).?;
        if (winner == i) {
            keep.appendAssumeCapacity(f);
        } else {
            std.log.debug(
                "extension: '{s}' from {s} shadowed by {s}",
                .{ f.name, f.source.label(), list.items[winner].source.label() },
            );
            drop.appendAssumeCapacity(f);
        }
    }

    // The map's keys still alias `drop`'s entries; deinit it first so we
    // can then free those entries' strings without leaving dangling keys.
    latest.deinit();
    for (drop.items) |*f| f.deinit(allocator);
    drop.deinit();

    // Replace list contents without re-freeing the entries we kept.
    list.clearRetainingCapacity();
    list.appendSlice(keep.items) catch unreachable;
    keep.deinit();
}

// ---------------------------------------------------------------------------
// Loading
// ---------------------------------------------------------------------------

/// Load one discovered extension and register its tools. Returns the
/// number of tools registered.
fn loadOne(
    allocator: Allocator,
    registry: *panto.ToolRegistry,
    found: Found,
) !usize {
    if (found.package_root) |root| {
        return lua_tool.loadExtensionWithPackageRoot(
            allocator,
            registry,
            found.script_path,
            root,
        );
    }
    return lua_tool.loadExtension(allocator, registry, found.script_path);
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

const testing = std.testing;

/// Helper: write a single file inside `dir` at `sub_path`.
fn writeFile(dir: Io.Dir, sub_path: []const u8, content: []const u8) !void {
    try dir.writeFile(testing.io, .{ .sub_path = sub_path, .data = content });
}

/// Helper: create a directory (path may contain separators).
fn makeDir(dir: Io.Dir, sub_path: []const u8) !void {
    try dir.createDirPath(testing.io, sub_path);
}

test "scanDir picks up single-file and directory-style extensions" {
    var tmp = testing.tmpDir(.{ .iterate = true });
    defer tmp.cleanup();

    // Layout:
    //   ext_root/
    //     alpha.lua
    //     beta/
    //       init.lua
    //       helper.lua
    //     .ignored.lua        (dotfile -> skipped)
    //     readme.txt          (non-.lua -> skipped)
    try makeDir(tmp.dir, "ext_root");
    try makeDir(tmp.dir, "ext_root/beta");
    try writeFile(tmp.dir, "ext_root/alpha.lua", "-- alpha\n");
    try writeFile(tmp.dir, "ext_root/beta/init.lua", "-- beta init\n");
    try writeFile(tmp.dir, "ext_root/beta/helper.lua", "-- helper\n");
    try writeFile(tmp.dir, "ext_root/.ignored.lua", "-- hidden\n");
    try writeFile(tmp.dir, "ext_root/readme.txt", "noise\n");

    // Absolute path to ext_root.
    var path_buf: [std.fs.max_path_bytes]u8 = undefined;
    const ext_root_len = try tmp.dir.realPathFile(testing.io, "ext_root", &path_buf);
    const ext_root = path_buf[0..ext_root_len];

    var list: std.array_list.Managed(Found) = .init(testing.allocator);
    defer {
        for (list.items) |*f| f.deinit(testing.allocator);
        list.deinit();
    }
    try scanDir(testing.allocator, testing.io, ext_root, .user, &list);

    try testing.expectEqual(@as(usize, 2), list.items.len);

    // Order is filesystem-dependent; sort by name for stable assertions.
    std.mem.sort(Found, list.items, {}, struct {
        fn lt(_: void, a: Found, b: Found) bool {
            return std.mem.lessThan(u8, a.name, b.name);
        }
    }.lt);

    try testing.expectEqualStrings("alpha", list.items[0].name);
    try testing.expect(list.items[0].package_root == null);
    try testing.expect(std.mem.endsWith(u8, list.items[0].script_path, "alpha.lua"));

    try testing.expectEqualStrings("beta", list.items[1].name);
    try testing.expect(list.items[1].package_root != null);
    try testing.expect(std.mem.endsWith(u8, list.items[1].script_path, "init.lua"));
}

test "duplicate name in same directory is an error" {
    var tmp = testing.tmpDir(.{ .iterate = true });
    defer tmp.cleanup();

    // Both `foo.lua` and `foo/init.lua` exist.
    try makeDir(tmp.dir, "ext_root/foo");
    try writeFile(tmp.dir, "ext_root/foo.lua", "-- single\n");
    try writeFile(tmp.dir, "ext_root/foo/init.lua", "-- dir\n");

    var path_buf: [std.fs.max_path_bytes]u8 = undefined;
    const n = try tmp.dir.realPathFile(testing.io, "ext_root", &path_buf);
    const ext_root = path_buf[0..n];

    var list: std.array_list.Managed(Found) = .init(testing.allocator);
    defer {
        for (list.items) |*f| f.deinit(testing.allocator);
        list.deinit();
    }

    const result = scanDir(testing.allocator, testing.io, ext_root, .project, &list);
    try testing.expectError(error.DuplicateExtensionInDirectory, result);
}

test "applyShadowing keeps the latest occurrence" {
    var list: std.array_list.Managed(Found) = .init(testing.allocator);
    defer {
        for (list.items) |*f| f.deinit(testing.allocator);
        list.deinit();
    }

    inline for (.{
        .{ "shared", "/u/shared.lua", Source.user },
        .{ "only_user", "/u/only_user.lua", Source.user },
        .{ "shared", "/p/shared.lua", Source.project },
        .{ "only_project", "/p/only_project.lua", Source.project },
    }) |row| {
        try list.append(.{
            .name = try testing.allocator.dupe(u8, row[0]),
            .script_path = try testing.allocator.dupe(u8, row[1]),
            .package_root = null,
            .source = row[2],
        });
    }

    try applyShadowing(testing.allocator, &list);
    try testing.expectEqual(@as(usize, 3), list.items.len);

    // "shared" survives once, from the project source.
    var shared_count: usize = 0;
    var shared_source: ?Source = null;
    for (list.items) |f| {
        if (std.mem.eql(u8, f.name, "shared")) {
            shared_count += 1;
            shared_source = f.source;
        }
    }
    try testing.expectEqual(@as(usize, 1), shared_count);
    try testing.expectEqual(@as(?Source, .project), shared_source);
}

test "loadFromDirs: project shadows user end-to-end" {
    var tmp = testing.tmpDir(.{ .iterate = true });
    defer tmp.cleanup();

    try makeDir(tmp.dir, "user_ext");
    try makeDir(tmp.dir, "project_ext");

    try writeFile(tmp.dir, "user_ext/greet.lua",
        \\panto.register_tool {
        \\  name = "greet", description = "user version",
        \\  schema = { type = "object" },
        \\  handler = function(input) return "USER" end,
        \\}
    );
    try writeFile(tmp.dir, "project_ext/greet.lua",
        \\panto.register_tool {
        \\  name = "greet", description = "project version",
        \\  schema = { type = "object" },
        \\  handler = function(input) return "PROJECT" end,
        \\}
    );

    var path_buf: [std.fs.max_path_bytes]u8 = undefined;
    const user_len = try tmp.dir.realPathFile(testing.io, "user_ext", &path_buf);
    const user_path = try testing.allocator.dupe(u8, path_buf[0..user_len]);
    defer testing.allocator.free(user_path);

    const proj_len = try tmp.dir.realPathFile(testing.io, "project_ext", &path_buf);
    const proj_path = try testing.allocator.dupe(u8, path_buf[0..proj_len]);
    defer testing.allocator.free(proj_path);

    var registry = panto.ToolRegistry.init(testing.allocator);
    defer registry.deinit();

    const n_tools = try loadFromDirs(
        testing.allocator,
        testing.io,
        &registry,
        user_path,
        proj_path,
    );
    try testing.expectEqual(@as(usize, 1), n_tools);

    const tool = registry.lookup("greet") orelse return error.NotRegistered;
    try testing.expectEqualStrings("project version", tool.description);

    const out = try tool.vtable.invoke(tool.ctx, "{}", testing.allocator);
    defer testing.allocator.free(out);
    try testing.expectEqualStrings("PROJECT", out);
}

test "loadFromDirs: directory-style extension can require siblings" {
    var tmp = testing.tmpDir(.{ .iterate = true });
    defer tmp.cleanup();

    try makeDir(tmp.dir, "ext/composer");
    try writeFile(tmp.dir, "ext/composer/util.lua",
        \\local M = {}
        \\function M.shout(s) return s:upper() .. "!" end
        \\return M
    );
    try writeFile(tmp.dir, "ext/composer/init.lua",
        \\local util = require("util")
        \\panto.register_tool {
        \\  name = "shout", description = "uppercase + bang",
        \\  schema = { type = "object", properties = { text = { type = "string" } } },
        \\  handler = function(input) return util.shout(input.text) end,
        \\}
    );

    var path_buf: [std.fs.max_path_bytes]u8 = undefined;
    const ext_len = try tmp.dir.realPathFile(testing.io, "ext", &path_buf);
    const ext_path = try testing.allocator.dupe(u8, path_buf[0..ext_len]);
    defer testing.allocator.free(ext_path);

    var registry = panto.ToolRegistry.init(testing.allocator);
    defer registry.deinit();

    const n = try loadFromDirs(
        testing.allocator,
        testing.io,
        &registry,
        null,
        ext_path,
    );
    try testing.expectEqual(@as(usize, 1), n);

    const tool = registry.lookup("shout") orelse return error.NotRegistered;
    const out = try tool.vtable.invoke(tool.ctx, "{\"text\":\"hi\"}", testing.allocator);
    defer testing.allocator.free(out);
    try testing.expectEqualStrings("HI!", out);
}

test "loadFromDirs: tool-name collision between extensions errors" {
    var tmp = testing.tmpDir(.{ .iterate = true });
    defer tmp.cleanup();

    try makeDir(tmp.dir, "ext");
    try writeFile(tmp.dir, "ext/alpha.lua",
        \\panto.register_tool {
        \\  name = "clash", description = "a",
        \\  schema = { type = "object" },
        \\  handler = function(input) return "a" end,
        \\}
    );
    try writeFile(tmp.dir, "ext/beta.lua",
        \\panto.register_tool {
        \\  name = "clash", description = "b",
        \\  schema = { type = "object" },
        \\  handler = function(input) return "b" end,
        \\}
    );

    var path_buf: [std.fs.max_path_bytes]u8 = undefined;
    const n = try tmp.dir.realPathFile(testing.io, "ext", &path_buf);
    const ext_path = try testing.allocator.dupe(u8, path_buf[0..n]);
    defer testing.allocator.free(ext_path);

    var registry = panto.ToolRegistry.init(testing.allocator);
    defer registry.deinit();

    const result = loadFromDirs(
        testing.allocator,
        testing.io,
        &registry,
        null,
        ext_path,
    );
    try testing.expectError(error.DuplicateTool, result);
}