odidere/internal/llm/tool.go

package llm

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"text/template"
	"time"

	"code.chimeric.al/chimerical/odidere/internal/config"
)

// fn provides template functions available to argument templates.
var fn = template.FuncMap{
	"json": func(v any) string {
		b, _ := json.Marshal(v)
		return string(b)
	},
}

// Tool represents an external tool that can be invoked by LLMs.
// Tools are executed as subprocesses with templated arguments,
// and described to the LLM via their JSON Schema parameters.
type Tool struct {
	// Name uniquely identifies the tool within a registry.
	Name string
	// Description explains the tool's purpose for the LLM.
	Description string
	// Command is the executable path or name.
	Command string
	// Arguments are Go templates expanded with LLM-provided parameters.
	Arguments []string
	// Parameters is a JSON Schema describing expected input from the LLM.
	Parameters map[string]any
	// timeout is the parsed execution duration limit.
	timeout time.Duration
}

// NewTool creates a Tool from a ToolConfig.
func NewTool(cfg config.ToolConfig) (*Tool, error) {
	var timeout time.Duration
	if cfg.Timeout != "" {
		var err error
		timeout, err = time.ParseDuration(cfg.Timeout)
		if err != nil {
			return nil, fmt.Errorf(
				"invalid timeout %q: %w", cfg.Timeout, err,
			)
		}
	}

	return &Tool{
		Name:        cfg.Name,
		Description: cfg.Description,
		Command:     cfg.Command,
		Arguments:   cfg.Arguments,
		Parameters:  cfg.Parameters,
		timeout:     timeout,
	}, nil
}

// ParseArguments expands argument templates with the provided JSON data.
// The args parameter should be a JSON object string; empty string or "{}"
// results in an empty data map. Templates producing empty strings are
// filtered from the result, allowing conditional arguments.
func (t *Tool) ParseArguments(args string) ([]string, error) {
	var data = map[string]any{}
	if args != "" && args != "{}" {
		if err := json.Unmarshal([]byte(args), &data); err != nil {
			return nil, fmt.Errorf(
				"invalid arguments JSON: %w", err,
			)
		}
	}

	var result []string
	for _, v := range t.Arguments {
		tmpl, err := template.New("").Funcs(fn).Parse(v)
		if err != nil {
			return nil, fmt.Errorf(
				"invalid template %q: %w", v, err,
			)
		}

		var buf bytes.Buffer
		if err := tmpl.Execute(&buf, data); err != nil {
			return nil, fmt.Errorf(
				"execute template %q: %w", v, err,
			)
		}

		// Filter out empty strings (unused conditional arguments).
		if s := buf.String(); s != "" {
			result = append(result, s)
		}
	}

	return result, nil
}

// ToAPI converts the tool definition into the API format suitable for
// inclusion in a ChatRequest.
func (t *Tool) ToAPI() APITool {
	return APITool{
		Type: "function",
		Function: &FunctionDef{
			Name:        t.Name,
			Description: t.Description,
			Parameters:  t.Parameters,
		},
	}
}

// Registry holds tools indexed by name and handles their execution.
// It validates all tool definitions at construction time to fail fast on
// configuration errors.
type Registry struct {
	tools map[string]*Tool
}

// NewRegistry creates a registry from the provided tool definitions.
// Returns an error if any tool fails validation or if duplicate names exist.
func NewRegistry(tools []config.ToolConfig) (*Registry, error) {
	var r = &Registry{
		tools: make(map[string]*Tool),
	}
	for _, tc := range tools {
		t, err := NewTool(tc)
		if err != nil {
			return nil, fmt.Errorf(
				"invalid tool %q: %w", tc.Name, err,
			)
		}

		if _, exists := r.tools[t.Name]; exists {
			return nil, fmt.Errorf(
				"duplicate tool name: %s", t.Name,
			)
		}
		r.tools[t.Name] = t
	}

	return r, nil
}

// Get returns a tool by name and a boolean indicating if it was found.
func (r *Registry) Get(name string) (*Tool, bool) {
	tool, ok := r.tools[name]
	return tool, ok
}

// List returns all registered tool names in arbitrary order.
func (r *Registry) List() []string {
	names := make([]string, 0, len(r.tools))
	for name := range r.tools {
		names = append(names, name)
	}
	return names
}

// ToAPI converts all registered tools into API format suitable for
// inclusion in a ChatRequest.
func (r *Registry) ToAPI() []APITool {
	if r == nil {
		return nil
	}
	tools := make([]APITool, 0, len(r.tools))
	for _, t := range r.tools {
		tools = append(tools, t.ToAPI())
	}
	return tools
}

// binaryOutputPrefix is the message shown when tool output is detected as binary.
const binaryOutputPrefix = "error: binary output detected: "

// ToolResult holds the structured output from a tool execution.
type ToolResult struct {
	Stdout       []byte
	Stderr       []byte
	ProcessState *os.ProcessState
}

// FormatResult formats the tool result using the structured label-based
// format with stdout, stderr, and exit code sections.
//
// http.DetectContentType implements the WHATWG MIME-sniffing spec and
// returns text/plain for any content consisting entirely of printable
// characters and whitespace. Text-based formats like JSON, JavaScript,
// and XML will be detected as text/plain; charset=utf-8. We assume
// checking for the "text/" prefix is sufficient.
func (r *ToolResult) FormatResult() string {
	var buf strings.Builder

	fmt.Fprintln(&buf, "▸ stdout:")
	if len(r.Stdout) > 0 {
		ctype := http.DetectContentType(r.Stdout)
		if !strings.HasPrefix(ctype, "text/") {
			fmt.Fprintf(&buf, binaryOutputPrefix+"%s\n", ctype)
		} else {
			buf.Write(r.Stdout)
			if r.Stdout[len(r.Stdout)-1] != '\n' {
				fmt.Fprint(&buf, "\n")
			}
		}
	}

	fmt.Fprintln(&buf, "▸ stderr:")
	if len(r.Stderr) > 0 {
		ctype := http.DetectContentType(r.Stderr)
		if !strings.HasPrefix(ctype, "text/") {
			fmt.Fprintf(&buf, binaryOutputPrefix+"%s\n", ctype)
		} else {
			buf.Write(r.Stderr)
			if r.Stderr[len(r.Stderr)-1] != '\n' {
				fmt.Fprint(&buf, "\n")
			}
		}
	}

	exitCode := r.ProcessState.ExitCode()
	fmt.Fprintf(&buf, "┗ exit: %d", exitCode)
	// Include signal info when the process was killed by a signal.
	if exitCode == -1 {
		fmt.Fprintf(&buf, " (%s)", r.ProcessState.String())
	}
	fmt.Fprintln(&buf)

	return buf.String()
}

// Execute runs a tool by name with the provided JSON arguments.
// It expands argument templates, executes the command as a subprocess, and
// returns a structured result with stdout, stderr, and exit code.
// The context can be used for cancellation; tool-specific timeouts are
// applied on top of any context deadline.
//
// Start failures (process never ran) return an error directly.
func (r *Registry) Execute(
	ctx context.Context, name string, args string,
) (*ToolResult, error) {
	tool, ok := r.tools[name]
	if !ok {
		return nil, fmt.Errorf("unknown tool: %q", name)
	}

	// Evaluate argument templates.
	cmdArgs, err := tool.ParseArguments(args)
	if err != nil {
		return nil, fmt.Errorf(
			"tool %q: parse arguments: %w", name, err,
		)
	}

	// If defined, use the timeout.
	if tool.timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, tool.timeout)
		defer cancel()
	}

	// Setup and run the command.
	var (
		stdout, stderr bytes.Buffer
		cmd            = exec.CommandContext(
			ctx, tool.Command, cmdArgs...,
		)
	)
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err = cmd.Run()

	// Extract exit code and handle errors.
	// After cmd.Run() returns, the possible error cases are:
	//
	// 1. Context error (ctx.Err() != nil): the context was cancelled or
	//    timed out. If the process had started, it was killed and
	//    ProcessState is set. If it had not, Process is nil.
	//
	// 2. Start failure (cmd.Process == nil): the executable was not found,
	//    the fork failed, or the context was already done before Start.
	//    No stdout/stderr to capture.
	//
	// 3. Process exited with error (err is *exec.ExitError): the process
	//    ran but exited with a non-zero status or was killed by a signal.
	//    ProcessState is set and contains the exit code.
	//
	// In all cases where cmd.Process != nil, ProcessState is guaranteed
	// to be set because Start() succeeded.
	if err != nil {
		switch {
		// Context was cancelled or timed out. The process was killed if
		// it had started. Distinguish tool timeout from other context
		// errors so we can return an appropriate message.
		case ctx.Err() != nil:
			if ctx.Err() == context.DeadlineExceeded &&
				tool.timeout > 0 {
				// Tool-specific timeout. Return a clear
				// message with the configured duration.
				return nil, fmt.Errorf(
					"tool %q: timed out after %v",
					name, tool.timeout,
				)
			}
			// Context cancelled or parent deadline exceeded.
			return nil, fmt.Errorf("tool %q: %w", name, ctx.Err())

		// Process never started. No ProcessState, no output.
		// The error is the raw start failure (e.g., *PathError).
		case cmd.Process == nil:
			return nil, fmt.Errorf("tool %q: %w", name, err)

		// Process ran and exited with non-zero status or was killed by a signal.
		// ProcessState is set; return it along with captured output.
		default:
			return &ToolResult{
				Stdout:       stdout.Bytes(),
				Stderr:       stderr.Bytes(),
				ProcessState: cmd.ProcessState,
			}, nil
		}
	}

	// No error. Process exited with status 0.
	return &ToolResult{
		Stdout:       stdout.Bytes(),
		Stderr:       stderr.Bytes(),
		ProcessState: cmd.ProcessState,
	}, nil
}