From e0823eff31d8ab6a82a0b5201c4b60bbced65a75 Mon Sep 17 00:00:00 2001 From: wizardchen Date: Fri, 29 May 2026 13:32:04 +0800 Subject: [PATCH] feat: enhance IM file service tests and refactor storage URL handling This commit introduces a new test suite for the IM file service, including a stub implementation for testing purposes. It adds tests for resolving file services based on storage providers and ensures proper fallback mechanisms for MinIO URLs. Additionally, the `rewriteStorageURLs` and `cleanIMContent` functions have been refactored to utilize a resolver for improved caching and efficiency. These changes enhance the robustness of file service handling and improve test coverage for various storage scenarios. --- .../service/file/cos_parse_test.go | 7 + internal/im/content_rewrite_test.go | 34 ++++- internal/im/im_file_service_test.go | 122 ++++++++++++++++++ internal/im/service.go | 112 ++++++++++++---- 4 files changed, 243 insertions(+), 32 deletions(-) create mode 100644 internal/im/im_file_service_test.go diff --git a/internal/application/service/file/cos_parse_test.go b/internal/application/service/file/cos_parse_test.go index 11420873..ab577836 100644 --- a/internal/application/service/file/cos_parse_test.go +++ b/internal/application/service/file/cos_parse_test.go @@ -20,3 +20,10 @@ func TestParseCosObjectName_CosScheme(t *testing.T) { require.NoError(t, err) assert.Equal(t, "weknora/10000/exports/a.png", key) } + +func TestParseCosObjectName_RejectsMinioScheme(t *testing.T) { + svc := &cosFileService{bucketURL: "https://b.cos.ap-shanghai.myqcloud.com/"} + _, err := svc.parseCosObjectName("minio://wizard-test/10000/exports/img.png") + require.Error(t, err) + assert.Contains(t, err.Error(), "minio") +} diff --git a/internal/im/content_rewrite_test.go b/internal/im/content_rewrite_test.go index 0450fd21..cf0749de 100644 --- a/internal/im/content_rewrite_test.go +++ b/internal/im/content_rewrite_test.go @@ -177,7 +177,7 @@ func TestResolveIMFileServiceForPath_LocalSchemeDespiteCOSDefault(t *testing.T) }, }, } - svc := resolveIMFileServiceForPath(tenant, "local://10000/exports/img.png") + svc := resolveIMFileServiceForPath(tenant, "local://10000/exports/img.png", nil) require.NotNil(t, svc) got, err := svc.GetFileURL(context.Background(), "local://10000/exports/img.png") require.NoError(t, err) @@ -201,7 +201,7 @@ func TestRewriteStorageURLs_LocalUsesPresignedAPI(t *testing.T) { } in := "![img](local://10000/exports/abc.png)" - out := rewriteStorageURLs(context.Background(), in, tenant) + out := rewriteStorageURLs(context.Background(), in, newIMFileServiceResolver(tenant, nil)) assert.Contains(t, out, "/api/v1/files/presigned") assert.NotContains(t, out, "myqcloud.com") } @@ -222,11 +222,11 @@ func TestRewriteStorageURLs_COSPathNotSignedAsLocalKey(t *testing.T) { }, } path := "cos://test-bucket/ap-shanghai/weknora/10000/exports/abc.png" - svc := resolveIMFileServiceForPath(tenant, path) + svc := resolveIMFileServiceForPath(tenant, path, nil) require.NotNil(t, svc) in := "![img](" + path + ")" - out := rewriteStorageURLs(context.Background(), in, tenant) + out := rewriteStorageURLs(context.Background(), in, newIMFileServiceResolver(tenant, nil)) if out != in { assert.False(t, strings.Contains(out, "local%3A"), "COS URL must not treat local:// as object key") } @@ -245,6 +245,7 @@ func TestFindIncompleteMarkdownImage(t *testing.T) { {"bare provider suffix without markdown", "text minio://wizard-test/10000/exp", -1}, {"two images complete", "![a](local://1/a.png) ![b](local://1/b.png)", -1}, {"first complete second incomplete", "![a](local://1/a.png) ![b](minio://part", 22}, + {"bracket inside alt text", "![a[b]](minio://wizard-test/10000/part", 0}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -436,11 +437,32 @@ func TestCleanIMContent_AfterStreamReassembly(t *testing.T) { sent, rem := simulateIMStreamFlush(parts) require.Empty(t, rem) joined := strings.Join(sent, "") - out := cleanIMContent(context.Background(), joined, tenant) + out := cleanIMContent(context.Background(), joined, tenant, nil) assert.Contains(t, out, "/api/v1/files/presigned") assert.NotContains(t, out, "local://10000") } +func TestHoldbackCutoff_BracketInAlt(t *testing.T) { + in := "text ![a[b]](minio://wizard-test/10000/exp" + want := strings.Index(in, "![a[b]]") + require.GreaterOrEqual(t, want, 0) + assert.Equal(t, want, holdbackCutoff(in)) +} + +func TestSimulateIMStreamFlush_BracketInAltMiddleImage(t *testing.T) { + part1 := "### 1\n![ok](minio://wizard-test/10000/exports/ok.png)\n\n### 2\n" + part2 := "![a[b]](minio://wizard-test/10000/exports/c91cf852" + part3 := "-suffix.png)\n" + + sent, rem := simulateIMStreamFlush([]string{part1, part2, part3}) + require.Empty(t, rem) + joined := strings.Join(sent, "") + assert.Contains(t, joined, "![a[b]](minio://wizard-test/10000/exports/c91cf852-suffix.png)") + for _, chunk := range sent { + assert.False(t, orphanMarkdownImageOpenRe.MatchString(chunk)) + } +} + func TestRewriteStorageURLs_MultipleImagesInOneChunk(t *testing.T) { t.Setenv("SYSTEM_AES_KEY", "weknora-test-aes-key-32bytes!!!") t.Setenv("APP_EXTERNAL_URL", "https://weknora.example.com") @@ -453,7 +475,7 @@ func TestRewriteStorageURLs_MultipleImagesInOneChunk(t *testing.T) { `![知识助理"知识库"管理视图界面](local://10000/exports/c91cf852.png)` + "\n\n### 3\n\n" + "![c](local://10000/exports/a0423e91.png)\n" - out := rewriteStorageURLs(context.Background(), doc, tenant) + out := rewriteStorageURLs(context.Background(), doc, newIMFileServiceResolver(tenant, nil)) assert.NotContains(t, out, "local://") assert.Equal(t, 3, strings.Count(out, "/api/v1/files/presigned")) } diff --git a/internal/im/im_file_service_test.go b/internal/im/im_file_service_test.go new file mode 100644 index 00000000..6524650b --- /dev/null +++ b/internal/im/im_file_service_test.go @@ -0,0 +1,122 @@ +package im + +import ( + "context" + "io" + "mime/multipart" + "testing" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// stubIMFileService implements interfaces.FileService for IM resolver tests. +type stubIMFileService struct { + getFileURL func(ctx context.Context, filePath string) (string, error) +} + +func (s *stubIMFileService) CheckConnectivity(context.Context) error { return nil } + +func (s *stubIMFileService) SaveFile(context.Context, *multipart.FileHeader, uint64, string) (string, error) { + return "", nil +} + +func (s *stubIMFileService) SaveBytes(context.Context, []byte, uint64, string, bool) (string, error) { + return "", nil +} + +func (s *stubIMFileService) GetFile(context.Context, string) (io.ReadCloser, error) { + return nil, nil +} + +func (s *stubIMFileService) GetFileURL(ctx context.Context, filePath string) (string, error) { + if s.getFileURL != nil { + return s.getFileURL(ctx, filePath) + } + return "https://global-storage.example/" + filePath, nil +} + +func (s *stubIMFileService) DeleteFile(context.Context, string) error { return nil } + +func TestBuildIMFileServiceForProvider_FallbackToGlobal(t *testing.T) { + stub := &stubIMFileService{} + tenant := &types.Tenant{ + StorageEngineConfig: &types.StorageEngineConfig{ + DefaultProvider: "cos", + COS: &types.COSEngineConfig{ + SecretID: "id", + SecretKey: "key", + BucketName: "bucket", + Region: "ap-shanghai", + }, + }, + } + + svc := buildIMFileServiceForProvider(tenant, "minio", stub) + require.NotNil(t, svc) + got, err := svc.GetFileURL(context.Background(), "minio://wizard-test/10000/exports/a.png") + require.NoError(t, err) + assert.Equal(t, "https://global-storage.example/minio://wizard-test/10000/exports/a.png", got) +} + +func TestIMFileServiceResolver_CachesPerProvider(t *testing.T) { + stub := &stubIMFileService{} + tenant := &types.Tenant{ + StorageEngineConfig: &types.StorageEngineConfig{ + DefaultProvider: "cos", + COS: &types.COSEngineConfig{ + SecretID: "id", + SecretKey: "key", + BucketName: "bucket", + Region: "ap-shanghai", + }, + }, + } + r := newIMFileServiceResolver(tenant, stub) + + svc1 := r.resolve("minio://wizard-test/10000/a.png") + svc2 := r.resolve("minio://wizard-test/10000/b.png") + assert.Same(t, svc1, svc2, "same provider should reuse cached FileService") + + svc3 := r.resolve("local://10000/c.png") + assert.NotSame(t, svc1, svc3, "different provider should use a different service") +} + +func TestRewriteStorageURLs_MinIOFallbackViaGlobal(t *testing.T) { + stub := &stubIMFileService{ + getFileURL: func(_ context.Context, filePath string) (string, error) { + return "https://minio.example/presigned?path=" + filePath, nil + }, + } + tenant := &types.Tenant{ + StorageEngineConfig: &types.StorageEngineConfig{ + DefaultProvider: "cos", + COS: &types.COSEngineConfig{ + SecretID: "id", + SecretKey: "key", + BucketName: "bucket", + Region: "ap-shanghai", + }, + }, + } + in := `![知识助理"知识库"管理视图界面](minio://wizard-test/10000/exports/c91cf852.png)` + resolver := newIMFileServiceResolver(tenant, stub) + out := rewriteStorageURLs(context.Background(), in, resolver) + assert.Contains(t, out, "https://minio.example/presigned") + assert.NotContains(t, out, "](minio://") +} + +func TestCleanIMContent_MinIOFallbackIntegration(t *testing.T) { + stub := &stubIMFileService{ + getFileURL: func(_ context.Context, _ string) (string, error) { + return "https://minio.example/img.png", nil + }, + } + tenant := &types.Tenant{ + StorageEngineConfig: &types.StorageEngineConfig{DefaultProvider: "cos"}, + } + in := "see ![x](minio://wizard-test/10000/exports/x.png) ok" + out := cleanIMContent(context.Background(), in, tenant, stub) + assert.Contains(t, out, "https://minio.example/img.png") +} diff --git a/internal/im/service.go b/internal/im/service.go index 1fe3f4ff..4c1dff03 100644 --- a/internal/im/service.go +++ b/internal/im/service.go @@ -90,9 +90,12 @@ 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, tenant *types.Tenant) string { +func rewriteStorageURLs(ctx context.Context, content string, resolver *imFileServiceResolver) string { + if resolver == nil { + return content + } return storageSchemeRe.ReplaceAllStringFunc(content, func(match string) string { - fileSvc := resolveIMFileServiceForPath(tenant, match) + fileSvc := resolver.resolve(match) if fileSvc == nil { logger.Warnf(ctx, "[IM] rewriteStorageURLs: no file service for src=%s", match) return match @@ -144,6 +147,15 @@ var incompleteMarkdownImageSuffixRe = regexp.MustCompile(`!\[[^\]]*\]\([^)]*$`) // findIncompleteMarkdownImage returns the byte offset of an unclosed ![alt](url // suffix at the end of s, or -1 if none. func findIncompleteMarkdownImage(s string) int { + // Prefer pairing a trailing provider:// fragment with the nearest preceding ![…]( + // so alt text may contain ']' (e.g. ![a[b]](minio://part). + if urlIdx := findIncompleteStorageURL(s); urlIdx >= 0 { + if imgIdx := strings.LastIndex(s[:urlIdx], "!["); imgIdx >= 0 { + if strings.Contains(s[imgIdx:urlIdx], "](") { + return imgIdx + } + } + } loc := incompleteMarkdownImageSuffixRe.FindStringIndex(s) if loc == nil { return -1 @@ -186,10 +198,11 @@ func holdbackCutoff(chunk string) int { // 1. Collapse XML blocks back to plain markdown // 2. Strip and citation tags // 3. Rewrite provider:// URLs to HTTP URLs (scheme-aware per tenant config) -func cleanIMContent(ctx context.Context, content string, tenant *types.Tenant) string { +func cleanIMContent(ctx context.Context, content string, tenant *types.Tenant, defaultFileSvc interfaces.FileService) string { content = stripImageXMLTags(content) content = stripIMCitationTags(content) - content = rewriteStorageURLs(ctx, content, tenant) + resolver := newIMFileServiceResolver(tenant, defaultFileSvc) + content = rewriteStorageURLs(ctx, content, resolver) return content } @@ -201,36 +214,77 @@ func imLocalStorageBaseDir() string { 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 +// imFileServiceResolver resolves and caches FileService instances per storage provider +// for the lifetime of one cleanIMContent / outbound message (avoids re-creating SDK clients +// for every URL in a long answer). +type imFileServiceResolver struct { + tenant *types.Tenant + defaultSvc interfaces.FileService + cache map[string]interfaces.FileService +} + +func newIMFileServiceResolver(tenant *types.Tenant, defaultSvc interfaces.FileService) *imFileServiceResolver { + return &imFileServiceResolver{ + tenant: tenant, + defaultSvc: defaultSvc, + cache: make(map[string]interfaces.FileService), } +} + +func (r *imFileServiceResolver) resolve(filePath string) interfaces.FileService { + provider := types.ParseProviderScheme(filePath) if provider == "" { - if sec != nil { - provider = strings.ToLower(strings.TrimSpace(sec.DefaultProvider)) + if r.tenant != nil && r.tenant.StorageEngineConfig != nil { + provider = strings.ToLower(strings.TrimSpace(r.tenant.StorageEngineConfig.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 + if svc, ok := r.cache[provider]; ok { + return svc + } + svc := buildIMFileServiceForProvider(r.tenant, provider, r.defaultSvc) + if svc != nil { + r.cache[provider] = svc } return svc } +// buildIMFileServiceForProvider selects the FileService for a storage provider. +// filePath scheme wins over tenant DefaultProvider. Falls back to the process-wide +// default FileService (STORAGE_TYPE / env) when tenant config is missing — mirrors +// ImageMultimodalService.resolveFileServiceForPayload (issue #1282). +func buildIMFileServiceForProvider( + tenant *types.Tenant, + provider string, + defaultSvc interfaces.FileService, +) interfaces.FileService { + baseDir := imLocalStorageBaseDir() + var sec *types.StorageEngineConfig + if tenant != nil { + sec = tenant.StorageEngineConfig + } + + svc, _, err := filesvc.NewFileServiceFromStorageConfig(provider, sec, baseDir) + if err == nil { + return svc + } + if provider == "local" { + externalURL := strings.TrimSpace(os.Getenv("APP_EXTERNAL_URL")) + return filesvc.NewLocalFileService(baseDir, externalURL) + } + if defaultSvc != nil { + return defaultSvc + } + return nil +} + +// resolveIMFileServiceForPath is a test/helper entry point without caching. +func resolveIMFileServiceForPath(tenant *types.Tenant, filePath string, defaultSvc interfaces.FileService) interfaces.FileService { + return newIMFileServiceResolver(tenant, defaultSvc).resolve(filePath) +} + const ( // wsLeaderTTL is the TTL for the Redis key used for WebSocket leader election. wsLeaderTTL = 15 * time.Second @@ -303,6 +357,10 @@ type Service struct { // (but NewStreamManager always returns at least a memory implementation). streamManager interfaces.StreamManager + // defaultFileSvc is the process-wide storage backend (STORAGE_TYPE / env). + // Used when tenant StorageEngineConfig cannot build a service for the URL scheme. + defaultFileSvc interfaces.FileService + // cmdRegistry holds all registered slash-commands. cmdRegistry *CommandRegistry @@ -466,6 +524,7 @@ func NewService( kbService interfaces.KnowledgeBaseService, modelService interfaces.ModelService, streamManager interfaces.StreamManager, + defaultFileSvc interfaces.FileService, redisClient *redis.Client, appCfg *config.Config, ) *Service { @@ -491,6 +550,7 @@ func NewService( kbService: kbService, modelService: modelService, streamManager: streamManager, + defaultFileSvc: defaultFileSvc, cmdRegistry: registry, channels: make(map[string]*channelState), adapterFactories: make(map[string]AdapterFactory), @@ -1230,7 +1290,7 @@ func (s *Service) executeQARequest(req *qaRequest) { } reply := &ReplyMessage{ - Content: cleanIMContent(ctx, answer, req.tenant), + Content: cleanIMContent(ctx, answer, req.tenant, s.defaultFileSvc), IsFinal: true, } if err := req.adapter.SendReply(ctx, req.msg, reply); err != nil { @@ -1932,7 +1992,7 @@ func (s *Service) handleMessageStream(ctx context.Context, msg *IncomingMessage, } if chunk != "" { - if err := streamer.SendStreamChunk(ctx, msg, streamID, cleanIMContent(ctx, chunk, tenant)); err != nil { + if err := streamer.SendStreamChunk(ctx, msg, streamID, cleanIMContent(ctx, chunk, tenant, s.defaultFileSvc)); err != nil { logger.Warnf(ctx, "[IM] SendStreamChunk failed: %v", err) } } @@ -2002,7 +2062,7 @@ func (s *Service) fallbackNonStream(ctx context.Context, msg *IncomingMessage, s answer = "抱歉，处理您的问题时出现了异常，请稍后再试。" } - return adapter.SendReply(ctx, msg, &ReplyMessage{Content: cleanIMContent(ctx, answer, tenant), IsFinal: true}) + return adapter.SendReply(ctx, msg, &ReplyMessage{Content: cleanIMContent(ctx, answer, tenant, s.defaultFileSvc), IsFinal: true}) } // runQA executes the WeKnora QA pipeline and returns the full answer text.