feat(wecom): implement shared HTTP clients for WeCom image handling and response URL posting

pkg/channels/wecom/aibot.go

@@ -22,6 +22,10 @@ import (
 	"github.com/sipeed/picoclaw/pkg/utils"
 )
 
+// responseURLHTTPClient is a shared HTTP client for posting to WeCom response_url.
+// Reusing it enables connection pooling across replies.
+var responseURLHTTPClient = &http.Client{Timeout: 15 * time.Second}
+
 // WeComAIBotChannel implements the Channel interface for WeCom AI Bot (企业微信智能机器人)
 type WeComAIBotChannel struct {
 	*channels.BaseChannel
@@ -794,8 +798,7 @@ func (c *WeComAIBotChannel) sendViaResponseURL(responseURL, content string) erro
 	}
 	req.Header.Set("Content-Type", "application/json; charset=utf-8")
 
-	client := &http.Client{Timeout: 15 * time.Second}
-	resp, err := client.Do(req)
+	resp, err := responseURLHTTPClient.Do(req)
 	if err != nil {
 		return fmt.Errorf("post to response_url failed: %w: %w", channels.ErrTemporary, err)
 	}
@@ -805,7 +808,8 @@ func (c *WeComAIBotChannel) sendViaResponseURL(responseURL, content string) erro
 		return nil
 	}
 
-	respBody, err := io.ReadAll(resp.Body)
+	const maxErrBody = 64 << 10 // 64 KB is more than enough for any error response
+	respBody, err := io.ReadAll(io.LimitReader(resp.Body, maxErrBody))
 	if err != nil {
 		return fmt.Errorf("reading response_url body: %w: %w", channels.ErrTemporary, err)
 	}

pkg/channels/wecom/aibot_test.go

@@ -3,6 +3,7 @@ package wecom
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/sipeed/picoclaw/pkg/bus"
 	"github.com/sipeed/picoclaw/pkg/channels"
@@ -335,3 +336,120 @@ func TestWSGenerateID(t *testing.T) {
 		ids[id] = true
 	}
 }
+
+// ---- Webhook streaming fallback tests ----
+
+// makeWebhookChannel creates a started WeComAIBotChannel for testing.
+func makeWebhookChannel(t *testing.T) *WeComAIBotChannel {
+	t.Helper()
+	cfg := config.WeComAIBotConfig{
+		Enabled:        true,
+		Token:          "test_token",
+		EncodingAESKey: "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFG",
+	}
+	ch, err := NewWeComAIBotChannel(cfg, bus.NewMessageBus())
+	if err != nil {
+		t.Fatalf("create channel: %v", err)
+	}
+	wc := ch.(*WeComAIBotChannel)
+	wc.ctx, wc.cancel = context.WithCancel(context.Background())
+	return wc
+}
+
+// makeStreamTask creates and registers a streamTask for testing.
+func makeStreamTask(t *testing.T, ch *WeComAIBotChannel, streamID, chatID string, deadline time.Time) *streamTask {
+	t.Helper()
+	task := &streamTask{
+		StreamID: streamID,
+		ChatID:   chatID,
+		Deadline: deadline,
+		answerCh: make(chan string, 1),
+	}
+	task.ctx, task.cancel = context.WithCancel(ch.ctx)
+	ch.taskMu.Lock()
+	ch.streamTasks[streamID] = task
+	ch.chatTasks[chatID] = append(ch.chatTasks[chatID], task)
+	ch.taskMu.Unlock()
+	return task
+}
+
+// TestGetStreamResponse_ImmediateAnswer verifies that when the agent has already
+// placed its answer in answerCh, getStreamResponse returns a finish=true response
+// and fully removes the task.
+func TestGetStreamResponse_ImmediateAnswer(t *testing.T) {
+	ch := makeWebhookChannel(t)
+	defer ch.cancel()
+
+	task := makeStreamTask(t, ch, "stream-1", "chat-1", time.Now().Add(30*time.Second))
+	task.answerCh <- "hello from agent"
+
+	result := ch.getStreamResponse(task, "ts123", "nonce123")
+	if result == "" {
+		t.Fatal("expected non-empty encrypted response")
+	}
+
+	ch.taskMu.RLock()
+	_, exists := ch.streamTasks["stream-1"]
+	ch.taskMu.RUnlock()
+	if exists {
+		t.Error("task should have been removed from streamTasks after normal finish")
+	}
+	if !task.Finished {
+		t.Error("task.Finished should be true after normal finish")
+	}
+}
+
+// TestGetStreamResponse_DeadlinePassed verifies that when the stream deadline has
+// elapsed (no agent reply yet), getStreamResponse closes the stream but keeps the
+// task alive so the response_url fallback can still deliver the answer.
+func TestGetStreamResponse_DeadlinePassed(t *testing.T) {
+	ch := makeWebhookChannel(t)
+	defer ch.cancel()
+
+	task := makeStreamTask(t, ch, "stream-2", "chat-2", time.Now().Add(-time.Millisecond))
+
+	result := ch.getStreamResponse(task, "ts456", "nonce456")
+	if result == "" {
+		t.Fatal("expected non-empty encrypted response")
+	}
+
+	ch.taskMu.RLock()
+	_, stillStreaming := ch.streamTasks["stream-2"]
+	ch.taskMu.RUnlock()
+	if stillStreaming {
+		t.Error("task should have been removed from streamTasks after deadline")
+	}
+	if !task.StreamClosed {
+		t.Error("task.StreamClosed should be true after deadline")
+	}
+	if task.Finished {
+		t.Error("task.Finished must remain false: agent reply still expected via response_url")
+	}
+}
+
+// TestGetStreamResponse_StillPending verifies that when neither the agent has
+// replied nor the deadline has passed, getStreamResponse returns without altering
+// task state (client should poll again).
+func TestGetStreamResponse_StillPending(t *testing.T) {
+	ch := makeWebhookChannel(t)
+	defer ch.cancel()
+
+	task := makeStreamTask(t, ch, "stream-3", "chat-3", time.Now().Add(30*time.Second))
+
+	result := ch.getStreamResponse(task, "ts789", "nonce789")
+	if result == "" {
+		t.Fatal("expected non-empty encrypted response")
+	}
+
+	ch.taskMu.RLock()
+	_, exists := ch.streamTasks["stream-3"]
+	ch.taskMu.RUnlock()
+	if !exists {
+		t.Error("pending task should still be in streamTasks")
+	}
+	if task.Finished || task.StreamClosed {
+		t.Error("pending task should not be finished or stream-closed")
+	}
+	// Cleanup.
+	ch.removeTask(task)
+}

pkg/channels/wecom/aibot_ws.go

@@ -1,6 +1,7 @@
 package wecom
 
 import (
+	"bytes"
 	"context"
 	"encoding/base64"
 	"encoding/json"
@@ -49,6 +50,10 @@ const (
 	wsLateReplyRouteTTL = 30 * time.Minute
 )
 
+// wsImageHTTPClient is a shared HTTP client for downloading inbound images.
+// Reusing it enables connection pooling across multiple image downloads.
+var wsImageHTTPClient = &http.Client{Timeout: wsImageDownloadTimeout}
+
 // WeComAIBotWSChannel implements channels.Channel for WeCom AI Bot using the
 // WebSocket long-connection API.
 // Unlike the webhook counterpart it does NOT implement WebhookHandler, so the
@@ -1040,49 +1045,10 @@ func wsGenerateID() string {
 
 // ---- Inbound image download helpers ----
 
-// downloadWSImage fetches and optionally decrypts a WeCom WS image resource.
-// aesKey is the per-resource AES key provided in the callback (may be empty).
-func (c *WeComAIBotWSChannel) downloadWSImage(ctx context.Context, imageURL, aesKey string) ([]byte, error) {
-	const maxSize = 20 << 20 // 20 MB
-	req, err := http.NewRequestWithContext(ctx, http.MethodGet, imageURL, nil)
-	if err != nil {
-		return nil, fmt.Errorf("create request: %w", err)
-	}
-	client := &http.Client{Timeout: wsImageDownloadTimeout}
-	resp, err := client.Do(req)
-	if err != nil {
-		return nil, fmt.Errorf("download: %w", err)
-	}
-	defer resp.Body.Close()
-	if resp.StatusCode != http.StatusOK {
-		return nil, fmt.Errorf("download HTTP %d", resp.StatusCode)
-	}
-	data, err := io.ReadAll(io.LimitReader(resp.Body, maxSize+1))
-	if err != nil {
-		return nil, fmt.Errorf("read body: %w", err)
-	}
-	if len(data) > maxSize {
-		return nil, fmt.Errorf("image too large (> %d MB)", maxSize>>20)
-	}
-	if aesKey == "" {
-		return data, nil
-	}
-	// WeCom per-image AES key: try standard base64 first, then WeCom 43-char format.
-	key, decErr := base64.StdEncoding.DecodeString(aesKey)
-	if decErr != nil || len(key) != 32 {
-		key, decErr = decodeWeComAESKey(aesKey)
-		if decErr != nil {
-			return nil, fmt.Errorf("decode image AES key: %w", decErr)
-		}
-	}
-	decrypted, err := decryptAESCBC(key, data)
-	if err != nil {
-		return nil, fmt.Errorf("decrypt image: %w", err)
-	}
-	return decrypted, nil
-}
-
 // storeWSImage downloads, optionally decrypts, and stores an inbound image.
+// It streams the HTTP response body directly into the MediaStore, avoiding any
+// caller-managed temp files. File lifecycle (creation and deletion) is owned
+// entirely by the store and cleaned up via store.ReleaseAll.
 func (c *WeComAIBotWSChannel) storeWSImage(
 	ctx context.Context,
 	chatID, msgID, imageURL, aesKey string,
@@ -1091,26 +1057,121 @@ func (c *WeComAIBotWSChannel) storeWSImage(
 	if store == nil {
 		return "", fmt.Errorf("no media store available")
 	}
-	data, err := c.downloadWSImage(ctx, imageURL, aesKey)
+
+	const maxSize = 20 << 20 // 20 MB
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, imageURL, nil)
 	if err != nil {
-		return "", err
+		return "", fmt.Errorf("create request: %w", err)
 	}
+	resp, err := wsImageHTTPClient.Do(req)
+	if err != nil {
+		return "", fmt.Errorf("download: %w", err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		return "", fmt.Errorf("download HTTP %d", resp.StatusCode)
+	}
+
 	mediaDir := filepath.Join(os.TempDir(), "picoclaw_media")
+	scope := channels.BuildMediaScope("wecom_aibot", chatID, msgID)
+	meta := media.MediaMeta{Filename: msgID + ".jpg", Source: "wecom_aibot"}
+
+	// lr wraps the response body with a +1 limit: if lr.N reaches 0 after the
+	// transfer, the body was at least maxSize+1 bytes and must be rejected.
+	lr := &io.LimitedReader{R: resp.Body, N: int64(maxSize) + 1}
+	var r io.Reader = lr
+
+	if aesKey != "" {
+		// AES-CBC decryption requires full ciphertext. Buffer the bounded encrypted
+		// bytes, decrypt, then pass the plaintext as a reader.
+		encrypted, readErr := io.ReadAll(lr)
+		if readErr != nil {
+			return "", fmt.Errorf("read for decrypt: %w", readErr)
+		}
+		if lr.N == 0 {
+			return "", fmt.Errorf("image too large (> %d MB)", maxSize>>20)
+		}
+		key, decErr := base64.StdEncoding.DecodeString(aesKey)
+		if decErr != nil || len(key) != 32 {
+			key, decErr = decodeWeComAESKey(aesKey)
+			if decErr != nil {
+				return "", fmt.Errorf("decode image AES key: %w", decErr)
+			}
+		}
+		decrypted, decErr := decryptAESCBC(key, encrypted)
+		if decErr != nil {
+			return "", fmt.Errorf("decrypt image: %w", decErr)
+		}
+		r = bytes.NewReader(decrypted)
+	}
+
+	// Fast path: FileMediaStore supports StoreFromReader which manages the file
+	// lifecycle entirely (temp-rename-fsync pattern, deleted by ReleaseAll).
+	if fsStore, ok := store.(*media.FileMediaStore); ok {
+		ref, storeErr := fsStore.StoreFromReader(r, meta, scope, mediaDir)
+		if storeErr != nil {
+			return "", storeErr
+		}
+		// For the no-AES path, check whether the response body hit the size limit.
+		if aesKey == "" && lr.N == 0 {
+			_ = store.ReleaseAll(scope) // remove the oversized file
+			return "", fmt.Errorf("image too large (> %d MB)", maxSize>>20)
+		}
+		return ref, nil
+	}
+
+	// Fallback path for non-FileMediaStore implementations: stream into a temp
+	// file using the temp-rename-fsync pattern, then register the path.
+	// The file lifecycle is managed by store.ReleaseAll().
 	if mkErr := os.MkdirAll(mediaDir, 0o700); mkErr != nil {
 		return "", fmt.Errorf("mkdir: %w", mkErr)
 	}
-	filename := msgID + ".jpg"
-	localPath := filepath.Join(mediaDir, utils.SanitizeFilename(filename))
-	if writeErr := os.WriteFile(localPath, data, 0o600); writeErr != nil {
-		return "", fmt.Errorf("write: %w", writeErr)
-	}
-	scope := channels.BuildMediaScope("wecom_aibot", chatID, msgID)
-	ref, err := store.Store(localPath, media.MediaMeta{
-		Filename: filename,
-		Source:   "wecom_aibot",
-	}, scope)
+	tmpPath := filepath.Join(mediaDir, ".tmp-"+wsGenerateID())
+	tmpFile, err := os.OpenFile(tmpPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
 	if err != nil {
-		os.Remove(localPath)
+		return "", fmt.Errorf("create temp file: %w", err)
+	}
+	cleanup := true
+	defer func() {
+		if cleanup {
+			tmpFile.Close()
+			os.Remove(tmpPath)
+		}
+	}()
+	n, cpErr := io.Copy(tmpFile, r)
+	syncErr := tmpFile.Sync()
+	closeErr := tmpFile.Close()
+	cleanup = false
+	if cpErr != nil {
+		os.Remove(tmpPath)
+		return "", fmt.Errorf("write image: %w", cpErr)
+	}
+	if syncErr != nil {
+		os.Remove(tmpPath)
+		return "", fmt.Errorf("sync image: %w", syncErr)
+	}
+	if closeErr != nil {
+		os.Remove(tmpPath)
+		return "", fmt.Errorf("close image: %w", closeErr)
+	}
+	if aesKey == "" && lr.N == 0 {
+		os.Remove(tmpPath)
+		return "", fmt.Errorf("image too large (> %d MB)", maxSize>>20)
+	}
+	_ = n
+	finalPath := filepath.Join(mediaDir, wsGenerateID()+".jpg")
+	if err = os.Rename(tmpPath, finalPath); err != nil {
+		os.Remove(tmpPath)
+		return "", fmt.Errorf("rename: %w", err)
+	}
+	if d, openErr := os.Open(mediaDir); openErr == nil {
+		_ = d.Sync()
+		d.Close()
+	}
+	ref, err := store.Store(finalPath, meta, scope)
+	if err != nil {
+		os.Remove(finalPath)
 		return "", fmt.Errorf("store: %w", err)
 	}
 	return ref, nil

pkg/channels/wecom/aibot_ws_test.go

@@ -0,0 +1,165 @@
+package wecom
+
+import (
+	"bytes"
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"testing"
+
+	"github.com/sipeed/picoclaw/pkg/bus"
+	"github.com/sipeed/picoclaw/pkg/channels"
+	"github.com/sipeed/picoclaw/pkg/config"
+	"github.com/sipeed/picoclaw/pkg/media"
+)
+
+// newTestWSChannel creates a WeComAIBotWSChannel ready for unit testing.
+func newTestWSChannel(t *testing.T) *WeComAIBotWSChannel {
+	t.Helper()
+	cfg := config.WeComAIBotConfig{
+		Enabled: true,
+		BotID:   "test_bot_id",
+		Secret:  "test_secret",
+	}
+	ch, err := newWeComAIBotWSChannel(cfg, bus.NewMessageBus())
+	if err != nil {
+		t.Fatalf("create WS channel: %v", err)
+	}
+	return ch
+}
+
+// TestStoreWSImage_NilStore verifies that storeWSImage returns an error when no
+// MediaStore has been injected.
+func TestStoreWSImage_NilStore(t *testing.T) {
+	ch := newTestWSChannel(t)
+	_, err := ch.storeWSImage(context.Background(), "chat1", "msg1", "http://any", "")
+	if err == nil {
+		t.Fatal("expected error when no MediaStore is set")
+	}
+}
+
+// TestStoreWSImage_HTTPError verifies that storeWSImage propagates HTTP errors
+// from the image server.
+func TestStoreWSImage_HTTPError(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		http.Error(w, "not found", http.StatusNotFound)
+	}))
+	defer srv.Close()
+
+	ch := newTestWSChannel(t)
+	ch.SetMediaStore(media.NewFileMediaStore())
+
+	_, err := ch.storeWSImage(context.Background(), "chat1", "msg1", srv.URL, "")
+	if err == nil {
+		t.Fatal("expected error for HTTP 404")
+	}
+}
+
+// TestStoreWSImage_ServerUnavailable verifies that storeWSImage returns a clear
+// error when the image server cannot be reached.
+func TestStoreWSImage_ServerUnavailable(t *testing.T) {
+	ch := newTestWSChannel(t)
+	ch.SetMediaStore(media.NewFileMediaStore())
+
+	// Port 1 is reserved and will refuse the connection immediately.
+	_, err := ch.storeWSImage(context.Background(), "chat1", "msg1", "http://127.0.0.1:1", "")
+	if err == nil {
+		t.Fatal("expected error for unreachable server")
+	}
+}
+
+// TestStoreWSImage_Success_NoAES verifies the happy path: the image is downloaded,
+// a media ref is returned, and the file persists and is readable via Resolve until
+// ReleaseAll is called.
+func TestStoreWSImage_Success_NoAES(t *testing.T) {
+	imageData := bytes.Repeat([]byte("x"), 256)
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write(imageData)
+	}))
+	defer srv.Close()
+
+	ch := newTestWSChannel(t)
+	store := media.NewFileMediaStore()
+	ch.SetMediaStore(store)
+
+	ref, err := ch.storeWSImage(context.Background(), "chat1", "msg1", srv.URL, "")
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+	if ref == "" {
+		t.Fatal("expected non-empty ref")
+	}
+
+	// File must be accessible after storeWSImage returns (no premature deletion).
+	path, err := store.Resolve(ref)
+	if err != nil {
+		t.Fatalf("ref should resolve: %v", err)
+	}
+	got, err := os.ReadFile(path)
+	if err != nil {
+		t.Fatalf("file should exist at %s: %v", path, err)
+	}
+	if !bytes.Equal(got, imageData) {
+		t.Errorf("content mismatch: got len=%d, want len=%d", len(got), len(imageData))
+	}
+
+	// ReleaseAll must delete the file (store owns lifecycle).
+	scope := channels.BuildMediaScope("wecom_aibot", "chat1", "msg1")
+	if err := store.ReleaseAll(scope); err != nil {
+		t.Fatalf("ReleaseAll failed: %v", err)
+	}
+	if _, err := os.Stat(path); !os.IsNotExist(err) {
+		t.Errorf("file should have been deleted by ReleaseAll, stat err: %v", err)
+	}
+}
+
+// TestStoreWSImage_MultipleMessages verifies that concurrent image messages with
+// different msgIDs do not collide and each resolve to distinct files.
+func TestStoreWSImage_MultipleMessages(t *testing.T) {
+	imageA := bytes.Repeat([]byte("a"), 64)
+	imageB := bytes.Repeat([]byte("b"), 64)
+
+	srvA := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write(imageA)
+	}))
+	defer srvA.Close()
+	srvB := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write(imageB)
+	}))
+	defer srvB.Close()
+
+	ch := newTestWSChannel(t)
+	store := media.NewFileMediaStore()
+	ch.SetMediaStore(store)
+
+	refA, err := ch.storeWSImage(context.Background(), "chat1", "msgA", srvA.URL, "")
+	if err != nil {
+		t.Fatalf("storeWSImage A: %v", err)
+	}
+	refB, err := ch.storeWSImage(context.Background(), "chat1", "msgB", srvB.URL, "")
+	if err != nil {
+		t.Fatalf("storeWSImage B: %v", err)
+	}
+	if refA == refB {
+		t.Fatal("distinct messages must produce distinct refs")
+	}
+
+	pathA, _ := store.Resolve(refA)
+	pathB, _ := store.Resolve(refB)
+	if pathA == pathB {
+		t.Fatal("distinct messages must be stored at distinct paths")
+	}
+
+	gotA, _ := os.ReadFile(pathA)
+	gotB, _ := os.ReadFile(pathB)
+	if !bytes.Equal(gotA, imageA) {
+		t.Errorf("content mismatch for message A")
+	}
+	if !bytes.Equal(gotB, imageB) {
+		t.Errorf("content mismatch for message B")
+	}
+}

pkg/media/store.go

@@ -2,7 +2,9 @@ package media
 
 import (
 	"fmt"
+	"io"
 	"os"
+	"path/filepath"
 	"sync"
 	"time"
 
@@ -108,6 +110,64 @@ func (s *FileMediaStore) Store(localPath string, meta MediaMeta, scope string) (
 	return ref, nil
 }
 
+// StoreFromReader creates a temp file in dir by streaming r into it, then
+// atomically renames it to a UUID-based final filename and registers it under
+// scope. The file is owned by the store and deleted by ReleaseAll.
+// This uses the same temp-write + sync + rename pattern as fileutil.WriteFileAtomic
+// to guarantee flash-storage durability.
+func (s *FileMediaStore) StoreFromReader(r io.Reader, meta MediaMeta, scope, dir string) (string, error) {
+	if err := os.MkdirAll(dir, 0o755); err != nil {
+		return "", fmt.Errorf("media store: mkdir: %w", err)
+	}
+	tmpPath := filepath.Join(dir, ".tmp-picoclaw-"+uuid.New().String())
+	tmpFile, err := os.OpenFile(tmpPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
+	if err != nil {
+		return "", fmt.Errorf("media store: create temp: %w", err)
+	}
+	cleanup := true
+	defer func() {
+		if cleanup {
+			tmpFile.Close()
+			os.Remove(tmpPath)
+		}
+	}()
+
+	if _, err = io.Copy(tmpFile, r); err != nil {
+		return "", fmt.Errorf("media store: write: %w", err)
+	}
+	if err = tmpFile.Sync(); err != nil {
+		return "", fmt.Errorf("media store: sync: %w", err)
+	}
+	if err = tmpFile.Close(); err != nil {
+		return "", fmt.Errorf("media store: close: %w", err)
+	}
+	cleanup = false
+
+	// Atomic rename to a UUID-based final name (preserving the original extension).
+	ext := filepath.Ext(meta.Filename)
+	finalPath := filepath.Join(dir, uuid.New().String()+ext)
+	if err = os.Rename(tmpPath, finalPath); err != nil {
+		os.Remove(tmpPath)
+		return "", fmt.Errorf("media store: rename: %w", err)
+	}
+	// Sync the directory so the rename is durable.
+	if d, openErr := os.Open(dir); openErr == nil {
+		_ = d.Sync()
+		d.Close()
+	}
+
+	ref := "media://" + uuid.New().String()
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.refs[ref] = mediaEntry{path: finalPath, meta: meta, storedAt: s.nowFunc()}
+	if s.scopeToRefs[scope] == nil {
+		s.scopeToRefs[scope] = make(map[string]struct{})
+	}
+	s.scopeToRefs[scope][ref] = struct{}{}
+	s.refToScope[ref] = scope
+	return ref, nil
+}
+
 // Resolve returns the local path for the given ref.
 func (s *FileMediaStore) Resolve(ref string) (string, error) {
 	s.mu.RLock()