feat(files): implement file storage functionality with save, load, and list operations
This commit is contained in:
276
internal/tools/files.go
Normal file
276
internal/tools/files.go
Normal file
@@ -0,0 +1,276 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
|
||||
"git.warky.dev/wdevs/amcs/internal/session"
|
||||
"git.warky.dev/wdevs/amcs/internal/store"
|
||||
thoughttypes "git.warky.dev/wdevs/amcs/internal/types"
|
||||
)
|
||||
|
||||
type FilesTool struct {
|
||||
store *store.DB
|
||||
sessions *session.ActiveProjects
|
||||
}
|
||||
|
||||
type SaveFileInput struct {
|
||||
Name string `json:"name" jsonschema:"file name including extension, for example photo.png or note.pdf"`
|
||||
ContentBase64 string `json:"content_base64" jsonschema:"file contents encoded as base64"`
|
||||
MediaType string `json:"media_type,omitempty" jsonschema:"optional MIME type such as image/png, application/pdf, or audio/mpeg"`
|
||||
Kind string `json:"kind,omitempty" jsonschema:"optional logical type such as image, document, audio, or file"`
|
||||
ThoughtID string `json:"thought_id,omitempty" jsonschema:"optional thought id to link this file to"`
|
||||
Project string `json:"project,omitempty" jsonschema:"optional project name or id when saving outside a linked thought"`
|
||||
}
|
||||
|
||||
type SaveFileOutput struct {
|
||||
File thoughttypes.StoredFile `json:"file"`
|
||||
}
|
||||
|
||||
type LoadFileInput struct {
|
||||
ID string `json:"id" jsonschema:"the stored file id"`
|
||||
}
|
||||
|
||||
type LoadFileOutput struct {
|
||||
File thoughttypes.StoredFile `json:"file"`
|
||||
ContentBase64 string `json:"content_base64"`
|
||||
}
|
||||
|
||||
type ListFilesInput struct {
|
||||
Limit int `json:"limit,omitempty" jsonschema:"maximum number of files to return"`
|
||||
ThoughtID string `json:"thought_id,omitempty" jsonschema:"optional thought id to list files for"`
|
||||
Project string `json:"project,omitempty" jsonschema:"optional project name or id to scope the listing"`
|
||||
Kind string `json:"kind,omitempty" jsonschema:"optional kind filter such as image, document, audio, or file"`
|
||||
}
|
||||
|
||||
type ListFilesOutput struct {
|
||||
Files []thoughttypes.StoredFile `json:"files"`
|
||||
}
|
||||
|
||||
func NewFilesTool(db *store.DB, sessions *session.ActiveProjects) *FilesTool {
|
||||
return &FilesTool{store: db, sessions: sessions}
|
||||
}
|
||||
|
||||
func (t *FilesTool) Save(ctx context.Context, req *mcp.CallToolRequest, in SaveFileInput) (*mcp.CallToolResult, SaveFileOutput, error) {
|
||||
name := strings.TrimSpace(in.Name)
|
||||
if name == "" {
|
||||
return nil, SaveFileOutput{}, errInvalidInput("name is required")
|
||||
}
|
||||
|
||||
contentBase64, mediaTypeFromDataURL := splitDataURL(strings.TrimSpace(in.ContentBase64))
|
||||
if contentBase64 == "" {
|
||||
return nil, SaveFileOutput{}, errInvalidInput("content_base64 is required")
|
||||
}
|
||||
|
||||
content, err := decodeBase64(contentBase64)
|
||||
if err != nil {
|
||||
return nil, SaveFileOutput{}, errInvalidInput("content_base64 must be valid base64")
|
||||
}
|
||||
if len(content) == 0 {
|
||||
return nil, SaveFileOutput{}, errInvalidInput("decoded file content must not be empty")
|
||||
}
|
||||
|
||||
project, err := resolveProject(ctx, t.store, t.sessions, req, in.Project, false)
|
||||
if err != nil {
|
||||
return nil, SaveFileOutput{}, err
|
||||
}
|
||||
|
||||
var thoughtID *uuid.UUID
|
||||
var projectID = projectIDPtr(project)
|
||||
if rawThoughtID := strings.TrimSpace(in.ThoughtID); rawThoughtID != "" {
|
||||
parsedThoughtID, err := parseUUID(rawThoughtID)
|
||||
if err != nil {
|
||||
return nil, SaveFileOutput{}, err
|
||||
}
|
||||
thought, err := t.store.GetThought(ctx, parsedThoughtID)
|
||||
if err != nil {
|
||||
return nil, SaveFileOutput{}, err
|
||||
}
|
||||
thoughtID = &parsedThoughtID
|
||||
projectID = thought.ProjectID
|
||||
if project != nil && thought.ProjectID != nil && *thought.ProjectID != project.ID {
|
||||
return nil, SaveFileOutput{}, errInvalidInput("project does not match the linked thought's project")
|
||||
}
|
||||
}
|
||||
|
||||
mediaType := normalizeMediaType(strings.TrimSpace(in.MediaType), mediaTypeFromDataURL, content)
|
||||
kind := normalizeFileKind(strings.TrimSpace(in.Kind), mediaType)
|
||||
sum := sha256.Sum256(content)
|
||||
|
||||
file := thoughttypes.StoredFile{
|
||||
Name: name,
|
||||
MediaType: mediaType,
|
||||
Kind: kind,
|
||||
Encoding: "base64",
|
||||
SizeBytes: int64(len(content)),
|
||||
SHA256: hex.EncodeToString(sum[:]),
|
||||
Content: content,
|
||||
ProjectID: projectID,
|
||||
}
|
||||
if thoughtID != nil {
|
||||
file.ThoughtID = thoughtID
|
||||
}
|
||||
|
||||
created, err := t.store.InsertStoredFile(ctx, file)
|
||||
if err != nil {
|
||||
return nil, SaveFileOutput{}, err
|
||||
}
|
||||
|
||||
if created.ThoughtID != nil {
|
||||
if err := t.store.AddThoughtAttachment(ctx, *created.ThoughtID, thoughtAttachmentFromFile(created)); err != nil {
|
||||
return nil, SaveFileOutput{}, err
|
||||
}
|
||||
}
|
||||
if created.ProjectID != nil {
|
||||
_ = t.store.TouchProject(ctx, *created.ProjectID)
|
||||
}
|
||||
|
||||
return nil, SaveFileOutput{File: created}, nil
|
||||
}
|
||||
|
||||
func (t *FilesTool) Load(ctx context.Context, _ *mcp.CallToolRequest, in LoadFileInput) (*mcp.CallToolResult, LoadFileOutput, error) {
|
||||
id, err := parseUUID(in.ID)
|
||||
if err != nil {
|
||||
return nil, LoadFileOutput{}, err
|
||||
}
|
||||
|
||||
file, err := t.store.GetStoredFile(ctx, id)
|
||||
if err != nil {
|
||||
return nil, LoadFileOutput{}, err
|
||||
}
|
||||
|
||||
return nil, LoadFileOutput{
|
||||
File: file,
|
||||
ContentBase64: base64.StdEncoding.EncodeToString(file.Content),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (t *FilesTool) List(ctx context.Context, req *mcp.CallToolRequest, in ListFilesInput) (*mcp.CallToolResult, ListFilesOutput, error) {
|
||||
project, err := resolveProject(ctx, t.store, t.sessions, req, in.Project, false)
|
||||
if err != nil {
|
||||
return nil, ListFilesOutput{}, err
|
||||
}
|
||||
|
||||
var thoughtID *uuid.UUID
|
||||
if rawThoughtID := strings.TrimSpace(in.ThoughtID); rawThoughtID != "" {
|
||||
parsedThoughtID, err := parseUUID(rawThoughtID)
|
||||
if err != nil {
|
||||
return nil, ListFilesOutput{}, err
|
||||
}
|
||||
thought, err := t.store.GetThought(ctx, parsedThoughtID)
|
||||
if err != nil {
|
||||
return nil, ListFilesOutput{}, err
|
||||
}
|
||||
thoughtID = &parsedThoughtID
|
||||
if project != nil && thought.ProjectID != nil && *thought.ProjectID != project.ID {
|
||||
return nil, ListFilesOutput{}, errInvalidInput("project does not match the linked thought's project")
|
||||
}
|
||||
if project == nil && thought.ProjectID != nil {
|
||||
project = &thoughttypes.Project{ID: *thought.ProjectID}
|
||||
}
|
||||
}
|
||||
|
||||
files, err := t.store.ListStoredFiles(ctx, thoughttypes.StoredFileFilter{
|
||||
Limit: normalizeFileLimit(in.Limit),
|
||||
ThoughtID: thoughtID,
|
||||
ProjectID: projectIDPtr(project),
|
||||
Kind: strings.TrimSpace(in.Kind),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, ListFilesOutput{}, err
|
||||
}
|
||||
if project != nil {
|
||||
_ = t.store.TouchProject(ctx, project.ID)
|
||||
}
|
||||
|
||||
return nil, ListFilesOutput{Files: files}, nil
|
||||
}
|
||||
|
||||
func thoughtAttachmentFromFile(file thoughttypes.StoredFile) thoughttypes.ThoughtAttachment {
|
||||
return thoughttypes.ThoughtAttachment{
|
||||
FileID: file.ID,
|
||||
Name: file.Name,
|
||||
MediaType: file.MediaType,
|
||||
Kind: file.Kind,
|
||||
SizeBytes: file.SizeBytes,
|
||||
SHA256: file.SHA256,
|
||||
}
|
||||
}
|
||||
|
||||
func splitDataURL(value string) (contentBase64 string, mediaType string) {
|
||||
const marker = ";base64,"
|
||||
if !strings.HasPrefix(value, "data:") {
|
||||
return value, ""
|
||||
}
|
||||
|
||||
prefix, payload, ok := strings.Cut(value, marker)
|
||||
if !ok {
|
||||
return value, ""
|
||||
}
|
||||
|
||||
mediaType = strings.TrimPrefix(prefix, "data:")
|
||||
return payload, strings.TrimSpace(mediaType)
|
||||
}
|
||||
|
||||
func decodeBase64(value string) ([]byte, error) {
|
||||
decoded, err := base64.StdEncoding.DecodeString(value)
|
||||
if err == nil {
|
||||
return decoded, nil
|
||||
}
|
||||
return base64.RawStdEncoding.DecodeString(value)
|
||||
}
|
||||
|
||||
func normalizeMediaType(explicit string, fromDataURL string, content []byte) string {
|
||||
switch {
|
||||
case explicit != "":
|
||||
return explicit
|
||||
case fromDataURL != "":
|
||||
return fromDataURL
|
||||
default:
|
||||
return http.DetectContentType(content)
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeFileKind(explicit string, mediaType string) string {
|
||||
if explicit != "" {
|
||||
return explicit
|
||||
}
|
||||
|
||||
switch {
|
||||
case strings.HasPrefix(mediaType, "image/"):
|
||||
return "image"
|
||||
case strings.HasPrefix(mediaType, "audio/"):
|
||||
return "audio"
|
||||
case strings.HasPrefix(mediaType, "video/"):
|
||||
return "video"
|
||||
case mediaType == "application/pdf" || strings.HasPrefix(mediaType, "text/") || strings.Contains(mediaType, "document"):
|
||||
return "document"
|
||||
default:
|
||||
return "file"
|
||||
}
|
||||
}
|
||||
|
||||
func projectIDPtr(project *thoughttypes.Project) *uuid.UUID {
|
||||
if project == nil {
|
||||
return nil
|
||||
}
|
||||
return &project.ID
|
||||
}
|
||||
|
||||
func normalizeFileLimit(limit int) int {
|
||||
switch {
|
||||
case limit <= 0:
|
||||
return 20
|
||||
case limit > 100:
|
||||
return 100
|
||||
default:
|
||||
return limit
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user