hexis/internal/provider/workspace_mcp.go

package provider

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/kami/hexis/internal/domain"
	"gopkg.in/yaml.v3"
)

func LoadToolAllowlist(path string) (ToolAllowlist, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return ToolAllowlist{}, fmt.Errorf("read allowlist: %w", err)
	}
	var allowlist ToolAllowlist
	if err := yaml.Unmarshal(data, &allowlist); err != nil {
		return ToolAllowlist{}, fmt.Errorf("parse allowlist: %w", err)
	}
	if allowlist.Tools == nil {
		allowlist.Tools = map[string]ToolMapping{}
	}
	return allowlist, nil
}

type WorkspaceMCPProvider struct {
	mu         sync.RWMutex
	baseURL    string
	httpClient *http.Client
	tools      []WorkspaceTool
	allowlist  ToolAllowlist
}

type WorkspaceTool struct {
	Name        string `json:"name"`
	Description string `json:"description"`
	InputSchema any    `json:"inputSchema"`
}

type ToolAllowlist struct {
	Tools map[string]ToolMapping `yaml:"tools" json:"tools"`
}

type ToolMapping struct {
	Capability   string   `yaml:"capability" json:"capability"`
	Risk         string   `yaml:"risk" json:"risk"`
	TargetType   string   `yaml:"target_type,omitempty" json:"target_type,omitempty"`
	ReadOnly     bool     `yaml:"read_only" json:"read_only"`
	SideEffects  string   `yaml:"side_effects,omitempty" json:"side_effects,omitempty"`
	Enabled      *bool    `yaml:"enabled,omitempty" json:"enabled,omitempty"`
	AllowParams  []string `yaml:"allow_params,omitempty" json:"allow_params,omitempty"`
}

func (r *ToolAllowlist) IsEnabled(name string) bool {
	m, ok := r.Tools[name]
	if !ok {
		return false
	}
	if m.Enabled != nil && !*m.Enabled {
		return false
	}
	return true
}

func (r *ToolAllowlist) Mapping(name string) (ToolMapping, bool) {
	m, ok := r.Tools[name]
	return m, ok
}

func NewWorkspaceMCPProvider(baseURL string, allowlist ToolAllowlist) *WorkspaceMCPProvider {
	return &WorkspaceMCPProvider{
		baseURL: strings.TrimRight(baseURL, "/"),
		httpClient: &http.Client{
			Timeout: 60 * time.Second,
		},
		allowlist: allowlist,
	}
}

func (p *WorkspaceMCPProvider) Name() string {
	return "workspace_mcp"
}

func (p *WorkspaceMCPProvider) DiscoverTools() ([]WorkspaceTool, error) {
	resp, err := p.httpClient.Get(fmt.Sprintf("%s/api/tools", p.baseURL))
	if err != nil {
		return nil, fmt.Errorf("discover workspace tools: %w", err)
	}
	defer resp.Body.Close()

	var result struct {
		Tools []WorkspaceTool `json:"tools"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return nil, fmt.Errorf("decode workspace tools: %w", err)
	}

	p.mu.Lock()
	p.tools = result.Tools
	p.mu.Unlock()

	return result.Tools, nil
}

func (p *WorkspaceMCPProvider) DiscoveredTools() []WorkspaceTool {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return p.tools
}

func (p *WorkspaceMCPProvider) Execute(capability *domain.Capability, req *domain.ExecuteRequest) (map[string]any, error) {
	toolName := p.capabilityToTool(capability.Name)
	if toolName == "" {
		return nil, fmt.Errorf("no workspace tool mapped for capability %q", capability.Name)
	}

	mapping, ok := p.allowlist.Mapping(toolName)
	if !ok {
		return nil, fmt.Errorf("tool %q not in allowlist", toolName)
	}
	if !p.allowlist.IsEnabled(toolName) {
		return nil, fmt.Errorf("tool %q is disabled in allowlist", toolName)
	}

	args := map[string]any{}
	if req.Arguments != nil {
		if len(mapping.AllowParams) > 0 {
			for k, v := range req.Arguments {
				for _, allowed := range mapping.AllowParams {
					if k == allowed {
						args[k] = v
						break
					}
				}
			}
		} else {
			args = req.Arguments
		}
	}

	body, _ := json.Marshal(args)
	resp, err := p.httpClient.Post(
		fmt.Sprintf("%s/api/tool/%s", p.baseURL, toolName),
		"application/json",
		bytes.NewReader(body),
	)
	if err != nil {
		return nil, fmt.Errorf("call workspace tool %q: %w", toolName, err)
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("read workspace tool response: %w", err)
	}

	if resp.StatusCode >= 400 {
		return map[string]any{
			"error":      string(respBody),
			"http_status": resp.StatusCode,
		}, fmt.Errorf("workspace tool %q returned %d: %s", toolName, resp.StatusCode, string(respBody))
	}

	var result any
	if err := json.Unmarshal(respBody, &result); err != nil {
		return map[string]any{
			"raw": string(respBody),
		}, nil
	}

	return map[string]any{
		"result": result,
	}, nil
}

func (p *WorkspaceMCPProvider) capabilityToTool(capName string) string {
	for toolName, mapping := range p.allowlist.Tools {
		if mapping.Capability == capName {
			return toolName
		}
	}
	return ""
}

func (p *WorkspaceMCPProvider) BuildCapabilities() []domain.Capability {
	var caps []domain.Capability
	for toolName, mapping := range p.allowlist.Tools {
		if !p.allowlist.IsEnabled(toolName) {
			continue
		}
		readOnly := mapping.ReadOnly
		if mapping.Risk == "read" {
			readOnly = true
		}
		now := time.Now().UTC()

		caps = append(caps, domain.Capability{
			ID:                  fmt.Sprintf("cap_ws_%s", strings.ReplaceAll(mapping.Capability, ".", "_")),
			Name:                mapping.Capability,
			Description:         fmt.Sprintf("Workspace tool: %s", toolName),
			TargetTypes:         ifString(mapping.TargetType != "", []string{mapping.TargetType}, nil),
			TargetEntityID:      "",
			Provider:            "workspace_mcp",
			Operation:           toolName,
			Risk:                mapping.Risk,
			ReadOnly:            readOnly,
			ExpectedSideEffects: mapping.SideEffects,
			CreatedAt:           now,
			UpdatedAt:           now,
			Version:             1,
		})
	}
	return caps
}

func ifString(cond bool, a, b []string) []string {
	if cond {
		return a
	}
	return b
}