dingtalk-workspace-cli/internal/shortcut/userdef/loader.go

// Copyright 2026 Alibaba Group
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package userdef loads user-defined shortcuts from ~/.dws/shortcuts/*.yaml and
// compiles them into the same shortcut.Shortcut model as the built-ins, so a
// distilled high-frequency operation behaves exactly like a native `+command`.
// This is the runtime half of the "high-frequency auto-distillation" story
// (see docs/shortcut-p2-design.md); the `dws shortcut add` command writes these
// YAML files.
package userdef

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/shortcut"
	"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/pkg/config"
	"gopkg.in/yaml.v3"
)

var commandSegmentRE = regexp.MustCompile(`^[\p{L}\p{N}_-][\p{L}\p{N}._-]*$`)

// Spec is the on-disk YAML form of a user-defined shortcut. It mirrors
// shortcut.Shortcut but expresses the Execute step declaratively via Exec.Bind.
type Spec struct {
	Version     int        `yaml:"version"`
	Service     string     `yaml:"service"`
	Command     string     `yaml:"command"`
	Product     string     `yaml:"product"`
	Description string     `yaml:"description"`
	Intent      string     `yaml:"intent"`
	Risk        string     `yaml:"risk"`
	Source      string     `yaml:"source"` // auto | manual
	Flags       []FlagSpec `yaml:"flags"`
	Exec        ExecSpec   `yaml:"execute"`
}

// FlagSpec is the YAML form of shortcut.Flag.
type FlagSpec struct {
	Name     string   `yaml:"name"`
	Type     string   `yaml:"type"`
	Default  string   `yaml:"default"`
	Desc     string   `yaml:"desc"`
	Required bool     `yaml:"required"`
	Enum     []string `yaml:"enum"`
}

// ExecSpec declares the single MCP call. Bind maps each tool parameter key to
// either a constant string, or a "${flag}" reference resolved from a flag value.
type ExecSpec struct {
	Tool string            `yaml:"tool"`
	Bind map[string]string `yaml:"bind"`
}

// Dir returns the directory holding user shortcut YAML files (~/.dws/shortcuts).
func Dir() string { return filepath.Join(config.DefaultConfigDir(), "shortcuts") }

// Load reads every ~/.dws/shortcuts/*.yaml, compiles valid specs, and registers
// them (skipping any that collide with an already-registered service+command,
// so built-ins always win). Returns the number registered and any per-file
// errors. It never fails hard — a bad file is reported and skipped.
func Load() (registered int, errs []error) {
	files, _ := filepath.Glob(filepath.Join(Dir(), "*.yaml"))
	existing := registeredKeys()
	for _, f := range files {
		s, err := parseFile(f)
		if err != nil {
			errs = append(errs, fmt.Errorf("%s: %w", filepath.Base(f), err))
			continue
		}
		key := s.Service + " " + s.Command
		if existing[key] {
			errs = append(errs, fmt.Errorf("%s: 与已有 shortcut %q 冲突，已跳过", filepath.Base(f), key))
			continue
		}
		shortcut.Register(Compile(s))
		existing[key] = true
		registered++
	}
	return registered, errs
}

func registeredKeys() map[string]bool {
	set := map[string]bool{}
	for _, s := range shortcut.All() {
		set[s.Service+" "+s.Command] = true
	}
	return set
}

func parseFile(path string) (Spec, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return Spec{}, err
	}
	var s Spec
	if err := yaml.Unmarshal(data, &s); err != nil {
		return Spec{}, err
	}
	return s, Validate(s)
}

// Validate checks the minimum required fields of a spec.
func Validate(s Spec) error {
	if s.Service == "" || s.Command == "" {
		return fmt.Errorf("service 和 command 为必填")
	}
	if !strings.HasPrefix(s.Command, "+") {
		return fmt.Errorf("command 必须以 + 开头，当前 %q", s.Command)
	}
	if !commandSegmentRE.MatchString(s.Service) ||
		!commandSegmentRE.MatchString(strings.TrimPrefix(s.Command, "+")) {
		return fmt.Errorf("service 和 command 只能包含字母、数字、点、下划线和连字符，且不能以点开头")
	}
	if s.Exec.Tool == "" {
		return fmt.Errorf("execute.tool 为必填")
	}
	return nil
}

// FilePath returns the safe on-disk path for a user-defined shortcut. Validate
// already constrains both components, while the Rel check is a final defense
// against future validation changes reintroducing directory traversal.
func FilePath(service, command string) (string, error) {
	probe := Spec{Service: service, Command: command, Exec: ExecSpec{Tool: "path-check"}}
	if err := Validate(probe); err != nil {
		return "", err
	}
	dir := Dir()
	name := service + "." + strings.TrimPrefix(command, "+") + ".yaml"
	path := filepath.Join(dir, name)
	rel, err := filepath.Rel(dir, path)
	if err != nil || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) {
		return "", fmt.Errorf("shortcut 文件路径超出配置目录")
	}
	return path, nil
}

// Compile turns a validated Spec into a runnable shortcut.Shortcut. Its Execute
// builds the MCP params from Exec.Bind: "${flag}" values are resolved from the
// flag (respecting its type), everything else is a literal constant. Optional
// flag references that the user did not set are omitted.
func Compile(s Spec) shortcut.Shortcut {
	flags := make([]shortcut.Flag, 0, len(s.Flags))
	flagType := map[string]shortcut.FlagType{}
	flagRequired := map[string]bool{}
	flagHasDefault := map[string]bool{}
	for _, f := range s.Flags {
		t := shortcut.FlagType(f.Type)
		if t == "" {
			t = shortcut.FlagString
		}
		flagType[f.Name] = t
		flagRequired[f.Name] = f.Required
		flagHasDefault[f.Name] = f.Default != ""
		flags = append(flags, shortcut.Flag{
			Name: f.Name, Type: t, Default: f.Default,
			Desc: f.Desc, Required: f.Required, Enum: f.Enum,
		})
	}

	risk := shortcut.Risk(s.Risk)
	if risk == "" {
		risk = shortcut.RiskRead
	}
	desc := s.Description
	if desc == "" {
		desc = s.Command + "（自定义）"
	}
	intent := s.Intent
	if intent == "" {
		intent = "用户自定义 shortcut（沉淀自高频操作）：" + desc
	}

	bind := s.Exec.Bind
	tool := s.Exec.Tool

	return shortcut.Shortcut{
		Service:     s.Service,
		Command:     s.Command,
		Product:     s.Product,
		Description: desc,
		Intent:      intent,
		Risk:        risk,
		Flags:       flags,
		UserDefined: true,
		Execute: func(rt *shortcut.RuntimeContext) error {
			params := map[string]any{}
			for key, tmpl := range bind {
				if name, ok := flagRef(tmpl); ok {
					if !rt.Changed(name) && !flagRequired[name] && !flagHasDefault[name] {
						continue // optional flag not provided → omit param
					}
					params[key] = readFlag(rt, name, flagType[name])
				} else {
					params[key] = tmpl // constant
				}
			}
			return rt.CallMCP(tool, params)
		},
	}
}

// flagRef reports whether v is a "${name}" reference and returns name.
func flagRef(v string) (string, bool) {
	if strings.HasPrefix(v, "${") && strings.HasSuffix(v, "}") {
		return strings.TrimSpace(v[2 : len(v)-1]), true
	}
	return "", false
}

func readFlag(rt *shortcut.RuntimeContext, name string, t shortcut.FlagType) any {
	switch t {
	case shortcut.FlagBool:
		return rt.Bool(name)
	case shortcut.FlagInt:
		return rt.Int(name)
	case shortcut.FlagStringSlice:
		return rt.StrSlice(name)
	default:
		return rt.Str(name)
	}
}