dingtalk-workspace-cli/internal/app/plugin_stdio_overlay_test.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 app

import (
	"encoding/json"
	"testing"
	"time"

	"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/cache"
	"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/executor"
	"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/market"
	"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/plugin"
	"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/transport"
	"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/pkg/config"
	"github.com/spf13/cobra"
)

// withCleanStdioRegistry snapshots and restores the package-level stdio
// client registry so tests that call RegisterStdioClient don't leak state
// across cases.
func withCleanStdioRegistry(t *testing.T) {
	t.Helper()
	stdioMu.Lock()
	prev := stdioClients
	stdioClients = make(map[string]*transport.StdioClient)
	stdioMu.Unlock()
	t.Cleanup(func() {
		stdioMu.Lock()
		stdioClients = prev
		stdioMu.Unlock()
	})
}

// newOverlayFixture constructs a plugin + stdio entry carrying an inline
// CLIOverlay with the given tool-override map. The stdio client is created
// but never started, since the overlay-first path does not require the
// subprocess to be running for command registration.
func newOverlayFixture(t *testing.T, pluginName, serverKey string, overlay market.CLIOverlay) (*plugin.Plugin, plugin.StdioServerClient) {
	t.Helper()
	raw, err := json.Marshal(overlay)
	if err != nil {
		t.Fatalf("marshal overlay: %v", err)
	}
	p := &plugin.Plugin{
		Manifest: plugin.Manifest{
			Name:        pluginName,
			Version:     "1.0.0",
			Description: pluginName + " plugin",
			MCPServers: map[string]*plugin.MCPServer{
				serverKey: {
					Type:    "stdio",
					Command: "/usr/bin/true", // never executed by overlay-first path
					CLI:     raw,
				},
			},
		},
		Root: t.TempDir(),
	}
	sc := plugin.StdioServerClient{
		Key:    serverKey,
		Client: transport.NewStdioClient("/usr/bin/true", nil, nil),
	}
	return p, sc
}

// TestRegisterStdioServerFromOverlay_NoDiscoveryStillBuildsCommands verifies
// the core promise of the overlay-first path: when overlay.json ships
// ToolOverrides, commands appear immediately — no subprocess probe.
func TestRegisterStdioServerFromOverlay_NoDiscoveryStillBuildsCommands(t *testing.T) {
	withCleanDynamicRegistry(t)
	withCleanStdioRegistry(t)

	overlay := market.CLIOverlay{
		ID:      "conference-local",
		Command: "conference-local",
		Groups: map[string]market.CLIGroupDef{
			"meeting": {Description: "会议控制"},
			"member":  {Description: "成员管理"},
		},
		ToolOverrides: map[string]market.CLIToolOverride{
			"create_meeting": {CLIName: "create", Group: "meeting", Description: "Create a meeting"},
			"end_meeting":    {CLIName: "end", Group: "meeting", Description: "End a meeting"},
			"mute_member":    {CLIName: "mute", Group: "member", Description: "Mute a member"},
		},
	}
	p, sc := newOverlayFixture(t, "conference-local", "conference-local", overlay)

	store := cache.NewStore(t.TempDir())
	cmds, desc, ok := registerStdioServerFromOverlay(p, sc, executor.EchoRunner{}, store)
	if !ok {
		t.Fatal("registerStdioServerFromOverlay returned ok=false, want true")
	}
	if len(cmds) == 0 {
		t.Fatal("registerStdioServerFromOverlay returned 0 commands, want >=1")
	}

	var root *struct{ name, path string }
	_ = root
	found := false
	for _, c := range cmds {
		if c.Name() == "conference-local" {
			found = true
			// Groups must be attached as sub-commands.
			groups := map[string]bool{}
			for _, sub := range c.Commands() {
				groups[sub.Name()] = true
			}
			if !groups["meeting"] {
				t.Errorf("missing 'meeting' group sub-command, children = %v", groups)
			}
			if !groups["member"] {
				t.Errorf("missing 'member' group sub-command, children = %v", groups)
			}
		}
	}
	if !found {
		names := []string{}
		for _, c := range cmds {
			names = append(names, c.Name())
		}
		t.Fatalf("missing top-level 'conference-local' command, got %v", names)
	}

	// AppendDynamicServer registration: product ID should land in
	// DirectRuntimeProductIDs so hideNonDirectRuntimeCommands keeps it
	// visible even under a restrictive VisibleProducts hook.
	if !DirectRuntimeProductIDs()["conference-local"] {
		t.Error("DirectRuntimeProductIDs missing 'conference-local'")
	}

	// RegisterStdioClient side-effect: the runtime must be able to look up
	// the StdioClient when the endpoint is invoked later.
	if _, ok := LookupStdioClient("conference-local/conference-local"); !ok {
		t.Error("LookupStdioClient missing conference-local/conference-local")
	}

	if desc.Endpoint != StdioEndpoint("conference-local", "conference-local") {
		t.Errorf("descriptor.Endpoint = %q, want %q", desc.Endpoint, StdioEndpoint("conference-local", "conference-local"))
	}
}

// TestRegisterStdioServerFromOverlay_WarmCacheEnrichesFlags pre-populates the
// tools cache with a schema-bearing DetailTool and asserts the resulting
// leaf command picks up the typed flag derived from InputSchema.
func TestRegisterStdioServerFromOverlay_WarmCacheEnrichesFlags(t *testing.T) {
	withCleanDynamicRegistry(t)
	withCleanStdioRegistry(t)

	overlay := market.CLIOverlay{
		ID:      "cache-plugin",
		Command: "cache-plugin",
		ToolOverrides: map[string]market.CLIToolOverride{
			"echo": {CLIName: "echo", Description: "Echo input"},
		},
	}
	p, sc := newOverlayFixture(t, "cache-plugin", "cache-plugin", overlay)

	store := cache.NewStore(t.TempDir())
	cacheKey := pluginCacheKey(p.Manifest.Name, sc.Key)
	if err := store.SaveTools(config.DefaultPartition, cacheKey, cache.ToolsSnapshot{
		SavedAt:   time.Now().UTC(),
		ServerKey: cacheKey,
		Tools: []transport.ToolDescriptor{
			{
				Name:        "echo",
				Description: "Echo the input",
				InputSchema: map[string]any{
					"type": "object",
					"properties": map[string]any{
						"message": map[string]any{"type": "string"},
					},
					"required": []any{"message"},
				},
			},
		},
	}); err != nil {
		t.Fatalf("SaveTools: %v", err)
	}

	cmds, _, ok := registerStdioServerFromOverlay(p, sc, executor.EchoRunner{}, store)
	if !ok || len(cmds) == 0 {
		t.Fatalf("overlay registration failed: ok=%v cmds=%d", ok, len(cmds))
	}

	var echoLeaf *leafMatch
	for _, top := range cmds {
		if top.Name() != "cache-plugin" {
			continue
		}
		for _, sub := range top.Commands() {
			if sub.Name() == "echo" {
				echoLeaf = &leafMatch{name: sub.Name(), hasFlag: sub.Flags().Lookup("message") != nil}
			}
		}
	}
	if echoLeaf == nil {
		t.Fatal("missing 'echo' leaf command under 'cache-plugin'")
	}
	if !echoLeaf.hasFlag {
		t.Error("warm-cache enrichment did not wire --message flag from InputSchema")
	}
}

type leafMatch struct {
	name    string
	hasFlag bool
}

// TestRegisterStdioServerFromOverlay_OverlayWithoutOverridesReturnsFalse
// asserts the fallback contract: when overlay.json is missing toolOverrides,
// the overlay-first path declines so the caller can route the entry through
// the legacy discovery-first registerStdioServer.
func TestRegisterStdioServerFromOverlay_OverlayWithoutOverridesReturnsFalse(t *testing.T) {
	withCleanDynamicRegistry(t)
	withCleanStdioRegistry(t)

	// Overlay with no ToolOverrides (simulates a plugin that relies entirely
	// on runtime discovery for its tool list).
	overlay := market.CLIOverlay{
		ID:      "legacy-plugin",
		Command: "legacy-plugin",
	}
	p, sc := newOverlayFixture(t, "legacy-plugin", "legacy-plugin", overlay)

	store := cache.NewStore(t.TempDir())
	cmds, _, ok := registerStdioServerFromOverlay(p, sc, executor.EchoRunner{}, store)
	if ok {
		t.Errorf("registerStdioServerFromOverlay ok=true for empty toolOverrides; want false")
	}
	if cmds != nil {
		t.Errorf("cmds = %v, want nil", cmds)
	}
	if DirectRuntimeProductIDs()["legacy-plugin"] {
		t.Error("legacy-plugin must NOT be appended to dynamic registry in fallback case")
	}
	if _, found := LookupStdioClient("legacy-plugin/legacy-plugin"); found {
		t.Error("stdio client must NOT be registered in fallback case")
	}
}

// TestRefreshStdioToolsCache_FailurePreservesCache guards against the
// "negative cache poisoning" bug: if discovery fails (subprocess not ready,
// timeout, empty tool list), the existing warm cache must remain intact so
// the next startup still enriches flags from the last good snapshot.
func TestRefreshStdioToolsCache_FailurePreservesCache(t *testing.T) {
	withCleanDynamicRegistry(t)
	withCleanStdioRegistry(t)

	p, sc := newOverlayFixture(t, "refresh-plugin", "refresh-plugin", market.CLIOverlay{
		ID:      "refresh-plugin",
		Command: "refresh-plugin",
	})

	store := cache.NewStore(t.TempDir())
	cacheKey := pluginCacheKey(p.Manifest.Name, sc.Key)
	goodSnapshot := cache.ToolsSnapshot{
		SavedAt:   time.Now().UTC(),
		ServerKey: cacheKey,
		Tools: []transport.ToolDescriptor{
			{
				Name:        "ping",
				Description: "Health check",
				InputSchema: map[string]any{"type": "object"},
			},
		},
	}
	if err := store.SaveTools(config.DefaultPartition, cacheKey, goodSnapshot); err != nil {
		t.Fatalf("seed SaveTools: %v", err)
	}

	// /usr/bin/true exits immediately, so Initialize + ListTools will fail
	// (no MCP handshake). discoverStdioTools returns nil → refresh must be
	// a no-op and must NOT overwrite the good cache with a null snapshot.
	refreshStdioToolsCache(p, sc, store, pluginColdTimeouts{stdio: 200 * time.Millisecond})

	got, _, err := store.LoadTools(config.DefaultPartition, cacheKey)
	if err != nil {
		t.Fatalf("LoadTools after failed refresh: %v", err)
	}
	if len(got.Tools) != 1 || got.Tools[0].Name != "ping" {
		t.Errorf("warm cache was overwritten by failed refresh: %+v", got.Tools)
	}
}

// TestLoadPlugins_OverlayFirstVisibleBeforeDiscovery is an integration-style
// test for the loadPlugins split decision: stdio plugins whose overlay ships
// ToolOverrides must have their commands visible on the root immediately,
// WITHOUT waiting on any discovery handshake. It drives the same sequence
// loadPlugins uses (registerStdioServerFromOverlay → root.AddCommand →
// hideNonDirectRuntimeCommands) and asserts the plugin command survives the
// visibility filter even when no discovery has run.
func TestLoadPlugins_OverlayFirstVisibleBeforeDiscovery(t *testing.T) {
	withCleanDynamicRegistry(t)
	withCleanStdioRegistry(t)

	// Simulate a wukong-like edition that declares a static VisibleProducts
	// whitelist NOT containing our plugin. This is the exact scenario where
	// the original bug surfaced.
	overrideVisibleProducts(t, []string{"calendar", "doc"})

	overlay := market.CLIOverlay{
		ID:      "conference-local",
		Command: "conference-local",
		ToolOverrides: map[string]market.CLIToolOverride{
			"create_meeting": {CLIName: "create", Description: "Create a meeting"},
		},
	}
	p, sc := newOverlayFixture(t, "conference-local", "conference-local", overlay)

	// No discovery runs — no cache seeded. This mirrors a cold-start where
	// the subprocess is unavailable (or just slow) yet the user expects
	// `dws --help` to still list the plugin.
	store := cache.NewStore(t.TempDir())
	cmds, _, ok := registerStdioServerFromOverlay(p, sc, executor.EchoRunner{}, store)
	if !ok {
		t.Fatal("registerStdioServerFromOverlay returned ok=false")
	}

	root := &cobra.Command{Use: "dws"}
	// Also add a sibling command that is NOT a registered product so we can
	// prove the visibility filter still hides non-product commands.
	bogus := &cobra.Command{Use: "bogus-not-a-product"}
	root.AddCommand(bogus)
	for _, c := range cmds {
		root.AddCommand(c)
	}

	hideNonDirectRuntimeCommands(root)

	var pluginCmd *cobra.Command
	for _, c := range root.Commands() {
		if c.Name() == "conference-local" {
			pluginCmd = c
		}
	}
	if pluginCmd == nil {
		t.Fatal("conference-local missing from root after overlay-first registration")
	}
	if pluginCmd.Hidden {
		t.Error("conference-local must stay visible (Hidden=false) after hideNonDirectRuntimeCommands")
	}
	if !bogus.Hidden {
		t.Error("bogus-not-a-product must be hidden by the visibility filter")
	}

	services := visibleMCPRootCommands(root)
	if !containsCommand(services, "conference-local") {
		t.Errorf("visibleMCPRootCommands missing conference-local: %v", commandNames(services))
	}
}

// TestHasOverlayToolOverrides exercises the split-decision helper used by
// loadPlugins to route stdio entries to overlay-first vs. legacy buckets.
func TestHasOverlayToolOverrides(t *testing.T) {
	cases := []struct {
		name    string
		overlay market.CLIOverlay
		want    bool
	}{
		{
			name:    "empty overlay",
			overlay: market.CLIOverlay{ID: "x", Command: "x"},
			want:    false,
		},
		{
			name: "overlay with overrides",
			overlay: market.CLIOverlay{
				ID:      "x",
				Command: "x",
				ToolOverrides: map[string]market.CLIToolOverride{
					"foo": {CLIName: "foo"},
				},
			},
			want: true,
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			p, sc := newOverlayFixture(t, "x", "x", tc.overlay)
			got := hasOverlayToolOverrides(p, sc)
			if got != tc.want {
				t.Errorf("hasOverlayToolOverrides = %v, want %v", got, tc.want)
			}
		})
	}
}