Skip to content

Commit fbdb2bc

Browse files
hengyunabchengyunabc
committed
feat(mcp): add view tool
1 parent ccbc933 commit fbdb2bc

File tree

2 files changed

+528
-0
lines changed

2 files changed

+528
-0
lines changed

‎core/src/main/java/com/taobao/arthas/core/mcp/tool/function/basic1000/ViewFileTool.java‎

Lines changed: 349 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,349 @@
1+
package com.taobao.arthas.core.mcp.tool.function.basic1000;
2+
3+
import com.fasterxml.jackson.core.type.TypeReference;
4+
import com.taobao.arthas.core.mcp.tool.function.AbstractArthasTool;
5+
import com.taobao.arthas.mcp.server.tool.ToolContext;
6+
import com.taobao.arthas.mcp.server.tool.annotation.Tool;
7+
import com.taobao.arthas.mcp.server.tool.annotation.ToolParam;
8+
import com.taobao.arthas.mcp.server.util.JsonParser;
9+
10+
import java.io.RandomAccessFile;
11+
import java.nio.charset.StandardCharsets;
12+
import java.nio.file.Files;
13+
import java.nio.file.Path;
14+
import java.nio.file.Paths;
15+
import java.util.*;
16+
17+
import static com.taobao.arthas.core.mcp.tool.function.StreamableToolUtils.createCompletedResponse;
18+
import static com.taobao.arthas.core.mcp.tool.function.StreamableToolUtils.createErrorResponse;
19+
20+
/**
21+
* ViewFile MCP Tool: 在允许目录内分段查看文件内容
22+
*/
23+
public class ViewFileTool extends AbstractArthasTool {
24+
25+
static final String ALLOWED_DIRS_ENV = "ARTHAS_MCP_VIEWFILE_ALLOWED_DIRS";
26+
27+
static final int DEFAULT_MAX_BYTES = 8192;
28+
static final int MAX_MAX_BYTES = 65536;
29+
30+
@Tool(
31+
name = "viewfile",
32+
description = "查看文件内容（仅允许在配置的目录白名单内查看），并支持 cursor/offset 分段读取，避免一次性返回大量内容。\n" +
33+
"默认允许目录：当前工作目录下的 arthas-output（若存在）、用户目录下的 ~/logs/（若存在）。\n" +
34+
"配置白名单目录：\n" +
35+
"- 环境变量: " + ALLOWED_DIRS_ENV + "=/path/a,/path/b\n" +
36+
"使用方式：\n" +
37+
"- 首次读取：传 path（可传 offset/maxBytes）\n" +
38+
"- 继续读取：传 cursor（由上一次返回结果提供）"
39+
)
40+
public String viewFile(
41+
@ToolParam(description = "文件路径（绝对路径或相对路径；相对路径会在允许目录下解析）。当提供 cursor 时可不传。", required = false)
42+
String path,
43+
44+
@ToolParam(description = "游标（上一段返回的 nextCursor），用于继续读取。提供 cursor 时会忽略 path/offset。", required = false)
45+
String cursor,
46+
47+
@ToolParam(description = "起始字节偏移量（默认 0）。", required = false)
48+
Long offset,
49+
50+
@ToolParam(description = "本次最多读取字节数（默认 8192，最大 65536）。", required = false)
51+
Integer maxBytes,
52+
53+
ToolContext toolContext
54+
) {
55+
try {
56+
List<Path> allowedRoots = loadAllowedRoots();
57+
if (allowedRoots.isEmpty()) {
58+
return JsonParser.toJson(createErrorResponse("viewfile 未配置允许目录白名单，且默认目录 arthas-output、~/logs/ 不可用。" +
59+
"请通过环境变量 " + ALLOWED_DIRS_ENV + "=/path/a,/path/b 进行配置。"));
60+
}
61+
62+
CursorRequest cursorRequest = parseCursorOrArgs(path, cursor, offset);
63+
Path targetFile = resolveAllowedFile(cursorRequest.path, allowedRoots);
64+
65+
int readMaxBytes = clampMaxBytes(maxBytes);
66+
long fileSize = Files.size(targetFile);
67+
68+
long requestedOffset = cursorRequest.offset;
69+
long effectiveOffset = adjustOffset(cursorRequest.cursorUsed, requestedOffset, fileSize);
70+
71+
byte[] bytes = readBytes(targetFile, effectiveOffset, readMaxBytes, fileSize);
72+
int safeLen = utf8SafeLength(bytes, bytes.length);
73+
String content = new String(bytes, 0, safeLen, StandardCharsets.UTF_8);
74+
75+
long nextOffset = effectiveOffset + safeLen;
76+
boolean eof = nextOffset >= fileSize;
77+
78+
Map<String, Object> result = new LinkedHashMap<>();
79+
result.put("path", targetFile.toString());
80+
result.put("fileSize", fileSize);
81+
result.put("requestedOffset", requestedOffset);
82+
result.put("startOffset", effectiveOffset);
83+
result.put("maxBytes", readMaxBytes);
84+
result.put("readBytes", safeLen);
85+
result.put("nextOffset", nextOffset);
86+
result.put("eof", eof);
87+
result.put("nextCursor", encodeCursor(targetFile.toString(), nextOffset));
88+
result.put("content", content);
89+
90+
if (cursorRequest.cursorUsed && requestedOffset > fileSize) {
91+
result.put("cursorReset", true);
92+
result.put("cursorResetReason", "offsetGreaterThanFileSize");
93+
}
94+
95+
return JsonParser.toJson(createCompletedResponse("ok", result));
96+
} catch (Exception e) {
97+
logger.error("viewfile error", e);
98+
return JsonParser.toJson(createErrorResponse("viewfile 执行失败: " + e.getMessage()));
99+
}
100+
}
101+
102+
private static final class CursorRequest {
103+
private final String path;
104+
private final long offset;
105+
private final boolean cursorUsed;
106+
107+
private CursorRequest(String path, long offset, boolean cursorUsed) {
108+
this.path = path;
109+
this.offset = offset;
110+
this.cursorUsed = cursorUsed;
111+
}
112+
}
113+
114+
private CursorRequest parseCursorOrArgs(String path, String cursor, Long offset) {
115+
if (cursor != null && !cursor.trim().isEmpty()) {
116+
CursorValue decoded = decodeCursor(cursor.trim());
117+
return new CursorRequest(decoded.path, decoded.offset, true);
118+
}
119+
if (path == null || path.trim().isEmpty()) {
120+
throw new IllegalArgumentException("必须提供 path 或 cursor");
121+
}
122+
if (offset != null && offset < 0) {
123+
throw new IllegalArgumentException("offset 不允许为负数");
124+
}
125+
long resolvedOffset = (offset != null) ? offset : 0L;
126+
return new CursorRequest(path.trim(), resolvedOffset, false);
127+
}
128+
129+
private static final class CursorValue {
130+
private final String path;
131+
private final long offset;
132+
133+
private CursorValue(String path, long offset) {
134+
this.path = path;
135+
this.offset = offset;
136+
}
137+
}
138+
139+
private CursorValue decodeCursor(String cursor) {
140+
try {
141+
byte[] jsonBytes = Base64.getUrlDecoder().decode(cursor);
142+
String json = new String(jsonBytes, StandardCharsets.UTF_8);
143+
144+
Map<String, Object> map = JsonParser.fromJson(json, new TypeReference<Map<String, Object>>() {});
145+
Object pathObj = map.get("path");
146+
Object offsetObj = map.get("offset");
147+
if (!(pathObj instanceof String) || ((String) pathObj).trim().isEmpty()) {
148+
throw new IllegalArgumentException("cursor 缺少 path");
149+
}
150+
if (!(offsetObj instanceof Number)) {
151+
throw new IllegalArgumentException("cursor 缺少 offset");
152+
}
153+
long offset = ((Number) offsetObj).longValue();
154+
if (offset < 0) {
155+
throw new IllegalArgumentException("cursor offset 不允许为负数");
156+
}
157+
return new CursorValue(((String) pathObj).trim(), offset);
158+
} catch (IllegalArgumentException e) {
159+
throw new IllegalArgumentException("cursor 解析失败: " + e.getMessage(), e);
160+
}
161+
}
162+
163+
private String encodeCursor(String path, long offset) {
164+
Map<String, Object> cursor = new LinkedHashMap<>();
165+
cursor.put("v", 1);
166+
cursor.put("path", path);
167+
cursor.put("offset", offset);
168+
String json = JsonParser.toJson(cursor);
169+
return Base64.getUrlEncoder().withoutPadding().encodeToString(json.getBytes(StandardCharsets.UTF_8));
170+
}
171+
172+
private List<Path> loadAllowedRoots() {
173+
String config = System.getenv(ALLOWED_DIRS_ENV);
174+
175+
List<Path> roots = new ArrayList<>();
176+
if (config != null && !config.trim().isEmpty()) {
177+
String[] parts = config.split(",");
178+
for (String part : parts) {
179+
String p = (part != null) ? part.trim() : "";
180+
if (p.isEmpty()) {
181+
continue;
182+
}
183+
try {
184+
Path root = Paths.get(p).toAbsolutePath().normalize();
185+
if (!Files.isDirectory(root)) {
186+
logger.warn("viewfile allowed dir ignored (not a directory): {}", root);
187+
continue;
188+
}
189+
roots.add(root.toRealPath());
190+
} catch (Exception e) {
191+
logger.warn("viewfile allowed dir ignored (invalid): {}", p, e);
192+
}
193+
}
194+
}
195+
196+
// 默认目录：arthas-output
197+
try {
198+
Path defaultRoot = Paths.get("arthas-output").toAbsolutePath().normalize();
199+
if (Files.isDirectory(defaultRoot)) {
200+
roots.add(defaultRoot.toRealPath());
201+
}
202+
} catch (Exception e) {
203+
logger.debug("viewfile default root ignored: arthas-output", e);
204+
}
205+
206+
// 默认目录：~/logs/
207+
try {
208+
Path userLogsRoot = Paths.get(System.getProperty("user.home"), "logs").toAbsolutePath().normalize();
209+
if (Files.isDirectory(userLogsRoot)) {
210+
roots.add(userLogsRoot.toRealPath());
211+
}
212+
} catch (Exception e) {
213+
logger.debug("viewfile default root ignored: ~/logs/", e);
214+
}
215+
216+
return deduplicate(roots);
217+
}
218+
219+
private static List<Path> deduplicate(List<Path> roots) {
220+
if (roots == null || roots.isEmpty()) {
221+
return Collections.emptyList();
222+
}
223+
LinkedHashSet<Path> set = new LinkedHashSet<>(roots);
224+
return new ArrayList<>(set);
225+
}
226+
227+
private Path resolveAllowedFile(String requestedPath, List<Path> allowedRoots) throws Exception {
228+
Path req = Paths.get(requestedPath);
229+
if (req.isAbsolute()) {
230+
Path real = req.toRealPath();
231+
assertRegularFile(real);
232+
if (!isUnderAllowedRoot(real, allowedRoots)) {
233+
throw new IllegalArgumentException("文件不在允许目录白名单内: " + requestedPath);
234+
}
235+
return real;
236+
}
237+
238+
for (Path root : allowedRoots) {
239+
Path candidate = root.resolve(req).normalize();
240+
if (!candidate.startsWith(root)) {
241+
continue;
242+
}
243+
if (!Files.exists(candidate)) {
244+
continue;
245+
}
246+
Path real = candidate.toRealPath();
247+
if (!real.startsWith(root)) {
248+
continue;
249+
}
250+
assertRegularFile(real);
251+
return real;
252+
}
253+
throw new IllegalArgumentException("文件不存在或不在允许目录白名单内: " + requestedPath);
254+
}
255+
256+
private static void assertRegularFile(Path file) {
257+
if (!Files.isRegularFile(file)) {
258+
throw new IllegalArgumentException("不是普通文件: " + file);
259+
}
260+
}
261+
262+
private static boolean isUnderAllowedRoot(Path file, List<Path> allowedRoots) {
263+
for (Path root : allowedRoots) {
264+
if (file.startsWith(root)) {
265+
return true;
266+
}
267+
}
268+
return false;
269+
}
270+
271+
private static int clampMaxBytes(Integer maxBytes) {
272+
int value = (maxBytes != null && maxBytes > 0) ? maxBytes : DEFAULT_MAX_BYTES;
273+
return Math.min(value, MAX_MAX_BYTES);
274+
}
275+
276+
private static long adjustOffset(boolean cursorUsed, long requestedOffset, long fileSize) {
277+
if (requestedOffset < 0) {
278+
throw new IllegalArgumentException("offset 不允许为负数");
279+
}
280+
if (requestedOffset <= fileSize) {
281+
return requestedOffset;
282+
}
283+
return cursorUsed ? 0L : fileSize;
284+
}
285+
286+
private static byte[] readBytes(Path file, long offset, int maxBytes, long fileSize) throws Exception {
287+
if (offset < 0 || offset > fileSize) {
288+
return new byte[0];
289+
}
290+
long remaining = fileSize - offset;
291+
int toRead = (int) Math.min(maxBytes, Math.max(0, remaining));
292+
if (toRead <= 0) {
293+
return new byte[0];
294+
}
295+
296+
byte[] buf = new byte[toRead];
297+
int read;
298+
try (RandomAccessFile raf = new RandomAccessFile(file.toFile(), "r")) {
299+
raf.seek(offset);
300+
read = raf.read(buf);
301+
}
302+
if (read <= 0) {
303+
return new byte[0];
304+
}
305+
return Arrays.copyOf(buf, read);
306+
}
307+
308+
/**
309+
* 避免把 UTF-8 多字节字符截断在末尾，导致展示出现大量 �。
310+
*/
311+
static int utf8SafeLength(byte[] bytes, int length) {
312+
if (bytes == null || length <= 0) {
313+
return 0;
314+
}
315+
int lastIndex = length - 1;
316+
int lastByte = bytes[lastIndex] & 0xFF;
317+
if ((lastByte & 0x80) == 0) {
318+
return length;
319+
}
320+
321+
int i = lastIndex;
322+
int continuation = 0;
323+
while (i >= 0 && (bytes[i] & 0xC0) == 0x80) {
324+
continuation++;
325+
i--;
326+
}
327+
if (i < 0) {
328+
return Math.max(0, length - continuation);
329+
}
330+
331+
int lead = bytes[i] & 0xFF;
332+
int expectedLen;
333+
if ((lead & 0xE0) == 0xC0) {
334+
expectedLen = 2;
335+
} else if ((lead & 0xF0) == 0xE0) {
336+
expectedLen = 3;
337+
} else if ((lead & 0xF8) == 0xF0) {
338+
expectedLen = 4;
339+
} else {
340+
return length;
341+
}
342+
343+
int actualLen = continuation + 1;
344+
if (actualLen < expectedLen) {
345+
return i;
346+
}
347+
return length;
348+
}
349+
}

0 commit comments

Comments
 (0)