odidere/internal/tool/tool.go

// Package tool provides configured subprocess tools.
package tool

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

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

// errToolTimeout marks a timeout caused by the tool's configured limit.
var errToolTimeout = errors.New("tool timeout")

// fn provides template functions available to argument templates.
var fn = template.FuncMap{
	"json": func(v any) string {
		b, _ := jsonv2.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 contains the templates used to construct the subprocess
	// arguments. See config.ToolConfig.Arguments for examples, behavior, and
	// limitations.
	Arguments []string
	// tmpl holds Arguments compiled to executable templates, one per
	// element, in the same order. Only NewTool populates it; on a Tool
	// built any other way it is nil, and ParseArguments returns no
	// arguments rather than an error.
	tmpl []*template.Template
	// Parameters describes the tool's accepted input to the LLM. It is not
	// used to validate a tool call before its arguments are rendered. See
	// config.ToolConfig.Parameters.
	Parameters map[string]any
	// timeout is the parsed execution duration limit.
	timeout time.Duration
}

// NewTool creates a Tool from a ToolConfig. It compiles the argument
// templates so that a template syntax error is reported here, at
// construction, rather than on the tool's first execution.
func NewTool(cfg config.ToolConfig) (*Tool, error) {
	templates := make([]*template.Template, len(cfg.Arguments))
	for i, arg := range cfg.Arguments {
		tmpl, err := template.New("").Funcs(fn).Parse(arg)
		if err != nil {
			return nil, fmt.Errorf(
				"invalid template %q: %w", arg, err,
			)
		}
		templates[i] = tmpl
	}

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

// ParseArguments converts tool-call arguments into program arguments. The
// input is the JSON object supplied by the LLM when it calls the tool. The
// returned slice is passed to the external program after the tool's Command;
// it does not include the command name itself.
//
// An empty string or JSON null is treated as an empty object. ParseArguments
// returns an error for any other input that is not a single valid JSON
// object, or if an argument template cannot be evaluated. A template that
// fails to parse is an error of NewTool, reported when the tool is
// constructed.
//
// JSON numbers are decoded as float64, which is encoding/json/v2's default
// for values decoded into an interface. Integers outside the IEEE 754 safe
// range (±2^53) do not survive this round trip: they are rounded and
// rendered in floating-point form (9007199254740993 becomes
// 9.007199254740992e+15). Tool parameters whose values may exceed that range
// should be declared as strings in the tool's Parameters.
//
// See config.ToolConfig.Arguments for template behavior and limitations.
func (t *Tool) ParseArguments(args string) ([]string, error) {
	var data = map[string]any{}
	if args != "" && args != "{}" {
		if err := jsonv2.Unmarshal([]byte(args), &data); err != nil {
			return nil, fmt.Errorf(
				"invalid arguments JSON: %w", err,
			)
		}
	}

	var result []string
	for i, tmpl := range t.tmpl {
		var buf bytes.Buffer
		if err := tmpl.Execute(&buf, data); err != nil {
			return nil, fmt.Errorf(
				"execute template %q: %w", t.Arguments[i], 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() llm.APITool {
	return llm.APITool{
		Type: llm.ToolTypeFunction,
		Function: &llm.FunctionDef{
			Name:        t.Name,
			Description: t.Description,
			Parameters:  t.Parameters,
		},
	}
}

// Registry holds tools, indexed by name, and handles their execution.
type Registry struct {
	tools map[string]*Tool
}

// NewRegistry creates a registry from the provided tool definitions.
// It returns an error if a tool cannot be constructed or a name is duplicated.
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() []llm.APITool {
	if r == nil {
		return nil
	}
	tools := make([]llm.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 {
	// FinishedAt is the wall-clock time at which the command stopped
	// running, including the wait for its output pipes to drain after it
	// exits.
	FinishedAt time.Time
	// ProcessState describes how the process ended: its exit code, or the
	// signal that killed it.
	ProcessState *os.ProcessState
	// StartedAt is the wall-clock time at which the command began running.
	StartedAt time.Time
	// Stderr holds the bytes the command wrote to standard error.
	Stderr []byte
	// Stdout holds the bytes the command wrote to standard output.
	Stdout []byte
}

// FormatResult formats the tool result into separate sections for stdout and
// stderr, and other metadata.
//
// Output sections appear even when empty. A missing trailing newline is
// added after captured output. The exit section contains the process exit
// code; if it is -1, the section also includes the process-state
// description.
//
// A section contains the captured bytes only when FormatResult classifies
// them as text; otherwise, it contains a single notice line naming the
// detected type. Classification uses a MIME type, inferred from the
// stream's leading bytes: any type beginning with "text/" counts as text.
// This classification can be wrong in both directions; see
// [http.DetectContentType] for the rules it follows and
// https://mimesniff.spec.whatwg.org/ for the algorithm it implements.
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")
			}
		}
	}

	fmt.Fprintf(&buf, "┗ start: %s\n", r.StartedAt.Format(time.RFC3339Nano))
	fmt.Fprintf(&buf, "┗ end: %s\n", r.FinishedAt.Format(time.RFC3339Nano))
	fmt.Fprintf(&buf, "┗ duration: %s\n",
		r.FinishedAt.Sub(r.StartedAt).Round(time.Millisecond))

	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.
//
// On timeout, context cancellation, or [exec.ErrWaitDelay], Execute may
// return both a non-nil ToolResult with partial output and a non-nil error.
// The ToolResult is nil only for start failures where the process never ran.
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. WithTimeoutCause attaches errToolTimeout
	// as the context cause so we can distinguish it from a parent deadline.
	if tool.timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeoutCause(
			ctx, tool.timeout, errToolTimeout,
		)
		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
	// Bound pipe draining after cancellation or process exit; descendants
	// may retain output descriptors. This does not terminate descendants.
	// https://pkg.go.dev/os/exec#Cmd.WaitDelay
	cmd.WaitDelay = time.Second

	// 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, ProcessState is nil.
	//
	// 2. Start failure (cmd.ProcessState == 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 the process started, ProcessState is guaranteed
	// to be set because Start() succeeded.
	started := time.Now()
	err = cmd.Run()
	finished := time.Now()
	if err != nil {
		// Build the result only if the process started. A non-nil result
		// always has a non-nil ProcessState.
		var result *ToolResult
		if cmd.ProcessState != nil {
			result = &ToolResult{
				FinishedAt:   finished,
				ProcessState: cmd.ProcessState,
				StartedAt:    started,
				Stderr:       stderr.Bytes(),
				Stdout:       stdout.Bytes(),
			}
		}

		switch {
		// Context was cancelled or timed out. If the process started,
		// result contains partial output. If it never started, result
		// is nil. Distinguish tool timeout from other context errors
		// so we can return an appropriate message.
		case ctx.Err() != nil:
			if errors.Is(context.Cause(ctx), errToolTimeout) {
				return result, fmt.Errorf(
					"tool %q: timed out after %v: %w",
					name, tool.timeout, errToolTimeout,
				)
			}
			return result, fmt.Errorf("tool %q: %w", name, ctx.Err())

		// Process never started, and it wasn't a context error.
		// The error is the raw start failure (e.g., *PathError).
		case result == nil:
			return nil, fmt.Errorf("tool %q: %w", name, err)

		// Output may be incomplete even when the process exited
		// successfully, if the output pipe drained slowly.
		case errors.Is(err, exec.ErrWaitDelay):
			return result, fmt.Errorf(
				"tool %q: output incomplete: %w", name, err,
			)

		// Process ran and exited with non-zero status or was killed
		// by a signal. ProcessState is set; return it with output.
		default:
			if _, ok := errors.AsType[*exec.ExitError](err); !ok {
				return nil, fmt.Errorf("tool %q: %w", name, err)
			}
			return result, nil
		}
	}

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