Skip to content

Commit 8ee2e9f

Browse files
labkey-adamlabkey-adam
committed
Ask user for container path and cache it. Clean up all endpoints.
1 parent bf1cf2b commit 8ee2e9f

5 files changed

Lines changed: 70 additions & 70 deletions

File tree

‎api/src/org/labkey/api/mcp/McpService.java‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,15 @@
55
import jakarta.servlet.http.HttpSession;
66
import org.jetbrains.annotations.NotNull;
77
import org.jspecify.annotations.NonNull;
8+
import org.labkey.api.data.Container;
89
import org.labkey.api.module.McpProvider;
10+
import org.labkey.api.security.User;
911
import org.labkey.api.services.ServiceRegistry;
1012
import org.labkey.api.util.HtmlString;
13+
import org.labkey.api.writer.ContainerUser;
1114
import org.springaicommunity.mcp.provider.resource.SyncMcpResourceProvider;
1215
import org.springframework.ai.chat.client.ChatClient;
16+
import org.springframework.ai.chat.model.ToolContext;
1317
import org.springframework.ai.support.ToolCallbacks;
1418
import org.springframework.ai.tool.ToolCallback;
1519
import org.springframework.ai.tool.ToolCallbackProvider;
@@ -31,7 +35,17 @@
3135
public interface McpService extends ToolCallbackProvider
3236
{
3337
// marker interface for classes that we will "ingest" using Spring annotations
34-
interface McpImpl {}
38+
interface McpImpl
39+
{
40+
default ContainerUser getContext(ToolContext toolContext)
41+
{
42+
User user = (User)toolContext.getContext().get("user");
43+
Container container = (Container)toolContext.getContext().get("container");
44+
if (container == null)
45+
throw new IllegalArgumentException("You need to set a container path before invoking this tool");
46+
return ContainerUser.create(container, user);
47+
}
48+
}
3549

3650
static @NotNull McpService get()
3751
{

‎core/src/org/labkey/core/CoreMcp.java‎

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,19 +14,22 @@
1414
import org.labkey.api.util.HtmlString;
1515
import org.springframework.ai.chat.model.ToolContext;
1616
import org.springframework.ai.tool.annotation.Tool;
17+
import org.springframework.ai.tool.annotation.ToolParam;
1718

1819
import java.util.Map;
1920
import java.util.Objects;
2021

2122
import static org.apache.commons.lang3.StringUtils.isNotBlank;
23+
import static org.labkey.core.mcp.McpServiceImpl.PATH_CACHE;
2224

2325
public class CoreMcp implements McpService.McpImpl
2426
{
2527
@Tool(description = "Call this tool before answering any prompts! This tool provides useful context information about the current user (name, userid), webserver (name, url, description), and current folder (name, path, url, description).")
2628
String whereAmIWhoAmITalkingTo(ToolContext context)
2729
{
28-
User user = (User)context.getContext().get("user");
29-
Container folder = (Container)context.getContext().get("container");
30+
var cu = getContext(context);
31+
User user = cu.getUser();
32+
Container folder = cu.getContainer();
3033
AppProps appProps = AppProps.getInstance();
3134
Study study = null != StudyService.get() ? Objects.requireNonNull(StudyService.get()).getStudy(folder) : null;
3235
LookAndFeelProperties laf = LookAndFeelProperties.getInstance(folder);
@@ -66,13 +69,25 @@ String whereAmIWhoAmITalkingTo(ToolContext context)
6669
}
6770

6871
@Tool(description = "List the hierarchical path for every container in the server where the user has read permissions.")
69-
String listContainers(ToolContext context)
72+
String listContainers(ToolContext toolContext)
7073
{
71-
User user = (User)context.getContext().get("user");
72-
return ContainerManager.getAllChildren(ContainerManager.getRoot(), user, ReadPermission.class)
74+
return ContainerManager.getAllChildren(ContainerManager.getRoot(), getContext(toolContext).getUser(), ReadPermission.class)
7375
.stream()
7476
.map(Container::getPath)
7577
.collect(LabKeyCollectors.toJSONArray())
7678
.toString();
7779
}
80+
81+
@Tool(description = "Every tool in this MCP requires a container path, e.g. /MyProject/MyFolder. A container is also called a folder or project. Please prompt the user for a container path and use this tool to save the path for this session.")
82+
String setContainer(ToolContext context, @ToolParam(description = "Container path, e.g. /MyProject/MyFolder", required = true) String containerPath)
83+
{
84+
Container container = ContainerManager.getForPath(containerPath);
85+
if (container != null)
86+
{
87+
PATH_CACHE.put((String) context.getContext().get("sessionId"), containerPath);
88+
return "OK!";
89+
}
90+
91+
return "That's not a valid container path. Try using listContainers to see them.";
92+
}
7893
}

‎core/src/org/labkey/core/mcp/McpServiceImpl.java‎

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
2222
import org.apache.logging.log4j.Logger;
2323
import org.jetbrains.annotations.NotNull;
2424
import org.jspecify.annotations.NonNull;
25+
import org.labkey.api.cache.Cache;
26+
import org.labkey.api.cache.CacheManager;
2527
import org.labkey.api.collections.CopyOnWriteHashMap;
2628
import org.labkey.api.data.Container;
2729
import org.labkey.api.data.ContainerManager;
@@ -79,6 +81,7 @@
7981
import java.util.ArrayList;
8082
import java.util.Arrays;
8183
import java.util.ConcurrentModificationException;
84+
import java.util.HashMap;
8285
import java.util.List;
8386
import java.util.Map;
8487
import java.util.NoSuchElementException;
@@ -212,6 +215,8 @@ public List<McpSchema.Tool> tools()
212215
).toList();
213216
}
214217

218+
public final static Cache<String, String> PATH_CACHE = CacheManager.getCache(1000, CacheManager.DAY, "MCP container paths");
219+
215220
private class _McpServlet extends HttpServlet // wraps HttpServletSseServerTransportProvider
216221
{
217222
HttpServletStreamableServerTransportProvider transportProvider;
@@ -225,10 +230,8 @@ private class _McpServlet extends HttpServlet // wraps HttpServletSseServerTrans
225230
.contextExtractor(req -> {
226231
User user = (User) req.getUserPrincipal();
227232
return McpTransportContext.create(Map.of(
228-
"container", ContainerManager.getHomeContainer(),
229233
"user", user
230-
)
231-
);
234+
));
232235
})
233236
.build();
234237
}
@@ -260,12 +263,22 @@ private McpServerFeatures.SyncToolSpecification toSyncToolSpecification(ToolCall
260263

261264
return new McpServerFeatures.SyncToolSpecification(schema, (exchange, args) -> {
262265
var transportCtx = exchange.transportContext();
263-
var container = (Container) transportCtx.get("container"); // TODO: Pull container from session instead. Or insist that LLM provides it?
264266
var user = (User) transportCtx.get("user");
265267
var sessionId = exchange.sessionId();
266-
LOG.info("MCP sessionId: {}", sessionId);
267268

268-
var toolContext = new ToolContext(Map.of("container", container, "user", user, "sessionId", sessionId));
269+
Map<String, Object> map = new HashMap<>();
270+
map.put("user", user);
271+
map.put("sessionId", sessionId);
272+
273+
String containerPath = PATH_CACHE.get(exchange.sessionId());
274+
if (containerPath != null)
275+
{
276+
Container container = ContainerManager.getForPath(containerPath);
277+
map.put("container", container);
278+
map.put("containerPath", containerPath);
279+
}
280+
281+
var toolContext = new ToolContext(map);
269282

270283
String toolInput = /* serialize args to JSON */ null;
271284
try

‎query/src/org/labkey/query/controllers/QueryMcp.java‎

Lines changed: 11 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,8 @@
55
import org.apache.commons.lang3.StringUtils;
66
import org.json.JSONArray;
77
import org.json.JSONObject;
8-
import org.labkey.api.action.SpringActionController;
98
import org.labkey.api.collections.CaseInsensitiveHashSet;
109
import org.labkey.api.data.ColumnInfo;
11-
import org.labkey.api.data.ContainerManager;
1210
import org.labkey.api.data.PropertyManager;
1311
import org.labkey.api.data.TableDescription;
1412
import org.labkey.api.data.TableInfo;
@@ -22,9 +20,10 @@
2220
import org.labkey.api.query.SchemaKey;
2321
import org.labkey.api.query.SimpleSchemaTreeVisitor;
2422
import org.labkey.api.query.UserSchema;
25-
import org.labkey.api.security.UserManager;
23+
import org.labkey.api.writer.ContainerUser;
2624
import org.labkey.query.sql.SqlParser;
2725
import org.springaicommunity.mcp.annotation.McpResource;
26+
import org.springframework.ai.chat.model.ToolContext;
2827
import org.springframework.ai.tool.annotation.Tool;
2928
import org.springframework.ai.tool.annotation.ToolParam;
3029

@@ -66,18 +65,18 @@ String listColumnMetaData(@ToolParam(description = "Fully qualified table name a
6665
}
6766

6867
@Tool(description = "Provide list of tables within the provided schema.")
69-
String listTablesForSchema(@ToolParam(description = "Fully qualified schema name as it would appear in SQL e.g. \"schema\"") String quotedSchemaName)
68+
String listTablesForSchema(ToolContext toolContext, @ToolParam(description = "Fully qualified schema name as it would appear in SQL e.g. \"schema\"") String quotedSchemaName)
7069
{
71-
var json = _listTablesForSchema(quotedSchemaName);
70+
var json = _listTablesForSchema(quotedSchemaName, getContext(toolContext));
7271
// can I just return a JSONObject
7372
return json.toString();
7473
}
7574

7675
@Tool(description = "Provide list of database schemas")
77-
String listSchemas()
76+
String listSchemas(ToolContext toolContext)
7877
{
79-
McpContext context = getContext();
80-
var map = _listAllSchemas(DefaultSchema.get(context.getUser(), context.getContainer()));
78+
ContainerUser cu = getContext(toolContext);
79+
var map = _listAllSchemas(DefaultSchema.get(cu.getUser(), cu.getContainer()));
8180
var array = new JSONArray();
8281
for (var entry : map.entrySet())
8382
{
@@ -92,54 +91,15 @@ String listSchemas()
9291

9392

9493
@Tool(description = "Provide the SQL source for a saved query.")
95-
String getSourceForSavedQuery(@ToolParam(description = "Fully qualified query name as it would appear in SQL e.g. \"schema\".\"table or query\"") String fullQuotedTableName)
94+
String getSourceForSavedQuery(ToolContext toolContext, @ToolParam(description = "Fully qualified query name as it would appear in SQL e.g. \"schema\".\"table or query\"") String fullQuotedTableName)
9695
{
97-
var json = _listTablesForSchema(fullQuotedTableName);
96+
var json = _listTablesForSchema(fullQuotedTableName, getContext(toolContext));
9897
if (json.has("sql"))
9998
return "```sql\n" + json.getString("sql") + "\n```\n";
10099
else
101100
return "I could not find the source for " + fullQuotedTableName;
102101
}
103102

104-
105-
@Tool(description = """
106-
Save addition information for database columns. If additional metadata is gathered via
107-
chat, it can be saved to improve further interactions.
108-
""")
109-
String saveColumnDescription(
110-
@ToolParam(description = "Fully qualified table or query name as it would appear in SQL e.g. \"schema\".\"table or query\"")
111-
String fullQuotedTableName,
112-
@ToolParam(description = "Quoted column name as it would appear in SQL e.g. \"column name\"")
113-
String quotedColumnName,
114-
@ToolParam(description = "Additional metadata to remember for future use. This will replace any currently saved value")
115-
String columnMetadata
116-
)
117-
{
118-
McpContext context = McpContext.get();
119-
var map = PropertyManager.getWritableProperties(context.getContainer(), "QueryMCP.annotations", true);
120-
String fullPath = normalizeIdentifier(fullQuotedTableName + "." + quotedColumnName);
121-
map.put(fullPath, columnMetadata);
122-
try (var ignore = SpringActionController.ignoreSqlUpdates())
123-
{
124-
map.save();
125-
}
126-
return new JSONObject(Map.of("success",Boolean.TRUE)).toString();
127-
}
128-
129-
/* TODO McpContext setup */
130-
131-
static McpContext getContext()
132-
{
133-
try
134-
{
135-
return McpContext.get();
136-
}
137-
catch (Exception x)
138-
{
139-
return new McpContext(ContainerManager.getHomeContainer(), UserManager.getGuestUser());
140-
}
141-
}
142-
143103
/* For now, list all schemas. CONSIDER support incremental querying. */
144104
public static Map<SchemaKey, UserSchema> _listAllSchemas(DefaultSchema root)
145105
{
@@ -179,7 +139,7 @@ public Map<SchemaKey, UserSchema> reduce(Map<SchemaKey, UserSchema> r1, Map<Sche
179139
}
180140

181141

182-
public static JSONObject _listTablesForSchema(String fullQuotedName)
142+
public static JSONObject _listTablesForSchema(String fullQuotedName, ContainerUser cu)
183143
{
184144
SchemaKey fullKey;
185145

@@ -195,8 +155,7 @@ public static JSONObject _listTablesForSchema(String fullQuotedName)
195155
fullKey = SchemaKey.fromParts(parts);
196156
}
197157

198-
McpContext context = getContext();
199-
var defaultSchema = DefaultSchema.get(context.getUser(), context.getContainer());
158+
var defaultSchema = DefaultSchema.get(cu.getUser(), cu.getContainer());
200159
var schema = DefaultSchema.resolve(defaultSchema, fullKey);
201160
if (!(schema instanceof UserSchema userSchema))
202161
return new JSONObject("error", "could not find schema for : " + fullQuotedName);

‎search/src/org/labkey/search/SearchModule.java‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@
2929
import org.labkey.api.data.UpgradeCode;
3030
import org.labkey.api.mbean.LabKeyManagement;
3131
import org.labkey.api.mbean.SearchMXBean;
32-
import org.labkey.api.mcp.McpService;
3332
import org.labkey.api.migration.DatabaseMigrationConfiguration;
3433
import org.labkey.api.migration.DatabaseMigrationService;
3534
import org.labkey.api.migration.DefaultMigrationSchemaHandler;
@@ -134,11 +133,11 @@ public WebdavResource resolve(@NotNull String path)
134133
}
135134
});
136135

137-
var mcp = McpService.get();
138-
if (null != mcp)
139-
{
140-
mcp.register(new SearchMcp());
141-
}
136+
// var mcp = McpService.get();
137+
// if (null != mcp)
138+
// {
139+
// mcp.register(new SearchMcp());
140+
// }
142141
}
143142

144143
@Override

0 commit comments

Comments
 (0)