package im import ( "bytes" "context" "encoding/base64" "encoding/json" "errors" "fmt" "io" "mime/multipart" "net/http" "net/textproto" "os" "path/filepath" "reflect" "regexp" "strings" "sync" "sync/atomic" "time" "unicode/utf8" agenttools "github.com/Tencent/WeKnora/internal/agent/tools" "github.com/Tencent/WeKnora/internal/config" apperrors "github.com/Tencent/WeKnora/internal/errors" "github.com/Tencent/WeKnora/internal/event" "github.com/Tencent/WeKnora/internal/infrastructure/docparser" "github.com/Tencent/WeKnora/internal/logger" mcppkg "github.com/Tencent/WeKnora/internal/mcp" "github.com/Tencent/WeKnora/internal/ratelimit" "github.com/Tencent/WeKnora/internal/storageurl" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" "github.com/google/uuid" "github.com/redis/go-redis/v9" "gorm.io/gorm" ) const ( imNoAnswerFallback = "抱歉,我暂时无法回答这个问题。" imErrorFallback = "抱歉,处理您的问题时出现了异常,请稍后再试。" imCancelledFallback = "抱歉,回答已被取消。" // dedupTTL is how long processed message IDs are retained. dedupTTL = 5 * time.Minute // dedupCleanupInterval is how often the dedup map is cleaned. dedupCleanupInterval = 1 * time.Minute // maxContentLength is the maximum allowed message content length. maxContentLength = 4096 // maxQuoteContentLength is the max runes to include from a quoted message. maxQuoteContentLength = 500 // maxIMAttachmentBytes bounds an attachment buffered for IM Q&A. maxIMAttachmentBytes = 32 << 20 // 32 MiB // maxIMVisionAttachmentBytes prevents large images from expanding into an oversized data URI. maxIMVisionAttachmentBytes = 8 << 20 // 8 MiB // maxIMAttachmentLines matches the legacy attachment prompt limit. maxIMAttachmentLines = 500 // maxIMAttachmentContentBytes bounds parsed text persisted in the IM message and injected into QA. maxIMAttachmentContentBytes = 32 << 10 // 32 KiB // imAttachmentReadTimeout bounds downloading and parsing before the QA request runs. imAttachmentReadTimeout = time.Minute // streamFlushInterval is how often buffered stream content is flushed to the IM platform. // This prevents API rate-limiting while keeping perceived latency low. streamFlushInterval = 300 * time.Millisecond // agentCompleteWaitTimeout bounds how long IM waits for EventAgentComplete after // the final answer stream finishes, preventing indefinite hangs. agentCompleteWaitTimeout = 10 * time.Second ) // imCitationTagRe matches inline citation tags produced by the agent pipeline. // These tags are rendered as interactive UI in the web frontend but are meaningless // in IM platforms, so they must be stripped before sending. var imCitationTagRe = regexp.MustCompile(`<(?:kb|web)\b[^>]*/?>`) // stripIMCitationTags removes and inline citation tags from s. func stripIMCitationTags(s string) string { return imCitationTagRe.ReplaceAllString(s, "") } // imageXMLBlockRe matches ... blocks produced by // EnrichContentWithImageInfo in the RAG context pipeline. These blocks contain // metadata for the LLM and must be stripped before sending to IM platforms. var imageXMLBlockRe = regexp.MustCompile(`(?s)]*>.*?`) // imageOriginalRe extracts the original markdown image syntax from tags. var imageOriginalRe = regexp.MustCompile(`(.*?)`) // stripImageXMLTags collapses blocks back to plain markdown. // Extracts the original ![alt](url) from when present, // otherwise drops the block entirely. func stripImageXMLTags(s string) string { return imageXMLBlockRe.ReplaceAllStringFunc(s, func(block string) string { if m := imageOriginalRe.FindStringSubmatch(block); len(m) > 1 { return m[1] } return "" }) } // rewriteStorageURLs replaces all storage references in content with HTTP URLs // so IM clients — which cannot attach WeKnora credentials to an image fetch — // can render them. See internal/storageurl for the shared implementation. func rewriteStorageURLs(ctx context.Context, content string, resolver *storageurl.FileServiceResolver) string { if resolver == nil { return content } return storageurl.Rewrite(ctx, content, resolver, "IM") } // ── Streaming holdback helpers ── // During streaming, content is flushed in 300ms batches. A storage reference or // an XML tag may be split across two batches. These helpers detect incomplete // patterns at the end of a chunk so the caller can hold them back until the // next flush completes them. // incompleteXMLTagRe matches the opening of an , , or tag // that reaches the end of the string without a closing '>'. var incompleteXMLTagRe = regexp.MustCompile( `<(?:image|image_original|image_caption|image_ocr|kb|web)[^>]*$`, ) // findIncompleteXMLTag returns the byte offset of a potentially truncated XML // tag at the tail of s, or -1 if none. func findIncompleteXMLTag(s string) int { loc := incompleteXMLTagRe.FindStringIndex(s) if loc == nil { return -1 } return loc[0] } // holdbackCutoff returns the earliest incomplete-pattern offset at the tail of // chunk, or len(chunk) if the chunk is safe to flush entirely. func holdbackCutoff(chunk string) int { cutoff := storageurl.HoldbackCutoff(chunk) if idx := findIncompleteXMLTag(chunk); idx >= 0 && idx < cutoff { cutoff = idx } return cutoff } // formatIMOutboundAnswer strips thinking/tool blocks and applies IM content cleanup. func formatIMOutboundAnswer(ctx context.Context, raw string, tenant *types.Tenant, defaultFileSvc interfaces.FileService, storageResolvers ...interfaces.StorageBackendResolver) string { return cleanIMContent(ctx, FormatIMDisplayContent(raw, StreamDisplayFinal), tenant, defaultFileSvc, storageResolvers...) } // formatIMOutboundAnswerOrFallback guarantees that cleanup cannot turn a // non-empty model payload (for example, a think-only response) into an empty IM // message. Callers may pass imErrorFallback or imCancelledFallback as raw when // QA itself failed or was stopped. func formatIMOutboundAnswerOrFallback( ctx context.Context, raw string, tenant *types.Tenant, defaultFileSvc interfaces.FileService, storageResolvers ...interfaces.StorageBackendResolver, ) string { content := formatIMOutboundAnswer(ctx, raw, tenant, defaultFileSvc, storageResolvers...) if strings.TrimSpace(content) == "" { return imNoAnswerFallback } return content } // imOutboundContext keeps values from ctx but drops cancellation. /stop cancels // the QA request; Feishu CardKit calls honor that deadline and would otherwise // leave the thinking placeholder on screen. func imOutboundContext(ctx context.Context) context.Context { return context.WithoutCancel(ctx) } // imQAFailureReply maps a QA error onto the user-visible IM fallback text. func imQAFailureReply(err error) string { if err == nil { return imNoAnswerFallback } if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { return imCancelledFallback } return imErrorFallback } // 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 (scheme-aware per tenant config) func cleanIMContent(ctx context.Context, content string, tenant *types.Tenant, defaultFileSvc interfaces.FileService, storageResolvers ...interfaces.StorageBackendResolver) string { content = stripImageXMLTags(content) content = stripIMCitationTags(content) resolver := newIMFileServiceResolver(tenant, defaultFileSvc, storageResolvers...).WithContext(ctx) content = rewriteStorageURLs(ctx, content, resolver) return content } func imLocalStorageBaseDir() string { return storageurl.LocalStorageBaseDir() } // newIMFileServiceResolver builds a per-message storage backend resolver. The // cache lives for one cleanIMContent / outbound message so a long answer does // not re-create an SDK client for every reference. func newIMFileServiceResolver( tenant *types.Tenant, defaultSvc interfaces.FileService, storageResolvers ...interfaces.StorageBackendResolver, ) *storageurl.FileServiceResolver { return storageurl.NewFileServiceResolver(tenant, defaultSvc, storageResolvers...) } func buildIMFileServiceForProvider( tenant *types.Tenant, provider string, defaultSvc interfaces.FileService, ) interfaces.FileService { return storageurl.BuildFileServiceForProvider(tenant, provider, defaultSvc) } // 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).ResolveFileService(filePath) } const ( // wsLeaderTTL is the TTL for the Redis key used for WebSocket leader election. wsLeaderTTL = 15 * time.Second // wsLeaderRenewInterval is how often the leader renews its lock. wsLeaderRenewInterval = 5 * time.Second // wsLeaderRetryInterval is how often non-leader instances try to acquire the lock. wsLeaderRetryInterval = 10 * time.Second // stopMarkerTTL is the TTL for cross-instance /stop markers in Redis. stopMarkerTTL = 30 * time.Second // stopPollInterval is how often in-flight workers check for remote /stop signals. stopPollInterval = 500 * time.Millisecond ) // ── Redis key prefixes ────────────────────────────────────────────────────── // All IM-related Redis keys are defined here for discoverability and to avoid // scattered string literals across multiple files. const ( RedisKeyLeader = "im:ws:leader:" // + channelID — WebSocket leader election RedisKeyDedup = "im:dedup:" // + messageID — message deduplication RedisKeyStop = "im:stop:" // + userKey — cross-instance /stop marker (pre-execution) RedisKeyInflight = "im:inflight:" // + userKey — maps userKey → sessionID:messageID for cross-instance /stop RedisKeyQueueUser = "im:queue:user:" // + userKey — global per-user queue counter RedisKeyRateLimit = "im:ratelimit:" // + key — sliding-window rate limiting RedisKeyGlobalGate = "im:global:active" // global concurrent worker counter // RedisChannelConfig broadcasts durable im_channels mutations so every // application replica invalidates its local adapter/config snapshot. RedisChannelConfig = "im:channel:config" defaultRateLimitWindow = 60 * time.Second defaultRateLimitMaxRequests = 10 ) // ErrChannelDisabled reports that a channel row exists but is disabled, so // callers can tell it apart from a missing channel or a transient failure. var ErrChannelDisabled = errors.New("channel is disabled") // channelState holds runtime state for a running IM channel. type channelState struct { Channel *IMChannel Adapter Adapter Cancel context.CancelFunc // for stopping websocket goroutines leaderCancel context.CancelFunc // stops the leader renewal goroutine (nil if not leader) } type leaderRetryState struct { cancel context.CancelFunc } type channelConfigEvent struct { ChannelID string `json:"channel_id"` SourceInstance string `json:"source_instance"` } // AdapterFactory creates an Adapter from an IMChannel configuration. // The second return value is an optional cleanup function (e.g., for stopping websocket connections). type AdapterFactory func(ctx context.Context, channel *IMChannel, msgHandler func(ctx context.Context, msg *IncomingMessage) error) (Adapter, context.CancelFunc, error) // inflightEntry tracks a running QA request, keyed by userKey in the inflight map. type inflightEntry struct { cancel context.CancelFunc sessionID string // set after assistant message is created assistantMessageID string // set after assistant message is created } // Service orchestrates IM message handling: // 1. Receives a unified IncomingMessage from an Adapter // 2. Resolves or creates a WeKnora session for the IM channel // 3. Dispatches slash-commands (/help, /kb, /clear, etc.) without entering QA // 4. Calls the WeKnora QA pipeline for normal messages // 5. Collects the streaming answer and sends it back via the Adapter type Service struct { db *gorm.DB sessionService interfaces.SessionService messageService interfaces.MessageService tenantService interfaces.TenantService agentService interfaces.CustomAgentService // knowledgeService is used for saving IM file messages to knowledge bases. knowledgeService interfaces.KnowledgeService // kbService is used by slash-commands (/info) to list and inspect knowledge bases. kbService interfaces.KnowledgeBaseService // oauthManager builds MCP OAuth authorization URLs so IM users can authorize // OAuth-enabled MCP services out-of-band (IM cannot resolve the in-conversation // prompt). May be nil, in which case a generic console hint is shown instead. oauthManager *mcppkg.OAuthManager // streamManager writes/reads QA events for distributed stop detection, // consistent with the web StopSession mechanism. May be nil in Lite mode // (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 documentReader interfaces.DocumentReader storageResolver interfaces.StorageBackendResolver // cmdRegistry holds all registered slash-commands. cmdRegistry *CommandRegistry // channels maps channel ID -> running channel state channels map[string]*channelState leaderRetries map[string]*leaderRetryState mu sync.RWMutex // adapterFactories maps platform name -> factory function adapterFactories map[string]AdapterFactory // processedMsgs tracks recently processed message IDs to prevent duplicate handling. processedMsgs sync.Map // rateLimiter enforces per-user sliding window rate limiting. // Uses Redis ZSET when available, falls back to local sliding window. rateLimiter *ratelimit.Limiter rateLimitMax int // inflight tracks in-progress QA requests, keyed by userKey // ("channelID:userID:chatID"). Allows /stop to abort a running request // on this instance and look up (sessionID, messageID) for StreamManager. inflight sync.Map // userKey -> *inflightEntry // qaQueue manages bounded queuing and worker-pool execution of QA requests, // providing backpressure to protect downstream LLM resources. qaQueue *qaQueue // redis is the optional Redis client for distributed state (dedup, rate // limiting, leader election, cross-instance /stop). When nil the service // falls back to local in-memory state (single-instance / Lite mode). redis *redis.Client // instanceID uniquely identifies this service instance for leader election. instanceID string stopCh chan struct{} stopOnce sync.Once subscriberOnce sync.Once stopped atomic.Bool } // makeUserKey builds the canonical key used to identify a user's request // across the queue, inflight map, and /stop command. // threadID should only be non-empty when channel.SessionMode == "thread"; // callers must guard this to avoid leaking thread scope into user-mode keys. func makeUserKey(channelID, userID, chatID, threadID string) string { if threadID != "" { return fmt.Sprintf("%s:%s:%s:%s", channelID, userID, chatID, threadID) } return fmt.Sprintf("%s:%s:%s", channelID, userID, chatID) } // nonTextTypeLabel maps a message type to a Chinese label for LLM instructions. var nonTextTypeLabel = map[string]string{ "image": "图片", "file": "文件", "video": "视频", "voice": "语音", } // formatQuotedContext formats a QuotedMessage into a labeled string for LLM context. // Returns empty string if quote is nil. // For non-text quotes, generates an instruction telling the LLM to acknowledge // the unprocessable content instead of a placeholder that causes hallucination. func formatQuotedContext(quote *QuotedMessage) string { if quote == nil { return "" } // Non-text quote: generate instruction, not content placeholder. if quote.NonTextType != "" { label := nonTextTypeLabel[quote.NonTextType] if label == "" { label = "该类型的" } return "用户引用了一条" + label + "消息,但你无法查看该内容。请直接告知用户你目前无法处理" + label + "消息,建议用户用文字描述问题。不要猜测该消息的内容。" } if quote.Content == "" { return "" } content := quote.Content runes := []rune(content) if len(runes) > maxQuoteContentLength { content = string(runes[:maxQuoteContentLength]) + "..." } // Prevent quoted content from escaping the XML tag boundary. content = strings.ReplaceAll(content, "", "") label := "以下是用户引用的一条历史消息,仅作为上下文参考:" if quote.IsBotMessage { label = "以下是用户引用的你(机器人)之前的回复,仅作为上下文参考:" } return label + "\n\n" + content + "\n" } // withIMIdentity injects a synthetic caller identity into the context for IM // callbacks. IM platforms verify their own signatures and bypass the auth // middleware, so the downstream QA pipeline would otherwise see an empty // UserID/TenantRole. Mirroring the API-key path's "system-" synthetic // user (recognised by types.IsSyntheticUserID) lets Organization-shared // knowledge bases be merged and resolved correctly, since the shared-KB code // gates on a non-empty UserID. Viewer is the least privilege sufficient to // retrieve shared KBs. func withIMIdentity(ctx context.Context, tenantID uint64, channelID string, msg *IncomingMessage) context.Context { ctx = context.WithValue(ctx, types.TenantIDContextKey, tenantID) ctx = context.WithValue(ctx, types.UserIDContextKey, fmt.Sprintf("system-%d", tenantID)) if msg != nil { principalID := fmt.Sprintf("%d:%s:%s:%s", tenantID, channelID, msg.Platform, msg.UserID) ctx = types.WithPrincipal(ctx, types.Principal{Type: types.PrincipalIMUser, ID: principalID}) } ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleViewer) // IM bots have no live client that can complete an in-conversation MCP OAuth // prompt, so mark the context non-interactive: the agent emits a one-shot // authorization notice (surfaced in the reply) instead of blocking until the // OAuth wait times out for every unauthorized service. ctx = types.WithMCPOAuthNonInteractive(ctx) return ctx } // imMCPAuthService identifies an OAuth-enabled MCP service that the IM user has // not authorized yet, collected during a turn so an authorization notice can be // appended to the reply. type imMCPAuthService struct { ID string Name string } // mcpOAuthCallbackURL returns the absolute backend callback URL registered with // the authorization server, derived from APP_EXTERNAL_URL. Empty when unset. func mcpOAuthCallbackURL() string { base := strings.TrimRight(strings.TrimSpace(os.Getenv("APP_EXTERNAL_URL")), "/") if base == "" { return "" } return base + "/api/v1/mcp-oauth/callback" } // buildIMMCPAuthNotice builds a user-facing authorization notice for the given // OAuth-enabled MCP services. When the OAuth manager and APP_EXTERNAL_URL are // available it generates a per-service authorization URL (StartAuthorization) // the user can open; otherwise it falls back to a console hint. Returns "" when // there is nothing to report. The user authorizes out-of-band, then re-sends // their message to use the service (IM cannot resolve the prompt inline). func (s *Service) buildIMMCPAuthNotice(ctx context.Context, services []imMCPAuthService) string { // Deduplicate by service ID, preserving order. seen := make(map[string]bool, len(services)) uniq := make([]imMCPAuthService, 0, len(services)) for _, svc := range services { if svc.ID == "" || seen[svc.ID] { continue } seen[svc.ID] = true uniq = append(uniq, svc) } if len(uniq) == 0 { return "" } tenantID, _ := types.TenantIDFromContext(ctx) principal := types.MCPOAuthPrincipalFromContext(ctx) redirectURI := mcpOAuthCallbackURL() frontendRedirect := strings.TrimSpace(os.Getenv("APP_EXTERNAL_URL")) if frontendRedirect == "" { frontendRedirect = "/" } var lines []string for _, svc := range uniq { name := strings.TrimSpace(svc.Name) if name == "" { name = svc.ID } var authURL string if s.oauthManager != nil && redirectURI != "" && tenantID != 0 && principal.Valid() { url, err := s.oauthManager.StartAuthorizationForService( ctx, tenantID, principal, svc.ID, redirectURI, frontendRedirect, ) if err != nil { logger.Warnf(ctx, "[IM] Failed to build MCP OAuth URL for service %s: %v", svc.ID, err) } else { authURL = url } } if authURL != "" { lines = append(lines, fmt.Sprintf("• %s:%s", name, authURL)) } else { lines = append(lines, fmt.Sprintf("• %s(请在 WeKnora 管理后台完成 OAuth 授权)", name)) } } return "⚠️ 以下 MCP 服务需要授权后才能使用,请点击链接完成授权,然后重新发送你的消息:\n" + strings.Join(lines, "\n") } // appendIMAuthNotice appends an authorization notice to an existing reply body, // separated by a blank line. When the body is empty the notice becomes the body. func appendIMAuthNotice(body, notice string) string { if notice == "" { return body } if strings.TrimSpace(body) == "" { return notice } return body + "\n\n" + notice } func buildIMQARequest( session *types.Session, query string, assistantMessageID string, userMessageID string, customAgent *types.CustomAgent, kbIDs []string, quote *QuotedMessage, attachments ...types.MessageAttachments, ) *types.QARequest { // WebSearchEnabled: the web handler passes this per-request from the // frontend toggle; for IM channels the user has no per-message toggle, // so we derive it from the agent config (the single source of truth). webSearchEnabled := customAgent != nil && customAgent.Config.WebSearchEnabled quotedContext := formatQuotedContext(quote) var requestAttachments types.MessageAttachments if len(attachments) > 0 { requestAttachments = attachments[0] } return &types.QARequest{ Session: session, Query: query, AssistantMessageID: assistantMessageID, CustomAgent: customAgent, KnowledgeBaseIDs: kbIDs, UserMessageID: userMessageID, WebSearchEnabled: webSearchEnabled, QuotedContext: quotedContext, Attachments: requestAttachments, } } func buildIMLastRequestState(agentID string, customAgent *types.CustomAgent, kbIDs []string) *types.SessionLastRequestState { state := &types.SessionLastRequestState{ AgentID: agentID, KnowledgeBaseIDs: append([]string(nil), kbIDs...), } if customAgent == nil { return state } if state.AgentID == "" { state.AgentID = customAgent.ID } state.AgentEnabled = customAgent.IsAgentMode() state.ModelID = customAgent.Config.ModelID state.WebSearchEnabled = customAgent.Config.WebSearchEnabled if len(state.KnowledgeBaseIDs) == 0 && len(customAgent.Config.KnowledgeBases) > 0 { state.KnowledgeBaseIDs = append([]string(nil), customAgent.Config.KnowledgeBases...) } return state } func createIMUserMessagePayload(sessionID, content, requestID string, attachments ...types.MessageAttachments) *types.Message { var messageAttachments types.MessageAttachments if len(attachments) > 0 { messageAttachments = attachments[0] } return &types.Message{ SessionID: sessionID, Role: "user", Content: content, RequestID: requestID, CreatedAt: time.Now(), IsCompleted: true, Channel: "im", Attachments: messageAttachments, } } type imDownloadedAttachment struct { fileName string content []byte } // prepareIMAttachments downloads an IM attachment and exposes its parsed text // (and, for images, a bounded data URI) to the QA pipeline. This is separate // from the optional background knowledge-base save. func (s *Service) prepareIMAttachments(ctx context.Context, msg *IncomingMessage, adapter Adapter) (types.MessageAttachments, []string, *imDownloadedAttachment, error) { if msg.MessageType != MessageTypeFile && msg.MessageType != MessageTypeImage { return nil, nil, nil, nil } if msg.FileSize > maxIMAttachmentBytes { return nil, nil, nil, fmt.Errorf("attachment exceeds the %d MiB limit", maxIMAttachmentBytes>>20) } downloader, ok := adapter.(FileDownloader) if !ok { return nil, nil, nil, fmt.Errorf("platform %s does not support attachment download", msg.Platform) } attachmentCtx, cancel := context.WithTimeout(ctx, imAttachmentReadTimeout) defer cancel() reader, fileName, err := downloader.DownloadFile(attachmentCtx, msg) if err != nil { return nil, nil, nil, err } defer reader.Close() content, err := io.ReadAll(io.LimitReader(reader, maxIMAttachmentBytes+1)) if err != nil { return nil, nil, nil, err } if len(content) > maxIMAttachmentBytes { return nil, nil, nil, fmt.Errorf("attachment exceeds the %d MiB limit", maxIMAttachmentBytes>>20) } if fileName == "" { fileName = msg.FileName } if msg.MessageType == MessageTypeImage && filepath.Ext(fileName) == "" { fileName += ".png" } ext := strings.ToLower(strings.TrimPrefix(filepath.Ext(fileName), ".")) if ext == "" { return nil, nil, nil, fmt.Errorf("attachment has no file extension") } attachment := types.MessageAttachment{FileName: fileName, FileType: "." + ext, FileSize: int64(len(content))} request := &types.ReadRequest{FileContent: content, FileName: fileName, FileType: ext} var result *types.ReadResult isImage := msg.MessageType == MessageTypeImage || docparser.IsImageFormat(ext) if isImage && s.documentReader != nil { result, err = s.documentReader.Read(attachmentCtx, request) if err != nil { logger.Warnf(ctx, "[IM] image OCR/document parsing failed, continuing with vision input: %v", err) result, err = nil, nil } } if result == nil && docparser.IsSimpleFormat(attachment.FileType) { result, err = (&docparser.SimpleFormatReader{}).Read(attachmentCtx, request) } else if result == nil && !isImage && s.documentReader != nil { result, err = s.documentReader.Read(attachmentCtx, request) } if err != nil { logger.Warnf(ctx, "[IM] attachment parsing failed, continuing with attachment metadata: %v", err) } if result != nil { applyIMAttachmentTruncation(result.MarkdownContent, &attachment) } var imageURLs []string if isImage && len(content) <= maxIMVisionAttachmentBytes { mediaType := http.DetectContentType(content) if !strings.HasPrefix(mediaType, "image/") { return nil, nil, nil, fmt.Errorf("invalid image content type: %s", mediaType) } imageURLs = []string{"data:" + mediaType + ";base64," + base64.StdEncoding.EncodeToString(content)} } else if isImage { logger.Warnf(ctx, "[IM] image is too large for direct vision input: size=%d limit=%d", len(content), maxIMVisionAttachmentBytes) } return types.MessageAttachments{attachment}, imageURLs, &imDownloadedAttachment{fileName: fileName, content: content}, nil } func applyIMAttachmentTruncation(content string, attachment *types.MessageAttachment) { attachment.LineCount = strings.Count(content, "\n") + 1 limited := truncateUTF8ByBytes(content, maxIMAttachmentContentBytes) lines := strings.SplitN(limited, "\n", maxIMAttachmentLines+1) if len(lines) > maxIMAttachmentLines { limited = strings.Join(lines[:maxIMAttachmentLines], "\n") } attachment.Content = limited attachment.IsTruncated = len(limited) < len(content) } func truncateUTF8ByBytes(content string, maxBytes int) string { if len(content) <= maxBytes { return content } end := maxBytes for end > 0 && !utf8.RuneStart(content[end]) { end-- } return content[:end] } func createIMAssistantMessagePayload(sessionID, requestID string) *types.Message { return &types.Message{ SessionID: sessionID, Role: "assistant", RequestID: requestID, CreatedAt: time.Now(), IsCompleted: false, Channel: "im", } } func collectIMKnowledgeReferences(dst *[]*types.SearchResult, refs interface{}) { switch v := refs.(type) { case []*types.SearchResult: *dst = append(*dst, v...) case []interface{}: for _, ref := range v { if sr, ok := ref.(*types.SearchResult); ok { *dst = append(*dst, sr) } } } } func sanitizeIMAgentSteps(raw interface{}) types.AgentSteps { switch steps := raw.(type) { case []types.AgentStep: return types.AgentSteps(agenttools.SanitizeAgentStepsForStorage(steps)) case types.AgentSteps: return types.AgentSteps(agenttools.SanitizeAgentStepsForStorage([]types.AgentStep(steps))) default: return nil } } func applyIMCompleteDataToMessage(msg *types.Message, data event.AgentCompleteData) { if msg == nil { return } if data.MessageID != "" && data.MessageID != msg.ID { return } msg.IsCompleted = true msg.AgentDurationMs = data.TotalDurationMs if usage, ok := data.Usage.(*types.TokenUsage); ok && usage != nil { msg.Usage = usage } if len(data.KnowledgeRefs) > 0 { refs := make([]*types.SearchResult, 0, len(data.KnowledgeRefs)) collectIMKnowledgeReferences(&refs, data.KnowledgeRefs) if len(refs) > 0 { msg.KnowledgeReferences = types.References(refs) } } if steps := sanitizeIMAgentSteps(data.AgentSteps); len(steps) > 0 { msg.AgentSteps = steps } } // waitForIMAgentComplete blocks until EventAgentComplete, ctx cancellation, or timeout. func waitForIMAgentComplete(ctx context.Context, completeDone <-chan struct{}, sessionID string) { timer := time.NewTimer(agentCompleteWaitTimeout) defer timer.Stop() select { case <-completeDone: case <-ctx.Done(): logger.Warnf(ctx, "[IM] QA context ended before agent complete event: session=%s", sessionID) case <-timer.C: logger.Warnf(ctx, "[IM] Timed out waiting for agent complete event: session=%s", sessionID) } } // pickIMStoredAnswer returns the best available answer text from IM stream buffers. func pickIMStoredAnswer(candidates ...string) string { for _, s := range candidates { if strings.TrimSpace(s) != "" { return s } } return "" } // mergeIMAgentAnswerBuffers copies optimistic/live answers into persistence buffers // when EventAgentFinalAnswer did not populate answerBuilder (e.g. cancel before complete). func mergeIMAgentAnswerBuffers(answerBuilder, answerOuter, agentLiveAnswer *strings.Builder, completeFinal string) { if answerBuilder.Len() > 0 { return } switch { case agentLiveAnswer.Len() > 0: live := agentLiveAnswer.String() answerBuilder.WriteString(live) if answerOuter.Len() == 0 { answerOuter.WriteString(live) } case answerOuter.Len() > 0: answerBuilder.WriteString(answerOuter.String()) case strings.TrimSpace(completeFinal) != "": answerBuilder.WriteString(completeFinal) answerOuter.WriteString(completeFinal) } } // resolveIMConfig extracts IM tuning parameters from the application config, // falling back to built-in defaults for any zero/nil values. func resolveIMConfig(appCfg *config.Config) (workers, maxQueue, maxPerUser, globalMaxWorkers int, rlWindow time.Duration, rlMax int) { workers = defaultWorkers maxQueue = defaultMaxQueueSize maxPerUser = defaultMaxPerUser rlWindow = defaultRateLimitWindow rlMax = defaultRateLimitMaxRequests if appCfg == nil || appCfg.IM == nil { return } im := appCfg.IM if im.Workers > 0 { workers = im.Workers } if im.MaxQueueSize > 0 { maxQueue = im.MaxQueueSize } if im.MaxPerUser > 0 { maxPerUser = im.MaxPerUser } if im.GlobalMaxWorkers > 0 { globalMaxWorkers = im.GlobalMaxWorkers } if im.RateLimitWindow > 0 { rlWindow = im.RateLimitWindow } if im.RateLimitMax > 0 { rlMax = im.RateLimitMax } return } // NewService creates a new IM service. // redisClient may be nil — in that case the service falls back to local // in-memory state (Lite / single-instance mode). // cfg may be nil — in that case built-in defaults are used. func NewService( db *gorm.DB, sessionService interfaces.SessionService, messageService interfaces.MessageService, tenantService interfaces.TenantService, agentService interfaces.CustomAgentService, knowledgeService interfaces.KnowledgeService, kbService interfaces.KnowledgeBaseService, streamManager interfaces.StreamManager, defaultFileSvc interfaces.FileService, documentReader interfaces.DocumentReader, oauthManager *mcppkg.OAuthManager, redisClient *redis.Client, appCfg *config.Config, storageResolver interfaces.StorageBackendResolver, ) *Service { // Resolve IM configuration with defaults. workers, maxQueue, maxPerUser, globalMaxWorkers, rlWindow, rlMax := resolveIMConfig(appCfg) // Build command registry. registry := NewCommandRegistry() registry.Register(newHelpCommand(registry)) registry.Register(newInfoCommand(kbService)) registry.Register(newSearchCommand(sessionService, kbService)) registry.Register(newStopCommand()) registry.Register(newClearCommand()) instanceID := uuid.New().String() s := &Service{ db: db, sessionService: sessionService, messageService: messageService, tenantService: tenantService, agentService: agentService, knowledgeService: knowledgeService, kbService: kbService, streamManager: streamManager, defaultFileSvc: defaultFileSvc, documentReader: documentReader, storageResolver: storageResolver, oauthManager: oauthManager, cmdRegistry: registry, channels: make(map[string]*channelState), leaderRetries: make(map[string]*leaderRetryState), adapterFactories: make(map[string]AdapterFactory), rateLimiter: ratelimit.New(redisClient, RedisKeyRateLimit, rlWindow, instanceID), rateLimitMax: rlMax, redis: redisClient, instanceID: instanceID, stopCh: make(chan struct{}), } // Initialize the QA worker pool and bounded queue. s.qaQueue = newQAQueue(workers, maxQueue, maxPerUser, globalMaxWorkers, s.executeQARequest, redisClient) s.qaQueue.Start(s.stopCh) // Start periodic cleanup loops. // Dedup cleanup is only needed in single-instance mode (local sync.Map); // when Redis handles dedup, the TTL on Redis keys handles expiry automatically. if redisClient == nil { go s.dedupCleanupLoop() } go s.rateLimiter.StartCleanup(s.stopCh) if redisClient != nil { globalInfo := "unlimited" if globalMaxWorkers > 0 { globalInfo = fmt.Sprintf("%d", globalMaxWorkers) } logger.Infof(context.Background(), "[IM] Multi-instance mode enabled (instance=%s, workers=%d, queue=%d, global_max=%s)", s.instanceID[:8], workers, maxQueue, globalInfo) } else { logger.Infof(context.Background(), "[IM] Single-instance mode (no Redis, workers=%d, queue=%d)", workers, maxQueue) } return s } // RegisterAdapterFactory registers a factory for creating adapters for a given platform. func (s *Service) RegisterAdapterFactory(platform string, factory AdapterFactory) { s.mu.Lock() defer s.mu.Unlock() s.adapterFactories[platform] = factory } // Stop gracefully shuts down the service, stopping all channels and background goroutines. func (s *Service) Stop() { s.stopOnce.Do(func() { s.stopped.Store(true) close(s.stopCh) if s.qaQueue != nil { s.qaQueue.Stop() } s.mu.Lock() defer s.mu.Unlock() for id, cs := range s.channels { s.stopChannelLocked(id, cs) } for id, retry := range s.leaderRetries { retry.cancel() delete(s.leaderRetries, id) } }) } // dedupCleanupLoop periodically cleans up expired entries from the dedup map. func (s *Service) dedupCleanupLoop() { ticker := time.NewTicker(dedupCleanupInterval) defer ticker.Stop() for { select { case <-ticker.C: cutoff := time.Now().Add(-dedupTTL) s.processedMsgs.Range(func(key, value interface{}) bool { if t, ok := value.(time.Time); ok && t.Before(cutoff) { s.processedMsgs.Delete(key) } return true }) case <-s.stopCh: return } } } // imImageConfigWarning returns an operator warning when IM channels are active // but APP_EXTERNAL_URL is unset. resource:// images then render only if the // storage backend is itself publicly reachable (e.g. a cloud bucket with a // public endpoint); on the default MinIO (internal minio:9000) or local // deployment it is not, so images silently break. Advisory only; returns "" // when there is nothing to warn about. Pure (no receiver/closures) so it is // unit-testable. func imImageConfigWarning(activeChannels int, externalURL string) string { if activeChannels == 0 || strings.TrimSpace(externalURL) != "" { return "" } return fmt.Sprintf("[IM] %d IM channel(s) active but APP_EXTERNAL_URL is unset; "+ "resource:// images render only if the storage backend is publicly reachable — "+ "otherwise set APP_EXTERNAL_URL so they route through nginx /r/", activeChannels) } // LoadAndStartChannels loads all enabled channels from the database and starts them. func (s *Service) LoadAndStartChannels() error { s.startChannelConfigSubscriber() ctx := context.Background() var channels []IMChannel if err := s.db.Where("enabled = ? AND deleted_at IS NULL", true).Find(&channels).Error; err != nil { return fmt.Errorf("load im channels: %w", err) } if msg := imImageConfigWarning(len(channels), os.Getenv("APP_EXTERNAL_URL")); msg != "" { logger.Warnf(ctx, "%s", msg) } for i := range channels { ch := channels[i] if err := s.StartChannel(&ch); err != nil { logger.Warnf(ctx, "[IM] Failed to start channel %s (%s/%s): %v", ch.ID, ch.Platform, ch.Name, err) } else if _, _, active := s.GetChannelAdapter(ch.ID); active { // Adapter initialization can launch an asynchronous connection attempt. // Do not claim network readiness here; platform connection logs report it. logger.Infof(ctx, "[IM] Initialized channel runtime: id=%s platform=%s name=%s mode=%s agent=%s", ch.ID, ch.Platform, ch.Name, ch.Mode, ch.AgentID) } else { logger.Infof(ctx, "[IM] Channel runtime is on standby: id=%s platform=%s name=%s mode=%s agent=%s", ch.ID, ch.Platform, ch.Name, ch.Mode, ch.AgentID) } } logger.Infof(ctx, "[IM] Loaded %d enabled channels", len(channels)) return nil } // startChannelConfigSubscriber listens for durable channel mutations made by // other application replicas. Redis Pub/Sub provides the fast path; callback // freshness checks and the leader renewal DB check remain the durable fallback // if an event is missed while a replica is disconnected. func (s *Service) startChannelConfigSubscriber() { if s.redis == nil { return } s.subscriberOnce.Do(func() { go s.channelConfigSubscriberLoop() }) } func (s *Service) channelConfigSubscriberLoop() { pubsub := s.redis.Subscribe(context.Background(), RedisChannelConfig) defer pubsub.Close() messages := pubsub.Channel() for { select { case <-s.stopCh: return case msg, ok := <-messages: if !ok { return } var event channelConfigEvent if err := json.Unmarshal([]byte(msg.Payload), &event); err != nil { logger.Warnf(context.Background(), "[IM] Ignore invalid channel config event: %v", err) continue } if event.ChannelID == "" || event.SourceInstance == s.instanceID { continue } s.reloadChannelFromDB(event.ChannelID, "config event") } } } func (s *Service) publishChannelConfigChange(channelID string) { if s.redis == nil || channelID == "" || s.stopped.Load() { return } payload, err := json.Marshal(channelConfigEvent{ ChannelID: channelID, SourceInstance: s.instanceID, }) if err != nil { return } if err := s.redis.Publish(context.Background(), RedisChannelConfig, payload).Err(); err != nil { // The database is the source of truth. Subscribers also verify freshness // on webhook callbacks / leader renewal, so publication is best-effort. logger.Warnf(context.Background(), "[IM] Publish channel config event failed for %s: %v", channelID, err) } } func (s *Service) reloadChannelFromDB(channelID, reason string) { fresh, err := s.GetChannelByID(channelID) if errors.Is(err, gorm.ErrRecordNotFound) { s.StopChannel(channelID) return } if err != nil { logger.Warnf(context.Background(), "[IM] Reload channel %s after %s failed: %v", channelID, reason, err) return } if !fresh.Enabled { s.StopChannel(channelID) return } if _, cached, running := s.GetChannelAdapter(channelID); running && sameChannelRuntimeConfig(cached, fresh) { return } logger.Infof(context.Background(), "[IM] Reloading channel %s after %s", channelID, reason) if err := s.StartChannel(fresh); err != nil && !s.stopped.Load() { logger.Warnf(context.Background(), "[IM] Reload channel %s after %s failed: %v", channelID, reason, err) } } // StartChannel creates and registers an adapter for the given channel. // For WebSocket channels with Redis available, only one instance acquires // the leader lock and opens the connection; other instances periodically // retry so they can take over if the leader dies. func (s *Service) StartChannel(channel *IMChannel) error { if s.stopped.Load() { return fmt.Errorf("im service is stopped") } s.mu.Lock() s.stopLeaderRetryLocked(channel.ID) factory, ok := s.adapterFactories[channel.Platform] if !ok { s.mu.Unlock() return fmt.Errorf("no adapter factory for platform: %s", channel.Platform) } // Stop existing channel if running if existing, ok := s.channels[channel.ID]; ok { s.stopChannelLocked(channel.ID, existing) } s.mu.Unlock() // For WebSocket / long-poll channels, try leader election to avoid // duplicate connections. Only one instance should actively poll or // maintain a persistent connection for each channel. if (channel.Mode == "websocket" || channel.Mode == "longpoll") && s.redis != nil { acquired := s.tryAcquireWSLeader(channel.ID) if !acquired { logger.Infof(context.Background(), "[IM] Channel %s %s owned by another instance, will retry", channel.ID, channel.Mode) s.scheduleWSLeaderRetry(channel) return nil } } return s.startChannelInternal(channel, factory) } // startChannelInternal does the actual adapter creation and registration. func (s *Service) startChannelInternal(channel *IMChannel, factory AdapterFactory) error { // Build the message handler that delegates to HandleMessage with this channel's config msgHandler := func(msgCtx context.Context, msg *IncomingMessage) error { return s.HandleMessage(msgCtx, msg, channel.ID) } ctx := context.Background() adapter, cancelFn, err := factory(ctx, channel, msgHandler) if err != nil { s.releaseWSLeader(channel.ID) // release lock on failure return fmt.Errorf("create adapter: %w", err) } // Start leader renewal goroutine for WebSocket / long-poll channels. var leaderCancel context.CancelFunc if (channel.Mode == "websocket" || channel.Mode == "longpoll") && s.redis != nil { leaderCtx, lCancel := context.WithCancel(context.Background()) leaderCancel = lCancel go s.wsLeaderRenewLoop(leaderCtx, channel.ID) } s.mu.Lock() // Stop() may have drained the channel map while the factory was connecting // above, since the factory runs unlocked. Re-check under the lock, otherwise // this adapter's long connection would outlive process shutdown. if s.stopped.Load() { s.mu.Unlock() if leaderCancel != nil { leaderCancel() } if cancelFn != nil { cancelFn() } s.releaseWSLeader(channel.ID) return fmt.Errorf("im service is stopped") } // The factory runs outside the lock and may finish after an Agent deletion // has already invalidated this channel. Fence that late startup using the // durable row while holding the same lock used by StopChannel. fresh, freshErr := s.GetChannelByID(channel.ID) if freshErr != nil || !fresh.Enabled || !sameChannelRuntimeConfig(channel, fresh) { s.mu.Unlock() if leaderCancel != nil { leaderCancel() } if cancelFn != nil { cancelFn() } s.releaseWSLeader(channel.ID) if freshErr != nil { return freshErr } return fmt.Errorf("channel configuration changed during startup") } // Idempotency: another goroutine may have started this channel while // factory was running above (factory is called unlocked). Stop the old // state before overwriting so its adapter / long connection doesn't leak. if existing, ok := s.channels[channel.ID]; ok { s.stopChannelLocked(channel.ID, existing) } s.channels[channel.ID] = &channelState{ Channel: channel, Adapter: adapter, Cancel: cancelFn, leaderCancel: leaderCancel, } s.mu.Unlock() return nil } // StopChannel stops and removes a running channel. func (s *Service) StopChannel(channelID string) { s.mu.Lock() s.stopLeaderRetryLocked(channelID) if cs, ok := s.channels[channelID]; ok { s.stopChannelLocked(channelID, cs) } s.mu.Unlock() s.cancelChannelRequests(channelID) } func (s *Service) stopLeaderRetryLocked(channelID string) { if retry, ok := s.leaderRetries[channelID]; ok { retry.cancel() delete(s.leaderRetries, channelID) } } // stopChannelLocked stops a channel and removes it from the map. // Caller must hold s.mu. func (s *Service) stopChannelLocked(channelID string, cs *channelState) { if cs.leaderCancel != nil { cs.leaderCancel() } if cs.Cancel != nil { cs.Cancel() } delete(s.channels, channelID) // For long-poll channels, do NOT release the leader lock immediately. // Let it expire naturally via TTL so the old poll goroutine has time to // fully drain before another instance takes over. This prevents a brief // dual-writer window where both old and new instances process messages. // For websocket channels, the connection closes synchronously, so // immediate release is safe. if cs.Channel != nil && cs.Channel.Mode == "longpoll" { logger.Infof(context.Background(), "[IM] Stopped longpoll channel: id=%s (leader lock will expire via TTL)", channelID) } else { s.releaseWSLeader(channelID) logger.Infof(context.Background(), "[IM] Stopped channel: id=%s", channelID) } } // ── WebSocket leader election ─────────────────────────────────────────────── // tryAcquireWSLeader attempts to acquire the Redis lock for a WebSocket channel. // Returns true if this instance is now the leader. func (s *Service) tryAcquireWSLeader(channelID string) bool { if s.redis == nil { return true // single-instance mode: always leader } key := RedisKeyLeader + channelID ok, err := s.redis.SetNX(context.Background(), key, s.instanceID, wsLeaderTTL).Result() if err != nil { logger.Warnf(context.Background(), "[IM] Redis leader election failed for %s: %v; connection will retry without taking leadership", channelID, err) return false } return ok } // releaseWSLeader releases the Redis leader lock for a WebSocket channel, // but only if this instance owns it. func (s *Service) releaseWSLeader(channelID string) { if s.redis == nil { return } key := RedisKeyLeader + channelID // Only delete if we own it (compare-and-delete via Lua). script := redis.NewScript(` if redis.call('GET', KEYS[1]) == ARGV[1] then return redis.call('DEL', KEYS[1]) end return 0 `) script.Run(context.Background(), s.redis, []string{key}, s.instanceID) } // wsLeaderRenewLoop periodically refreshes the leader lock TTL. // Stops when ctx is cancelled (channel stopped) or if the lock is lost. func (s *Service) wsLeaderRenewLoop(ctx context.Context, channelID string) { key := RedisKeyLeader + channelID ticker := time.NewTicker(wsLeaderRenewInterval) defer ticker.Stop() for { select { case <-ticker.C: // Only renew if we still own the lock. script := redis.NewScript(` if redis.call('GET', KEYS[1]) == ARGV[1] then redis.call('PEXPIRE', KEYS[1], ARGV[2]) return 1 end return 0 `) result, err := script.Run(ctx, s.redis, []string{key}, s.instanceID, wsLeaderTTL.Milliseconds()).Int64() if err != nil || result == 0 { logger.Warnf(context.Background(), "[IM] Lost leadership for channel %s, stopping adapter and scheduling recovery", channelID) s.handleWSLeadershipLoss(channelID) return } // Still the leader — verify the channel is still active. A // delete/disable is served by whichever instance got the HTTP // request; without this check the leader would keep the long // connection open until process restart. The renew interval // bounds the worst-case lag. ch, err := s.GetChannelByID(channelID) if errors.Is(err, gorm.ErrRecordNotFound) { logger.Infof(context.Background(), "[IM] Channel %s deleted; leader stepping down", channelID) s.StopChannel(channelID) return } if err != nil { // Transient DB error — don't stop a possibly-healthy channel // on a DB hiccup. Skip this round; the next renewal re-checks. logger.Warnf(context.Background(), "[IM] DB check failed for channel %s during leader renewal: %v (skipping this round)", channelID, err) continue } if !ch.Enabled { logger.Infof(context.Background(), "[IM] Channel %s disabled; leader stepping down", channelID) s.StopChannel(channelID) return } _, cached, running := s.GetChannelAdapter(channelID) if !running { return } if !sameChannelRuntimeConfig(cached, ch) { logger.Infof(context.Background(), "[IM] Channel %s config changed; rebuilding runtime", channelID) // StartChannel synchronously stops the old runtime, releases its // lease, and then competes to start the fresh configuration. Return // because the old renewal context has been cancelled. if err := s.StartChannel(ch); err != nil && !s.stopped.Load() { logger.Warnf(context.Background(), "[IM] Rebuild changed channel %s failed: %v", channelID, err) } return } case <-ctx.Done(): return } } } // handleWSLeadershipLoss tears down the adapter that no longer owns its lease // and puts the enabled channel back into the existing takeover loop. The retry // loop re-reads the durable channel row before reconnecting, so a concurrent // delete, disable, or config update cannot resurrect a stale runtime. func (s *Service) handleWSLeadershipLoss(channelID string) { _, channel, running := s.GetChannelAdapter(channelID) if !running || channel == nil { return } s.StopChannel(channelID) if s.stopped.Load() { return } s.scheduleWSLeaderRetry(channel) } func (s *Service) scheduleWSLeaderRetry(channel *IMChannel) { retryCtx, cancel := context.WithCancel(context.Background()) state := &leaderRetryState{cancel: cancel} s.mu.Lock() if existing, ok := s.leaderRetries[channel.ID]; ok { existing.cancel() } if s.stopped.Load() { s.mu.Unlock() cancel() return } s.leaderRetries[channel.ID] = state s.mu.Unlock() go s.wsLeaderRetryLoop(retryCtx, channel, state) } // wsLeaderRetryLoop periodically tries to acquire the WebSocket leader lock. // When it succeeds, it starts the channel adapter. A per-channel retry state // ensures repeated config events cannot accumulate duplicate retry goroutines. func (s *Service) wsLeaderRetryLoop(ctx context.Context, channel *IMChannel, state *leaderRetryState) { ticker := time.NewTicker(wsLeaderRetryInterval) defer ticker.Stop() defer func() { s.mu.Lock() if s.leaderRetries[channel.ID] == state { delete(s.leaderRetries, channel.ID) } s.mu.Unlock() }() for { select { case <-ticker.C: // Check if channel is already running (another goroutine may have started it). if _, _, ok := s.GetChannelAdapter(channel.ID); ok { return } if s.tryAcquireWSLeader(channel.ID) { // Re-check the DB before starting: the channel may have been // deleted or disabled on another instance while we waited for // leadership. The in-memory `channel` is a startup snapshot and // won't reflect that, so without this guard we'd resurrect a // stopped channel (and reopen its long connection). fresh, err := s.GetChannelByID(channel.ID) if errors.Is(err, gorm.ErrRecordNotFound) { // Channel was actually deleted — give up for good. s.releaseWSLeader(channel.ID) logger.Infof(context.Background(), "[IM] Channel %s deleted while waiting for leadership; aborting leader takeover", channel.ID) return } if err != nil { // Transient DB error — don't make a destructive decision. // Release the lock for this round and retry on the next tick. s.releaseWSLeader(channel.ID) logger.Warnf(context.Background(), "[IM] DB check failed for channel %s during leader takeover: %v (will retry)", channel.ID, err) continue } if !fresh.Enabled { s.releaseWSLeader(channel.ID) logger.Infof(context.Background(), "[IM] Channel %s disabled while waiting for leadership; aborting leader takeover", channel.ID) return } channel = fresh // use latest config (credentials/mode may have changed) logger.Infof(context.Background(), "[IM] Acquired leadership for channel %s, starting adapter", channel.ID) s.mu.RLock() factory, ok := s.adapterFactories[channel.Platform] s.mu.RUnlock() if !ok { return } if err := s.startChannelInternal(channel, factory); err != nil { logger.Warnf(context.Background(), "[IM] Failed to start channel %s after acquiring leadership: %v", channel.ID, err) } return } case <-ctx.Done(): return case <-s.stopCh: return } } } // ── Cross-instance /stop via StreamManager ─────────────────────────────────── // // The mechanism mirrors the web StopSession flow: // 1. /stop writes a stop StreamEvent to StreamManager (keyed by sessionID + messageID) // 2. A per-request watcher polls StreamManager and cancels the context on detection // // A Redis marker (im:stop:{userKey}) is kept as a lightweight pre-execution // check for requests that haven't created an assistant message yet. // checkAndClearStopMarker checks if a pre-execution /stop marker exists for // the given userKey. If found, it deletes the marker and returns true. func (s *Service) checkAndClearStopMarker(ctx context.Context, userKey string) bool { if s.redis == nil { return false } stopKey := RedisKeyStop + userKey deleted, err := s.redis.Del(ctx, stopKey).Result() if err != nil { return false } return deleted > 0 } // storeInflightMapping writes the (sessionID, assistantMessageID) to Redis so // that /stop on any instance can look it up and write to StreamManager. func (s *Service) storeInflightMapping(ctx context.Context, userKey, sessionID, messageID string) { if s.redis == nil { return } val := sessionID + ":" + messageID if err := s.redis.Set(ctx, RedisKeyInflight+userKey, val, 10*time.Minute).Err(); err != nil { logger.Warnf(ctx, "[IM] Failed to store inflight mapping: %v", err) } } // clearInflightMapping removes the inflight mapping from Redis. func (s *Service) clearInflightMapping(ctx context.Context, userKey string) { if s.redis == nil { return } s.redis.Del(ctx, RedisKeyInflight+userKey) } // loadInflightMapping retrieves (sessionID, messageID) from Redis. func (s *Service) loadInflightMapping(ctx context.Context, userKey string) (sessionID, messageID string, ok bool) { if s.redis == nil { return "", "", false } val, err := s.redis.Get(ctx, RedisKeyInflight+userKey).Result() if err != nil { return "", "", false } parts := strings.SplitN(val, ":", 2) if len(parts) != 2 { return "", "", false } return parts[0], parts[1], true } // writeStopEvent writes a stop event to StreamManager, matching the web // StopSession pattern. The QA watcher goroutine detects it and cancels. func (s *Service) writeStopEvent(ctx context.Context, sessionID, messageID string) { stopEvt := interfaces.StreamEvent{ ID: fmt.Sprintf("stop-%d", time.Now().UnixNano()), Type: types.ResponseType(event.EventStop), Content: "", Done: true, Timestamp: time.Now(), Data: map[string]interface{}{ "session_id": sessionID, "message_id": messageID, "reason": "user_requested", "source": "im", }, } if err := s.streamManager.AppendEvent(ctx, sessionID, messageID, stopEvt); err != nil { logger.Warnf(ctx, "[IM] Failed to write stop event to StreamManager: %v", err) } } // watchStreamManagerStop polls StreamManager for stop events and cancels the // QA context when one is detected. This is the IM equivalent of the web SSE // handler's stop detection loop. Exits when ctx is done. func (s *Service) watchStreamManagerStop(ctx context.Context, sessionID, messageID string, cancel context.CancelFunc) { ticker := time.NewTicker(stopPollInterval) defer ticker.Stop() offset := 0 for { select { case <-ctx.Done(): return case <-ticker.C: events, newOffset, err := s.streamManager.GetEvents(ctx, sessionID, messageID, offset) if err != nil { continue } for _, evt := range events { if evt.Type == types.ResponseType(event.EventStop) { logger.Infof(ctx, "[IM] Stop event from StreamManager, cancelling: session=%s message=%s", sessionID, messageID) cancel() return } } offset = newOffset } } } // GetChannelAdapter returns the adapter and channel config for a given channel ID. func (s *Service) GetChannelAdapter(channelID string) (Adapter, *IMChannel, bool) { s.mu.RLock() defer s.mu.RUnlock() cs, ok := s.channels[channelID] if !ok { return nil, nil, false } return cs.Adapter, cs.Channel, true } // EnsureChannelAdapter loads the durable channel row before serving a webhook // callback. This prevents a replica that missed an invalidation event from // verifying or processing the callback with stale credentials/configuration. func (s *Service) EnsureChannelAdapter(channelID string) (Adapter, *IMChannel, error) { fresh, err := s.GetChannelByID(channelID) if err != nil { // A deleted channel may still exist in this replica's runtime map. Do // not tear down a healthy runtime on a transient database failure. if errors.Is(err, gorm.ErrRecordNotFound) { s.StopChannel(channelID) } return nil, nil, err } if !fresh.Enabled { s.StopChannel(channelID) return nil, fresh, ErrChannelDisabled } adapter, cached, ok := s.GetChannelAdapter(channelID) if !ok || !sameChannelRuntimeConfig(cached, fresh) { if err := s.StartChannel(fresh); err != nil { return nil, fresh, err } adapter, cached, ok = s.GetChannelAdapter(channelID) if !ok { return nil, fresh, fmt.Errorf("channel adapter is not active on this instance") } } return adapter, cached, nil } // sameChannelRuntimeConfig compares every field that affects adapter creation // or message routing. It intentionally does not compare UpdatedAt: PostgreSQL // timestamp precision can differ from the in-memory value assigned by GORM, // which would otherwise rebuild webhook adapters on every callback. func sameChannelRuntimeConfig(cached, fresh *IMChannel) bool { if cached == nil || fresh == nil { return false } return cached.ID == fresh.ID && cached.TenantID == fresh.TenantID && cached.AgentID == fresh.AgentID && cached.Platform == fresh.Platform && cached.Enabled == fresh.Enabled && cached.Mode == fresh.Mode && cached.OutputMode == fresh.OutputMode && cached.KnowledgeBaseID == fresh.KnowledgeBaseID && cached.SessionMode == fresh.SessionMode && equalChannelCredentials(cached.Credentials, fresh.Credentials) } func equalChannelCredentials(left, right types.JSON) bool { var leftValue, rightValue any if err := json.Unmarshal(left, &leftValue); err != nil { return string(left) == string(right) } if err := json.Unmarshal(right, &rightValue); err != nil { return false } return reflect.DeepEqual(leftValue, rightValue) } // GetChannelByID loads a channel from the database. func (s *Service) GetChannelByID(channelID string) (*IMChannel, error) { var ch IMChannel if err := s.db.Where("id = ? AND deleted_at IS NULL", channelID).First(&ch).Error; err != nil { return nil, err } return &ch, nil } // GetChannelByIDAndTenant loads a channel from the database, scoped to a specific tenant. func (s *Service) GetChannelByIDAndTenant(channelID string, tenantID uint64) (*IMChannel, error) { var ch IMChannel if err := s.db.Where("id = ? AND tenant_id = ? AND deleted_at IS NULL", channelID, tenantID).First(&ch).Error; err != nil { return nil, err } return &ch, nil } // isDuplicate checks if a message has already been processed. // // Multi-instance mode (Redis available): uses Redis SetNX for cross-instance // deduplication. If Redis fails, returns true (fail-closed) to prevent // duplicate processing across instances — a dropped message can be retried // by the user, but a duplicate LLM response wastes resources and confuses. // // Single-instance mode (no Redis): uses a local sync.Map, which is sufficient // when only one instance receives messages. func (s *Service) isDuplicate(ctx context.Context, messageID string) bool { if s.redis != nil { key := RedisKeyDedup + messageID ok, err := s.redis.SetNX(ctx, key, "1", dedupTTL).Result() if err == nil { return !ok // SetNX returns true when key was newly set (not a duplicate) } // Redis is configured but failed — fail-closed to avoid cross-instance // duplicate processing. The user can simply resend the message. logger.Errorf(ctx, "[IM] Redis dedup failed (fail-closed, message dropped): %v", err) return true } // Single-instance mode: local dedup is sufficient. _, loaded := s.processedMsgs.LoadOrStore(messageID, time.Now()) return loaded } // HandleMessage processes an incoming IM message end-to-end using channel config. func (s *Service) HandleMessage(ctx context.Context, msg *IncomingMessage, channelID string) error { // Dedup: skip if this message was already processed (IM platforms may retry) if msg.MessageID != "" { if s.isDuplicate(ctx, msg.MessageID) { logger.Infof(ctx, "[IM] Skipping duplicate message: %s", msg.MessageID) return nil } } // Reject overly long messages to protect the QA pipeline contentRunes := []rune(msg.Content) if len(contentRunes) > maxContentLength { logger.Warnf(ctx, "[IM] Message too long (%d runes), truncating to %d", len(contentRunes), maxContentLength) msg.Content = string(contentRunes[:maxContentLength]) } // Get channel config (moved before rate limit so we can reply to the user) adapter, channel, err := s.EnsureChannelAdapter(channelID) if err != nil { return fmt.Errorf("channel unavailable: %w", err) } // Resolve threadID for key building — only include in thread mode to avoid // leaking thread scope into user-mode rate limit / inflight keys. threadID := "" if channel.SessionMode == string(SessionModeThread) { threadID = msg.ThreadID } // Rate limit: enforce per-user sliding window to prevent abuse. // Slash-commands (/stop, /clear, etc.) bypass rate limiting so the user // always retains control over the bot even under heavy messaging. isCommand := s.cmdRegistry.IsRegistered(msg.Content) if !isCommand { rateLimitKey := makeUserKey(channelID, msg.UserID, msg.ChatID, threadID) if !s.rateLimiter.Allow(ctx, rateLimitKey, s.rateLimitMax) { logger.Warnf(ctx, "[IM] Rate limited: channel=%s user=%s chat=%s", channelID, msg.UserID, msg.ChatID) _ = adapter.SendReply(ctx, msg, &ReplyMessage{ Content: "您的消息发送过于频繁,请稍后再试。", IsFinal: true, }) return nil } } tenantID := channel.TenantID agentID := channel.AgentID logger.Infof(ctx, "[IM] HandleMessage: channel=%s platform=%s user=%s chat=%s msgtype=%s content_len=%d", channelID, msg.Platform, msg.UserID, msg.ChatID, msg.MessageType, len(msg.Content)) logger.Debugf(ctx, "[IM] HandleMessage detail: msgid=%s raw_msgtype=%s filekey=%s filename=%s", msg.MessageID, msg.Extra["raw_msgtype"], msg.FileKey, msg.FileName) // ── File/Image message handling ── // File messages use the normal QA path as well. A configured knowledge base // only adds a best-effort, asynchronous save; it must never replace or block // the reply to this message. With no configured knowledge base, simply skip // the save rather than rejecting the message. if msg.MessageType == MessageTypeFile || msg.MessageType == MessageTypeImage { msg.Content = fileMessageQAContent(msg) } // Never send an empty query into rewrite/intent classification. Some IM // platforms deliver unsupported rich message shapes with no normalized text; // allowing those through makes the model infer an unrelated intent from an // empty query (for example, rewriting it as a greeting). if hint, empty := emptyIncomingMessageReply(msg); empty { logger.Infof(ctx, "[IM] Skipping QA for message without content: type=%s raw_type=%s", msg.MessageType, msg.Extra["raw_msgtype"]) if err := adapter.SendReply(ctx, msg, &ReplyMessage{ Content: hint, IsFinal: true, }); err != nil { logger.Warnf(ctx, "[IM] Failed to send empty-message hint reply: %v", err) } return nil } // 1. Get tenant tenant, err := s.tenantService.GetTenantByID(ctx, tenantID) if err != nil { return fmt.Errorf("get tenant: %w", err) } sessionCtx := context.WithValue(ctx, types.TenantInfoContextKey, tenant) sessionCtx = withIMIdentity(sessionCtx, tenantID, channelID, msg) // A configured channel authorizes this exact Agent. Never fall back to a // different Agent if the original was deleted or cannot be loaded. sessionCtx, customAgent, err := s.resolveConfiguredChannelExecution(sessionCtx, channel) if err != nil { return fmt.Errorf("configured channel agent unavailable: %w", err) } // 2. Resolve or create a WeKnora session channelSession, err := s.resolveSession(sessionCtx, msg, tenantID, agentID, channelID, channel.SessionMode) if err != nil { return fmt.Errorf("resolve session: %w", err) } sessionCtx = types.WithSessionID(sessionCtx, channelSession.SessionID) // ── Slash-command dispatch ── // Commands are handled before the QA pipeline so they respond instantly. if cmd, args, ok := s.cmdRegistry.Parse(msg.Content); ok { return s.handleCommand(sessionCtx, cmd, args, msg, adapter, channel, channelSession, customAgent) } // Unrecognised slash-word: show help hint instead of sending to QA. if LooksLikeCommand(msg.Content) { _ = adapter.SendReply(ctx, msg, &ReplyMessage{ Content: "未知指令,发送 `/help` 查看所有可用指令。", IsFinal: true, }) return nil } // 4. Get the WeKnora session session, err := s.sessionService.GetSession(sessionCtx, channelSession.SessionID) if err != nil { // The underlying session may have been deleted from the UI while the // ChannelSession mapping still exists (GORM soft-delete does not trigger // SQL ON DELETE CASCADE). Recover by soft-deleting the stale mapping and // re-creating a fresh session so the IM bot doesn't become permanently // unresponsive. (fixes #1046, #1499) if isSessionNotFound(err) { logger.Warnf(ctx, "[IM] Session %s not found (deleted?), recycling stale channel session %s", channelSession.SessionID, channelSession.ID) if delErr := s.db.Delete(&ChannelSession{}, "id = ?", channelSession.ID).Error; delErr != nil { logger.Warnf(ctx, "[IM] Failed to delete stale channel session %s: %v", channelSession.ID, delErr) } channelSession, err = s.resolveSession(sessionCtx, msg, tenantID, agentID, channelID, channel.SessionMode) if err != nil { return fmt.Errorf("resolve session (retry): %w", err) } sessionCtx = types.WithSessionID(sessionCtx, channelSession.SessionID) session, err = s.sessionService.GetSession(sessionCtx, channelSession.SessionID) if err != nil { return fmt.Errorf("get session (retry): %w", err) } } else { return fmt.Errorf("get session: %w", err) } } // Title an untitled IM session from its first text message, like web chats. // GenerateTitleAsync self-guards on a non-empty title and persists to the DB; // nil eventBus is fine (IM has no live stream — the sidebar reloads it). if session.Title == "" && strings.TrimSpace(msg.Content) != "" { // Copy the session: the async title goroutine writes Title while the QA // worker below shares the same *session. sessionForTitle := *session titleModelID := "" if customAgent != nil && customAgent.Config.ModelID != "" { titleModelID = customAgent.Config.ModelID } s.sessionService.GenerateTitleAsync(sessionCtx, &sessionForTitle, msg.Content, titleModelID, nil) } s.persistIMLastRequestState(sessionCtx, session.ID, agentID, customAgent, nil) // 5. Enqueue the QA request into the bounded worker pool. // The worker pool controls LLM concurrency and provides backpressure. qaCtx, qaCancel := context.WithCancel(sessionCtx) userKey := makeUserKey(channelID, msg.UserID, msg.ChatID, threadID) req := &qaRequest{ ctx: qaCtx, cancel: qaCancel, msg: msg, session: session, agent: customAgent, adapter: adapter, channel: channel, channelID: channelID, tenant: tenant, userKey: userKey, } pos, enqueueErr := s.qaQueue.Enqueue(req) if enqueueErr != nil { qaCancel() logger.Warnf(ctx, "[IM] Queue rejected: user=%s reason=%v", msg.UserID, enqueueErr) _ = adapter.SendReply(ctx, msg, &ReplyMessage{ Content: "当前排队人数较多,请稍后再试。", IsFinal: true, }) return nil } if pos > 0 { logger.Infof(ctx, "[IM] Enqueued: user=%s pos=%d depth=%d", msg.UserID, pos, s.qaQueue.Metrics().Depth) // In multi-instance mode the local queue position does not reflect global // depth, so use a generic "queued" hint instead of an exact number. queueMsg := fmt.Sprintf("收到,前面还有 %d 条消息在处理,请稍候 ⏳", pos) if s.redis != nil { queueMsg = "收到,当前排队中,请稍候 ⏳" } _ = adapter.SendReply(ctx, msg, &ReplyMessage{ Content: queueMsg, IsFinal: true, }) } else { logger.Infof(ctx, "[IM] Enqueued: user=%s pos=0 (immediate)", msg.UserID) } return nil } func emptyIncomingMessageReply(msg *IncomingMessage) (string, bool) { if msg == nil || strings.TrimSpace(msg.Content) != "" { return "", false } // Image/file events are allowed to arrive without a caption. Do not depend // on fileMessageQAContent having already filled Content — a later reorder // of HandleMessage must not reject attachments. hasAttachment := msg.MessageType == MessageTypeFile || msg.MessageType == MessageTypeImage || strings.TrimSpace(msg.FileKey) != "" if hasAttachment { return "", false } rawType := "" if msg.Extra != nil { rawType = strings.ToLower(strings.TrimSpace(msg.Extra["raw_msgtype"])) } switch rawType { case "audio": return "未能识别这条语音中的文字内容。请改用纯文本发送,或再说一遍。", true case "video": return "暂不支持视频消息。请改用纯文本发送;图片或文件请单独发送。", true default: return "未能识别这条消息中的文字内容。请改用纯文本发送;图片或文件请单独发送。", true } } func (s *Service) persistIMLastRequestState(ctx context.Context, sessionID, agentID string, customAgent *types.CustomAgent, kbIDs []string) { state := buildIMLastRequestState(agentID, customAgent, kbIDs) if err := s.sessionService.UpdateSessionLastRequestState(logger.CloneContext(context.WithoutCancel(ctx)), sessionID, state); err != nil { logger.Warnf(ctx, "[IM] persist last_request_state failed for session %s: %v", sessionID, err) } } // executeQARequest is the worker handler that runs the QA pipeline for a queued request. // It is called by qaQueue workers and must not block indefinitely. func (s *Service) executeQARequest(req *qaRequest) { ctx := req.ctx defer req.cancel() // Track in-flight request so /stop can cancel it. entry := &inflightEntry{cancel: req.cancel} s.inflight.Store(req.userKey, entry) defer s.inflight.Delete(req.userKey) verifiedCtx, freshAgent, err := s.resolveConfiguredChannelExecution(ctx, req.channel) if err != nil { logger.Errorf(ctx, "[IM] Queued request rejected because channel or agent is unavailable: channel=%s error=%v", req.channelID, err) return } ctx, req.agent = verifiedCtx, freshAgent // Check if a pre-execution /stop was issued while this request was queued. if s.checkAndClearStopMarker(ctx, req.userKey) { logger.Infof(ctx, "[IM] Request cancelled by remote /stop before execution: %s", req.userKey) return } // NOTE: StreamManager-based stop detection is started inside handleMessageStream / // runQA after the assistant message is created (that's when we have the // sessionID + messageID needed to poll StreamManager). // kbIDs is left empty so the QA pipeline resolves them from the agent config. var kbIDs []string attachments, imageURLs, downloaded, err := s.prepareIMAttachments(ctx, req.msg, req.adapter) if err != nil { logger.Warnf(ctx, "[IM] attachment preparation failed: %v", err) if sendErr := req.adapter.SendReply(ctx, req.msg, &ReplyMessage{Content: "❌ 无法读取此附件,请重试或改用文字描述。", IsFinal: true}); sendErr != nil { logger.Warnf(ctx, "[IM] Failed to send attachment error reply: %v", sendErr) } return } if req.channel.KnowledgeBaseID != "" && downloaded != nil { go s.processDownloadedFileToKnowledgeBase( context.WithoutCancel(ctx), req.channel, downloaded, ) } // Determine output mode from channel config. streamDisabled := req.channel.OutputMode == "full" if streamDisabled { if progressSender, ok := req.adapter.(FullOutputProgressSender); ok && progressSender.SupportsFullOutputProgress() { // Full output still starts the platform stream so users immediately see // its thinking placeholder. No intermediate reasoning/tool content is // sent; the placeholder is replaced only after QA completes. if err := s.handleMessageFullOutput( ctx, req.msg, req.session, req.agent, kbIDs, attachments, imageURLs, progressSender, req.adapter, req.userKey, req.tenant, ); err != nil { logger.Errorf(ctx, "[IM] Full-output QA failed: %v", err) } return } } else if streamer, ok := req.adapter.(StreamSender); ok { // Stream mode sends intermediate reasoning and answer updates. if err := s.handleMessageStream( ctx, req.msg, req.session, req.agent, kbIDs, attachments, imageURLs, streamer, req.adapter, req.userKey, req.tenant, ); err != nil { logger.Errorf(ctx, "[IM] Stream QA failed: %v", err) } return } // Non-streaming fallback: collect full answer then send. answer, err := s.runQA(ctx, req.session, req.msg.Content, req.agent, kbIDs, attachments, imageURLs, req.userKey, req.msg.Quote) if err != nil { logger.Errorf(ctx, "[IM] QA failed: %v, sending fallback reply", err) answer = imQAFailureReply(err) } outCtx := imOutboundContext(ctx) reply := &ReplyMessage{ Content: formatIMOutboundAnswerOrFallback(outCtx, answer, req.tenant, s.defaultFileSvc, s.storageResolver), IsFinal: true, } if err := req.adapter.SendReply(outCtx, req.msg, reply); err != nil { logger.Errorf(ctx, "[IM] Send reply failed: %v", err) return } logger.Infof(ctx, "[IM] Reply sent: channel=%s platform=%s user=%s answer_len=%d", req.channelID, req.msg.Platform, req.msg.UserID, len(answer)) } // handleMessageFullOutput keeps the channel's full-output semantics while // providing immediate progress feedback on adapters that support replaceable // stream messages. StartStream creates the platform placeholder; no intermediate // updates are sent, and the final answer replaces it exactly once. func (s *Service) handleMessageFullOutput( ctx context.Context, msg *IncomingMessage, session *types.Session, customAgent *types.CustomAgent, kbIDs []string, attachments types.MessageAttachments, imageURLs []string, streamer FullOutputProgressSender, adapter Adapter, userKey string, tenant *types.Tenant, ) error { streamID, err := streamer.StartStream(ctx, msg) if err != nil { logger.Warnf(ctx, "[IM] StartStream failed for full output, falling back to plain reply: %v", err) return s.fallbackNonStream( ctx, msg, session, customAgent, kbIDs, attachments, imageURLs, adapter, userKey, tenant, ) } answer, qaErr := s.runQA( ctx, session, msg.Content, customAgent, kbIDs, attachments, imageURLs, userKey, msg.Quote, ) if qaErr != nil { logger.Errorf(ctx, "[IM] Full-output QA failed: %v, sending fallback reply", qaErr) answer = imQAFailureReply(qaErr) } // QA (and /stop) may have cancelled ctx. Platform updates must still run so // the thinking card is replaced instead of hanging forever. outCtx := imOutboundContext(ctx) finalContent := formatIMOutboundAnswerOrFallback(outCtx, answer, tenant, s.defaultFileSvc, s.storageResolver) finalizeErr := streamer.FinalizeStream(outCtx, msg, streamID, finalContent) if finalizeErr != nil { logger.Warnf(ctx, "[IM] FinalizeStream failed for full output: %v", finalizeErr) } endErr := streamer.EndStream(outCtx, msg, streamID) if endErr != nil { logger.Warnf(ctx, "[IM] EndStream failed for full output: %v", endErr) } // If the placeholder could not be replaced, send a plain final reply so the // user still receives the answer instead of being left on "thinking". var fallbackErr error if finalizeErr != nil { fallbackErr = adapter.SendReply(outCtx, msg, &ReplyMessage{Content: finalContent, IsFinal: true}) if fallbackErr != nil { logger.Errorf(ctx, "[IM] Plain reply fallback after full-output finalize failure failed: %v", fallbackErr) } } if finalizeErr == nil || fallbackErr == nil { logger.Infof( ctx, "[IM] Full-output reply sent: platform=%s user=%s answer_len=%d", msg.Platform, msg.UserID, len(answer), ) return nil } return errors.Join(finalizeErr, endErr, fallbackErr) } // handleCommand executes a slash-command and sends the result back to the user. // It also handles side effects (ActionClear, ActionStop). func (s *Service) handleCommand( ctx context.Context, cmd Command, args []string, msg *IncomingMessage, adapter Adapter, channel *IMChannel, channelSession *ChannelSession, customAgent *types.CustomAgent, ) error { agentName := "" if customAgent != nil { agentName = customAgent.Name } cmdCtx := &CommandContext{ Incoming: msg, Session: channelSession, TenantID: channel.TenantID, AgentName: agentName, CustomAgent: customAgent, ChannelOutputMode: channel.OutputMode, } result, err := cmd.Execute(ctx, cmdCtx, args) if err != nil { logger.Errorf(ctx, "[IM] Command /%s error: %v", cmd.Name(), err) _ = adapter.SendReply(ctx, msg, &ReplyMessage{ Content: "抱歉,执行指令时出现了异常,请稍后再试。", IsFinal: true, }) return err } // Handle service-level side effects. switch result.Action { case ActionClear: // Soft-delete the current ChannelSession so the next IM message // starts a completely fresh WeKnora session. Conversation history // is keyed by session ID and rebuilt from DB on demand, so no // separate cache invalidation step is needed. if err := s.db.Model(&ChannelSession{}). Where("id = ?", channelSession.ID). Update("deleted_at", time.Now()).Error; err != nil { logger.Warnf(ctx, "[IM] Failed to soft-delete channel session: %v", err) } case ActionStop: stopThreadID := "" if channel.SessionMode == string(SessionModeThread) { stopThreadID = msg.ThreadID } inflightKey := makeUserKey(channel.ID, msg.UserID, msg.ChatID, stopThreadID) // 1. Try local cancel: remove from queue or cancel in-flight. var localSessionID, localMessageID string localStopped := s.qaQueue.Remove(inflightKey) if localStopped { logger.Infof(ctx, "[IM] Cancelled queued QA: key=%s", inflightKey) } else if raw, loaded := s.inflight.LoadAndDelete(inflightKey); loaded { e := raw.(*inflightEntry) e.cancel() localStopped = true localSessionID = e.sessionID localMessageID = e.assistantMessageID logger.Infof(ctx, "[IM] Cancelled in-flight QA: key=%s", inflightKey) } // 2. Write stop event to StreamManager (same as web StopSession). // For local stop with known IDs, write directly. // For cross-instance, look up Redis inflight mapping to get IDs. sessionID, messageID := localSessionID, localMessageID if sessionID == "" || messageID == "" { // Try cross-instance lookup. sessionID, messageID, _ = s.loadInflightMapping(ctx, inflightKey) } if sessionID != "" && messageID != "" { s.writeStopEvent(ctx, sessionID, messageID) logger.Infof(ctx, "[IM] Wrote stop event to StreamManager: session=%s message=%s", sessionID, messageID) } // 3. Set Redis marker as fallback for requests not yet executing // (no assistant message yet → no StreamManager entry to poll). if s.redis != nil { s.redis.Set(ctx, RedisKeyStop+inflightKey, "1", stopMarkerTTL) } if !localStopped && sessionID == "" { logger.Infof(ctx, "[IM] Set cross-instance stop marker (no inflight found): key=%s", inflightKey) } } // Send the command reply, respecting the configured output mode. sent := false if channel.OutputMode != "full" { if streamer, ok := adapter.(StreamSender); ok { if err := s.sendStreamReply(ctx, msg, streamer, result.Content); err != nil { logger.Warnf(ctx, "[IM] Stream reply for command /%s failed, falling back: %v", cmd.Name(), err) } else { sent = true } } } if !sent { _ = adapter.SendReply(ctx, msg, &ReplyMessage{ Content: result.Content, IsFinal: true, }) } logger.Infof(ctx, "[IM] Command /%s executed: channel=%s user=%s action=%d", cmd.Name(), channel.ID, msg.UserID, result.Action) return nil } // sendStreamReply sends a complete content string via the streaming interface. func (s *Service) sendStreamReply(ctx context.Context, msg *IncomingMessage, streamer StreamSender, content string) error { streamID, err := streamer.StartStream(ctx, msg) if err != nil { return fmt.Errorf("start stream: %w", err) } if err := streamer.UpdateStreamContent(ctx, msg, streamID, content); err != nil { return fmt.Errorf("update stream content: %w", err) } if err := streamer.FinalizeStream(ctx, msg, streamID, content); err != nil { return fmt.Errorf("finalize stream: %w", err) } if err := streamer.EndStream(ctx, msg, streamID); err != nil { return fmt.Errorf("end stream: %w", err) } return nil } // isSessionNotFound reports whether err indicates the underlying WeKnora // session no longer exists. The session repository translates GORM's // ErrRecordNotFound into apperrors.ErrSessionNotFound, so the application // sentinel is what GetSession returns today; the GORM check is kept as a // safety net in case a future repository revert bypasses the translation. func isSessionNotFound(err error) bool { return errors.Is(err, apperrors.ErrSessionNotFound) || errors.Is(err, gorm.ErrRecordNotFound) } // resolveSession dispatches to the appropriate session resolution strategy // based on the channel's session mode. func (s *Service) resolveSession(ctx context.Context, msg *IncomingMessage, tenantID uint64, agentID string, imChannelID string, sessionMode string) (*ChannelSession, error) { switch SessionMode(sessionMode) { case SessionModeThread: return s.resolveThreadSession(ctx, msg, tenantID, agentID, imChannelID) default: // SessionModeUser return s.resolveUserSession(ctx, msg, tenantID, agentID, imChannelID) } } // buildUserSessionTitle produces a human-distinguishable title for a user-mode // IM session. Platform adapters only surface ChatID, not a readable chat name, // so we fall back to short ID suffixes to keep group/DM sessions visually distinct. // Platform prefix is intentionally omitted — the UI renders a platform icon badge // alongside the title, so the `[feishu]` prefix would be redundant clutter. func buildUserSessionTitle(msg *IncomingMessage) string { var b strings.Builder if msg.UserName != "" { b.WriteString(msg.UserName) } else if msg.UserID != "" { b.WriteString("user ") b.WriteString(shortID(msg.UserID)) } else { b.WriteString("user") } if msg.ChatType == ChatTypeGroup && msg.ChatID != "" { fmt.Fprintf(&b, " · group %s", shortID(msg.ChatID)) } else if msg.ChatType == ChatTypeDirect { b.WriteString(" · dm") } return b.String() } // buildThreadSessionTitle produces a title for a thread-mode IM session. // In thread mode different users can share one session, so the user name is // omitted and chat/thread IDs carry the distinguishing information. // Platform prefix is omitted for the same reason as buildUserSessionTitle. func buildThreadSessionTitle(msg *IncomingMessage) string { var b strings.Builder if msg.ChatID != "" { fmt.Fprintf(&b, "chat %s · ", shortID(msg.ChatID)) } b.WriteString("thread ") b.WriteString(shortID(msg.ThreadID)) return b.String() } // shortID returns the last 8 characters of id, or id itself when shorter. // Used to keep long platform IDs readable inside titles without losing uniqueness. func shortID(id string) string { if len(id) > 8 { return id[len(id)-8:] } return id } // imInitialSessionTitle picks a new IM session's starting title: "" when the // message has text to summarise (so it gets a content-based title later, like // web chats), otherwise the IM identity title so the row is never blank. func imInitialSessionTitle(msg *IncomingMessage, identityTitle func(*IncomingMessage) string) string { if strings.TrimSpace(msg.Content) != "" { return "" } return identityTitle(msg) } // resolveUserSession finds or creates a ChannelSession keyed by (platform, user_id, chat_id, tenant_id, agent_id). // This is the original session resolution strategy. // // Invariant: a cache miss creates a brand-new session — we never attach a second // mapping to an existing session. The session-list source filter (repository // QueryPaged) relies on this one-mapping-per-session property; if this ever // re-maps an existing session, that JOIN needs a one-row-per-session guard. func (s *Service) resolveUserSession(ctx context.Context, msg *IncomingMessage, tenantID uint64, agentID string, imChannelID string) (*ChannelSession, error) { var cs ChannelSession result := s.db.Where("platform = ? AND user_id = ? AND chat_id = ? AND tenant_id = ? AND agent_id = ? AND im_channel_id = ? AND deleted_at IS NULL", string(msg.Platform), msg.UserID, msg.ChatID, tenantID, agentID, imChannelID). First(&cs) if result.Error == nil { return &cs, nil } if result.Error != gorm.ErrRecordNotFound { return nil, fmt.Errorf("query channel session: %w", result.Error) } // Create a new WeKnora session. Start untitled when there's text to summarise // so it gets a content-based title after the first message (see HandleMessage); // fall back to the IM identity title otherwise. title := imInitialSessionTitle(msg, buildUserSessionTitle) newSession := &types.Session{ TenantID: tenantID, Title: title, Description: fmt.Sprintf("Auto-created from %s IM integration", msg.Platform), } createdSession, err := s.sessionService.CreateSession(ctx, newSession) if err != nil { return nil, fmt.Errorf("create session: %w", err) } // Create the channel-session mapping; use a unique constraint fallback // to handle concurrent creation attempts for the same channel. cs = ChannelSession{ Platform: string(msg.Platform), UserID: msg.UserID, ChatID: msg.ChatID, SessionID: createdSession.ID, TenantID: tenantID, AgentID: agentID, IMChannelID: imChannelID, } if err := s.db.Create(&cs).Error; err != nil { if delErr := s.db.Where("id = ?", createdSession.ID).Delete(createdSession).Error; delErr != nil { logger.Warnf(ctx, "[IM] Failed to clean up orphaned session %s: %v", createdSession.ID, delErr) } var existing ChannelSession if findErr := s.db.Where("platform = ? AND user_id = ? AND chat_id = ? AND tenant_id = ? AND agent_id = ? AND im_channel_id = ? AND deleted_at IS NULL", string(msg.Platform), msg.UserID, msg.ChatID, tenantID, agentID, imChannelID). First(&existing).Error; findErr != nil { return nil, fmt.Errorf("create channel session: %w (lookup fallback: %v)", err, findErr) } return &existing, nil } logger.Infof(ctx, "[IM] Created new session mapping: channel=%s/%s/%s -> session=%s", msg.Platform, msg.UserID, msg.ChatID, createdSession.ID) return &cs, nil } // resolveThreadSession finds or creates a ChannelSession keyed by (platform, chat_id, thread_id, tenant_id, agent_id). // In thread mode, each message thread gets its own session. Multiple users in the // same thread share the same session. Top-level messages use their own ID as // ThreadID, creating a new session per top-level message. func (s *Service) resolveThreadSession(ctx context.Context, msg *IncomingMessage, tenantID uint64, agentID string, imChannelID string) (*ChannelSession, error) { threadID := msg.ThreadID if threadID == "" { // Defense-in-depth: frontend blocks thread mode for unsupported platforms, // but if ThreadID is somehow empty, fall back to user-mode resolution // to avoid creating a shared session for all empty-thread messages. logger.Warnf(ctx, "[IM] Thread mode but ThreadID is empty (platform=%s chat=%s), falling back to user session", msg.Platform, msg.ChatID) return s.resolveUserSession(ctx, msg, tenantID, agentID, imChannelID) } var cs ChannelSession result := s.db.Where( "platform = ? AND chat_id = ? AND thread_id = ? AND tenant_id = ? AND agent_id = ? AND im_channel_id = ? AND deleted_at IS NULL", string(msg.Platform), msg.ChatID, threadID, tenantID, agentID, imChannelID, ).First(&cs) if result.Error == nil { return &cs, nil } if result.Error != gorm.ErrRecordNotFound { return nil, fmt.Errorf("query thread session: %w", result.Error) } // Start untitled when there's text to summarise so it gets a content-based // title after the first message; fall back to the chat/thread identity title. title := imInitialSessionTitle(msg, buildThreadSessionTitle) newSession := &types.Session{ TenantID: tenantID, Title: title, Description: fmt.Sprintf("Thread-based session from %s IM", msg.Platform), } createdSession, err := s.sessionService.CreateSession(ctx, newSession) if err != nil { return nil, fmt.Errorf("create session: %w", err) } cs = ChannelSession{ Platform: string(msg.Platform), UserID: msg.UserID, // record the first creator ChatID: msg.ChatID, ThreadID: threadID, SessionID: createdSession.ID, TenantID: tenantID, AgentID: agentID, IMChannelID: imChannelID, } if err := s.db.Create(&cs).Error; err != nil { // Unique constraint fallback for concurrent creation. if delErr := s.db.Where("id = ?", createdSession.ID).Delete(createdSession).Error; delErr != nil { logger.Warnf(ctx, "[IM] Failed to clean up orphaned session %s: %v", createdSession.ID, delErr) } var existing ChannelSession if findErr := s.db.Where( "platform = ? AND chat_id = ? AND thread_id = ? AND tenant_id = ? AND agent_id = ? AND im_channel_id = ? AND deleted_at IS NULL", string(msg.Platform), msg.ChatID, threadID, tenantID, agentID, imChannelID, ).First(&existing).Error; findErr != nil { return nil, fmt.Errorf("create thread session: %w (lookup fallback: %v)", err, findErr) } return &existing, nil } logger.Infof(ctx, "[IM] Created new thread session: platform=%s thread=%s chat=%s -> session=%s", msg.Platform, threadID, msg.ChatID, createdSession.ID) return &cs, nil } // ── Agent tool call progress formatting ────────────────────────────── // These helpers format tool-call / tool-result events as Markdown text // that is injected into the streaming reply so IM users can see the // agent's reasoning process in real-time. // ───────────────────────────────────────────────────────────────────── // internalToolNames lists tools whose execution should NOT be displayed in IM // messages because they are internal reasoning aids (thinking, planning) rather // than user-facing actions. var internalToolNames = map[string]bool{ "thinking": true, "todo_write": true, } // isToolVisibleToUser returns true if the tool's execution progress should be // displayed to the IM user. Internal reasoning tools (thinking, planning) are // hidden. func isToolVisibleToUser(toolName string) bool { return !internalToolNames[toolName] } // briefToolSummary extracts a short human-readable summary from tool output. // Returns empty string if no suitable summary can be extracted. func briefToolSummary(output string) string { const maxRunes = 40 if output == "" { return "" } output = strings.TrimSpace(output) if output == "" { return "" } // Skip structured data (JSON, XML, etc.) if output[0] == '{' || output[0] == '[' || output[0] == '<' { return "" } // Take first non-empty line if idx := strings.IndexByte(output, '\n'); idx >= 0 { output = strings.TrimSpace(output[:idx]) } if output == "" { return "" } runes := []rune(output) if len(runes) > maxRunes { return string(runes[:maxRunes]) + "..." } return output } // 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, attachments types.MessageAttachments, imageURLs []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, attachments, imageURLs, adapter, userKey, tenant) } // Prepare the QA pipeline // No total deadline: each agent round has its own LLMCallTimeout (default 120s). // A hard pipeline deadline would kill multi-round agent reasoning prematurely. qaCtx, qaCancel := context.WithCancel(ctx) defer qaCancel() useAgent := customAgent != nil && customAgent.IsAgentMode() eventBus := event.NewEventBus() var ( bufMu sync.Mutex reasoningInner streamSection // quick QA: model reasoning_content agentInner streamSection // agent: retracted text + thoughts agentLiveAnswer strings.Builder // agent: optimistic answer before tool retract answerOuter strings.Builder // final answer for display persistence answerBuilder strings.Builder // answer persisted to DB qaErr error done = make(chan struct{}) completeDone = make(chan struct{}) closeOnce sync.Once completeOnce sync.Once agentDone bool assistantMsg *types.Message seenToolCalls = make(map[string]bool) agentToolIdx = make(map[string]int) pipelineIdx = make(map[string]int) agentToolSteps []IMToolStep pipelineToolSteps []IMToolStep agentCompleteFinalAnswer string streamedAny bool // mcpAuthServices collects OAuth services that need out-of-band // authorization (IM cannot resolve the in-conversation prompt). mcpAuthServices []imMCPAuthService mcpAuthSeen = make(map[string]bool) ) closeDone := func() { closeOnce.Do(func() { close(done) }) } closeComplete := func() { completeOnce.Do(func() { close(completeDone) }) } agentWrite := func(s string) { if s == "" { return } agentInner.write(s) streamedAny = true } reasoningWrite := func(s string) { if s == "" { return } reasoningInner.write(s) streamedAny = true } // retractAgentLiveAnswer moves the optimistic answer into the think block (Web: superseded preamble). retractAgentLiveAnswer := func() { if agentLiveAnswer.Len() == 0 { return } if agentInner.text.Len() > 0 { agentInner.ensureNewlineBefore() } agentInner.write(agentLiveAnswer.String()) agentLiveAnswer.Reset() } getStreamParts := func() IMStreamParts { mode := IMStreamModeQuickQA if useAgent { mode = IMStreamModeAgent } return IMStreamParts{ Mode: mode, PipelineToolSteps: pipelineToolSteps, ReasoningInner: reasoningInner.text.String(), AgentInner: agentInner.text.String(), AgentToolSteps: agentToolSteps, LiveAnswer: agentLiveAnswer.String(), Answer: answerOuter.String(), } } // Subscribe to answer chunks. eventBus.On(event.EventAgentFinalAnswer, func(_ context.Context, evt event.Event) error { data, ok := evt.Data.(event.AgentFinalAnswerData) if !ok { return nil } bufMu.Lock() if useAgent && !agentDone { if data.Content != "" { agentLiveAnswer.WriteString(data.Content) streamedAny = true } } else { answerOuter.WriteString(data.Content) answerBuilder.WriteString(data.Content) streamedAny = true } bufMu.Unlock() if data.Done { closeDone() } return nil }) eventBus.On(event.EventError, func(_ context.Context, evt event.Event) error { data, ok := evt.Data.(event.ErrorData) if !ok { return nil } logger.Errorf(ctx, "[IM] QA stream error: %s", data.Error) bufMu.Lock() qaErr = fmt.Errorf("QA pipeline error: %s", data.Error) bufMu.Unlock() closeDone() closeComplete() return nil }) eventBus.On(event.EventAgentReferences, func(_ context.Context, evt event.Event) error { data, ok := evt.Data.(event.AgentReferencesData) if !ok { return nil } bufMu.Lock() if assistantMsg != nil { refs := []*types.SearchResult(assistantMsg.KnowledgeReferences) collectIMKnowledgeReferences(&refs, data.References) assistantMsg.KnowledgeReferences = types.References(refs) } bufMu.Unlock() return nil }) eventBus.On(event.EventAgentComplete, func(_ context.Context, evt event.Event) error { data, ok := evt.Data.(event.AgentCompleteData) if !ok { return nil } bufMu.Lock() agentDone = true agentCompleteFinalAnswer = data.FinalAnswer applyIMCompleteDataToMessage(assistantMsg, data) mergeIMAgentAnswerBuffers(&answerBuilder, &answerOuter, &agentLiveAnswer, data.FinalAnswer) bufMu.Unlock() closeComplete() return nil }) // Subscribe to agent thought events — stream thinking content into block eventBus.On(event.EventAgentThought, func(_ context.Context, evt event.Event) error { data, ok := evt.Data.(event.AgentThoughtData) if !ok { return nil } bufMu.Lock() if useAgent { agentWrite(data.Content) } else { reasoningWrite(data.Content) } bufMu.Unlock() return nil }) // Subscribe to agent tool call events — write status line into the think block. eventBus.On(event.EventAgentToolCall, func(_ context.Context, evt event.Event) error { data, ok := evt.Data.(event.AgentToolCallData) if !ok { return nil } if !isToolVisibleToUser(data.ToolName) { return nil } bufMu.Lock() if seenToolCalls[data.ToolCallID] { if useAgent { upsertIMToolStep(&agentToolSteps, agentToolIdx, data.ToolCallID, func(step *IMToolStep) { if data.Arguments != nil { step.Arguments = data.Arguments } }) streamedAny = true } bufMu.Unlock() return nil } seenToolCalls[data.ToolCallID] = true if !useAgent && IsRAGPipelineToolName(data.ToolName) { upsertIMToolStep(&pipelineToolSteps, pipelineIdx, data.ToolCallID, func(step *IMToolStep) { step.ToolName = data.ToolName step.Pending = true step.Arguments = data.Arguments }) streamedAny = true } else if useAgent { retractAgentLiveAnswer() upsertIMToolStep(&agentToolSteps, agentToolIdx, data.ToolCallID, func(step *IMToolStep) { step.ToolName = data.ToolName step.Pending = true step.Arguments = data.Arguments }) streamedAny = true } bufMu.Unlock() logger.Debugf(ctx, "[IM] Tool call streamed to IM: tool=%s id=%s", data.ToolName, data.ToolCallID) return nil }) // Subscribe to agent tool result events eventBus.On(event.EventAgentToolResult, func(_ context.Context, evt event.Event) error { data, ok := evt.Data.(event.AgentToolResultData) if !ok { return nil } if !isToolVisibleToUser(data.ToolName) { return nil } bufMu.Lock() if !useAgent && IsRAGPipelineToolName(data.ToolName) { upsertIMToolStep(&pipelineToolSteps, pipelineIdx, data.ToolCallID, func(step *IMToolStep) { step.ToolName = data.ToolName step.Pending = false step.Success = data.Success step.Data = data.Data step.Output = data.Output }) streamedAny = true } else if useAgent { upsertIMToolStep(&agentToolSteps, agentToolIdx, data.ToolCallID, func(step *IMToolStep) { step.ToolName = data.ToolName step.Pending = false step.Success = data.Success step.Data = data.Data step.Output = data.Output }) streamedAny = true } bufMu.Unlock() logger.Debugf(ctx, "[IM] Tool result streamed to IM: tool=%s success=%v duration=%dms", data.ToolName, data.Success, data.Duration) return nil }) // An OAuth-enabled MCP service the IM user has not authorized yet. IM cannot // resolve the in-conversation prompt, so collect the service name and append // an authorization notice to the final reply (deduped per service). eventBus.On(event.EventMCPOAuthRequired, func(_ context.Context, evt event.Event) error { data, ok := evt.Data.(event.MCPOAuthRequiredData) if !ok { return nil } bufMu.Lock() if !mcpAuthSeen[data.ServiceID] { mcpAuthSeen[data.ServiceID] = true mcpAuthServices = append(mcpAuthServices, imMCPAuthService{ID: data.ServiceID, Name: data.ServiceName}) } bufMu.Unlock() return nil }) // Determine whether to use agent mode (already set above for event handlers). requestID := uuid.New().String() // Create user message userMsg, err := s.messageService.CreateMessage(qaCtx, createIMUserMessagePayload(session.ID, msg.Content, requestID, attachments)) if err != nil { return fmt.Errorf("create user message: %w", err) } // Create placeholder assistant message assistantMsg, err = s.messageService.CreateMessage(qaCtx, createIMAssistantMessagePayload(session.ID, requestID)) if err != nil { return fmt.Errorf("create assistant message: %w", err) } // Register inflight mapping so cross-instance /stop can find this request // and write a stop event to StreamManager. if raw, ok := s.inflight.Load(userKey); ok { e := raw.(*inflightEntry) e.sessionID = session.ID e.assistantMessageID = assistantMsg.ID } s.storeInflightMapping(qaCtx, userKey, session.ID, assistantMsg.ID) defer s.clearInflightMapping(ctx, userKey) // Start StreamManager stop watcher — mirrors web's handleAgentEventsForSSE // stop detection. Cancels qaCtx if a stop event is written by any instance. go s.watchStreamManagerStop(qaCtx, session.ID, assistantMsg.ID, qaCancel) // Run QA async go func() { var err error req := buildIMQARequest(session, msg.Content, assistantMsg.ID, userMsg.ID, customAgent, kbIDs, msg.Quote, attachments) req.ImageURLs = imageURLs if req.QuotedContext != "" { logger.Debugf(qaCtx, "[IM] QuotedContext set: length=%d", len(req.QuotedContext)) } if useAgent { err = s.sessionService.AgentQA(qaCtx, req, eventBus) } else { err = s.sessionService.KnowledgeQA(qaCtx, req, eventBus) } if err != nil { logger.Errorf(ctx, "[IM] QA stream execution error: %v", err) bufMu.Lock() qaErr = fmt.Errorf("QA execution error: %w", err) bufMu.Unlock() closeDone() closeComplete() } }() // Flush loop: periodically send buffered content to the IM platform. // A holdback mechanism prevents flushing incomplete provider:// URLs or // XML tags that straddle a chunk boundary (see holdbackCutoff). ticker := time.NewTicker(streamFlushInterval) defer ticker.Stop() flush := func() { bufMu.Lock() parts := getStreamParts() agentRunning := useAgent && !agentDone bufMu.Unlock() displaySource := FormatIMIntermediateFromParts(parts, agentRunning) if displaySource == "" { return } if cut := holdbackCutoff(displaySource); cut < len(displaySource) { displaySource = displaySource[:cut] } display := cleanIMContent(ctx, displaySource, tenant, s.defaultFileSvc, s.storageResolver) if err := streamer.UpdateStreamContent(ctx, msg, streamID, display); err != nil { logger.Warnf(ctx, "[IM] UpdateStreamContent failed: %v", err) } } loop: for { select { case <-ticker.C: flush() case <-done: break loop case <-qaCtx.Done(): break loop } } if useAgent { waitForIMAgentComplete(qaCtx, completeDone, session.ID) } bufMu.Lock() parts := getStreamParts() resolvedAnswer := pickIMStoredAnswer( answerBuilder.String(), answerOuter.String(), agentLiveAnswer.String(), agentCompleteFinalAnswer, ) if parts.Answer == "" { parts.Answer = resolvedAnswer } answer := resolvedAnswer finalErr := qaErr noVisibleContent := !streamedAny && strings.TrimSpace(resolvedAnswer) == "" authServices := append([]imMCPAuthService(nil), mcpAuthServices...) bufMu.Unlock() finalDisplay := cleanIMContent(ctx, FormatIMFinalFromParts(parts), tenant, s.defaultFileSvc, s.storageResolver) if noVisibleContent || finalDisplay == "" { fallback := imNoAnswerFallback if finalErr != nil { fallback = imQAFailureReply(finalErr) } finalDisplay = fallback if answer == "" { answer = fallback } } if notice := s.buildIMMCPAuthNotice(ctx, authServices); notice != "" { finalDisplay = appendIMAuthNotice(finalDisplay, notice) answer = appendIMAuthNotice(answer, notice) } if err := streamer.FinalizeStream(ctx, msg, streamID, finalDisplay); err != nil { logger.Warnf(ctx, "[IM] FinalizeStream failed: %v", err) } // End the stream if err := streamer.EndStream(ctx, msg, streamID); err != nil { logger.Warnf(ctx, "[IM] EndStream failed: %v", err) } if answer == "" { answer = imNoAnswerFallback } assistantMsg.Content = answer assistantMsg.IsCompleted = true if err := s.messageService.UpdateMessage(ctx, assistantMsg); err != nil { logger.Warnf(ctx, "[IM] Failed to update assistant message: %v", err) } logger.Infof(ctx, "[IM] Stream reply sent: platform=%s user=%s answer_len=%d", msg.Platform, msg.UserID, len(answer)) return nil } // fallbackNonStream is used when streaming initialization fails. func (s *Service) fallbackNonStream(ctx context.Context, msg *IncomingMessage, session *types.Session, customAgent *types.CustomAgent, kbIDs []string, attachments types.MessageAttachments, imageURLs []string, adapter Adapter, userKey string, tenant *types.Tenant) error { answer, err := s.runQA(ctx, session, msg.Content, customAgent, kbIDs, attachments, imageURLs, userKey, msg.Quote) if err != nil { logger.Errorf(ctx, "[IM] QA fallback failed: %v", err) answer = imQAFailureReply(err) } outCtx := imOutboundContext(ctx) return adapter.SendReply(outCtx, msg, &ReplyMessage{ Content: formatIMOutboundAnswerOrFallback(outCtx, answer, tenant, s.defaultFileSvc, s.storageResolver), IsFinal: true, }) } // runQA executes the WeKnora QA pipeline and returns the full answer text. func (s *Service) runQA(ctx context.Context, session *types.Session, query string, customAgent *types.CustomAgent, kbIDs []string, attachments types.MessageAttachments, imageURLs []string, userKey string, quote *QuotedMessage) (string, error) { // Cancellable context (no hard deadline): each agent round has its own // LLMCallTimeout. The context can still be cancelled by /stop. ctx, cancel := context.WithCancel(ctx) defer cancel() eventBus := event.NewEventBus() // Thread-safe answer collection var answerMu sync.Mutex var answerBuilder strings.Builder var qaErr error var mcpAuthServices []imMCPAuthService mcpAuthSeen := make(map[string]bool) done := make(chan struct{}) completeDone := make(chan struct{}) var closeOnce sync.Once var completeOnce sync.Once closeDone := func() { closeOnce.Do(func() { close(done) }) } closeComplete := func() { completeOnce.Do(func() { close(completeDone) }) } eventBus.On(event.EventAgentFinalAnswer, func(ctx context.Context, evt event.Event) error { data, ok := evt.Data.(event.AgentFinalAnswerData) if !ok { return nil } answerMu.Lock() answerBuilder.WriteString(data.Content) answerMu.Unlock() if data.Done { closeDone() } return nil }) eventBus.On(event.EventError, func(ctx context.Context, evt event.Event) error { data, ok := evt.Data.(event.ErrorData) if !ok { return nil } logger.Errorf(ctx, "[IM] QA error: %s", data.Error) answerMu.Lock() qaErr = fmt.Errorf("QA pipeline error: %s", data.Error) answerMu.Unlock() closeDone() closeComplete() return nil }) // Collect OAuth services that need out-of-band authorization (IM cannot // resolve the in-conversation prompt); appended to the answer below. eventBus.On(event.EventMCPOAuthRequired, func(_ context.Context, evt event.Event) error { data, ok := evt.Data.(event.MCPOAuthRequiredData) if !ok { return nil } answerMu.Lock() if !mcpAuthSeen[data.ServiceID] { mcpAuthSeen[data.ServiceID] = true mcpAuthServices = append(mcpAuthServices, imMCPAuthService{ID: data.ServiceID, Name: data.ServiceName}) } answerMu.Unlock() return nil }) // Determine whether to use agent mode useAgent := customAgent != nil && customAgent.IsAgentMode() // Generate a shared RequestID to pair user and assistant messages for history requestID := uuid.New().String() // Create user message so it appears in conversation history userMsg, err := s.messageService.CreateMessage(ctx, createIMUserMessagePayload(session.ID, query, requestID, attachments)) if err != nil { return "", fmt.Errorf("create user message: %w", err) } // Create a placeholder assistant message assistantMsg, err := s.messageService.CreateMessage(ctx, createIMAssistantMessagePayload(session.ID, requestID)) if err != nil { return "", fmt.Errorf("create assistant message: %w", err) } eventBus.On(event.EventAgentReferences, func(ctx context.Context, evt event.Event) error { data, ok := evt.Data.(event.AgentReferencesData) if !ok { return nil } answerMu.Lock() refs := []*types.SearchResult(assistantMsg.KnowledgeReferences) collectIMKnowledgeReferences(&refs, data.References) assistantMsg.KnowledgeReferences = types.References(refs) answerMu.Unlock() return nil }) eventBus.On(event.EventAgentComplete, func(ctx context.Context, evt event.Event) error { data, ok := evt.Data.(event.AgentCompleteData) if !ok { return nil } answerMu.Lock() applyIMCompleteDataToMessage(assistantMsg, data) if answerBuilder.Len() == 0 && strings.TrimSpace(data.FinalAnswer) != "" { answerBuilder.WriteString(data.FinalAnswer) } answerMu.Unlock() closeComplete() return nil }) // Register inflight mapping for cross-instance /stop via StreamManager. if raw, ok := s.inflight.Load(userKey); ok { e := raw.(*inflightEntry) e.sessionID = session.ID e.assistantMessageID = assistantMsg.ID } s.storeInflightMapping(ctx, userKey, session.ID, assistantMsg.ID) defer s.clearInflightMapping(ctx, userKey) // Start StreamManager stop watcher. go s.watchStreamManagerStop(ctx, session.ID, assistantMsg.ID, cancel) // Run QA async go func() { var err error req := buildIMQARequest(session, query, assistantMsg.ID, userMsg.ID, customAgent, kbIDs, quote, attachments) req.ImageURLs = imageURLs if req.QuotedContext != "" { logger.Debugf(ctx, "[IM] QuotedContext set: length=%d", len(req.QuotedContext)) } if useAgent { err = s.sessionService.AgentQA(ctx, req, eventBus) } else { err = s.sessionService.KnowledgeQA(ctx, req, eventBus) } if err != nil { logger.Errorf(ctx, "[IM] QA execution error: %v", err) answerMu.Lock() qaErr = fmt.Errorf("QA execution error: %w", err) answerMu.Unlock() closeDone() closeComplete() } }() // Wait for completion or cancellation (e.g., /stop) select { case <-done: if useAgent { waitForIMAgentComplete(ctx, completeDone, session.ID) } case <-ctx.Done(): // Mark assistant message as completed to avoid dangling incomplete records assistantMsg.Content = imCancelledFallback assistantMsg.IsCompleted = true // Use a fresh context since the original is cancelled if updateErr := s.messageService.UpdateMessage(context.WithoutCancel(ctx), assistantMsg); updateErr != nil { logger.Warnf(ctx, "[IM] Failed to update cancelled assistant message: %v", updateErr) } return "", fmt.Errorf("QA cancelled: %w", ctx.Err()) } answerMu.Lock() answer := answerBuilder.String() qaError := qaErr authServices := append([]imMCPAuthService(nil), mcpAuthServices...) answerMu.Unlock() if answer == "" && qaError != nil { return "", qaError } if answer == "" { answer = imNoAnswerFallback } if notice := s.buildIMMCPAuthNotice(ctx, authServices); notice != "" { answer = appendIMAuthNotice(answer, notice) } // Update assistant message with the full answer (including citation tags for web rendering). assistantMsg.Content = answer assistantMsg.IsCompleted = true if err := s.messageService.UpdateMessage(ctx, assistantMsg); err != nil { logger.Warnf(ctx, "[IM] Failed to update assistant message: %v", err) } // Return raw answer — callers apply cleanIMContent with the appropriate FileService. return answer, nil } // ── CRUD operations for IM channels ── // ListChannelsByAgent returns all channels for a given agent within a tenant. func (s *Service) ListChannelsByAgent(agentID string, tenantID uint64) ([]IMChannel, error) { var channels []IMChannel if err := s.db.Where("agent_id = ? AND tenant_id = ? AND deleted_at IS NULL", agentID, tenantID). Order("created_at DESC").Find(&channels).Error; err != nil { return nil, err } return channels, nil } // ChannelWithAgent augments an IMChannel summary with its owning agent's display name. // Credentials are intentionally omitted so this type is safe to return from a // tenant-scoped list endpoint; callers that need credentials must use the // per-agent endpoint which enforces the same tenant scope anyway. type ChannelWithAgent struct { ID string `json:"id"` TenantID uint64 `json:"tenant_id"` AgentID string `json:"agent_id"` AgentName string `json:"agent_name"` Platform string `json:"platform"` Name string `json:"name"` Enabled bool `json:"enabled"` Mode string `json:"mode"` OutputMode string `json:"output_mode"` SessionMode string `json:"session_mode"` BotIdentity string `json:"bot_identity"` CreatedAt time.Time `json:"created_at"` UpdatedAt time.Time `json:"updated_at"` } // ListChannelsByTenant returns all non-deleted IM channels in the given tenant, // joined with custom_agents.name. Built-in agent names are re-localized from // YAML i18n using ctx (the DB column may be empty or frozen in the writer's // language). Channels whose custom agent was soft-deleted are excluded so // overview lists stay consistent after agent removal. func (s *Service) ListChannelsByTenant(ctx context.Context, tenantID uint64) ([]ChannelWithAgent, error) { builtinIDs := types.GetBuiltinAgentIDs() var rows []ChannelWithAgent q := s.db.Table("im_channels AS c"). Select(`c.id, c.tenant_id, c.agent_id, COALESCE(a.name, '') AS agent_name, c.platform, c.name, c.enabled, c.mode, c.output_mode, c.session_mode, c.bot_identity, c.created_at, c.updated_at`). Joins(`LEFT JOIN custom_agents AS a ON a.id = c.agent_id AND a.tenant_id = c.tenant_id AND a.deleted_at IS NULL`). Where("c.tenant_id = ? AND c.deleted_at IS NULL", tenantID) if len(builtinIDs) > 0 { q = q.Where("a.id IS NOT NULL OR c.agent_id IN ?", builtinIDs) } else { q = q.Where("a.id IS NOT NULL") } err := excludeOtherPersonalChannels(q, ctx).Order("c.created_at DESC").Scan(&rows).Error if err != nil { return nil, err } relocalizeBuiltinChannelAgentNames(ctx, rows) return rows, nil } // relocalizeBuiltinChannelAgentNames overlays YAML i18n names onto overview // rows for built-in agents. Custom agents are left unchanged. func relocalizeBuiltinChannelAgentNames(ctx context.Context, rows []ChannelWithAgent) { for i := range rows { if a := types.GetBuiltinAgentWithContext(ctx, rows[i].AgentID, rows[i].TenantID); a != nil && a.Name != "" { rows[i].AgentName = a.Name } } } // CreateChannel creates a new IM channel and optionally starts it. // Returns a duplicate_bot error if the bot identity is already used by another channel. func (s *Service) CreateChannel(channel *IMChannel) error { if err := s.checkDuplicateBot(channel, ""); err != nil { return err } if err := s.db.Transaction(func(tx *gorm.DB) error { if err := lockChannelAgent(tx, channel); err != nil { return err } return tx.Create(channel).Error }); err != nil { return err } if channel.Enabled { if err := s.StartChannel(channel); err != nil { logger.Warnf(context.Background(), "[IM] Created channel %s but failed to start: %v", channel.ID, err) } } s.publishChannelConfigChange(channel.ID) return nil } // SetChannelAgentID validates and assigns a new agent for an existing channel. func (s *Service) SetChannelAgentID(ctx context.Context, channel *IMChannel, agentID string) error { agentID = strings.TrimSpace(agentID) if agentID == "" { return fmt.Errorf("agent_id is required") } agent, err := s.agentService.GetAgentByID(ctx, agentID) if err != nil { return err } if agent == nil || agent.TenantID != channel.TenantID { return fmt.Errorf("agent not found") } channel.AgentID = agentID return nil } // UpdateChannel updates a channel and restarts it if needed. // Returns a duplicate_bot error if the bot identity is already used by another channel. func (s *Service) UpdateChannel(channel *IMChannel) error { if err := s.checkDuplicateBot(channel, channel.ID); err != nil { return err } if err := s.db.Transaction(func(tx *gorm.DB) error { if err := lockChannelAgent(tx, channel); err != nil { return err } // Updates cannot resurrect a concurrently deleted channel. result := tx.Model(&IMChannel{}).Where("id = ? AND tenant_id = ?", channel.ID, channel.TenantID).Select("*").Updates(channel) if result.Error != nil { return result.Error } if result.RowsAffected != 1 { return gorm.ErrRecordNotFound } return nil }); err != nil { return err } // Restart channel: stop old, start new if enabled s.StopChannel(channel.ID) if channel.Enabled { if err := s.StartChannel(channel); err != nil { logger.Warnf(context.Background(), "[IM] Updated channel %s but failed to restart: %v", channel.ID, err) } } s.publishChannelConfigChange(channel.ID) return nil } // DeleteChannelsByAgent stops and soft-deletes every IM channel bound to the // given agent within the tenant. Used when a custom agent is removed so // overview lists and running adapters do not outlive the agent. func (s *Service) DeleteChannelsByAgent(agentID string, tenantID uint64) error { var channels []IMChannel if err := s.db.Where("agent_id = ? AND tenant_id = ? AND deleted_at IS NULL", agentID, tenantID). Find(&channels).Error; err != nil { return err } if len(channels) == 0 { return nil } if err := s.db.Where("agent_id = ? AND tenant_id = ? AND deleted_at IS NULL", agentID, tenantID). Delete(&IMChannel{}).Error; err != nil { return err } for i := range channels { s.StopChannel(channels[i].ID) s.publishChannelConfigChange(channels[i].ID) } return nil } // DeleteChannel soft-deletes a channel and stops it. Only deletes if the channel belongs to the given tenant. func (s *Service) DeleteChannel(channelID string, tenantID uint64) error { result := s.db.Where("id = ? AND tenant_id = ?", channelID, tenantID).Delete(&IMChannel{}) if result.Error != nil { return result.Error } if result.RowsAffected == 0 { return fmt.Errorf("channel not found") } s.StopChannel(channelID) s.publishChannelConfigChange(channelID) return nil } // ToggleChannel enables or disables a channel. Only toggles if the channel belongs to the given tenant. func (s *Service) ToggleChannel(channelID string, tenantID uint64) (*IMChannel, error) { var ch IMChannel if err := s.db.Where("id = ? AND tenant_id = ? AND deleted_at IS NULL", channelID, tenantID).First(&ch).Error; err != nil { return nil, err } ch.Enabled = !ch.Enabled if err := s.UpdateChannel(&ch); err != nil { return nil, err } return &ch, nil } // checkDuplicateBot queries the bot_identity index to see if another active channel // already uses the same bot. This is an O(1) index lookup, not a full table scan. // The DB unique index on bot_identity serves as an additional safety net. // excludeID is the channel's own ID (for updates); pass "" for new channels. func (s *Service) checkDuplicateBot(channel *IMChannel, excludeID string) error { // Compute bot_identity the same way the BeforeSave hook will botKey := channel.computeBotIdentity() if botKey == "" { return nil } var existing IMChannel query := s.db.Where("bot_identity = ? AND deleted_at IS NULL", botKey) if excludeID != "" { query = query.Where("id != ?", excludeID) } if err := query.First(&existing).Error; err != nil { if err == gorm.ErrRecordNotFound { return nil // no conflict } return fmt.Errorf("check duplicate bot: %w", err) } return fmt.Errorf("duplicate_bot: this bot is already bound to channel %q (%s); each bot can only be connected to one channel", existing.Name, existing.ID) } // ── File message handling ────────────────────────────────────────────── // These methods handle file messages received via IM platforms. // Files are downloaded from the IM platform, validated, and saved to the // configured knowledge base asynchronously. The user receives a notification // at the start and end of processing. // ──────────────────────────────────────────────────────────────────────── // supportedKBFileExts is the set of file extensions that can be saved to a knowledge base. var supportedKBFileExts = map[string]bool{ "pdf": true, "txt": true, "docx": true, "doc": true, "md": true, "markdown": true, "png": true, "jpg": true, "jpeg": true, "gif": true, "csv": true, "xlsx": true, "xls": true, "pptx": true, "ppt": true, } // fileMessageQAContent turns a file-only platform event into a valid QA query. // IM adapters intentionally leave Content empty for file/image messages, while // the QA API requires a non-empty query. Preserve a caption when an adapter // provides one; otherwise identify the uploaded file without claiming that its // contents have already been read. func fileMessageQAContent(msg *IncomingMessage) string { if strings.TrimSpace(msg.Content) != "" { return msg.Content } fileName := strings.TrimSpace(msg.FileName) if fileName == "" { fileName = "未命名文件" } return fmt.Sprintf("我上传了文件「%s」。请确认已收到,并告知我接下来可以如何协助。", fileName) } // processDownloadedFileToKnowledgeBase stores bytes already downloaded for QA. // It deliberately has no user-facing notifications: the originating file message // receives exactly its normal QA reply, while persistence remains background work. func (s *Service) processDownloadedFileToKnowledgeBase(ctx context.Context, channel *IMChannel, file *imDownloadedAttachment) { kbID := channel.KnowledgeBaseID tenantID := channel.TenantID // Build context with tenant info for the knowledge service tenant, err := s.tenantService.GetTenantByID(ctx, tenantID) if err != nil { logger.Errorf(ctx, "[IM] Failed to get tenant %d for file processing: %v", tenantID, err) return } kbCtx := context.WithValue(ctx, types.TenantIDContextKey, tenantID) kbCtx = context.WithValue(kbCtx, types.TenantInfoContextKey, tenant) fileName := file.fileName ext := fileExtension(fileName) if !supportedKBFileExts[ext] { logger.Infof(ctx, "[IM] Unsupported file type after download: %s (file=%s)", ext, fileName) return } // Create a multipart.FileHeader compatible wrapper fh := newInMemoryFileHeader(fileName, file.content) // Create knowledge entry via the knowledge service knowledge, err := s.knowledgeService.CreateKnowledgeFromFile(kbCtx, kbID, fh, nil, nil, "", nil, imPlatformToChannel(channel.Platform), nil) if err != nil { errMsg := err.Error() // Check for duplicate file if strings.Contains(errMsg, "duplicate") || strings.Contains(errMsg, "already exists") { logger.Infof(ctx, "[IM] File already exists in knowledge base: %s", fileName) return } logger.Errorf(ctx, "[IM] Failed to create knowledge from file: %v", err) return } logger.Infof(ctx, "[IM] File saved to knowledge base: kb=%s knowledge=%s file=%s", kbID, knowledge.ID, fileName) } // fileExtension extracts the lowercase file extension from a filename. func fileExtension(filename string) string { parts := strings.Split(filename, ".") if len(parts) < 2 { return "" } return strings.ToLower(parts[len(parts)-1]) } // imPlatformToChannel maps an IM platform identifier to a Knowledge.Channel constant. func imPlatformToChannel(platform string) string { switch strings.ToLower(platform) { case "wechat": return types.ChannelWechat case "wecom", "wxwork": return types.ChannelWecom case "feishu", "lark": return types.ChannelFeishu case "dingtalk": return types.ChannelDingtalk case "slack": return types.ChannelSlack default: return types.ChannelIM } } // newInMemoryFileHeader wraps in-memory file content as a *multipart.FileHeader // so it can be passed to CreateKnowledgeFromFile which expects a multipart upload. func newInMemoryFileHeader(filename string, data []byte) *multipart.FileHeader { body := &bytes.Buffer{} writer := multipart.NewWriter(body) h := make(textproto.MIMEHeader) h.Set("Content-Disposition", fmt.Sprintf(`form-data; name="file"; filename="%s"`, filename)) h.Set("Content-Type", "application/octet-stream") part, err := writer.CreatePart(h) if err != nil { // Fallback: return a minimal FileHeader return &multipart.FileHeader{Filename: filename, Size: int64(len(data))} } _, _ = part.Write(data) _ = writer.Close() // Parse the multipart body to extract the FileHeader reader := multipart.NewReader(body, writer.Boundary()) form, err := reader.ReadForm(int64(len(data)) + 1024) if err != nil || form == nil { return &multipart.FileHeader{Filename: filename, Size: int64(len(data))} } files := form.File["file"] if len(files) == 0 { return &multipart.FileHeader{Filename: filename, Size: int64(len(data))} } return files[0] }