Cyrene_v2/internal/runtime/tool/tool.go

// Package tool implements singleshot tool dispatch for the agent runtime
// (DevDocs/design/04-运行时.md §4). A Dispatcher wraps the global capability
// registry, resolves a named Tool, validates the caller-supplied arguments
// against the tool's JSON Schema, and executes it under a per-tool timeout,
// turning every failure into a ToolError that is safe to feed back to the LLM.
//
// This package deliberately covers only a single invocation: the loop upper
// bound (8 rounds) and retry belong to the agent loop (internal/runtime/loop),
// not to dispatch.
package tool

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/restrict"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
)

// DefaultTimeout is the per-invoke deadline used when a tool has no explicit
// override and the Dispatcher was built without WithDefaultTimeout.
const DefaultTimeout = 30 * time.Second

// Execution modes a tool may declare to the loop. A tool that mutates shared
// state (file writes, device control) should be marked sequential so the loop
// does not run it concurrently with other calls.
const (
	// ExecutionModeParallel is the default: the loop may run such a tool
	// concurrently with other calls in a batch.
	ExecutionModeParallel = "parallel"
	// ExecutionModeSequential forces the loop to run the whole batch serially so
	// this tool is never invoked concurrently.
	ExecutionModeSequential = "sequential"
)

// Error-kind labels distinguishing why a tool call failed, fed back to the model
// as a structured {kind,message} so it can react differently to a bad argument
// (re-issue a corrected call) versus an execution failure.
const (
	// KindArgumentValidation marks a JSON-Schema argument-validation failure.
	KindArgumentValidation = "argument-validation"
	// KindExecution marks a runtime execution failure (timeout, panic, tool error).
	KindExecution = "execution"
)

// Sentinel errors returned by the Dispatcher. Use errors.Is to test for them;
// they are wrapped with %w so the underlying cause stays available.
var (
	// ErrUnknownTool is returned when the requested name is not a registered
	// tool, or resolves to a capability of a non-Tool kind.
	ErrUnknownTool = errors.New("tool: unknown tool")
	// ErrInvalidArgs is returned when the caller-supplied arguments fail the
	// tool's JSON Schema validation (missing required, wrong type).
	ErrInvalidArgs = errors.New("tool: invalid arguments")
	// ErrTimeout is returned when a tool does not complete before its deadline.
	ErrTimeout = errors.New("tool: execution timeout")
	// ErrNilRegistry is returned by New when the supplied registry is nil.
	ErrNilRegistry = errors.New("tool: nil registry")
)

// ToolError is the error the Dispatcher returns to the caller on any failed
// invocation. It carries the tool Name and a human/LLM-readable Message and
// wraps the underlying cause so sentinel classification survives errors.Is and
// errors.As.
type ToolError struct {
	// Name is the tool that failed.
	Name string
	// Message is a safe, human- and LLM-readable description of the failure.
	Message string
	// cause is the wrapped underlying error (may be nil).
	cause error
}

// Error implements the error interface.
func (e *ToolError) Error() string {
	if e.cause == nil {
		return fmt.Sprintf("tool %q failed: %s", e.Name, e.Message)
	}
	return fmt.Sprintf("tool %q failed: %s: %v", e.Name, e.Message, e.cause)
}

// Unwrap exposes the wrapped cause so errors.Is/errors.As traverse the chain.
func (e *ToolError) Unwrap() error { return e.cause }

// Kind classifies the failure as "argument-validation" (the caller-supplied
// arguments failed the JSON Schema) or "execution" (anything after validation:
// unknown tool, timeout, panic or an Execute error). The loop uses it to feed a
// structured {kind,message} tool result back to the model.
func (e *ToolError) Kind() string {
	if e.cause != nil && errors.Is(e.cause, ErrInvalidArgs) {
		return KindArgumentValidation
	}
	return KindExecution
}

// Option configures a Dispatcher at construction time.
type Option func(*Dispatcher)

// Dispatcher resolves, validates and executes tools registered in a capability
// registry. It is safe for concurrent use.
type Dispatcher struct {
	reg      *registry.Registry
	dft      time.Duration
	mu       sync.RWMutex // guards timeouts and modes
	timeouts map[string]time.Duration
	modes    map[string]string // tool name -> ExecutionMode* (default parallel)
}

// New builds a Dispatcher backed by reg, applying opts in order. It returns an
// error wrapping ErrNilRegistry when reg is nil. The default per-invoke timeout
// is DefaultTimeout unless overridden with WithDefaultTimeout.
func New(reg *registry.Registry, opts ...Option) (*Dispatcher, error) {
	if reg == nil {
		return nil, fmt.Errorf("tool: new: %w", ErrNilRegistry)
	}
	d := &Dispatcher{
		reg:      reg,
		dft:      DefaultTimeout,
		timeouts: make(map[string]time.Duration),
		modes:    make(map[string]string),
	}
	for _, opt := range opts {
		if opt != nil {
			opt(d)
		}
	}
	return d, nil
}

// WithDefaultTimeout sets the default per-invoke deadline for tools without an
// explicit override. A non-positive value disables the deadline.
func WithDefaultTimeout(timeout time.Duration) Option {
	return func(d *Dispatcher) { d.dft = timeout }
}

// Register registers t as a tool with the underlying capability registry,
// delegating duplicate/invalid handling to the registry. A nil tool is rejected
// by the registry's own validation.
func (d *Dispatcher) Register(t sdk.Tool) error {
	if d.reg == nil {
		return fmt.Errorf("tool: register: %w", ErrNilRegistry)
	}
	return d.reg.RegisterTool(t)
}

// SetToolTimeout overrides the per-invoke deadline for the named tool. A
// non-positive value disables the deadline for that tool. The override applies
// regardless of whether the tool is registered yet.
func (d *Dispatcher) SetToolTimeout(name string, timeout time.Duration) error {
	if name == "" {
		return fmt.Errorf("tool: set timeout: %w", ErrUnknownTool)
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	d.timeouts[name] = timeout
	return nil
}

// SetToolExecutionMode marks a tool as running "sequential" or "parallel" (the
// default). A sequential tool is never invoked concurrently by the loop, so
// shared-state tools (file writes, device control) can declare it. An empty mode
// clears the override back to the default. Only ExecutionModeSequential is
// treated as sequential.
func (d *Dispatcher) SetToolExecutionMode(name, mode string) error {
	if name == "" {
		return fmt.Errorf("tool: set execution mode: %w", ErrUnknownTool)
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	if mode == "" || mode == ExecutionModeParallel {
		delete(d.modes, name)
		return nil
	}
	d.modes[name] = mode
	return nil
}

// IsSequential reports whether the named tool has been declared sequential. An
// unknown name or one never marked sequential reports false (the default).
func (d *Dispatcher) IsSequential(name string) bool {
	d.mu.RLock()
	defer d.mu.RUnlock()
	return d.modes[name] == ExecutionModeSequential
}

// PolicyEngine returns a restrict.Policy scoped to env, reading the editable
// reachability markers from the underlying capability registry (design/11 §3).
// It is the loop's authoritative source for the group/stranger tool gate and is
// fail-closed: a nil receiver or nil registry denies every tool in a non-owner
// environment, and the owner environment always allows. The returned policy is
// immutable and safe for concurrent use.
func (d *Dispatcher) PolicyEngine(env restrict.ActorEnv) *restrict.Policy {
	if d == nil || d.reg == nil {
		return restrict.SetContext(nil, env)
	}
	return restrict.SetContext(d.reg, env)
}

// Invoke executes the tool registered as name with the given args. It resolves
// the tool from the registry, validates args against the tool's Parameters JSON
// Schema, and runs Execute under the tool's deadline (per-tool override, else
// default). Any failure — unknown tool, invalid args, timeout, panic or an
// Execute error — is returned as a *ToolError wrapping the matching sentinel,
// so it is safe to feed back to the LLM.
func (d *Dispatcher) Invoke(ctx context.Context, name string, args map[string]any) (any, error) {
	if name == "" {
		return nil, &ToolError{Name: name, Message: "tool name is empty", cause: ErrUnknownTool}
	}
	cap, err := d.reg.Get(name)
	if err != nil {
		return nil, &ToolError{Name: name, Message: "unknown tool", cause: ErrUnknownTool}
	}
	if cap.Kind != registry.KindTool {
		return nil, &ToolError{Name: name, Message: "not a callable tool", cause: ErrUnknownTool}
	}
	tool, ok := cap.Value.(sdk.Tool)
	if !ok {
		return nil, &ToolError{Name: name, Message: "tool value is not an sdk.Tool", cause: ErrUnknownTool}
	}

	params, perr := tool.Parameters()
	if perr != nil {
		return nil, &ToolError{
			Name:    name,
			Message: "invalid tool schema",
			cause:   fmt.Errorf("tool: parameters %q: %w", name, perr),
		}
	}
	if verr := validateArgs(params, args); verr != nil {
		return nil, &ToolError{
			Name:    name,
			Message: verr.Error(),
			cause:   fmt.Errorf("tool: validate %q: %w", name, ErrInvalidArgs),
		}
	}

	invokeCtx, cancel := d.invokeCtx(ctx, name)
	defer cancel()

	result, execErr := callTool(invokeCtx, tool, args)
	if execErr != nil {
		return nil, wrapInvokeErr(name, execErr)
	}
	return result, nil
}

// invokeCtx derives the execution context for a tool, applying the tool's
// deadline when it is positive.
func (d *Dispatcher) invokeCtx(ctx context.Context, name string) (context.Context, context.CancelFunc) {
	timeout := d.timeoutFor(name)
	if timeout > 0 {
		return context.WithTimeout(ctx, timeout)
	}
	return context.WithCancel(ctx)
}

// timeoutFor returns the per-tool deadline override, or the default.
func (d *Dispatcher) timeoutFor(name string) time.Duration {
	d.mu.RLock()
	defer d.mu.RUnlock()
	if t, ok := d.timeouts[name]; ok {
		return t
	}
	return d.dft
}

// wrapInvokeErr converts a raw execution error into a *ToolError, classifying a
// deadline exceeded as ErrTimeout and preserving every other cause so errors.As
// can recover the ToolError and errors.Is can reach the underlying error.
func wrapInvokeErr(name string, err error) error {
	switch {
	case errors.Is(err, context.DeadlineExceeded):
		return &ToolError{Name: name, Message: "execution timed out", cause: ErrTimeout}
	case errors.Is(err, context.Canceled):
		return &ToolError{Name: name, Message: "execution canceled", cause: err}
	default:
		return &ToolError{Name: name, Message: err.Error(), cause: err}
	}
}

// callTool runs tool.Execute in a goroutine so a tool that ignores its context
// is still interrupted by the deadline. The goroutine recovers panics, so a
// panicking tool never crashes the process. The result channel is buffered so a
// goroutine that outlives a timeout can still deliver without leaking a sender.
func callTool(ctx context.Context, tool sdk.Tool, args any) (any, error) {
	type result struct {
		val any
		err error
	}
	ch := make(chan result, 1)
	go func() {
		defer func() {
			if r := recover(); r != nil {
				ch <- result{err: fmt.Errorf("tool panicked: %v", r)}
			}
		}()
		val, err := tool.Execute(ctx, args)
		ch <- result{val: val, err: err}
	}()
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case r := <-ch:
		return r.val, r.err
	}
}

// ToolSchemas returns the registered function-call schemas in registration
// order, ready to be advertised to a model in a model.Request. A capability that
// is not a callable sdk.Tool, or whose Parameters() reports an error, is
// skipped. A nil registry returns nil. The returned slice is owned by the caller
// and safe to mutate.
func (d *Dispatcher) ToolSchemas() []model.ToolSpec {
	if d.reg == nil {
		return nil
	}
	caps := d.reg.List(registry.KindTool)
	out := make([]model.ToolSpec, 0, len(caps))
	for _, cap := range caps {
		t, ok := cap.Value.(sdk.Tool)
		if !ok {
			continue
		}
		params, err := t.Parameters()
		if err != nil {
			continue
		}
		out = append(out, model.ToolSpec{
			Name:        t.Name(),
			Description: t.Description(),
			Parameters:  params,
		})
	}
	return out
}