Render Mermaid diagrams and worktree images in agent Markdown

internal/server/images.go

@@ -0,0 +1,89 @@
+package server
+
+import (
+	"errors"
+	"fmt"
+	"io/fs"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"g.fc.run/misc/byre/internal/api"
+)
+
+var worktreeImageTypes = map[string]string{
+	".png":  "image/png",
+	".jpg":  "image/jpeg",
+	".jpeg": "image/jpeg",
+	".gif":  "image/gif",
+	".webp": "image/webp",
+	".svg":  "image/svg+xml",
+}
+
+func (s *Server) GetTaskImage(w http.ResponseWriter, r *http.Request, taskID string, params api.GetTaskImageParams) {
+	t, err := s.store.GetTask(r.Context(), taskID)
+	if err != nil {
+		writeStoreError(w, err)
+		return
+	}
+	if t.Status != api.TaskStatusActive {
+		writeError(w, http.StatusConflict, fmt.Errorf("task %s is %s and has no worktree", t.Id, t.Status))
+		return
+	}
+
+	root, err := filepath.EvalSymlinks(t.WorktreePath)
+	if err != nil {
+		writeError(w, http.StatusInternalServerError, err)
+		return
+	}
+
+	path := params.Path
+	if !filepath.IsAbs(path) {
+		path = filepath.Join(t.WorktreePath, path)
+	}
+	real, err := filepath.EvalSymlinks(path)
+	if errors.Is(err, fs.ErrNotExist) {
+		writeError(w, http.StatusNotFound, fmt.Errorf("%s does not exist", params.Path))
+		return
+	}
+	if err != nil {
+		writeError(w, http.StatusBadRequest, err)
+		return
+	}
+
+	rel, err := filepath.Rel(root, real)
+	if err != nil || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) {
+		writeError(w, http.StatusForbidden, fmt.Errorf("%s is outside the task's worktree", params.Path))
+		return
+	}
+
+	contentType := worktreeImageTypes[strings.ToLower(filepath.Ext(real))]
+	if contentType == "" {
+		writeError(w, http.StatusBadRequest, fmt.Errorf("%s is not a PNG, JPEG, GIF, WebP or SVG file", params.Path))
+		return
+	}
+
+	f, err := os.Open(real)
+	if err != nil {
+		writeError(w, http.StatusInternalServerError, err)
+		return
+	}
+	defer f.Close()
+
+	info, err := f.Stat()
+	if err != nil {
+		writeError(w, http.StatusInternalServerError, err)
+		return
+	}
+	if !info.Mode().IsRegular() {
+		writeError(w, http.StatusBadRequest, fmt.Errorf("%s is not a file", params.Path))
+		return
+	}
+
+	w.Header().Set("Content-Type", contentType)
+	w.Header().Set("X-Content-Type-Options", "nosniff")
+	w.Header().Set("Content-Security-Policy", "default-src 'none'; style-src 'unsafe-inline'; sandbox")
+	w.Header().Set("Cache-Control", "private, no-cache")
+	http.ServeContent(w, r, "", info.ModTime(), f)
+}