feat: add unit tests for COS object name parsing and enhance error handling

This commit introduces unit tests for the `parseCosObjectName` method in the `cosFileService`, ensuring it correctly rejects local scheme URLs and properly parses COS scheme URLs. Additionally, the `parseCosObjectName` method has been updated to return an error for unsupported schemes, improving error handling in the `GetFile` and `DeleteFile` methods. This enhancement ensures more robust handling of file paths in the application.
2026-05-29 12:21:17 +08:00
parent 83808cc5b7
commit 5d02404fdd
5 changed files with 167 additions and 36 deletions
@@ -123,7 +123,10 @@ func (s *cosFileService) SaveFile(ctx context.Context,

 // GetFile retrieves a file from COS storage by its path URL
 func (s *cosFileService) GetFile(ctx context.Context, filePathUrl string) (io.ReadCloser, error) {
-	objectName := s.parseCosObjectName(filePathUrl)
+	objectName, err := s.parseCosObjectName(filePathUrl)
+	if err != nil {
+		return nil, err
+	}
 	if err := utils.SafeObjectKey(objectName); err != nil {
 		return nil, fmt.Errorf("invalid file path: %w", err)
 	}
@@ -136,11 +139,14 @@ func (s *cosFileService) GetFile(ctx context.Context, filePathUrl string) (io.Re

 // DeleteFile removes a file from COS storage
 func (s *cosFileService) DeleteFile(ctx context.Context, filePath string) error {
-	objectName := s.parseCosObjectName(filePath)
+	objectName, err := s.parseCosObjectName(filePath)
+	if err != nil {
+		return err
+	}
 	if err := utils.SafeObjectKey(objectName); err != nil {
 		return fmt.Errorf("invalid file path: %w", err)
 	}
-	_, err := s.client.Object.Delete(ctx, objectName)
+	_, err = s.client.Object.Delete(ctx, objectName)
 	if err != nil {
 		return fmt.Errorf("failed to delete file: %w", err)
 	}
@@ -150,18 +156,23 @@ func (s *cosFileService) DeleteFile(ctx context.Context, filePath string) error
 // parseCosObjectName extracts the object name from:
 // - provider scheme: cos://{bucket}/{region}/{objectKey}
 // - legacy URL: https://bucket.cos.region.myqcloud.com/{objectKey}
-func (s *cosFileService) parseCosObjectName(filePath string) string {
+func (s *cosFileService) parseCosObjectName(filePath string) (string, error) {
+	for _, other := range []string{"local://", "minio://", "s3://", "tos://", "oss://", "ks3://", "obs://"} {
+		if strings.HasPrefix(filePath, other) {
+			return "", fmt.Errorf("cos file service cannot resolve %s path", strings.Split(other, "://")[0])
+		}
+	}
 	// Provider scheme format: cos://{bucket}/{region}/{objectKey}
 	if strings.HasPrefix(filePath, cosScheme) {
 		rest := strings.TrimPrefix(filePath, cosScheme)
 		parts := strings.SplitN(rest, "/", 3)
 		if len(parts) == 3 {
-			return parts[2]
+			return parts[2], nil
 		}
-		return rest
+		return rest, nil
 	}
 	// Legacy format: https://bucket.cos.region.myqcloud.com/{objectKey}
-	return strings.TrimPrefix(filePath, s.bucketURL)
+	return strings.TrimPrefix(filePath, s.bucketURL), nil
 }

 // SaveBytes saves bytes data to COS
@@ -212,7 +223,10 @@ func (s *cosFileService) GetFileURL(ctx context.Context, filePath string) (strin
 		return presignedURL.String(), nil
 	}

-	objectName := s.parseCosObjectName(filePath)
+	objectName, err := s.parseCosObjectName(filePath)
+	if err != nil {
+		return "", err
+	}
 	if err := utils.SafeObjectKey(objectName); err != nil {
 		return "", fmt.Errorf("invalid file path: %w", err)
 	}
@@ -0,0 +1,22 @@
+package file
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestParseCosObjectName_RejectsLocalScheme(t *testing.T) {
+	svc := &cosFileService{bucketURL: "https://b.cos.ap-shanghai.myqcloud.com/"}
+	_, err := svc.parseCosObjectName("local://10000/exports/img.png")
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "local")
+}
+
+func TestParseCosObjectName_CosScheme(t *testing.T) {
+	svc := &cosFileService{bucketURL: "https://b.cos.ap-shanghai.myqcloud.com/"}
+	key, err := svc.parseCosObjectName("cos://bucket/ap-shanghai/weknora/10000/exports/a.png")
+	require.NoError(t, err)
+	assert.Equal(t, "weknora/10000/exports/a.png", key)
+}
@@ -1,9 +1,13 @@
 package im

 import (
+	"context"
+	"strings"
 	"testing"

+	"github.com/Tencent/WeKnora/internal/types"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )

 func TestStripIMCitationTags(t *testing.T) {
@@ -157,6 +161,76 @@ func TestFindIncompleteXMLTag(t *testing.T) {
 	}
 }

+func TestResolveIMFileServiceForPath_LocalSchemeDespiteCOSDefault(t *testing.T) {
+	t.Setenv("SYSTEM_AES_KEY", "weknora-test-aes-key-32bytes!!!")
+	t.Setenv("APP_EXTERNAL_URL", "https://weknora.example.com")
+
+	tenant := &types.Tenant{
+		StorageEngineConfig: &types.StorageEngineConfig{
+			DefaultProvider: "cos",
+			COS: &types.COSEngineConfig{
+				SecretID:   "id",
+				SecretKey:  "key",
+				BucketName: "bucket",
+				Region:     "ap-shanghai",
+			},
+		},
+	}
+	svc := resolveIMFileServiceForPath(tenant, "local://10000/exports/img.png")
+	require.NotNil(t, svc)
+	got, err := svc.GetFileURL(context.Background(), "local://10000/exports/img.png")
+	require.NoError(t, err)
+	assert.Contains(t, got, "/api/v1/files/presigned")
+}
+
+func TestRewriteStorageURLs_LocalUsesPresignedAPI(t *testing.T) {
+	t.Setenv("SYSTEM_AES_KEY", "weknora-test-aes-key-32bytes!!!")
+	t.Setenv("APP_EXTERNAL_URL", "https://weknora.example.com")
+
+	tenant := &types.Tenant{
+		StorageEngineConfig: &types.StorageEngineConfig{
+			DefaultProvider: "cos",
+			COS: &types.COSEngineConfig{
+				SecretID:   "id",
+				SecretKey:  "key",
+				BucketName: "bucket",
+				Region:     "ap-shanghai",
+			},
+		},
+	}
+
+	in := "![img](local://10000/exports/abc.png)"
+	out := rewriteStorageURLs(context.Background(), in, tenant)
+	assert.Contains(t, out, "/api/v1/files/presigned")
+	assert.NotContains(t, out, "myqcloud.com")
+}
+
+func TestRewriteStorageURLs_COSPathNotSignedAsLocalKey(t *testing.T) {
+	// Without real COS credentials, GetFileURL may fail; ensure we never embed
+	// local:// as a COS object key when rewriting fails.
+	tenant := &types.Tenant{
+		StorageEngineConfig: &types.StorageEngineConfig{
+			DefaultProvider: "cos",
+			COS: &types.COSEngineConfig{
+				SecretID:   "id",
+				SecretKey:  "key",
+				BucketName: "test-bucket",
+				Region:     "ap-shanghai",
+				PathPrefix: "weknora",
+			},
+		},
+	}
+	path := "cos://test-bucket/ap-shanghai/weknora/10000/exports/abc.png"
+	svc := resolveIMFileServiceForPath(tenant, path)
+	require.NotNil(t, svc)
+
+	in := "![img](" + path + ")"
+	out := rewriteStorageURLs(context.Background(), in, tenant)
+	if out != in {
+		assert.False(t, strings.Contains(out, "local%3A"), "COS URL must not treat local:// as object key")
+	}
+}
+
 func TestHoldbackCutoff(t *testing.T) {
 	tests := []struct {
 		name string
@@ -9,7 +9,6 @@ import (

 	"github.com/Tencent/WeKnora/internal/logger"
 	"github.com/Tencent/WeKnora/internal/types"
-	"github.com/Tencent/WeKnora/internal/types/interfaces"
 	"github.com/redis/go-redis/v9"
 )

@@ -44,9 +43,8 @@ type qaRequest struct {
 	channel   *IMChannel
 	channelID string

-	// fileSvc resolves storage URLs to HTTP URLs for the tenant's storage backend.
-	// May be nil if the tenant has no storage config.
-	fileSvc interfaces.FileService
+	// tenant is used to resolve provider:// URLs in outbound replies (scheme-aware).
+	tenant *types.Tenant

 	// userKey is "channelID:userID:chatID", used for per-user limits and /stop.
 	userKey    string
@@ -90,11 +90,13 @@ var storageSchemeRe = regexp.MustCompile(`\b(local|minio|s3|cos|tos|oss)://[^\s)
 //   - Failure or no-op rewrite logs at WARN. The no-op case typically means
 //     APP_EXTERNAL_URL is not configured for the local backend, which is
 //     the most common cause of "image broken in IM" reports.
-func rewriteStorageURLs(ctx context.Context, content string, fileSvc interfaces.FileService) string {
-	if fileSvc == nil {
-		return content
-	}
+func rewriteStorageURLs(ctx context.Context, content string, tenant *types.Tenant) string {
 	return storageSchemeRe.ReplaceAllStringFunc(content, func(match string) string {
+		fileSvc := resolveIMFileServiceForPath(tenant, match)
+		if fileSvc == nil {
+			logger.Warnf(ctx, "[IM] rewriteStorageURLs: no file service for src=%s", match)
+			return match
+		}
 		httpURL, err := fileSvc.GetFileURL(ctx, match)
 		if err != nil {
 			logger.Warnf(ctx, "[IM] rewriteStorageURLs failed: src=%s err=%v", match, err)
@@ -165,29 +167,50 @@ func holdbackCutoff(chunk string) int {
 // cleanIMContent applies all IM-specific content transformations:
 //  1. Collapse <image> XML blocks back to plain markdown
 //  2. Strip <kb/> and <web/> citation tags
-//  3. Rewrite provider:// URLs to HTTP URLs (if fileSvc is available)
-func cleanIMContent(ctx context.Context, content string, fileSvc interfaces.FileService) string {
+//  3. Rewrite provider:// URLs to HTTP URLs (scheme-aware per tenant config)
+func cleanIMContent(ctx context.Context, content string, tenant *types.Tenant) string {
 	content = stripImageXMLTags(content)
 	content = stripIMCitationTags(content)
-	content = rewriteStorageURLs(ctx, content, fileSvc)
+	content = rewriteStorageURLs(ctx, content, tenant)
 	return content
 }

-// buildTenantFileService creates a FileService for the given tenant's storage config.
-// Returns nil if the tenant has no storage config or if creation fails.
-func buildTenantFileService(tenant *types.Tenant) interfaces.FileService {
-	if tenant == nil {
-		return nil
-	}
-	baseDir := os.Getenv("LOCAL_STORAGE_BASE_DIR")
+func imLocalStorageBaseDir() string {
+	baseDir := strings.TrimSpace(os.Getenv("LOCAL_STORAGE_BASE_DIR"))
 	if baseDir == "" {
 		baseDir = "/data/files"
 	}
-	fileSvc, _, err := filesvc.NewFileServiceFromStorageConfig("", tenant.StorageEngineConfig, baseDir)
+	return baseDir
+}
+
+// resolveIMFileServiceForPath selects the FileService that matches the URL scheme in
+// filePath (e.g. local:// → local backend), not the tenant's default provider.
+// Mirrors /api/v1/files/presigned and ImageMultimodalService.resolveFileServiceForPayload.
+func resolveIMFileServiceForPath(tenant *types.Tenant, filePath string) interfaces.FileService {
+	baseDir := imLocalStorageBaseDir()
+	provider := types.ParseProviderScheme(filePath)
+	var sec *types.StorageEngineConfig
+	if tenant != nil {
+		sec = tenant.StorageEngineConfig
+	}
+	if provider == "" {
+		if sec != nil {
+			provider = strings.ToLower(strings.TrimSpace(sec.DefaultProvider))
+		}
+		if provider == "" {
+			return nil
+		}
+	}
+
+	svc, _, err := filesvc.NewFileServiceFromStorageConfig(provider, sec, baseDir)
 	if err != nil {
+		if provider == "local" {
+			externalURL := strings.TrimSpace(os.Getenv("APP_EXTERNAL_URL"))
+			return filesvc.NewLocalFileService(baseDir, externalURL)
+		}
 		return nil
 	}
-	return fileSvc
+	return svc
 }

 const (
@@ -1100,7 +1123,7 @@ func (s *Service) HandleMessage(ctx context.Context, msg *IncomingMessage, chann
 		adapter:   adapter,
 		channel:   channel,
 		channelID: channelID,
-		fileSvc:   buildTenantFileService(tenant),
+		tenant:    tenant,
 		userKey:   userKey,
 	}

@@ -1172,7 +1195,7 @@ func (s *Service) executeQARequest(req *qaRequest) {
 	// If the adapter supports streaming and output is not "full", use streaming.
 	if !streamDisabled {
 		if streamer, ok := req.adapter.(StreamSender); ok {
-			if err := s.handleMessageStream(ctx, req.msg, req.session, req.agent, kbIDs, streamer, req.adapter, req.userKey, req.fileSvc); err != nil {
+			if err := s.handleMessageStream(ctx, req.msg, req.session, req.agent, kbIDs, streamer, req.adapter, req.userKey, req.tenant); err != nil {
 				span.SetStatus(codes.Error, err.Error())
 				logger.Errorf(ctx, "[IM] Stream QA failed: %v", err)
 			}
@@ -1189,7 +1212,7 @@ func (s *Service) executeQARequest(req *qaRequest) {
 	}

 	reply := &ReplyMessage{
-		Content: cleanIMContent(ctx, answer, req.fileSvc),
+		Content: cleanIMContent(ctx, answer, req.tenant),
 		IsFinal: true,
 	}
 	if err := req.adapter.SendReply(ctx, req.msg, reply); err != nil {
@@ -1641,12 +1664,12 @@ func briefToolSummary(output string) string {
 // handleMessageStream runs the QA pipeline and streams answer chunks to the IM platform
 // in real-time via the StreamSender interface. Chunks are batched at streamFlushInterval
 // to avoid API rate-limiting.
-func (s *Service) handleMessageStream(ctx context.Context, msg *IncomingMessage, session *types.Session, customAgent *types.CustomAgent, kbIDs []string, streamer StreamSender, adapter Adapter, userKey string, fileSvc interfaces.FileService) error {
+func (s *Service) handleMessageStream(ctx context.Context, msg *IncomingMessage, session *types.Session, customAgent *types.CustomAgent, kbIDs []string, streamer StreamSender, adapter Adapter, userKey string, tenant *types.Tenant) error {
 	// Start the stream on the IM platform (e.g., create Feishu streaming card)
 	streamID, err := streamer.StartStream(ctx, msg)
 	if err != nil {
 		logger.Warnf(ctx, "[IM] StartStream failed, falling back to non-streaming: %v", err)
-		return s.fallbackNonStream(ctx, msg, session, customAgent, kbIDs, adapter, userKey, fileSvc)
+		return s.fallbackNonStream(ctx, msg, session, customAgent, kbIDs, adapter, userKey, tenant)
 	}

 	// Prepare the QA pipeline
@@ -1891,7 +1914,7 @@ func (s *Service) handleMessageStream(ctx context.Context, msg *IncomingMessage,
 		}

 		if chunk != "" {
-			if err := streamer.SendStreamChunk(ctx, msg, streamID, cleanIMContent(ctx, chunk, fileSvc)); err != nil {
+			if err := streamer.SendStreamChunk(ctx, msg, streamID, cleanIMContent(ctx, chunk, tenant)); err != nil {
 				logger.Warnf(ctx, "[IM] SendStreamChunk failed: %v", err)
 			}
 		}
@@ -1954,14 +1977,14 @@ loop:
 }

 // fallbackNonStream is used when streaming initialization fails.
-func (s *Service) fallbackNonStream(ctx context.Context, msg *IncomingMessage, session *types.Session, customAgent *types.CustomAgent, kbIDs []string, adapter Adapter, userKey string, fileSvc interfaces.FileService) error {
+func (s *Service) fallbackNonStream(ctx context.Context, msg *IncomingMessage, session *types.Session, customAgent *types.CustomAgent, kbIDs []string, adapter Adapter, userKey string, tenant *types.Tenant) error {
 	answer, err := s.runQA(ctx, session, msg.Content, customAgent, kbIDs, userKey, msg.Quote)
 	if err != nil {
 		logger.Errorf(ctx, "[IM] QA fallback failed: %v", err)
 		answer = "抱歉，处理您的问题时出现了异常，请稍后再试。"
 	}

-	return adapter.SendReply(ctx, msg, &ReplyMessage{Content: cleanIMContent(ctx, answer, fileSvc), IsFinal: true})
+	return adapter.SendReply(ctx, msg, &ReplyMessage{Content: cleanIMContent(ctx, answer, tenant), IsFinal: true})
 }

 // runQA executes the WeKnora QA pipeline and returns the full answer text.