package vip.mate.tool.disclosure; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; import org.springframework.ai.support.ToolCallbacks; import org.springframework.ai.tool.ToolCallback; import org.springframework.ai.tool.annotation.Tool; import org.springframework.test.util.ReflectionTestUtils; import vip.mate.agent.AgentToolSet; import vip.mate.tool.ToolRegistry; import vip.mate.tool.mcp.model.McpServerEntity; import vip.mate.tool.mcp.service.McpServerService; import vip.mate.tool.model.AvailableToolDTO; import vip.mate.tool.model.ToolEntity; import vip.mate.tool.service.AvailableToolService; import vip.mate.tool.service.ToolService; import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.Set; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; class ToolDisclosureServiceTest { /** Fixture beans whose @Tool function names drive tier resolution. */ static class Tools { @Tool(description = "text to image") public String image_generate() { return ""; } @Tool(description = "a plain core tool") public String my_core_tool() { return ""; } } /** Fixture whose class simple name is {@code ImageGenerateTool} and function * name is {@code image_generate} — mirrors the real builtin's name skew so * the class-name → function-name bridge can be tested. */ static class ImageGenerateTool { @Tool(description = "text to image") public String image_generate() { return ""; } } /** Global tool set the bridge resolves DB class/bean names against. */ private static AgentToolSet globalSet() { Object t1 = new Tools(); Object t2 = new ImageGenerateTool(); List cbs = new ArrayList<>(); cbs.addAll(List.of(ToolCallbacks.from(t1))); cbs.addAll(List.of(ToolCallbacks.from(t2))); Map beanNames = Map.of(t1, "tools", t2, "imageGenerateTool"); return AgentToolSet.fromCallbacks(List.of(t1, t2), cbs, beanNames::get); } private static ToolEntity toolRow(String name, String type, String tier) { ToolEntity t = new ToolEntity(); t.setName(name); t.setToolType(type); t.setDisclosureTier(tier); return t; } private static McpServerEntity server(Long id, String name, String tier) { McpServerEntity s = new McpServerEntity(); s.setId(id); s.setName(name); s.setDisclosureTier(tier); return s; } private static AvailableToolDTO mcpDto(String name, Long serverId) { return AvailableToolDTO.builder().source("mcp").providerId(serverId).name(name).build(); } private DefaultToolDisclosureService service(List tools, List servers, List available) { ToolService ts = mock(ToolService.class); McpServerService ms = mock(McpServerService.class); AvailableToolService as = mock(AvailableToolService.class); ToolRegistry tr = mock(ToolRegistry.class); lenient().when(ts.listTools()).thenReturn(tools); lenient().when(ms.listAll()).thenReturn(servers); lenient().when(as.listAvailable()).thenReturn(available); lenient().when(tr.getEnabledToolSet()).thenReturn(globalSet()); return new DefaultToolDisclosureService(ts, ms, as, tr); } @Test @DisplayName("meta-tools enable_tool / load_skill are always core") void metaToolsAlwaysCore() { var svc = service(List.of(toolRow("enable_tool", "builtin", "extension")), List.of(), List.of()); assertEquals(DisclosureTier.CORE, svc.resolveTierByName("enable_tool")); assertEquals(DisclosureTier.CORE, svc.resolveTierByName("load_skill")); } @Test @DisplayName("generative tools default to extension even without a DB row") void generativeDefaultsExtension() { var svc = service(List.of(), List.of(), List.of()); assertEquals(DisclosureTier.EXTENSION, svc.resolveTierByName("image_generate")); assertEquals(DisclosureTier.EXTENSION, svc.resolveTierByName("browser_use")); } @Test @DisplayName("unknown tools default to core (conservative)") void unknownDefaultsCore() { var svc = service(List.of(), List.of(), List.of()); assertEquals(DisclosureTier.CORE, svc.resolveTierByName("memory_recall")); } @Test @DisplayName("mate_tool.disclosure_tier overrides the code default") void dbRowOverrides() { var svc = service(List.of(toolRow("my_core_tool", "builtin", "extension")), List.of(), List.of()); assertEquals(DisclosureTier.EXTENSION, svc.resolveTierByName("my_core_tool")); } @Test @DisplayName("DB tier stored by Java class name bridges to the runtime function name") void dbTierBridgesClassNameToFunctionName() { // mate_tool.name = class name; resolveTier is queried by function name. var hidden = service(List.of(toolRow("ImageGenerateTool", "builtin", "extension")), List.of(), List.of()); assertEquals(DisclosureTier.EXTENSION, hidden.resolveTierByName("image_generate")); // Admin un-hides it by setting the row to core; the DB value must win over // the code-level extension default. var unhidden = service(List.of(toolRow("ImageGenerateTool", "builtin", "core")), List.of(), List.of()); assertEquals(DisclosureTier.CORE, unhidden.resolveTierByName("image_generate")); } @Test @DisplayName("MCP tool tier follows its owning server") void mcpFollowsServer() { var extSvc = service(List.of(), List.of(server(7L, "github", "extension")), List.of(mcpDto("mcp_github_create_issue", 7L))); assertEquals(DisclosureTier.EXTENSION, extSvc.resolveTierByName("mcp_github_create_issue")); var coreSvc = service(List.of(), List.of(server(7L, "github", "core")), List.of(mcpDto("mcp_github_create_issue", 7L))); assertEquals(DisclosureTier.CORE, coreSvc.resolveTierByName("mcp_github_create_issue")); } @Test @DisplayName("MCP tool whose server has no tier set defaults to core (visible)") void mcpDefaultsCoreWhenServerTierUnset() { var svc = service(List.of(), List.of(server(7L, "github", null)), List.of(mcpDto("mcp_github_create_issue", 7L))); assertEquals(DisclosureTier.CORE, svc.resolveTierByName("mcp_github_create_issue")); } @Test @DisplayName("split partitions into active (core + enabled) and the full extension catalog") void splitPartitions() { var svc = service(List.of(), List.of(), List.of()); AgentToolSet set = AgentToolSet.fromCallbacks(List.of(new Tools()), List.of(ToolCallbacks.from(new Tools()))); var noneEnabled = svc.split(set, Set.of()); assertEquals(List.of("my_core_tool"), names(noneEnabled.activeCallbacks())); assertEquals(List.of("image_generate"), names(noneEnabled.extensionCatalog())); var imgEnabled = svc.split(set, Set.of("image_generate")); assertTrue(names(imgEnabled.activeCallbacks()).contains("image_generate")); assertTrue(names(imgEnabled.activeCallbacks()).contains("my_core_tool")); assertEquals(List.of("image_generate"), names(imgEnabled.extensionCatalog())); } @Test @DisplayName("legacy mode advertises everything and renders no catalog") void legacyMode() { var svc = service(List.of(), List.of(), List.of()); ReflectionTestUtils.setField(svc, "disclosureMode", "legacy"); AgentToolSet set = AgentToolSet.fromCallbacks(List.of(new Tools()), List.of(ToolCallbacks.from(new Tools()))); var split = svc.split(set, Set.of()); assertEquals(2, split.activeCallbacks().size()); assertTrue(split.extensionCatalog().isEmpty()); assertEquals("", svc.renderExtensionCatalog(set, 8192)); assertEquals(DisclosureTier.CORE, svc.resolveTierByName("image_generate")); } @Test @DisplayName("renderExtensionCatalog lists extension tools under a heading") void rendersCatalog() { var svc = service(List.of(), List.of(), List.of()); AgentToolSet set = AgentToolSet.fromCallbacks(List.of(new Tools()), List.of(ToolCallbacks.from(new Tools()))); String catalog = svc.renderExtensionCatalog(set, 8192); assertTrue(catalog.contains("## Extension Tools")); assertTrue(catalog.contains("image_generate")); assertTrue(catalog.contains("enable_tool")); assertFalse(catalog.contains("my_core_tool"), "core tools must not appear in the extension catalog"); } private static List names(List cbs) { return cbs.stream().map(c -> c.getToolDefinition().name()).toList(); } }