Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 91 additions & 29 deletions src/main/java/run/halo/mcpserver/tools/AttachmentTools.java
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,10 @@
import static org.springframework.data.domain.Sort.Order.asc;
import static org.springframework.data.domain.Sort.Order.desc;

import io.modelcontextprotocol.spec.McpSchema;
import java.net.MalformedURLException;
import java.net.URI;
import java.net.URL;
import java.util.Base64;
import java.util.List;
import java.util.Map;
Expand Down Expand Up @@ -38,10 +42,17 @@ class AttachmentTools extends ToolSupport implements ToolGroup {
static final String LIST = "halo_list_attachments";
static final String GET = "halo_get_attachment";
static final String UPLOAD = "halo_upload_attachment";
static final String UPLOAD_FROM_URL = "halo_upload_attachment_from_url";
static final String DELETE = "halo_delete_attachment";

private static final int MAX_CONTENT_BYTES = 7 * 1024 * 1024;
private static final int MAX_ENCODED_CHARS = (MAX_CONTENT_BYTES + 2) / 3 * 4;
private static final McpSchema.ToolAnnotations CREATE_OPEN_WORLD = McpSchema.ToolAnnotations.builder()
.readOnlyHint(false)
.destructiveHint(false)
.idempotentHint(false)
.openWorldHint(true)
.build();

private final ReactiveExtensionClient client;
private final AttachmentService attachmentService;
Expand All @@ -60,7 +71,7 @@ class AttachmentTools extends ToolSupport implements ToolGroup {

@Override
public List<BuiltInTool> tools() {
return List.of(listTool(), getTool(), uploadTool(), deleteTool());
return List.of(listTool(), getTool(), uploadTool(), uploadFromUrlTool(), deleteTool());
}

Mono<ToolPayload> list(Map<String, Object> arguments) {
Expand Down Expand Up @@ -100,37 +111,50 @@ Mono<ToolPayload> upload(Map<String, Object> arguments) {
.subscribeOn(Schedulers.boundedElastic())
.flatMap(bytes -> {
var buffer = DefaultDataBufferFactory.sharedInstance.wrap(bytes);
return attachmentService
.upload(
return uploadedAttachment(
attachmentService.upload(
config.policyName(),
config.groupName(),
filename,
Flux.just(buffer),
mediaType)
.switchIfEmpty(Mono.error(new McpToolException(
"ATTACHMENT_UNAVAILABLE", "Halo did not create the attachment")))
.flatMap(attachment -> attachmentService
.getPermalink(attachment)
.doOnNext(permalink -> {
if (attachment.getStatus() == null) {
attachment.setStatus(new Attachment.AttachmentStatus());
}
attachment.getStatus().setPermalink(permalink.toString());
})
.onErrorResume(error -> {
log.warn("Failed to resolve permalink for uploaded attachment {}",
filename, error);
return Mono.empty();
})
.thenReturn(attachment))
.map(attachment -> payload(
ContentPayloads.attachment(attachment),
"Uploaded attachment " + filename));
mediaType),
filename);
})
.doFinally(ignored -> reservation.close());
});
}

Mono<ToolPayload> uploadFromUrl(Map<String, Object> arguments) {
var url = httpUrl(requiredString(arguments, "url"));
var requestedFilename = optionalString(arguments, "filename", null);
var filename = requestedFilename == null ? null : safeFilename(requestedFilename);
return attachmentConfig().flatMap(config -> uploadedAttachment(
attachmentService.uploadFromUrl(url, config.policyName(), config.groupName(), filename),
filename));
}

private Mono<ToolPayload> uploadedAttachment(Mono<Attachment> upload, String filename) {
return upload.switchIfEmpty(Mono.error(new McpToolException(
"ATTACHMENT_UNAVAILABLE", "Halo did not create the attachment")))
.flatMap(attachment -> attachmentService
.getPermalink(attachment)
.doOnNext(permalink -> {
if (attachment.getStatus() == null) {
attachment.setStatus(new Attachment.AttachmentStatus());
}
attachment.getStatus().setPermalink(permalink.toString());
})
.onErrorResume(error -> {
log.warn("Failed to resolve permalink for uploaded attachment {}", filename, error);
return Mono.empty();
})
.thenReturn(attachment))
.map(attachment -> {
var resultName = filename == null ? attachment.getMetadata().getName() : filename;
return payload(ContentPayloads.attachment(attachment), "Uploaded attachment " + resultName);
});
}

private static byte[] decode(String encoded) {
final byte[] bytes;
try {
Expand Down Expand Up @@ -206,12 +230,7 @@ private BuiltInTool uploadTool() {
"ATTACHMENT",
objectSchema(
map(
"filename", map(
"type", "string",
"minLength", 1,
"maxLength", 255,
"pattern", "^[^/\\\\\\x00-\\x1F\\x7F]+$",
"description", "File name without a directory path."),
"filename", filenameSchema(),
"mediaType", stringSchema("IANA media type; defaults to application/octet-stream."),
"contentBase64",
map(
Expand All @@ -226,6 +245,24 @@ private BuiltInTool uploadTool() {
this::upload);
}

private BuiltInTool uploadFromUrlTool() {
return tool(
UPLOAD_FROM_URL,
"Upload Halo attachment from URL",
"Transfer an absolute HTTP or HTTPS URL using the Console attachment policy and group from Halo system settings.",
"从 URL 上传附件",
"使用 Halo 系统设置中的 Console 附件策略和分组,将 HTTP 或 HTTPS URL 转存为附件。",
"ATTACHMENT",
objectSchema(
map(
"url", stringSchema("Absolute HTTP or HTTPS URL."),
"filename", filenameSchema()),
List.of("url")),
ContentPayloads.attachmentSchema(),
CREATE_OPEN_WORLD,
this::uploadFromUrl);
}

private BuiltInTool deleteTool() {
return tool(
DELETE,
Expand Down Expand Up @@ -260,6 +297,31 @@ private static MediaType mediaType(Object value) {
}
}

private static URL httpUrl(String value) {
try {
var uri = URI.create(value);
var scheme = uri.getScheme();
if (!uri.isAbsolute()
|| uri.getHost() == null
|| !("http".equalsIgnoreCase(scheme) || "https".equalsIgnoreCase(scheme))) {
throw new IllegalArgumentException();
}
return uri.toURL();
} catch (IllegalArgumentException | MalformedURLException error) {
throw new McpToolException(
"INVALID_ARGUMENT", "url must be an absolute http or https URL", error);
}
}

private static Map<String, Object> filenameSchema() {
return map(
"type", "string",
"minLength", 1,
"maxLength", 255,
"pattern", "^[^/\\\\\\x00-\\x1F\\x7F]+$",
"description", "File name without a directory path.");
}

private static String safeFilename(String filename) {
if (filename.length() > 255 || filename.contains("/") || filename.contains("\\")
|| filename.chars().anyMatch(Character::isISOControl)) {
Expand Down
87 changes: 87 additions & 0 deletions src/test/java/run/halo/mcpserver/tools/AttachmentToolsTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,13 @@
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.ArgumentMatchers.isNull;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import java.net.URI;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.Map;
Expand Down Expand Up @@ -40,6 +42,87 @@ class AttachmentToolsTest {
@Mock
McpAuthorization authorization;

@Test
void advertisesUrlUploadInTheAttachmentCatalog() {
var tools = new AttachmentTools(
client, attachmentService, new AttachmentUploadLimiter(), authorization);

var upload = tools.tools().stream()
.filter(tool -> tool.specification().tool().name().equals("halo_upload_attachment_from_url"))
.findFirst()
.orElseThrow();
var protocolTool = upload.protocolTool();
var schema = protocolTool.inputSchema();
var properties = (Map<?, ?>) schema.get("properties");

assertThat(upload.category()).isEqualTo("ATTACHMENT");
assertThat(schema.get("required")).isEqualTo(java.util.List.of("url"));
assertThat(properties.get("url")).isNotNull();
assertThat(properties.get("filename")).isNotNull();
assertThat(protocolTool.annotations().readOnlyHint()).isFalse();
assertThat(protocolTool.annotations().destructiveHint()).isFalse();
assertThat(protocolTool.annotations().idempotentHint()).isFalse();
assertThat(protocolTool.annotations().openWorldHint()).isTrue();
}

@Test
void transfersUrlWithConsoleConfigAndReturnsResolvedPermalink() throws Exception {
var tools = new AttachmentTools(
client, attachmentService, new AttachmentUploadLimiter(), authorization);
stubAttachmentConfig();
var source = URI.create("https://example.com/source.png").toURL();
var attachment = new Attachment();
attachment.setMetadata(ToolSupport.metadata("stored.png"));
when(attachmentService.uploadFromUrl(source, "local", "default", "photo.png"))
.thenReturn(Mono.just(attachment));
when(attachmentService.getPermalink(attachment))
.thenReturn(Mono.just(URI.create("https://example.com/stored.png")));

StepVerifier.create(tools.uploadFromUrl(Map.of(
"url", source.toString(),
"filename", "photo.png")))
.assertNext(payload -> {
assertThat(payload.summary()).contains("photo.png");
assertThat(payload.data().toString()).contains("https://example.com/stored.png");
})
.verifyComplete();

verify(attachmentService).uploadFromUrl(source, "local", "default", "photo.png");
}

@Test
void omitsFilenameWhenTransferringUrl() throws Exception {
var tools = new AttachmentTools(
client, attachmentService, new AttachmentUploadLimiter(), authorization);
stubAttachmentConfig();
var source = URI.create("http://example.com/source.png").toURL();
var attachment = new Attachment();
attachment.setMetadata(ToolSupport.metadata("source.png"));
when(attachmentService.uploadFromUrl(eq(source), eq("local"), eq("default"), isNull()))
.thenReturn(Mono.just(attachment));
when(attachmentService.getPermalink(attachment))
.thenReturn(Mono.just(URI.create("https://example.com/source.png")));

StepVerifier.create(tools.uploadFromUrl(Map.of("url", source.toString())))
.assertNext(payload -> assertThat(payload.summary()).contains("source.png"))
.verifyComplete();

verify(attachmentService).uploadFromUrl(eq(source), eq("local"), eq("default"), isNull());
}

@Test
void rejectsNonHttpAndRelativeUrlsBeforeTransfer() {
var tools = new AttachmentTools(
client, attachmentService, new AttachmentUploadLimiter(), authorization);

for (var url : java.util.List.of("ftp://example.com/a.png", "/relative/a.png")) {
assertThatThrownBy(() -> tools.uploadFromUrl(Map.of("url", url)))
.isInstanceOf(McpToolException.class)
.hasMessage("url must be an absolute http or https URL");
}
verify(attachmentService, never()).uploadFromUrl(any(URL.class), any(), any(), any());
}

@Test
void validatesAttachmentBase64BeforeUpload() {
var tools = new AttachmentTools(
Expand Down Expand Up @@ -300,6 +383,10 @@ void deletionUsesExtensionLifecycleSoReconcilerCleansStorage() {

private void stubKeyId(String keyId) {
when(authorization.keyId()).thenReturn(Mono.just(keyId));
stubAttachmentConfig();
}

private void stubAttachmentConfig() {
var defaults = new ConfigMap();
defaults.setData(Map.of());
when(client.fetch(ConfigMap.class, SystemSetting.SYSTEM_CONFIG_DEFAULT))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ void organizesToolsByHaloDomainAndUsesChineseConsoleDescriptions() {
new AttachmentTools(client, mock(AttachmentService.class), new AttachmentUploadLimiter(), authorization));

var names = tools.names().toList();
assertThat(tools.tools()).hasSize(33);
assertThat(tools.tools()).hasSize(34);
assertThat(names)
.doesNotHaveDuplicates()
.contains(
Expand Down