From 5d02404fdd81c9bde3a77efb80cf6319cf309f84 Mon Sep 17 00:00:00 2001 From: wizardchen Date: Fri, 29 May 2026 12:21:17 +0800 Subject: [PATCH] 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. --- internal/application/service/file/cos.go | 30 ++++++-- .../service/file/cos_parse_test.go | 22 ++++++ internal/im/content_rewrite_test.go | 74 +++++++++++++++++++ internal/im/qaqueue.go | 6 +- internal/im/service.go | 71 ++++++++++++------ 5 files changed, 167 insertions(+), 36 deletions(-) create mode 100644 internal/application/service/file/cos_parse_test.go diff --git a/internal/application/service/file/cos.go b/internal/application/service/file/cos.go index 3f933040..94f45917 100644 --- a/internal/application/service/file/cos.go +++ b/internal/application/service/file/cos.go @@ -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) } diff --git a/internal/application/service/file/cos_parse_test.go b/internal/application/service/file/cos_parse_test.go new file mode 100644 index 00000000..11420873 --- /dev/null +++ b/internal/application/service/file/cos_parse_test.go @@ -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) +} diff --git a/internal/im/content_rewrite_test.go b/internal/im/content_rewrite_test.go index d7955e21..2470d42c 100644 --- a/internal/im/content_rewrite_test.go +++ b/internal/im/content_rewrite_test.go @@ -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 diff --git a/internal/im/qaqueue.go b/internal/im/qaqueue.go index cc691573..13ff57ce 100644 --- a/internal/im/qaqueue.go +++ b/internal/im/qaqueue.go @@ -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 diff --git a/internal/im/service.go b/internal/im/service.go index 1519cc1d..89dc7cb7 100644 --- a/internal/im/service.go +++ b/internal/im/service.go @@ -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 XML blocks back to plain markdown // 2. Strip and 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.