feat: tools add upload

Author: YangSen-qn

src/mcp_server/core/storage/resource.py

@@ -5,9 +5,12 @@
 from mcp import types
 from urllib.parse import unquote
 
+from mcp.server.lowlevel.helper_types import ReadResourceContents
+
 from .storage import StorageService
 from ...consts import consts
 from ...resource import resource
+from ...resource.resource import ResourceContents
 
 logger = logging.getLogger(consts.LOGGER_NAME)
 
@@ -88,7 +91,7 @@ async def process_bucket_with_semaphore(bucket):
         logger.info(f"Returning {len(resources)} resources")
         return resources
 
-    async def read_resource(self, uri: types.AnyUrl, **kwargs) -> str:
+    async def read_resource(self, uri: types.AnyUrl, **kwargs) -> ResourceContents:
         """
         Read content from an S3 resource and return structured response
 
@@ -120,7 +123,7 @@ async def read_resource(self, uri: types.AnyUrl, **kwargs) -> str:
         if content_type.startswith("image/"):
             file_content = base64.b64encode(file_content).decode("utf-8")
 
-        return file_content
+        return [ReadResourceContents(mime_type=content_type, content=file_content)]
 
 
 def register_resource_provider(storage: StorageService):

src/mcp_server/core/storage/storage.py

@@ -1,5 +1,4 @@
 import aioboto3
-import asyncio
 import logging
 import qiniu
 
@@ -199,21 +198,45 @@ async def get_object(
                     return response
 
             except Exception as e:
-                last_exception = e
-                if "NoSuchKey" in str(e):
-                    raise
-
-                attempt += 1
-                if attempt < max_retries:
-                    wait_time = 2 ** attempt
-                    logger.warning(
-                        f"Attempt {attempt} failed, retrying in {wait_time} seconds: {str(e)}"
-                    )
-                    await asyncio.sleep(wait_time)
-                continue
+                logger.warning(
+                    f"Attempt {attempt} failed, exception: {str(e)}"
+                )
+                raise e
 
         raise last_exception or Exception("Failed to get object after all retries")
 
+    def upload_text_data(self, bucket: str, key: str, data: str, overwrite: bool = False) -> list[dict[str:Any]]:
+        policy = {
+            "insertOnly": 1,
+        }
+
+        if overwrite:
+            policy["insertOnly"] = 0
+            policy["scope"] = f"{bucket}:{key}"
+
+        token = self.auth.upload_token(bucket=bucket, key=key, policy=policy)
+        ret, info = qiniu.put_data(up_token=token, key=key, data=bytes(data, encoding="utf-8"))
+        if info.status_code != 200:
+            raise Exception(f"Failed to upload object: {info.text}")
+
+        return self.get_object_url(bucket, key)
+
+    def upload_file(self, bucket: str, key: str, file_path: str, overwrite: bool = False) -> list[dict[str:Any]]:
+        policy = {
+            "insertOnly": 1,
+        }
+
+        if overwrite:
+            policy["insertOnly"] = 0
+            policy["scope"] = f"{bucket}:{key}"
+
+        token = self.auth.upload_token(bucket=bucket, key=key, policy=policy)
+        ret, info = qiniu.put_file(up_token=token, key=key, file_path=file_path)
+        if info.status_code != 200:
+            raise Exception(f"Failed to upload object: {info.text}")
+
+        return self.get_object_url(bucket, key)
+
     def is_text_file(self, key: str) -> bool:
         """Determine if a file is text-based by its extension"""
         text_extensions = {

src/mcp_server/core/storage/tools.py

@@ -109,6 +109,70 @@ async def get_object(self, **kwargs) -> list[ImageContent] | list[TextContent]:
             text_content = str(file_content)
         return [types.TextContent(type="text", text=text_content)]
 
+    @tools.tool_meta(
+        types.Tool(
+            name="UploadTextData",
+            description="Upload text data to Qiniu bucket. In the UploadData request, specify the full key name for the object. Path-style requests are not supported.",
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "bucket": {
+                        "type": "string",
+                        "description": _BUCKET_DESC,
+                    },
+                    "key": {
+                        "type": "string",
+                        "description": "Key of the object to upload. Length Constraints: Minimum length of 1.",
+                    },
+                    "data": {
+                        "type": "string",
+                        "description": "The data to upload.",
+                    },
+                    "overwrite": {
+                        "type": "boolean",
+                        "description": "Whether to overwrite the existing object if it already exists.",
+                    },
+                },
+                "required": ["bucket", "key", "data"],
+            }
+        )
+    )
+    def upload_text_data(self, **kwargs) -> list[types.TextContent]:
+        urls = self.storage.upload_text_data(**kwargs)
+        return [types.TextContent(type="text", text=str(urls))]
+
+    @tools.tool_meta(
+        types.Tool(
+            name="UploadFile",
+            description="Upload file to Qiniu bucket. In the UploadFile request, specify the full key name for the object. Path-style requests are not supported.",
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "bucket": {
+                        "type": "string",
+                        "description": _BUCKET_DESC,
+                    },
+                    "key": {
+                        "type": "string",
+                        "description": "Key of the object to upload. Length Constraints: Minimum length of 1.",
+                    },
+                    "file_path": {
+                        "type": "string",
+                        "description": "The file path of file to upload.",
+                    },
+                    "overwrite": {
+                        "type": "boolean",
+                        "description": "Whether to overwrite the existing object if it already exists.",
+                    },
+                },
+                "required": ["bucket", "key", "file_path"],
+            }
+        )
+    )
+    def upload_file(self, **kwargs) -> list[types.TextContent]:
+        urls = self.storage.upload_file(**kwargs)
+        return [types.TextContent(type="text", text=str(urls))]
+
     @tools.tool_meta(
         types.Tool(
             name="GetObjectURL",
@@ -149,6 +213,8 @@ def register_tools(storage: StorageService):
             tool_impl.list_buckets,
             tool_impl.list_objects,
             tool_impl.get_object,
+            tool_impl.upload_text_data,
+            tool_impl.upload_file,
             tool_impl.get_object_url,
         ]
     )

src/mcp_server/resource/resource.py

@@ -1,12 +1,15 @@
 import logging
 from abc import abstractmethod
-from typing import Dict, AsyncGenerator
+from typing import Dict, AsyncGenerator, Iterable
 
 from mcp import types
+from mcp.server.lowlevel.helper_types import ReadResourceContents
+
 from ..consts import consts
 
 logger = logging.getLogger(consts.LOGGER_NAME)
 
+ResourceContents = str | bytes | Iterable[ReadResourceContents]
 
 class ResourceProvider:
     def __init__(self, scheme: str):
@@ -17,7 +20,7 @@ async def list_resources(self, **kwargs) -> list[types.Resource]:
         pass
 
     @abstractmethod
-    async def read_resource(self, uri: types.AnyUrl, **kwargs) -> str:
+    async def read_resource(self, uri: types.AnyUrl, **kwargs) -> ResourceContents:
         pass
 
 
@@ -35,7 +38,7 @@ async def list_resources(**kwargs) -> AsyncGenerator[types.Resource, None]:
     return
 
 
-async def read_resource(uri: types.AnyUrl, **kwargs) -> str:
+async def read_resource(uri: types.AnyUrl, **kwargs) -> ResourceContents:
     if len(_all_resource_providers) == 0:
         return ""
 
@@ -52,6 +55,7 @@ def register_resource_provider(provider: ResourceProvider):
 
 
 __all__ = [
+    "ResourceContents",
     "ResourceProvider",
     "list_resources",
     "read_resource",