path: root/src/extension_loader.zig

//! Extension discovery: walk well-known directories, locate Lua extensions,
//! and load each one into a long-lived `LuaRuntime`.
//!
//! 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:
//!   - `<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 one 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.
//!
//! Symlinks: followed normally. Dotfiles: skipped.

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

const Allocator = std.mem.Allocator;
const Io = std.Io;
const LuaRuntime = lua_runtime.LuaRuntime;

/// 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.
    package_root: ?[]u8,
    /// Which search-path source this came from.
    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 load every extension found in the standard paths into
/// `runtime`. Returns the number of *tools* (not extensions) declared.
///
/// `environ_map` is consulted for `HOME` and `XDG_CONFIG_HOME`. The
/// project directory is always `cwd()/.panto/extensions`.
pub fn discoverAndLoad(
    allocator: Allocator,
    io: Io,
    environ_map: *const std.process.Environ.Map,
    runtime: *LuaRuntime,
) !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, runtime, user_dir, project_dir);
}

/// Lower-level entry point: load from explicit user/project paths.
/// Either path may be null; missing directories are silently skipped.
pub fn loadFromDirs(
    allocator: Allocator,
    io: Io,
    runtime: *LuaRuntime,
    user_dir: ?[]const u8,
    project_dir: ?[]const u8,
) !usize {
    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);

    try applyShadowing(allocator, &found);

    const before = runtime.toolCount();
    for (found.items) |f| {
        runtime.loadExtension(f.script_path, f.package_root) catch |err| {
            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 '{s}' ({s})",
            .{ f.name, f.source.label() },
        );
    }
    return runtime.toolCount() - before;
}

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

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;
}

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
// ---------------------------------------------------------------------------

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| {
        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;

        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 },
                );
            }
            var dup = f;
            dup.deinit(allocator);
            return error.DuplicateExtensionInDirectory;
        }
        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,
    };
}

fn classifyDirectory(
    allocator: Allocator,
    io: Io,
    parent: Io.Dir,
    dir_path: []const u8,
    entry_name: []const u8,
    source: Source,
) !?Found {
    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,
    };

    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
// ---------------------------------------------------------------------------

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);
        }
    }

    latest.deinit();
    for (drop.items) |*f| f.deinit(allocator);
    drop.deinit();

    list.clearRetainingCapacity();
    list.appendSlice(keep.items) catch unreachable;
    keep.deinit();
}

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

const testing = std.testing;

fn writeFile(dir: Io.Dir, sub_path: []const u8, content: []const u8) !void {
    try dir.writeFile(testing.io, .{ .sub_path = sub_path, .data = content });
}

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();

    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");

    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);

    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();

    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);

    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 (via long-lived runtime)" {
    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 rt = try LuaRuntime.create(testing.allocator);
    defer rt.deinit();

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

    // Invoke the tool through the source and verify the project handler ran.
    var src = rt.toolSource();
    const calls = [_]panto.ToolCall{.{ .tool_name = "greet", .input = "{}" }};
    var results: [1]panto.ToolCallResult = .{.{ .err = error.SourceDroppedCall }};
    try src.vtable.invoke_batch(src.ctx, &calls, &results, testing.allocator);
    defer testing.allocator.free(results[0].ok);
    try testing.expectEqualStrings("PROJECT", results[0].ok);
}

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 rt = try LuaRuntime.create(testing.allocator);
    defer rt.deinit();

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