diff --git a/frontend/src/components/doc-content.vue b/frontend/src/components/doc-content.vue index de189e97..ca4a0a71 100644 --- a/frontend/src/components/doc-content.vue +++ b/frontend/src/components/doc-content.vue @@ -54,6 +54,7 @@ marked.use({ }); const renderer = new marked.Renderer(); let page = 1; +let loadingChunks = false; let doc = null; let down = ref() let mdContentWrap = ref() @@ -136,9 +137,13 @@ onMounted(() => { doc.addEventListener('scroll', handleDetailsScroll); }) }) -onUpdated(() => { - page = 1 -}) +watch(() => props.details?.id, () => { + page = 1; + loadingChunks = false; +}); +watch(() => props.details?.md?.length, () => { + loadingChunks = false; +}); onUnmounted(() => { doc.removeEventListener('scroll', handleDetailsScroll); }) @@ -559,12 +564,13 @@ const downloadFile = () => { }); }; const handleDetailsScroll = () => { - if (doc) { - let pageNum = Math.ceil(props.details.total / 20); + if (doc && !loadingChunks) { + let pageNum = Math.ceil(props.details.total / 25); const { scrollTop, scrollHeight, clientHeight } = doc; if (scrollTop + clientHeight >= scrollHeight) { - page++; - if (props.details.md.length < props.details.total && page <= pageNum) { + if (props.details.md.length < props.details.total && page + 1 <= pageNum) { + page++; + loadingChunks = true; emit("getDoc", page); } } diff --git a/internal/application/service/agent_service.go b/internal/application/service/agent_service.go index 90219220..f7f50c28 100644 --- a/internal/application/service/agent_service.go +++ b/internal/application/service/agent_service.go @@ -109,7 +109,7 @@ func (s *agentService) CreateAgentEngine( // Register MCP tools from enabled services for this tenant tenantID := uint64(0) - if tid, ok := ctx.Value(types.TenantIDContextKey).(uint64); ok { + if tid, ok := types.TenantIDFromContext(ctx); ok { tenantID = tid } if tenantID > 0 && s.mcpServiceService != nil && s.mcpManager != nil { @@ -550,7 +550,7 @@ func (s *agentService) getSelectedDocumentInfos(ctx context.Context, knowledgeID // Get tenant ID from context tenantID := uint64(0) - if tid, ok := ctx.Value(types.TenantIDContextKey).(uint64); ok { + if tid, ok := types.TenantIDFromContext(ctx); ok { tenantID = tid } diff --git a/internal/application/service/chat_pipline/memory.go b/internal/application/service/chat_pipline/memory.go index 15bbf0f0..58ba2964 100644 --- a/internal/application/service/chat_pipline/memory.go +++ b/internal/application/service/chat_pipline/memory.go @@ -104,12 +104,12 @@ func (p *MemoryPlugin) handleStorage( {Role: "user", Content: chatManage.Query}, {Role: "assistant", Content: chatManage.ChatResponse.Content}, } - // Capture UserID and SessionID for goroutine userID := chatManage.UserID sessionID := chatManage.SessionID + bgCtx := context.WithoutCancel(ctx) go func() { - if err := p.memoryService.AddEpisode(ctx, userID, sessionID, messages); err != nil { - logger.Errorf(ctx, "failed to add episode: %v", err) + if err := p.memoryService.AddEpisode(bgCtx, userID, sessionID, messages); err != nil { + logger.Errorf(bgCtx, "failed to add episode: %v", err) } }() return nil @@ -118,11 +118,11 @@ func (p *MemoryPlugin) handleStorage( // If streaming, subscribe to events if chatManage.EventBus != nil { var fullResponse string - // Capture UserID and SessionID for closure userID := chatManage.UserID sessionID := chatManage.SessionID + bgCtx := context.WithoutCancel(ctx) - chatManage.EventBus.On(types.EventType(event.EventAgentFinalAnswer), func(ctx context.Context, evt types.Event) error { + chatManage.EventBus.On(types.EventType(event.EventAgentFinalAnswer), func(_ context.Context, evt types.Event) error { data, ok := evt.Data.(event.AgentFinalAnswerData) if !ok { return nil @@ -134,8 +134,8 @@ func (p *MemoryPlugin) handleStorage( {Role: "assistant", Content: fullResponse}, } go func() { - if err := p.memoryService.AddEpisode(ctx, userID, sessionID, messages); err != nil { - logger.Errorf(ctx, "failed to add episode: %v", err) + if err := p.memoryService.AddEpisode(bgCtx, userID, sessionID, messages); err != nil { + logger.Errorf(bgCtx, "failed to add episode: %v", err) } }() } diff --git a/internal/application/service/chat_pipline/merge.go b/internal/application/service/chat_pipline/merge.go index c1ed13ee..0ff3daca 100644 --- a/internal/application/service/chat_pipline/merge.go +++ b/internal/application/service/chat_pipline/merge.go @@ -220,7 +220,7 @@ func (p *PluginMerge) populateFAQAnswers( return results } - tenantID, _ := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, _ := types.TenantIDFromContext(ctx) if tenantID == 0 && chatManage != nil { tenantID = chatManage.TenantID } @@ -351,7 +351,7 @@ func (p *PluginMerge) expandShortContextWithNeighbors( return results } - tenantID, _ := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, _ := types.TenantIDFromContext(ctx) if tenantID == 0 && chatManage != nil { tenantID = chatManage.TenantID } diff --git a/internal/application/service/chat_pipline/search.go b/internal/application/service/chat_pipline/search.go index df55dd51..302ff414 100644 --- a/internal/application/service/chat_pipline/search.go +++ b/internal/application/service/chat_pipline/search.go @@ -503,7 +503,7 @@ func (p *PluginSearch) searchWebIfEnabled(ctx context.Context, chatManage *types if !chatManage.WebSearchEnabled || p.webSearchService == nil || p.tenantService == nil { return nil } - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil || tenant.WebSearchConfig == nil || tenant.WebSearchConfig.Provider == "" { pipelineWarn(ctx, "Search", "web_config_missing", map[string]interface{}{ "tenant_id": chatManage.TenantID, diff --git a/internal/application/service/chat_pipline/search_entity.go b/internal/application/service/chat_pipline/search_entity.go index b5542527..a33c2b16 100644 --- a/internal/application/service/chat_pipline/search_entity.go +++ b/internal/application/service/chat_pipline/search_entity.go @@ -138,7 +138,7 @@ func (p *PluginSearchEntity) OnEvent(ctx context.Context, logger.Infof(ctx, "No new chunk found") return next() } - chunks, err := p.chunkRepo.ListChunksByID(ctx, ctx.Value(types.TenantIDContextKey).(uint64), chunkIDs) + chunks, err := p.chunkRepo.ListChunksByID(ctx, types.MustTenantIDFromContext(ctx), chunkIDs) if err != nil { logger.Errorf(ctx, "Failed to list chunks, session_id: %s, error: %v", chatManage.SessionID, err) return next() @@ -149,7 +149,7 @@ func (p *PluginSearchEntity) OnEvent(ctx context.Context, } knowledges, err := p.knowledgeRepo.GetKnowledgeBatch( ctx, - ctx.Value(types.TenantIDContextKey).(uint64), + types.MustTenantIDFromContext(ctx), knowledgeIDs, ) if err != nil { diff --git a/internal/application/service/chunk.go b/internal/application/service/chunk.go index b322c428..9533c70f 100644 --- a/internal/application/service/chunk.go +++ b/internal/application/service/chunk.go @@ -86,7 +86,7 @@ func (s *chunkService) CreateChunks(ctx context.Context, chunks []*types.Chunk) // - *types.Chunk: Retrieved chunk if found // - error: Any error encountered during retrieval func (s *chunkService) GetChunkByID(ctx context.Context, id string) (*types.Chunk, error) { - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Getting chunk by ID, ID: %s, tenant ID: %d", id, tenantID) chunk, err := s.chunkRepository.GetChunkByID(ctx, tenantID, id) if err != nil { @@ -126,7 +126,7 @@ func (s *chunkService) ListChunksByKnowledgeID(ctx context.Context, knowledgeID logger.Info(ctx, "Start listing chunks by knowledge ID") logger.Infof(ctx, "Knowledge ID: %s", knowledgeID) - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Tenant ID: %d", tenantID) chunks, err := s.chunkRepository.ListChunksByKnowledgeID(ctx, tenantID, knowledgeID) @@ -155,7 +155,7 @@ func (s *chunkService) ListChunksByKnowledgeID(ctx context.Context, knowledgeID func (s *chunkService) ListPagedChunksByKnowledgeID(ctx context.Context, knowledgeID string, page *types.Pagination, chunkType []types.ChunkType, ) (*types.PageResult, error) { - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) chunks, total, err := s.chunkRepository.ListPagedChunksByKnowledgeID( ctx, tenantID, @@ -236,7 +236,7 @@ func (s *chunkService) UpdateChunks(ctx context.Context, chunks []*types.Chunk) // Returns: // - error: Any error encountered during deletion func (s *chunkService) DeleteChunk(ctx context.Context, id string) error { - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) err := s.chunkRepository.DeleteChunk(ctx, tenantID, id) if err != nil { logger.ErrorWithFields(ctx, err, map[string]interface{}{ @@ -263,7 +263,7 @@ func (s *chunkService) DeleteChunks(ctx context.Context, ids []string) error { logger.Info(ctx, "Start deleting chunks in batch") logger.Infof(ctx, "Deleting %d chunks", len(ids)) - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Tenant ID: %d", tenantID) err := s.chunkRepository.DeleteChunks(ctx, tenantID, ids) @@ -291,7 +291,7 @@ func (s *chunkService) DeleteChunksByKnowledgeID(ctx context.Context, knowledgeI logger.Info(ctx, "Start deleting all chunks by knowledge ID") logger.Infof(ctx, "Knowledge ID: %s", knowledgeID) - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Tenant ID: %d", tenantID) err := s.chunkRepository.DeleteChunksByKnowledgeID(ctx, tenantID, knowledgeID) @@ -311,7 +311,7 @@ func (s *chunkService) DeleteByKnowledgeList(ctx context.Context, ids []string) logger.Info(ctx, "Start deleting all chunks by knowledge IDs") logger.Infof(ctx, "Knowledge IDs: %v", ids) - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Tenant ID: %d", tenantID) err := s.chunkRepository.DeleteByKnowledgeList(ctx, tenantID, ids) @@ -353,7 +353,7 @@ func (s *chunkService) ListChunkByParentID( func (s *chunkService) DeleteGeneratedQuestion(ctx context.Context, chunkID string, questionID string) error { logger.Infof(ctx, "Deleting generated question, chunk ID: %s, question ID: %s", chunkID, questionID) - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // 1. Get the chunk chunk, err := s.chunkRepository.GetChunkByID(ctx, tenantID, chunkID) @@ -404,7 +404,7 @@ func (s *chunkService) DeleteGeneratedQuestion(ctx context.Context, chunkID stri // The source_id format is: {chunk_id}-{question_id} sourceID := fmt.Sprintf("%s-%s", chunkID, questionID) - tenantInfo := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenantInfo, _ := types.TenantInfoFromContext(ctx) retrieveEngine, err := retriever.NewCompositeRetrieveEngine(s.retrieveEngine, tenantInfo.GetEffectiveEngines()) if err != nil { logger.ErrorWithFields(ctx, err, map[string]interface{}{ diff --git a/internal/application/service/custom_agent.go b/internal/application/service/custom_agent.go index 65fb0b5c..433e73ba 100644 --- a/internal/application/service/custom_agent.go +++ b/internal/application/service/custom_agent.go @@ -46,7 +46,7 @@ func (s *customAgentService) CreateAgent(ctx context.Context, agent *types.Custo } // Get tenant ID from context - tenantID, ok := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, ok := types.TenantIDFromContext(ctx) if !ok { return nil, ErrInvalidTenantID } @@ -90,7 +90,7 @@ func (s *customAgentService) GetAgentByID(ctx context.Context, id string) (*type } // Get tenant ID from context - tenantID, ok := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, ok := types.TenantIDFromContext(ctx) if !ok { return nil, ErrInvalidTenantID } @@ -142,7 +142,7 @@ func (s *customAgentService) GetAgentByIDAndTenant(ctx context.Context, id strin // ListAgents lists all agents for the current tenant (including built-in agents) func (s *customAgentService) ListAgents(ctx context.Context) ([]*types.CustomAgent, error) { - tenantID, ok := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, ok := types.TenantIDFromContext(ctx) if !ok { return nil, ErrInvalidTenantID } @@ -204,7 +204,7 @@ func (s *customAgentService) UpdateAgent(ctx context.Context, agent *types.Custo } // Get tenant ID from context - tenantID, ok := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, ok := types.TenantIDFromContext(ctx) if !ok { return nil, ErrInvalidTenantID } @@ -330,7 +330,7 @@ func (s *customAgentService) DeleteAgent(ctx context.Context, id string) error { } // Get tenant ID from context - tenantID, ok := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, ok := types.TenantIDFromContext(ctx) if !ok { return ErrInvalidTenantID } @@ -370,7 +370,7 @@ func (s *customAgentService) CopyAgent(ctx context.Context, id string) (*types.C } // Get tenant ID from context - tenantID, ok := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, ok := types.TenantIDFromContext(ctx) if !ok { return nil, ErrInvalidTenantID } diff --git a/internal/application/service/evaluation.go b/internal/application/service/evaluation.go index d6d1cb92..407f462e 100644 --- a/internal/application/service/evaluation.go +++ b/internal/application/service/evaluation.go @@ -109,7 +109,7 @@ func (e *EvaluationService) EvaluationResult(ctx context.Context, taskID string) return nil, err } - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof( ctx, "Checking tenant ID match, task tenant ID: %d, current tenant ID: %d", @@ -138,7 +138,7 @@ func (e *EvaluationService) Evaluation(ctx context.Context, datasetID, knowledgeBaseID, chatModelID, rerankModelID) // Get tenant ID from context for multi-tenancy support - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Tenant ID: %d", tenantID) // Handle knowledge base creation if not provided diff --git a/internal/application/service/file/minio.go b/internal/application/service/file/minio.go index 0fcf2f42..2b3e617d 100644 --- a/internal/application/service/file/minio.go +++ b/internal/application/service/file/minio.go @@ -7,6 +7,7 @@ import ( "io" "mime/multipart" "path/filepath" + "strings" "time" "github.com/Tencent/WeKnora/internal/types/interfaces" @@ -89,6 +90,27 @@ func CheckMinioConnectivity(ctx context.Context, endpoint, accessKeyID, secretAc return svc.CheckConnectivity(ctx) } +// parseMinioFilePath extracts the object name from a "minio://bucket/key" path, +// validates the bucket matches this service instance, and checks for path traversal. +func (s *minioFileService) parseMinioFilePath(filePath string) (string, error) { + const prefix = "minio://" + if !strings.HasPrefix(filePath, prefix) { + return "", fmt.Errorf("invalid MinIO file path: %s", filePath) + } + rest := strings.TrimPrefix(filePath, prefix) + parts := strings.SplitN(rest, "/", 2) + if len(parts) != 2 || parts[0] == "" || parts[1] == "" { + return "", fmt.Errorf("invalid MinIO file path: %s", filePath) + } + if parts[0] != s.bucketName { + return "", fmt.Errorf("bucket mismatch in path: got %s, want %s", parts[0], s.bucketName) + } + if err := utils.SafeObjectKey(parts[1]); err != nil { + return "", fmt.Errorf("invalid file path: %w", err) + } + return parts[1], nil +} + // SaveFile saves a file to MinIO func (s *minioFileService) SaveFile(ctx context.Context, file *multipart.FileHeader, tenantID uint64, knowledgeID string, @@ -118,55 +140,28 @@ func (s *minioFileService) SaveFile(ctx context.Context, // GetFile gets a file from MinIO func (s *minioFileService) GetFile(ctx context.Context, filePath string) (io.ReadCloser, error) { - // Parse MinIO path - // Format: minio://bucketName/objectName - if len(filePath) < 9 || filePath[:8] != "minio://" { - return nil, fmt.Errorf("invalid MinIO file path: %s", filePath) + objectName, err := s.parseMinioFilePath(filePath) + if err != nil { + return nil, err } - - // Extract object name - objectName := filePath[9+len(s.bucketName):] - if len(objectName) > 0 && objectName[0] == '/' { - objectName = objectName[1:] - } - if err := utils.SafeObjectKey(objectName); err != nil { - return nil, fmt.Errorf("invalid file path: %w", err) - } - - // Get object obj, err := s.client.GetObject(ctx, s.bucketName, objectName, minio.GetObjectOptions{}) if err != nil { return nil, fmt.Errorf("failed to get file from MinIO: %w", err) } - return obj, nil } // DeleteFile deletes a file func (s *minioFileService) DeleteFile(ctx context.Context, filePath string) error { - // Parse MinIO path - // Format: minio://bucketName/objectName - if len(filePath) < 9 || filePath[:8] != "minio://" { - return fmt.Errorf("invalid MinIO file path: %s", filePath) - } - - // Extract object name - objectName := filePath[9+len(s.bucketName):] - if len(objectName) > 0 && objectName[0] == '/' { - objectName = objectName[1:] - } - if err := utils.SafeObjectKey(objectName); err != nil { - return fmt.Errorf("invalid file path: %w", err) - } - - // Delete object - err := s.client.RemoveObject(ctx, s.bucketName, objectName, minio.RemoveObjectOptions{ - GovernanceBypass: true, - }) + objectName, err := s.parseMinioFilePath(filePath) if err != nil { + return err + } + if err := s.client.RemoveObject(ctx, s.bucketName, objectName, minio.RemoveObjectOptions{ + GovernanceBypass: true, + }); err != nil { return fmt.Errorf("failed to delete file: %w", err) } - return nil } @@ -194,25 +189,13 @@ func (s *minioFileService) SaveBytes(ctx context.Context, data []byte, tenantID // GetFileURL returns a presigned download URL for the file func (s *minioFileService) GetFileURL(ctx context.Context, filePath string) (string, error) { - // Parse MinIO path - if len(filePath) < 9 || filePath[:8] != "minio://" { - return "", fmt.Errorf("invalid MinIO file path: %s", filePath) + objectName, err := s.parseMinioFilePath(filePath) + if err != nil { + return "", err } - - // Extract object name - objectName := filePath[9+len(s.bucketName):] - if len(objectName) > 0 && objectName[0] == '/' { - objectName = objectName[1:] - } - if err := utils.SafeObjectKey(objectName); err != nil { - return "", fmt.Errorf("invalid file path: %w", err) - } - - // Generate presigned URL (valid for 24 hours) presignedURL, err := s.client.PresignedGetObject(ctx, s.bucketName, objectName, 24*time.Hour, nil) if err != nil { return "", fmt.Errorf("failed to generate presigned URL: %w", err) } - return presignedURL.String(), nil } diff --git a/internal/application/service/image_multimodal.go b/internal/application/service/image_multimodal.go index 6fcad570..e5f3e1c6 100644 --- a/internal/application/service/image_multimodal.go +++ b/internal/application/service/image_multimodal.go @@ -4,8 +4,6 @@ import ( "context" "encoding/json" "fmt" - "io" - "net/http" "os" "strings" "time" @@ -16,6 +14,7 @@ import ( "github.com/Tencent/WeKnora/internal/models/vlm" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" + secutils "github.com/Tencent/WeKnora/internal/utils" "github.com/google/uuid" "github.com/hibiken/asynq" ) @@ -355,17 +354,5 @@ func (s *ImageMultimodalService) resolveVLM(ctx context.Context, kbID string) (v // downloadImageFromURL downloads image bytes from an HTTP(S) URL. func downloadImageFromURL(imageURL string) ([]byte, error) { - if !strings.HasPrefix(imageURL, "http://") && !strings.HasPrefix(imageURL, "https://") { - return nil, fmt.Errorf("unsupported URL scheme: %s", imageURL) - } - client := &http.Client{Timeout: 60 * time.Second} - resp, err := client.Get(imageURL) - if err != nil { - return nil, fmt.Errorf("HTTP GET: %w", err) - } - defer resp.Body.Close() - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("HTTP %d for %s", resp.StatusCode, imageURL) - } - return io.ReadAll(resp.Body) + return secutils.DownloadBytes(imageURL) } diff --git a/internal/application/service/knowledgebase.go b/internal/application/service/knowledgebase.go index afd650a7..15f4eaac 100644 --- a/internal/application/service/knowledgebase.go +++ b/internal/application/service/knowledgebase.go @@ -82,7 +82,7 @@ func (s *knowledgeBaseService) CreateKnowledgeBase(ctx context.Context, kb.ID = uuid.New().String() } kb.CreatedAt = time.Now() - kb.TenantID = ctx.Value(types.TenantIDContextKey).(uint64) + kb.TenantID = types.MustTenantIDFromContext(ctx) kb.UpdatedAt = time.Now() kb.EnsureDefaults() @@ -158,7 +158,7 @@ func (s *knowledgeBaseService) GetKnowledgeBasesByIDsOnly(ctx context.Context, i // ListKnowledgeBases returns all knowledge bases for a tenant func (s *knowledgeBaseService) ListKnowledgeBases(ctx context.Context) ([]*types.KnowledgeBase, error) { - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) kbs, err := s.repo.ListKnowledgeBasesByTenantID(ctx, tenantID) if err != nil { @@ -324,8 +324,8 @@ func (s *knowledgeBaseService) DeleteKnowledgeBase(ctx context.Context, id strin logger.Infof(ctx, "Deleting knowledge base, ID: %s", id) // Get tenant ID from context - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) - tenantInfo := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenantID := types.MustTenantIDFromContext(ctx) + tenantInfo, _ := types.TenantInfoFromContext(ctx) // Step 1: Delete the knowledge base record first (mark as deleted) logger.Infof(ctx, "Deleting knowledge base from database") @@ -543,7 +543,7 @@ func (s *knowledgeBaseService) SetEmbeddingModel(ctx context.Context, id string, func (s *knowledgeBaseService) CopyKnowledgeBase(ctx context.Context, srcKB string, dstKB string, ) (*types.KnowledgeBase, *types.KnowledgeBase, error) { - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // Load source KB with tenant scope to prevent cross-tenant cloning sourceKB, err := s.repo.GetKnowledgeBaseByIDAndTenant(ctx, srcKB, tenantID) if err != nil { @@ -593,8 +593,8 @@ func (s *knowledgeBaseService) HybridSearch(ctx context.Context, ) ([]*types.SearchResult, error) { logger.Infof(ctx, "Hybrid search parameters, knowledge base ID: %s, query text: %s", id, params.QueryText) - tenantInfo := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) - currentTenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantInfo, _ := types.TenantInfoFromContext(ctx) + currentTenantID := types.MustTenantIDFromContext(ctx) // Create a composite retrieval engine with tenant's configured retrievers retrieveEngine, err := retriever.NewCompositeRetrieveEngine(s.retrieveEngine, tenantInfo.GetEffectiveEngines()) @@ -880,7 +880,7 @@ func (s *knowledgeBaseService) iterativeRetrieveWithDeduplication(ctx context.Co filteredOutChunks := make(map[string]struct{}) queryTextLower := strings.ToLower(strings.TrimSpace(queryText)) - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) for i := 0; i < maxIterations; i++ { // Update TopK in retrieve params @@ -1017,7 +1017,7 @@ func (s *knowledgeBaseService) filterByNegativeQuestions(ctx context.Context, return chunks } - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // Collect chunk IDs chunkIDs := make([]string, 0, len(chunks)) @@ -1104,7 +1104,7 @@ func (s *knowledgeBaseService) processSearchResults(ctx context.Context, return nil, nil } - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // Prepare data structures for efficient processing var knowledgeIDs []string diff --git a/internal/application/service/message.go b/internal/application/service/message.go index 1c5792d3..ea44c518 100644 --- a/internal/application/service/message.go +++ b/internal/application/service/message.go @@ -60,7 +60,7 @@ func (s *messageService) CreateMessage(ctx context.Context, message *types.Messa logger.Infof(ctx, "Creating message for session ID: %s", message.SessionID) // Check if the session exists to validate the message belongs to a valid session - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Checking if session exists, tenant ID: %d, session ID: %s", tenantID, message.SessionID) _, err := s.sessionRepo.Get(ctx, tenantID, message.SessionID) if err != nil { @@ -94,7 +94,7 @@ func (s *messageService) GetMessage(ctx context.Context, sessionID string, messa logger.Infof(ctx, "Getting message, session ID: %s, message ID: %s", sessionID, messageID) // Verify the session exists before attempting to retrieve the message - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Checking if session exists, tenant ID: %d", tenantID) _, err := s.sessionRepo.Get(ctx, tenantID, sessionID) if err != nil { @@ -132,7 +132,7 @@ func (s *messageService) GetMessagesBySession(ctx context.Context, logger.Infof(ctx, "Getting messages for session ID: %s, page: %d, pageSize: %d", sessionID, page, pageSize) // Verify the session exists before retrieving messages - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Checking if session exists, tenant ID: %d", tenantID) _, err := s.sessionRepo.Get(ctx, tenantID, sessionID) if err != nil { @@ -253,7 +253,7 @@ func (s *messageService) UpdateMessage(ctx context.Context, message *types.Messa logger.Infof(ctx, "Updating message, ID: %s, session ID: %s", message.ID, message.SessionID) // Verify the session exists before updating the message - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Checking if session exists, tenant ID: %d", tenantID) _, err := s.sessionRepo.Get(ctx, tenantID, message.SessionID) if err != nil { @@ -288,7 +288,7 @@ func (s *messageService) DeleteMessage(ctx context.Context, sessionID string, me logger.Infof(ctx, "Deleting message, session ID: %s, message ID: %s", sessionID, messageID) // Verify the session exists before deleting the message - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Checking if session exists, tenant ID: %d", tenantID) _, err := s.sessionRepo.Get(ctx, tenantID, sessionID) if err != nil { diff --git a/internal/application/service/model.go b/internal/application/service/model.go index e4795224..dc553f9c 100644 --- a/internal/application/service/model.go +++ b/internal/application/service/model.go @@ -104,7 +104,7 @@ func (s *modelService) GetModelByID(ctx context.Context, id string) (*types.Mode return nil, errors.New("model ID cannot be empty") } - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // Fetch model from repository model, err := s.repo.GetByID(ctx, tenantID, id) @@ -148,7 +148,7 @@ func (s *modelService) GetModelByID(ctx context.Context, id string) (*types.Mode func (s *modelService) ListModels(ctx context.Context) ([]*types.Model, error) { logger.Info(ctx, "Start listing models") - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Listing models for tenant ID: %d", tenantID) // List models from repository with no additional filters @@ -170,7 +170,7 @@ func (s *modelService) UpdateModel(ctx context.Context, model *types.Model) erro logger.Infof(ctx, "Updating model ID: %s, name: %s", model.ID, model.Name) // Check if the model is builtin - builtin models cannot be updated - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) existingModel, err := s.repo.GetByID(ctx, tenantID, model.ID) if err != nil { logger.ErrorWithFields(ctx, err, map[string]interface{}{ @@ -202,7 +202,7 @@ func (s *modelService) DeleteModel(ctx context.Context, id string) error { logger.Info(ctx, "Start deleting model") logger.Infof(ctx, "Deleting model ID: %s", id) - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Tenant ID: %d", tenantID) // Check if the model is builtin - builtin models cannot be deleted @@ -368,7 +368,7 @@ func (s *modelService) GetChatModel(ctx context.Context, modelId string) (chat.C return nil, errors.New("model ID cannot be empty") } - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // Get the model directly from repository to avoid status checks model, err := s.repo.GetByID(ctx, tenantID, modelId) @@ -412,7 +412,7 @@ func (s *modelService) GetVLMModel(ctx context.Context, modelId string) (vlm.VLM return nil, errors.New("model ID cannot be empty") } - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) model, err := s.repo.GetByID(ctx, tenantID, modelId) if err != nil { diff --git a/internal/application/service/session.go b/internal/application/service/session.go index ae670a4b..8a09f0f6 100644 --- a/internal/application/service/session.go +++ b/internal/application/service/session.go @@ -114,7 +114,7 @@ func (s *sessionService) GetSession(ctx context.Context, id string) (*types.Sess } // Get tenant ID from context - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Retrieving session, ID: %s, tenant ID: %d", id, tenantID) // Get session from repository @@ -134,7 +134,7 @@ func (s *sessionService) GetSession(ctx context.Context, id string) (*types.Sess // GetSessionsByTenant retrieves all sessions for the current tenant func (s *sessionService) GetSessionsByTenant(ctx context.Context) ([]*types.Session, error) { // Get tenant ID from context - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Retrieving all sessions for tenant, tenant ID: %d", tenantID) // Get sessions from repository @@ -157,7 +157,7 @@ func (s *sessionService) GetPagedSessionsByTenant(ctx context.Context, pagination *types.Pagination, ) (*types.PageResult, error) { // Get tenant ID from context - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // Get paged sessions from repository sessions, total, err := s.sessionRepo.GetPagedByTenantID(ctx, tenantID, pagination) if err != nil { @@ -203,7 +203,7 @@ func (s *sessionService) DeleteSession(ctx context.Context, id string) error { } // Get tenant ID from context - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // Cleanup temporary KB stored in Redis for this session if err := s.webSearchStateRepo.DeleteWebSearchTempKBState(ctx, id); err != nil { @@ -236,7 +236,7 @@ func (s *sessionService) BatchDeleteSessions(ctx context.Context, ids []string) } // Get tenant ID from context - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // Cleanup associated resources for each session for _, id := range ids { @@ -643,7 +643,7 @@ func (s *sessionService) KnowledgeQA( ) // Get UserID from context - userID, _ := ctx.Value(types.UserIDContextKey).(string) + userID, _ := types.UserIDFromContext(ctx) chatManage := &types.ChatManage{ Query: query, @@ -790,7 +790,7 @@ func (s *sessionService) selectChatModelID( ) (string, error) { // If no knowledge base IDs but have knowledge IDs, derive KB IDs from knowledge IDs (include shared KB files) if len(knowledgeBaseIDs) == 0 && len(knowledgeIDs) > 0 { - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) knowledgeList, err := s.knowledgeService.GetKnowledgeBatchWithSharedAccess(ctx, tenantID, knowledgeIDs) if err != nil { logger.Warnf(ctx, "Failed to get knowledge batch for model selection: %v", err) @@ -899,7 +899,7 @@ func (s *sessionService) resolveKnowledgeBasesFromAgent( // unrelated KBs from other organisations into the agent's retrieval scope. isSharedAgent := sessionTenantID != 0 && sessionTenantID != customAgent.TenantID if !isSharedAgent { - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) userIDVal := ctx.Value(types.UserIDContextKey) if userIDVal != nil { if userID, ok := userIDVal.(string); ok && userID != "" && s.kbShareService != nil { @@ -1021,7 +1021,7 @@ func (s *sessionService) buildSearchTargets( kbByID[kb.ID] = kb } } - userID, _ := ctx.Value(types.UserIDContextKey).(string) + userID, _ := types.UserIDFromContext(ctx) for _, kbID := range knowledgeBaseIDs { fullKBSet[kbID] = true kb := kbByID[kbID] @@ -1114,7 +1114,7 @@ func (s *sessionService) KnowledgeQAByEvent(ctx context.Context, // Set up tracing attributes logger.Infof(ctx, "Trigger event list: %v", methods) span.SetAttributes( - attribute.String("request_id", ctx.Value(types.RequestIDContextKey).(string)), + attribute.String("request_id", func() string { id, _ := types.RequestIDFromContext(ctx); return id }()), attribute.String("query", chatManage.Query), attribute.String("method", strings.Join(methods, ",")), ) @@ -1163,7 +1163,7 @@ func (s *sessionService) SearchKnowledge(ctx context.Context, knowledgeBaseIDs, knowledgeIDs, query) // Get tenant ID from context - tenantID, ok := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, ok := types.TenantIDFromContext(ctx) if !ok { logger.Error(ctx, "Failed to get tenant ID from context") return nil, fmt.Errorf("tenant ID not found in context") @@ -1181,7 +1181,7 @@ func (s *sessionService) SearchKnowledge(ctx context.Context, } // Create default retrieval parameters - userID, _ := ctx.Value(types.UserIDContextKey).(string) + userID, _ := types.UserIDFromContext(ctx) chatManage := &types.ChatManage{ Query: query, RewriteQuery: query, @@ -1528,7 +1528,7 @@ func (s *sessionService) getContextManagerForSession( chatModel chat.Chat, ) interfaces.ContextManager { // Get tenant to access global context configuration - tenant, _ := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) // Determine which context config to use: tenant-level or default var contextConfig *types.ContextConfig if tenant != nil && tenant.ContextConfig != nil { diff --git a/internal/application/service/tag.go b/internal/application/service/tag.go index 51bb748f..7e9fb88a 100644 --- a/internal/application/service/tag.go +++ b/internal/application/service/tag.go @@ -73,7 +73,7 @@ func (s *knowledgeTagService) ListTags( } // Check access permission - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) if kb.TenantID != tenantID { // Get user ID from context userIDVal := ctx.Value(types.UserIDContextKey) @@ -192,7 +192,7 @@ func (s *knowledgeTagService) UpdateTag( if id == "" { return nil, werrors.NewBadRequestError("标签ID不能为空") } - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) tag, err := s.repo.GetByID(ctx, tenantID, id) if err != nil { return nil, err @@ -225,7 +225,7 @@ func (s *knowledgeTagService) DeleteTag(ctx context.Context, id string, force bo if id == "" { return werrors.NewBadRequestError("标签ID不能为空") } - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) tag, err := s.repo.GetByID(ctx, tenantID, id) if err != nil { return err @@ -243,7 +243,7 @@ func (s *knowledgeTagService) DeleteTag(ctx context.Context, id string, force bo } // Get tenant info for effective engines - tenantInfo := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenantInfo, _ := types.TenantInfoFromContext(ctx) // Helper function to delete chunks and enqueue index deletion task deleteChunksAndEnqueueIndexDelete := func() error { diff --git a/internal/handler/initialization.go b/internal/handler/initialization.go index e1d3395d..094d6fca 100644 --- a/internal/handler/initialization.go +++ b/internal/handler/initialization.go @@ -1912,7 +1912,7 @@ func (h *InitializationHandler) testMultimodalWithDocReader( return nil, fmt.Errorf("DocReader service not configured") } - requestID, _ := ctx.Value(types.RequestIDContextKey).(string) + requestID, _ := types.RequestIDFromContext(ctx) readResult, err := h.documentReader.Read(ctx, &types.ReadRequest{ FileContent: imageContent, diff --git a/internal/handler/knowledge.go b/internal/handler/knowledge.go index b68fcfac..c6e582b7 100644 --- a/internal/handler/knowledge.go +++ b/internal/handler/knowledge.go @@ -44,74 +44,10 @@ func NewKnowledgeHandler( } // validateKnowledgeBaseAccess validates access permissions to a knowledge base -// Returns the knowledge base, the knowledge base ID, effective tenant ID, permission level, and any errors encountered -// For owned KBs, effectiveTenantID is the caller's tenant ID -// For shared KBs, effectiveTenantID is the source tenant ID (owner's tenant) +// using the ":id" URL path parameter. It delegates to validateKnowledgeBaseAccessWithKBID. func (h *KnowledgeHandler) validateKnowledgeBaseAccess(c *gin.Context) (*types.KnowledgeBase, string, uint64, types.OrgMemberRole, error) { - ctx := c.Request.Context() - - // Get tenant ID from context - tenantID := c.GetUint64(types.TenantIDContextKey.String()) - if tenantID == 0 { - logger.Error(ctx, "Failed to get tenant ID") - return nil, "", 0, "", errors.NewUnauthorizedError("Unauthorized") - } - - // Get user ID from context (needed for shared KB permission check) - userID, userExists := c.Get(types.UserIDContextKey.String()) - - // Get knowledge base ID from URL path parameter kbID := secutils.SanitizeForLog(c.Param("id")) - if kbID == "" { - logger.Error(ctx, "Knowledge base ID is empty") - return nil, "", 0, "", errors.NewBadRequestError("Knowledge base ID cannot be empty") - } - - // Get knowledge base details - kb, err := h.kbService.GetKnowledgeBaseByID(ctx, kbID) - if err != nil { - logger.ErrorWithFields(ctx, err, nil) - return nil, kbID, 0, "", errors.NewInternalServerError(err.Error()) - } - - // Check 1: Verify tenant ownership (owner has full access) - if kb.TenantID == tenantID { - return kb, kbID, tenantID, types.OrgRoleAdmin, nil - } - - // Check 2: If not owner, check organization shared access - if userExists && h.kbShareService != nil { - // Check if user has shared access through organization - permission, isShared, permErr := h.kbShareService.CheckUserKBPermission(ctx, kbID, userID.(string)) - if permErr == nil && isShared { - // User has shared access, get the source tenant ID for queries - sourceTenantID, srcErr := h.kbShareService.GetKBSourceTenant(ctx, kbID) - if srcErr == nil { - logger.Infof(ctx, "User %s accessing shared KB %s with permission %s, source tenant: %d", - userID.(string), kbID, permission, sourceTenantID) - return kb, kbID, sourceTenantID, permission, nil - } - } - } - - // Check 3: If not owner and no direct share, allow if user has any shared agent that can access this KB (e.g. opened from "通过智能体可见" list) - if userExists && h.agentShareService != nil { - can, err := h.agentShareService.UserCanAccessKBViaSomeSharedAgent(ctx, userID.(string), tenantID, kb) - if err == nil && can { - logger.Infof(ctx, "User %s accessing KB %s via some shared agent", userID.(string), kbID) - return kb, kbID, kb.TenantID, types.OrgRoleViewer, nil - } - } - - // No permission: not owner and no shared access - logger.Warnf( - ctx, - "Permission denied to access this knowledge base, tenant ID mismatch, "+ - "requested tenant ID: %d, knowledge base tenant ID: %d", - tenantID, - kb.TenantID, - ) - return nil, kbID, 0, "", errors.NewForbiddenError("Permission denied to access this knowledge base") + return h.validateKnowledgeBaseAccessWithKBID(c, kbID) } // validateKnowledgeBaseAccessWithKBID validates access to the given knowledge base ID (e.g. from query or body). diff --git a/internal/handler/organization.go b/internal/handler/organization.go index 8d24d1e5..a4286819 100644 --- a/internal/handler/organization.go +++ b/internal/handler/organization.go @@ -1225,7 +1225,7 @@ func (h *OrganizationHandler) ListSharedKnowledgeBases(c *gin.Context) { ctx := c.Request.Context() userID := c.GetString(types.UserIDContextKey.String()) - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) sharedKBs, err := h.shareService.ListSharedKnowledgeBases(ctx, userID, tenantID) if err != nil { diff --git a/internal/handler/session/helpers.go b/internal/handler/session/helpers.go index e9c2da34..03abf522 100644 --- a/internal/handler/session/helpers.go +++ b/internal/handler/session/helpers.go @@ -215,7 +215,7 @@ func getFloat64(m map[string]interface{}, key string) float64 { // It prioritizes tenant-level ConversationConfig, then falls back to config.yaml defaults func (h *Handler) createDefaultSummaryConfig(ctx context.Context) *types.SummaryConfig { // Try to get tenant from context - tenant, _ := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) // Initialize with config.yaml defaults cfg := &types.SummaryConfig{ @@ -259,7 +259,7 @@ func (h *Handler) createDefaultSummaryConfig(ctx context.Context) *types.Summary // It prioritizes tenant-level ConversationConfig, then falls back to config.yaml defaults func (h *Handler) fillSummaryConfigDefaults(ctx context.Context, config *types.SummaryConfig) { // Try to get tenant from context - tenant, _ := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) // Determine default values: tenant config first, then config.yaml var defaultPrompt, defaultContextTemplate, defaultNoMatchPrefix string diff --git a/internal/handler/tag.go b/internal/handler/tag.go index 7db819de..25b3e7b2 100644 --- a/internal/handler/tag.go +++ b/internal/handler/tag.go @@ -94,7 +94,7 @@ func (h *TagHandler) resolveTagIDWithCtx(c *gin.Context, ctx context.Context) (s tagIDParam := secutils.SanitizeForLog(c.Param("tag_id")) if seqID, err := strconv.ParseInt(tagIDParam, 10, 64); err == nil { - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) tag, err := h.tagRepo.GetBySeqID(ctx, tenantID, seqID) if err != nil { return "", errors.NewNotFoundError("标签不存在") diff --git a/internal/handler/tenant.go b/internal/handler/tenant.go index 2e7db31f..a89e2e6d 100644 --- a/internal/handler/tenant.go +++ b/internal/handler/tenant.go @@ -439,7 +439,7 @@ type AgentConfigRequest struct { // @Router /tenants/kv/agent-config [get] func (h *TenantHandler) GetTenantAgentConfig(c *gin.Context) { ctx := c.Request.Context() - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -529,7 +529,7 @@ func (h *TenantHandler) updateTenantAgentConfigInternal(c *gin.Context) { } // Get existing tenant - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -670,7 +670,7 @@ func (h *TenantHandler) updateTenantWebSearchConfigInternal(c *gin.Context) { return } - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -711,7 +711,7 @@ func (h *TenantHandler) GetTenantWebSearchConfig(c *gin.Context) { ctx := c.Request.Context() logger.Info(ctx, "Start getting tenant web search config") // Get tenant - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -728,7 +728,7 @@ func (h *TenantHandler) GetTenantWebSearchConfig(c *gin.Context) { // GetTenantParserEngineConfig returns the tenant's parser engine config (MinerU endpoint, API key, etc.). func (h *TenantHandler) GetTenantParserEngineConfig(c *gin.Context) { ctx := c.Request.Context() - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -753,7 +753,7 @@ func (h *TenantHandler) updateTenantParserEngineConfigInternal(c *gin.Context) { c.Error(errors.NewValidationError("Invalid request data").WithDetails(err.Error())) return } - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -780,7 +780,7 @@ func (h *TenantHandler) updateTenantParserEngineConfigInternal(c *gin.Context) { // GetTenantStorageEngineConfig returns the tenant's storage engine config (Local, MinIO, COS parameters). func (h *TenantHandler) GetTenantStorageEngineConfig(c *gin.Context) { ctx := c.Request.Context() - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -805,7 +805,7 @@ func (h *TenantHandler) updateTenantStorageEngineConfigInternal(c *gin.Context) c.Error(errors.NewValidationError("Invalid request data").WithDetails(err.Error())) return } - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -897,7 +897,7 @@ func validateConversationConfig(req *types.ConversationConfig) error { // @Router /tenants/kv/conversation-config [get] func (h *TenantHandler) GetTenantConversationConfig(c *gin.Context) { ctx := c.Request.Context() - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -934,7 +934,7 @@ func (h *TenantHandler) updateTenantConversationInternal(c *gin.Context) { } // Get existing tenant - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) diff --git a/internal/infrastructure/docparser/helpers.go b/internal/infrastructure/docparser/helpers.go new file mode 100644 index 00000000..b0ba4d9b --- /dev/null +++ b/internal/infrastructure/docparser/helpers.go @@ -0,0 +1,110 @@ +package docparser + +import ( + "context" + "fmt" + "sort" + "strings" + "time" +) + +// stringOr returns val (trimmed) if non-empty, otherwise fallback. +func stringOr(val, fallback string) string { + val = strings.TrimSpace(val) + if val == "" { + return fallback + } + return val +} + +// parseBoolOr parses a truthy string ("true","1","yes"), returning fallback on empty. +func parseBoolOr(val string, fallback bool) bool { + val = strings.ToLower(strings.TrimSpace(val)) + if val == "" { + return fallback + } + return val == "true" || val == "1" || val == "yes" +} + +// firstNonEmpty returns the first non-empty string, or "" if all are empty. +func firstNonEmpty(vals ...string) string { + for _, v := range vals { + if v != "" { + return v + } + } + return "" +} + +// sleepCtx sleeps for d but returns early if ctx is cancelled. +func sleepCtx(ctx context.Context, d time.Duration) { + t := time.NewTimer(d) + defer t.Stop() + select { + case <-ctx.Done(): + case <-t.C: + } +} + +// logResponseStructure recursively logs the structure of an API response, +// truncating large string values. label identifies the subsystem (e.g. "MinerU"). +func logResponseStructure(label string, obj interface{}, prefix string) { + switch v := obj.(type) { + case map[string]interface{}: + keys := make([]string, 0, len(v)) + for k := range v { + keys = append(keys, k) + } + sort.Strings(keys) + logger.Printf("DEBUG: [%s] %s = {object with %d keys: %s}", label, prefix, len(v), strings.Join(keys, ", ")) + for _, key := range keys { + val := v[key] + path := prefix + "." + key + switch inner := val.(type) { + case map[string]interface{}: + logResponseStructure(label, inner, path) + case []interface{}: + logger.Printf("DEBUG: [%s] %s = [array with %d items]", label, path, len(inner)) + if len(inner) > 0 { + logger.Printf("DEBUG: [%s] %s[0] type=%T", label, path, inner[0]) + if len(inner) <= 3 { + for i, item := range inner { + logResponseStructure(label, item, fmt.Sprintf("%s[%d]", path, i)) + } + } else { + logResponseStructure(label, inner[0], path+"[0]") + logger.Printf("DEBUG: [%s] ... and %d more items in %s", label, len(inner)-1, path) + } + } + case string: + if len(inner) > 200 { + logger.Printf("DEBUG: [%s] %s = string(%d chars): %.200s...", label, path, len(inner), inner) + } else { + logger.Printf("DEBUG: [%s] %s = %q", label, path, inner) + } + case float64: + logger.Printf("DEBUG: [%s] %s = %v (number)", label, path, inner) + case bool: + logger.Printf("DEBUG: [%s] %s = %v (bool)", label, path, inner) + case nil: + logger.Printf("DEBUG: [%s] %s = null", label, path) + default: + logger.Printf("DEBUG: [%s] %s = %v (%T)", label, path, val, val) + } + } + case []interface{}: + logger.Printf("DEBUG: [%s] %s = [array with %d items]", label, prefix, len(v)) + if len(v) > 0 { + if len(v) <= 3 { + for i, item := range v { + logResponseStructure(label, item, fmt.Sprintf("%s[%d]", prefix, i)) + } + } else { + logResponseStructure(label, v[0], prefix+"[0]") + logger.Printf("DEBUG: [%s] ... and %d more items in %s", label, len(v)-1, prefix) + } + } + default: + logger.Printf("DEBUG: [%s] %s = %v (%T)", label, prefix, v, v) + } +} diff --git a/internal/infrastructure/docparser/image_resolver.go b/internal/infrastructure/docparser/image_resolver.go index a8a54eda..964a9349 100644 --- a/internal/infrastructure/docparser/image_resolver.go +++ b/internal/infrastructure/docparser/image_resolver.go @@ -2,20 +2,16 @@ package docparser import ( "fmt" - "io" - "net/http" "os" "path/filepath" "regexp" "strings" - "time" "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/utils" "github.com/google/uuid" ) -var httpImageClient = &http.Client{Timeout: 60 * time.Second} - // StoredImage describes an image that has been saved to local storage. type StoredImage struct { OriginalRef string // reference in the original markdown @@ -179,19 +175,7 @@ func (r *ImageResolver) readImageBytes( } func downloadImage(url string) ([]byte, error) { - resp, err := httpImageClient.Get(url) - if err != nil { - return nil, fmt.Errorf("HTTP GET: %w", err) - } - defer resp.Body.Close() - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("HTTP %d for %s", resp.StatusCode, url) - } - data, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("read body: %w", err) - } - return data, nil + return utils.DownloadBytes(url) } func extFromMime(mime string) string { diff --git a/internal/infrastructure/docparser/mineru_cloud_converter.go b/internal/infrastructure/docparser/mineru_cloud_converter.go index 44bd08c1..4171684e 100644 --- a/internal/infrastructure/docparser/mineru_cloud_converter.go +++ b/internal/infrastructure/docparser/mineru_cloud_converter.go @@ -297,7 +297,7 @@ func (c *MinerUCloudReader) fetchBatchStatus(ctx context.Context, batchID string // Pretty-print the structure to reveal all available fields var rawObj interface{} if err := json.Unmarshal(pollResp.Data.ExtractResult, &rawObj); err == nil { - logCloudResponseStructure(rawObj, "extract_result") + logResponseStructure("MinerUCloud", rawObj, "extract_result") } // The extract_result can be either a single object or an array @@ -472,68 +472,7 @@ func readZipEntryBytes(f *zip.File) ([]byte, error) { return io.ReadAll(rc) } -// logCloudResponseStructure recursively logs the structure of the Cloud API response, -// truncating large string values to reveal all available fields (e.g. caption, bbox, etc.). -func logCloudResponseStructure(obj interface{}, prefix string) { - switch v := obj.(type) { - case map[string]interface{}: - keys := make([]string, 0, len(v)) - for k := range v { - keys = append(keys, k) - } - sort.Strings(keys) - logger.Printf("DEBUG: [MinerUCloud] %s = {object with %d keys: %s}", prefix, len(v), strings.Join(keys, ", ")) - for _, key := range keys { - val := v[key] - path := prefix + "." + key - switch inner := val.(type) { - case map[string]interface{}: - logCloudResponseStructure(inner, path) - case []interface{}: - logger.Printf("DEBUG: [MinerUCloud] %s = [array with %d items]", path, len(inner)) - if len(inner) > 0 { - logger.Printf("DEBUG: [MinerUCloud] %s[0] type=%T", path, inner[0]) - if len(inner) <= 3 { - for i, item := range inner { - logCloudResponseStructure(item, fmt.Sprintf("%s[%d]", path, i)) - } - } else { - logCloudResponseStructure(inner[0], path+"[0]") - logger.Printf("DEBUG: [MinerUCloud] ... and %d more items in %s", len(inner)-1, path) - } - } - case string: - if len(inner) > 200 { - logger.Printf("DEBUG: [MinerUCloud] %s = string(%d chars): %.200s...", path, len(inner), inner) - } else { - logger.Printf("DEBUG: [MinerUCloud] %s = %q", path, inner) - } - case float64: - logger.Printf("DEBUG: [MinerUCloud] %s = %v (number)", path, inner) - case bool: - logger.Printf("DEBUG: [MinerUCloud] %s = %v (bool)", path, inner) - case nil: - logger.Printf("DEBUG: [MinerUCloud] %s = null", path) - default: - logger.Printf("DEBUG: [MinerUCloud] %s = %v (%T)", path, val, val) - } - } - case []interface{}: - logger.Printf("DEBUG: [MinerUCloud] %s = [array with %d items]", prefix, len(v)) - if len(v) > 0 { - if len(v) <= 3 { - for i, item := range v { - logCloudResponseStructure(item, fmt.Sprintf("%s[%d]", prefix, i)) - } - } else { - logCloudResponseStructure(v[0], prefix+"[0]") - logger.Printf("DEBUG: [MinerUCloud] ... and %d more items in %s", len(v)-1, prefix) - } - } - default: - logger.Printf("DEBUG: [MinerUCloud] %s = %v (%T)", prefix, v, v) - } -} +// logCloudResponseStructure is now unified in helpers.go as logResponseStructure("MinerUCloud", ...) // PingMinerUCloud checks if the MinerU Cloud API is reachable with the given API key. func PingMinerUCloud(apiKey, baseURL string) (bool, string) { @@ -568,22 +507,4 @@ func PingMinerUCloud(apiKey, baseURL string) (bool, string) { return true, "" } -// --- helpers --- - -func firstNonEmpty(vals ...string) string { - for _, v := range vals { - if v != "" { - return v - } - } - return "" -} - -func sleepCtx(ctx context.Context, d time.Duration) { - t := time.NewTimer(d) - defer t.Stop() - select { - case <-ctx.Done(): - case <-t.C: - } -} +// firstNonEmpty, sleepCtx are defined in helpers.go diff --git a/internal/infrastructure/docparser/mineru_converter.go b/internal/infrastructure/docparser/mineru_converter.go index a70fef0c..11394fb1 100644 --- a/internal/infrastructure/docparser/mineru_converter.go +++ b/internal/infrastructure/docparser/mineru_converter.go @@ -25,12 +25,12 @@ var b64DataURIPattern = regexp.MustCompile(`^data:image/(\w+);base64,(.+)$`) // MinerUReader calls a self-hosted MinerU API to read/convert documents. type MinerUReader struct { - endpoint string - backend string // "pipeline", "vlm-*", "hybrid-*" - formulaEnable bool - tableEnable bool - ocrEnable bool - language string + endpoint string + backend string // "pipeline", "vlm-*", "hybrid-*" + formulaEnable bool + tableEnable bool + ocrEnable bool + language string } // NewMinerUReader creates a reader from ParserEngineOverrides. @@ -95,18 +95,18 @@ func (c *MinerUReader) callFileParse(ctx context.Context, content []byte) (strin // Form fields fields := map[string]string{ - "return_md": "true", - "return_images": "true", - "table_enable": fmt.Sprintf("%v", c.tableEnable), - "formula_enable": fmt.Sprintf("%v", c.formulaEnable), - "parse_method": "ocr", - "start_page_id": "0", - "end_page_id": "99999", - "backend": c.backend, - "response_format_zip": "false", - "return_middle_json": "false", - "return_model_output": "false", - "return_content_list": "true", + "return_md": "true", + "return_images": "true", + "table_enable": fmt.Sprintf("%v", c.tableEnable), + "formula_enable": fmt.Sprintf("%v", c.formulaEnable), + "parse_method": "ocr", + "start_page_id": "0", + "end_page_id": "99999", + "backend": c.backend, + "response_format_zip": "false", + "return_middle_json": "false", + "return_model_output": "false", + "return_content_list": "true", } if !c.ocrEnable { fields["parse_method"] = "txt" @@ -162,7 +162,7 @@ func (c *MinerUReader) callFileParse(ctx context.Context, content []byte) (strin // Also pretty-print the top-level structure (without large base64 blobs) var rawMap map[string]interface{} if err := json.Unmarshal(respBody, &rawMap); err == nil { - c.logResponseStructure(rawMap, "") + c.logMinerUResponseStructure(rawMap, "") } var result mineruFileParseResponse @@ -225,60 +225,9 @@ func (c *MinerUReader) processImages(mdContent string, imagesB64 map[string]stri return refs, mdContent } -// logResponseStructure recursively logs the structure of the API response, -// truncating large string values to reveal all available fields. -func (c *MinerUReader) logResponseStructure(obj interface{}, prefix string) { - switch v := obj.(type) { - case map[string]interface{}: - for key, val := range v { - path := prefix + "." + key - if path[0] == '.' { - path = path[1:] - } - switch inner := val.(type) { - case map[string]interface{}: - logger.Printf("DEBUG: [MinerU] Response field: %s = {object with %d keys: %s}", path, len(inner), mapKeys(inner)) - c.logResponseStructure(inner, path) - case []interface{}: - logger.Printf("DEBUG: [MinerU] Response field: %s = [array with %d items]", path, len(inner)) - if len(inner) > 0 { - logger.Printf("DEBUG: [MinerU] Response field: %s[0] type=%T", path, inner[0]) - if len(inner) <= 3 { - for i, item := range inner { - c.logResponseStructure(item, fmt.Sprintf("%s[%d]", path, i)) - } - } else { - c.logResponseStructure(inner[0], path+"[0]") - logger.Printf("DEBUG: [MinerU] ... and %d more items in %s", len(inner)-1, path) - } - } - case string: - if len(inner) > 200 { - logger.Printf("DEBUG: [MinerU] Response field: %s = string(%d chars): %.200s...", path, len(inner), inner) - } else { - logger.Printf("DEBUG: [MinerU] Response field: %s = %q", path, inner) - } - case float64: - logger.Printf("DEBUG: [MinerU] Response field: %s = %v (number)", path, inner) - case bool: - logger.Printf("DEBUG: [MinerU] Response field: %s = %v (bool)", path, inner) - case nil: - logger.Printf("DEBUG: [MinerU] Response field: %s = null", path) - default: - logger.Printf("DEBUG: [MinerU] Response field: %s = %v (%T)", path, inner, inner) - } - } - default: - logger.Printf("DEBUG: [MinerU] Value at %s: %v (%T)", prefix, v, v) - } -} - -func mapKeys(m map[string]interface{}) string { - keys := make([]string, 0, len(m)) - for k := range m { - keys = append(keys, k) - } - return strings.Join(keys, ", ") +// logMinerUResponseStructure logs the structure of the MinerU API response. +func (c *MinerUReader) logMinerUResponseStructure(obj interface{}, prefix string) { + logResponseStructure("MinerU", obj, prefix) } // PingMinerU checks if the self-hosted MinerU service is reachable. @@ -309,21 +258,3 @@ func htmlToMarkdown(content string) string { } return md } - -// --- helpers --- - -func stringOr(val, fallback string) string { - val = strings.TrimSpace(val) - if val == "" { - return fallback - } - return val -} - -func parseBoolOr(val string, fallback bool) bool { - val = strings.ToLower(strings.TrimSpace(val)) - if val == "" { - return fallback - } - return val == "true" || val == "1" || val == "yes" -} diff --git a/internal/types/context_helpers.go b/internal/types/context_helpers.go new file mode 100644 index 00000000..0f383787 --- /dev/null +++ b/internal/types/context_helpers.go @@ -0,0 +1,47 @@ +package types + +import "context" + +// TenantIDFromContext extracts the tenant ID from ctx. +// Returns (0, false) when the key is absent or the value is not uint64. +func TenantIDFromContext(ctx context.Context) (uint64, bool) { + v, ok := ctx.Value(TenantIDContextKey).(uint64) + return v, ok +} + +// MustTenantIDFromContext extracts the tenant ID from ctx, panicking if missing. +func MustTenantIDFromContext(ctx context.Context) uint64 { + v, ok := TenantIDFromContext(ctx) + if !ok { + panic("types.TenantIDContextKey not set in context") + } + return v +} + +// TenantInfoFromContext extracts the *Tenant from ctx. +func TenantInfoFromContext(ctx context.Context) (*Tenant, bool) { + v, ok := ctx.Value(TenantInfoContextKey).(*Tenant) + return v, ok && v != nil +} + +// RequestIDFromContext extracts the request ID string from ctx. +func RequestIDFromContext(ctx context.Context) (string, bool) { + v, ok := ctx.Value(RequestIDContextKey).(string) + return v, ok && v != "" +} + +// UserIDFromContext extracts the user ID string from ctx. +func UserIDFromContext(ctx context.Context) (string, bool) { + v, ok := ctx.Value(UserIDContextKey).(string) + return v, ok && v != "" +} + +// SessionTenantIDFromContext extracts the session-owner tenant ID from ctx. +// Falls back to TenantIDFromContext when the session key is absent. +func SessionTenantIDFromContext(ctx context.Context) (uint64, bool) { + v, ok := ctx.Value(SessionTenantIDContextKey).(uint64) + if ok && v != 0 { + return v, true + } + return TenantIDFromContext(ctx) +} diff --git a/internal/utils/httputil.go b/internal/utils/httputil.go new file mode 100644 index 00000000..b1d3ed0a --- /dev/null +++ b/internal/utils/httputil.go @@ -0,0 +1,32 @@ +package utils + +import ( + "fmt" + "io" + "net/http" + "strings" + "time" +) + +var defaultHTTPClient = &http.Client{Timeout: 60 * time.Second} + +// DownloadBytes fetches the content at the given HTTP(S) URL and returns the +// raw bytes. It reuses a package-level http.Client with a 60-second timeout. +func DownloadBytes(url string) ([]byte, error) { + if !strings.HasPrefix(url, "http://") && !strings.HasPrefix(url, "https://") { + return nil, fmt.Errorf("unsupported URL scheme: %s", url) + } + resp, err := defaultHTTPClient.Get(url) + if err != nil { + return nil, fmt.Errorf("HTTP GET: %w", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("HTTP %d for %s", resp.StatusCode, url) + } + data, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("read body: %w", err) + } + return data, nil +}