package handlers import ( "context" "database/sql" "encoding/base64" "encoding/json" "fmt" "io" "log" "net/http" "strings" "github.com/gin-gonic/gin" capspkg "git.gobha.me/xcaliber/chat-switchboard/capabilities" "git.gobha.me/xcaliber/chat-switchboard/crypto" "git.gobha.me/xcaliber/chat-switchboard/database" "git.gobha.me/xcaliber/chat-switchboard/events" "git.gobha.me/xcaliber/chat-switchboard/models" "git.gobha.me/xcaliber/chat-switchboard/providers" "git.gobha.me/xcaliber/chat-switchboard/storage" "git.gobha.me/xcaliber/chat-switchboard/store" "git.gobha.me/xcaliber/chat-switchboard/tools" ) // ── Request Types ─────────────────────────── type completionRequest struct { ChannelID string `json:"channel_id"` // preferred; validated manually below ChatID string `json:"chat_id"` // deprecated alias — maps to channel_id Content string `json:"content" binding:"required"` Model string `json:"model,omitempty"` PresetID string `json:"preset_id,omitempty"` // if set, unwraps preset → base model + config APIConfigID string `json:"provider_config_id,omitempty"` MaxTokens int `json:"max_tokens,omitempty"` Temperature *float64 `json:"temperature,omitempty"` TopP *float64 `json:"top_p,omitempty"` Stream *bool `json:"stream,omitempty"` AttachmentIDs []string `json:"attachment_ids,omitempty"` // staged attachment UUIDs to include in request DisabledTools []string `json:"disabled_tools,omitempty"` // tool names to exclude from this request } // CompletionHandler proxies LLM requests through the backend. type CompletionHandler struct { vault *crypto.KeyResolver stores store.Stores hub *events.Hub // WebSocket hub for browser tool bridge objStore storage.ObjectStore // file storage for attachment content (nil = disabled) } // NewCompletionHandler creates a new handler. func NewCompletionHandler(vault *crypto.KeyResolver, stores store.Stores, hub *events.Hub, objStore storage.ObjectStore) *CompletionHandler { return &CompletionHandler{vault: vault, stores: stores, hub: hub, objStore: objStore} } // ── Chat Completion ───────────────────────── // POST /api/v1/chat/completions // // Flow: // 1. Validate request, verify chat ownership // 2. Resolve api_config (from request, chat, or user default) // 3. Load conversation history from messages // 4. Persist user message // 5. Attach tool definitions if model supports tools // 6. Call provider (stream or non-stream) // 7. Tool loop: if LLM requests tool calls, execute them and re-call // 8. Stream SSE to client / return JSON // 9. Persist assistant message with token counts func (h *CompletionHandler) Complete(c *gin.Context) { var req completionRequest if err := c.ShouldBindJSON(&req); err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) return } // Support chat_id as alias during frontend transition channelID := req.ChannelID if channelID == "" { channelID = req.ChatID } if channelID == "" { c.JSON(http.StatusBadRequest, gin.H{"error": "channel_id is required"}) return } userID := getUserID(c) // Verify channel ownership if !userOwnsChannel(c, channelID, userID) { return } // ── Preset unwrap: preset overrides defaults, explicit request fields win ── var presetSystemPrompt string var personaID string // tracks active persona for KB scoping if req.PresetID != "" { preset := ResolvePreset(h.stores, req.PresetID, userID) if preset == nil { c.JSON(http.StatusBadRequest, gin.H{"error": "preset not found or not accessible"}) return } personaID = preset.ID // Preset provides defaults; explicit request fields take priority if req.Model == "" { req.Model = preset.BaseModelID } if req.APIConfigID == "" && preset.ProviderConfigID != nil { req.APIConfigID = *preset.ProviderConfigID } if req.Temperature == nil && preset.Temperature != nil { req.Temperature = preset.Temperature } if req.MaxTokens == 0 && preset.MaxTokens != nil { req.MaxTokens = *preset.MaxTokens } if preset.SystemPrompt != "" { presetSystemPrompt = preset.SystemPrompt } } // Resolve provider config providerCfg, providerID, model, configID, providerScope, err := h.resolveConfig(userID, channelID, req) if err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) return } // ── Team policy: require_private_providers ── if err := enforcePrivateProviderPolicy(userID, configID); err != nil { c.JSON(http.StatusForbidden, gin.H{"error": err.Error()}) return } provider, err := providers.Get(providerID) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return } // Load conversation history messages, err := h.loadConversation(channelID, userID, presetSystemPrompt, personaID) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load conversation"}) return } // Resolve capabilities early — needed for vision gating below caps := h.getModelCapabilities(c, model, configID) // Build user message — multimodal if attachments are present userMsg := providers.Message{ Role: "user", Content: req.Content, } if len(req.AttachmentIDs) > 0 && h.objStore != nil { parts, augContent, _, err := h.buildMultimodalParts(c, channelID, req.Content, req.AttachmentIDs, caps) if err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) return } if len(parts) > 0 { // Multimodal (images present): use ContentParts userMsg.ContentParts = parts } else if augContent != req.Content { // Doc-only: use augmented text content userMsg.Content = augContent } } messages = append(messages, userMsg) // Persist user message (text-only for storage — multimodal parts are ephemeral) msgID, err := h.persistMessage(channelID, userID, "user", req.Content, "", 0, 0, nil, nil) if err != nil { log.Printf("Failed to persist user message: %v", err) } // Link attachments to the persisted message if msgID != "" && len(req.AttachmentIDs) > 0 { for _, attID := range req.AttachmentIDs { if err := h.stores.Attachments.SetMessageID(c.Request.Context(), attID, msgID); err != nil { log.Printf("Failed to link attachment %s to message %s: %v", attID, msgID, err) } } } // Build provider request provReq := providers.CompletionRequest{ Model: model, Messages: messages, } if req.MaxTokens > 0 { provReq.MaxTokens = req.MaxTokens } else { // ResolveMaxOutput checks: caps → known models → context/8 → 4096 provReq.MaxTokens = capspkg.ResolveMaxOutput(model, caps) } if req.Temperature != nil { provReq.Temperature = req.Temperature } if req.TopP != nil { provReq.TopP = req.TopP } // Attach tool definitions if model supports tool calling and tools are available hasBrowserTools := h.hub != nil && h.hub.IsConnected(userID) if caps.ToolCalling && (tools.HasTools() || hasBrowserTools) { provReq.Tools = h.buildToolDefs(c.Request.Context(), userID, hasBrowserTools, req.DisabledTools) } // Determine streaming stream := true // default if req.Stream != nil { stream = *req.Stream } if stream { h.streamCompletion(c, provider, providerCfg, provReq, channelID, userID, model, configID, providerScope, personaID) } else { h.syncCompletion(c, provider, providerCfg, provReq, channelID, userID, model, configID, providerScope, personaID) } } // buildToolDefs converts registered tools to provider-format tool definitions. // If includeBrowser is true, also fetches tool schemas from enabled browser extensions. // Tools whose names appear in disabledTools are excluded. func (h *CompletionHandler) buildToolDefs(ctx context.Context, userID string, includeBrowser bool, disabledTools []string) []providers.ToolDef { // Build disabled set for O(1) lookup disabled := make(map[string]bool, len(disabledTools)) for _, name := range disabledTools { disabled[name] = true } allTools := tools.AllDefinitionsFiltered(disabled) defs := make([]providers.ToolDef, 0, len(allTools)) for _, t := range allTools { defs = append(defs, providers.ToolDef{ Type: "function", Function: providers.FunctionDef{ Name: t.Name, Description: t.Description, Parameters: t.Parameters, }, }) } // Append browser-defined tool schemas from extensions if includeBrowser { exts, err := h.stores.Extensions.ListForUser(ctx, userID) if err != nil { log.Printf("⚠️ Failed to load extensions for tools: %v", err) return defs } for _, ext := range exts { if ext.Tier != "browser" { continue } var manifest struct { Tools []struct { Name string `json:"name"` Description string `json:"description"` Parameters json.RawMessage `json:"parameters"` Tier string `json:"tier"` } `json:"tools"` } if err := json.Unmarshal(ext.Manifest, &manifest); err != nil { continue } for _, t := range manifest.Tools { if disabled[t.Name] { continue } defs = append(defs, providers.ToolDef{ Type: "function", Function: providers.FunctionDef{ Name: t.Name, Description: t.Description, Parameters: t.Parameters, }, }) } } } return defs } // ── List Available Tools ──────────────────── // GET /api/v1/tools // // Returns all registered server-side tools plus browser extension tools // for the current user. Used by the frontend to build the tools toggle menu. type toolInfo struct { Name string `json:"name"` DisplayName string `json:"display_name"` Description string `json:"description"` Category string `json:"category"` } func (h *CompletionHandler) ListTools(c *gin.Context) { userID := c.GetString("user_id") // Server-side tools allDefs := tools.AllDefinitions() result := make([]toolInfo, 0, len(allDefs)) for _, d := range allDefs { result = append(result, toolInfo{ Name: d.Name, DisplayName: d.DisplayName, Description: d.Description, Category: d.Category, }) } // Browser extension tools if h.hub != nil && h.hub.IsConnected(userID) { exts, err := h.stores.Extensions.ListForUser(c.Request.Context(), userID) if err == nil { for _, ext := range exts { if ext.Tier != "browser" { continue } var manifest struct { Tools []struct { Name string `json:"name"` Description string `json:"description"` } `json:"tools"` } if err := json.Unmarshal(ext.Manifest, &manifest); err != nil { continue } for _, t := range manifest.Tools { result = append(result, toolInfo{ Name: t.Name, Description: t.Description, Category: "browser", }) } } } } c.JSON(http.StatusOK, gin.H{"data": result}) } // ── Streaming Completion (SSE) with Tool Loop ── const maxToolIterations = 10 // safety limit on tool call rounds func (h *CompletionHandler) streamCompletion( c *gin.Context, provider providers.Provider, cfg providers.ProviderConfig, req providers.CompletionRequest, channelID, userID, model, configID, providerScope, personaID string, ) { result := streamWithToolLoop(c, provider, cfg, &req, model, userID, channelID, personaID, h.hub) // Persist assistant response if result.Content != "" { if _, err := h.persistMessage(channelID, userID, "assistant", result.Content, model, result.InputTokens, result.OutputTokens, nil, toolActivityJSON(result.ToolActivity)); err != nil { log.Printf("Failed to persist assistant message: %v", err) } } // Log usage h.logUsage(c, channelID, userID, configID, providerScope, model, result.InputTokens, result.OutputTokens, result.CacheCreationTokens, result.CacheReadTokens) } // ── Non-Streaming Completion with Tool Loop ── func (h *CompletionHandler) syncCompletion( c *gin.Context, provider providers.Provider, cfg providers.ProviderConfig, req providers.CompletionRequest, channelID, userID, model, configID, providerScope, personaID string, ) { var totalInput, totalOutput int var totalCacheCreation, totalCacheRead int var finalContent string var allToolActivity []map[string]interface{} for iteration := 0; iteration < maxToolIterations; iteration++ { resp, err := provider.ChatCompletion(c.Request.Context(), cfg, req) if err != nil { c.JSON(http.StatusBadGateway, gin.H{"error": "provider error: " + err.Error()}) return } totalInput += resp.InputTokens totalOutput += resp.OutputTokens totalCacheCreation += resp.CacheCreationTokens totalCacheRead += resp.CacheReadTokens // No tool calls — normal response if len(resp.ToolCalls) == 0 || resp.FinishReason != "tool_calls" { finalContent = resp.Content // Persist assistant response if _, err := h.persistMessage(channelID, userID, "assistant", finalContent, model, totalInput, totalOutput, nil, toolActivityJSON(allToolActivity)); err != nil { log.Printf("Failed to persist assistant message: %v", err) } // Log usage h.logUsage(c, channelID, userID, configID, providerScope, model, totalInput, totalOutput, totalCacheCreation, totalCacheRead) // Return OpenAI-compatible response c.JSON(http.StatusOK, gin.H{ "model": resp.Model, "choices": []gin.H{ { "message": gin.H{ "role": "assistant", "content": finalContent, }, "finish_reason": resp.FinishReason, }, }, "usage": gin.H{ "prompt_tokens": totalInput, "completion_tokens": totalOutput, "total_tokens": totalInput + totalOutput, }, }) return } // ── Tool execution round ──────────────── assistantMsg := providers.Message{ Role: "assistant", Content: resp.Content, } for _, tc := range resp.ToolCalls { assistantMsg.ToolCalls = append(assistantMsg.ToolCalls, tc) } req.Messages = append(req.Messages, assistantMsg) execCtx := tools.ExecutionContext{ UserID: userID, ChannelID: channelID, PersonaID: personaID, } for _, tc := range resp.ToolCalls { call := tools.ToolCall{ ID: tc.ID, Name: tc.Function.Name, Arguments: tc.Function.Arguments, } log.Printf("🔧 Executing tool: %s (call %s)", call.Name, call.ID) result := tools.ExecuteCall(c.Request.Context(), execCtx, call) // Collect for persistence allToolActivity = append(allToolActivity, map[string]interface{}{ "id": tc.ID, "name": tc.Function.Name, "arguments": tc.Function.Arguments, "result": result.Content, "is_error": result.IsError, }) req.Messages = append(req.Messages, providers.Message{ Role: "tool", Content: result.Content, ToolCallID: result.ToolCallID, Name: result.Name, }) } } // Max iterations reached c.JSON(http.StatusOK, gin.H{ "model": model, "choices": []gin.H{ { "message": gin.H{ "role": "assistant", "content": "[Tool execution limit reached]", }, "finish_reason": "stop", }, }, }) } // ── Model Capabilities ────────────────────── // escapeJSON escapes a string for safe embedding in a JSON string value. func escapeJSON(s string) string { s = strings.ReplaceAll(s, `\`, `\\`) s = strings.ReplaceAll(s, `"`, `\"`) s = strings.ReplaceAll(s, "\n", `\n`) s = strings.ReplaceAll(s, "\r", `\r`) s = strings.ReplaceAll(s, "\t", `\t`) return s } // getModelCapabilities looks up capabilities from model_catalog DB, // then overlays with known model defaults and heuristic detection. func (h *CompletionHandler) getModelCapabilities(c *gin.Context, model, apiConfigID string) models.ModelCapabilities { return ResolveModelCaps(c, model, apiConfigID) } // ── Multimodal Assembly ───────────────────── // Builds content parts from attachments for the user message. // // Returns: // - parts: ContentParts array (non-nil only when images are present) // - augContent: enriched text content with document context (for doc-only case) // - validIDs: attachment IDs that were successfully processed // - error: if any attachment is invalid or vision is needed but missing // // Rules: // - Images → base64 data URI (requires vision capability) // - Documents with extracted_text → text injection // - Documents without extraction → filename placeholder // - Text is always the first part func (h *CompletionHandler) buildMultimodalParts( c *gin.Context, channelID, textContent string, attachmentIDs []string, caps models.ModelCapabilities, ) ([]providers.ContentPart, string, []string, error) { parts := []providers.ContentPart{ {Type: "text", Text: textContent}, } var docTexts []string var validIDs []string hasImage := false for _, attID := range attachmentIDs { att, err := h.stores.Attachments.GetByID(c.Request.Context(), attID) if err != nil { return nil, "", nil, fmt.Errorf("attachment %s not found", attID) } // Security: verify attachment belongs to this channel if att.ChannelID != channelID { return nil, "", nil, fmt.Errorf("attachment %s does not belong to this channel", attID) } if isImageContentType(att.ContentType) { // Vision gating: reject images if model lacks vision if !caps.Vision { return nil, "", nil, fmt.Errorf("model does not support image input; remove image attachments or choose a vision-capable model") } // Read image from storage, base64 encode, build data URI reader, _, _, err := h.objStore.Get(c.Request.Context(), att.StorageKey) if err != nil { log.Printf("Failed to read attachment %s from storage: %v", attID, err) parts = append(parts, providers.ContentPart{ Type: "text", Text: fmt.Sprintf("[Image: %s — failed to read from storage]", att.Filename), }) validIDs = append(validIDs, attID) continue } data, err := io.ReadAll(reader) reader.Close() if err != nil { log.Printf("Failed to read attachment %s bytes: %v", attID, err) validIDs = append(validIDs, attID) continue } b64 := base64.StdEncoding.EncodeToString(data) dataURI := fmt.Sprintf("data:%s;base64,%s", att.ContentType, b64) parts = append(parts, providers.ContentPart{ Type: "image_url", ImageURL: &providers.ImageURL{ URL: dataURI, Detail: "auto", }, }) hasImage = true } else if att.ExtractedText != nil && *att.ExtractedText != "" { // Document with extracted text → inject as context docText := fmt.Sprintf("[Document: %s]\n%s", att.Filename, *att.ExtractedText) parts = append(parts, providers.ContentPart{ Type: "text", Text: docText, }) docTexts = append(docTexts, docText) } else { // Document without extraction (pending, failed, or not extractable) status := "pending" if s, ok := att.Metadata["extraction_status"].(string); ok { status = s } placeholder := fmt.Sprintf("[Attached file: %s (extraction %s)]", att.Filename, status) parts = append(parts, providers.ContentPart{ Type: "text", Text: placeholder, }) docTexts = append(docTexts, placeholder) } validIDs = append(validIDs, attID) } // If images present → use ContentParts (multimodal array) if hasImage { return parts, textContent, validIDs, nil } // Doc-only: merge document context into a single text string (more efficient) augmented := textContent for _, dt := range docTexts { augmented += "\n\n" + dt } return nil, augmented, validIDs, nil } // isImageContentType returns true for MIME types that should be sent as // base64 image content parts rather than text extraction. func isImageContentType(ct string) bool { return strings.HasPrefix(ct, "image/") } // ── Config Resolution ─────────────────────── // Priority: request.provider_config_id → chat.provider_config_id → user's first active config func (h *CompletionHandler) resolveConfig(userID string, channelID string, req completionRequest) (providers.ProviderConfig, string, string, string, string, error) { var configID string // 1. Explicit config from request if req.APIConfigID != "" { configID = req.APIConfigID } // 2. Config from channel if configID == "" { var channelConfigID *string err := database.DB.QueryRow( database.Q(`SELECT provider_config_id FROM channels WHERE id = $1`), channelID, ).Scan(&channelConfigID) if err == nil && channelConfigID != nil { configID = *channelConfigID } } // 3. User's first active config (personal first, then global — excludes team providers) if configID == "" { err := database.DB.QueryRow(database.Q(` SELECT id FROM provider_configs WHERE is_active = true AND ( (scope = 'personal' AND owner_id = $1) OR scope = 'global' ) ORDER BY scope ASC, created_at ASC LIMIT 1 `), userID).Scan(&configID) if err != nil { return providers.ProviderConfig{}, "", "", "", "", fmt.Errorf("no API config found — add one at /api-configs") } } // Load the config — allow personal, global, OR team configs the user belongs to var providerID, endpoint string var providerScope string var modelDefault *string var apiKeyEnc, keyNonce []byte var keyScope string var customHeadersJSON, providerSettingsJSON []byte // $2/userID appears twice in the query; Postgres reuses positional params, // SQLite needs each ? bound separately. configArgs := []interface{}{configID, userID} if database.IsSQLite() { configArgs = append(configArgs, userID) } err := database.DB.QueryRow(database.Q(` SELECT provider, endpoint, scope, api_key_enc, key_nonce, key_scope, model_default, headers, settings FROM provider_configs WHERE id = $1 AND is_active = true AND (scope = 'global' OR (scope = 'personal' AND owner_id = $2) OR (scope = 'team' AND owner_id IN (SELECT team_id FROM team_members WHERE user_id = $2))) `), configArgs...).Scan(&providerID, &endpoint, &providerScope, &apiKeyEnc, &keyNonce, &keyScope, &modelDefault, &customHeadersJSON, &providerSettingsJSON) if err == sql.ErrNoRows { return providers.ProviderConfig{}, "", "", "", "", fmt.Errorf("API config not found or not accessible") } if err != nil { return providers.ProviderConfig{}, "", "", "", "", fmt.Errorf("failed to load API config: %w", err) } // Resolve model: request > config default model := req.Model if model == "" && modelDefault != nil { model = *modelDefault } if model == "" { return providers.ProviderConfig{}, "", "", "", "", fmt.Errorf("no model specified and no default model in config") } // Decrypt API key using the appropriate tier key := "" if len(apiKeyEnc) > 0 { if h.vault != nil { var err error key, err = h.vault.Decrypt(apiKeyEnc, keyNonce, keyScope, userID) if err != nil { if err == crypto.ErrVaultLocked { return providers.ProviderConfig{}, "", "", "", "", fmt.Errorf("personal vault is locked — please log in again") } return providers.ProviderConfig{}, "", "", "", "", fmt.Errorf("failed to decrypt API key: %w", err) } } else { // No vault — key stored as raw bytes (unencrypted fallback) key = string(apiKeyEnc) } } // Parse custom headers customHeaders := make(map[string]string) if customHeadersJSON != nil { _ = json.Unmarshal(customHeadersJSON, &customHeaders) } // Parse provider-specific settings providerSettings := make(map[string]interface{}) if providerSettingsJSON != nil { _ = json.Unmarshal(providerSettingsJSON, &providerSettings) } return providers.ProviderConfig{ Endpoint: endpoint, APIKey: key, CustomHeaders: customHeaders, Settings: providerSettings, }, providerID, model, configID, providerScope, nil } // ── Conversation Loader ───────────────────── // loadConversation builds the message history for the LLM by walking the // active path through the message tree. Only the current branch is sent — // sibling branches are never included in context. // // Summary-aware: if the path contains a summary node (metadata.type = "summary"), // messages before it are replaced by the summary content as a system message. func (h *CompletionHandler) loadConversation(channelID, userID, presetSystemPrompt, personaID string) ([]providers.Message, error) { messages := make([]providers.Message, 0) // ── Admin system prompt (always injected first, no opt out) ── adminPrompt, err := h.stores.GlobalConfig.Get(context.Background(), "system_prompt") if err == nil { if content, ok := adminPrompt["content"].(string); ok && content != "" { messages = append(messages, providers.Message{ Role: "system", Content: content, }) } } // ── User/preset system prompt (appended after admin prompt) ── var systemPrompt *string _ = database.DB.QueryRow( database.Q(`SELECT system_prompt FROM channels WHERE id = $1`), channelID, ).Scan(&systemPrompt) // Preset system prompt takes priority; channel system prompt is fallback if presetSystemPrompt != "" { messages = append(messages, providers.Message{ Role: "system", Content: presetSystemPrompt, }) } else if systemPrompt != nil && *systemPrompt != "" { messages = append(messages, providers.Message{ Role: "system", Content: *systemPrompt, }) } // ── KB hint (nudge LLM to use kb_search when KBs are active) ── if kbHint := BuildKBHint(context.Background(), h.stores, channelID, userID, personaID); kbHint != "" { messages = append(messages, providers.Message{ Role: "system", Content: kbHint, }) } // Walk the active path (root → leaf) instead of loading all messages path, err := getActivePath(channelID, userID) if err != nil { return nil, err } // Find the most recent summary node in the path summaryIdx := -1 for i, m := range path { if isSummaryMessage(&m) { summaryIdx = i } } // If we found a summary, inject it as a system message and skip prior messages startIdx := 0 if summaryIdx >= 0 { messages = append(messages, providers.Message{ Role: "system", Content: "Previous conversation summary:\n" + path[summaryIdx].Content, }) startIdx = summaryIdx + 1 } for _, m := range path[startIdx:] { if m.Role == "system" { continue // system prompts handled above } // Skip summary nodes that aren't the boundary (shouldn't happen, but defensive) if isSummaryMessage(&m) { continue } messages = append(messages, providers.Message{ Role: m.Role, Content: m.Content, }) } return messages, nil } // ── Message Persistence ───────────────────── // persistMessage inserts a message into the tree, updates the cursor, and // returns the new message's ID. // // Parent resolution: // - parentOverride (explicit parent for edit/regen operations) // - cursor's active_leaf_id (normal conversation flow) // - latest message by created_at (fallback for legacy/missing cursor) // toolActivityJSON marshals tool activity for persistence. Returns nil if empty. func toolActivityJSON(activity []map[string]interface{}) json.RawMessage { if len(activity) == 0 { return nil } b, _ := json.Marshal(activity) return b } func (h *CompletionHandler) persistMessage(channelID, userID, role, content, model string, inputTokens, outputTokens int, parentOverride *string, toolCallsJSON json.RawMessage) (string, error) { var tokensUsed *int if inputTokens > 0 || outputTokens > 0 { total := inputTokens + outputTokens tokensUsed = &total } // Determine parent var parentID *string if parentOverride != nil { parentID = parentOverride } else { parentID, _ = getActiveLeaf(channelID, userID) } // Compute sibling_index siblingIdx := nextSiblingIndex(channelID, parentID) // Determine participant participantType := "user" participantID := userID if role == "assistant" { participantType = "model" participantID = model } // Prepare tool_calls JSONB (nil → NULL) var tcVal interface{} if len(toolCallsJSON) > 0 { tcVal = string(toolCallsJSON) } // Insert with RETURNING id (Postgres) or pre-generated id (SQLite) var newID string if database.IsSQLite() { newID = store.NewID() _, err := database.DB.Exec(` INSERT INTO messages (id, channel_id, role, content, model, tokens_used, tool_calls, parent_id, participant_type, participant_id, sibling_index) VALUES (?, ?, ?, ?, NULLIF(?, ''), ?, ?, ?, ?, ?, ?) `, newID, channelID, role, content, model, tokensUsed, tcVal, parentID, participantType, participantID, siblingIdx) if err != nil { return "", err } } else { err := database.DB.QueryRow(` INSERT INTO messages (channel_id, role, content, model, tokens_used, tool_calls, parent_id, participant_type, participant_id, sibling_index) VALUES ($1, $2, $3, NULLIF($4, ''), $5, $6, $7, $8, $9, $10) RETURNING id `, channelID, role, content, model, tokensUsed, tcVal, parentID, participantType, participantID, siblingIdx, ).Scan(&newID) if err != nil { return "", err } } // Update cursor to point at new message if userID != "" { _ = updateCursor(channelID, userID, newID) } // Touch channel updated_at _, _ = database.DB.Exec(database.Q(`UPDATE channels SET updated_at = NOW() WHERE id = $1`), channelID) return newID, nil } // ── Usage Logging ───────────────────────── // logUsage records token usage and cost in the usage_log table. // Always logs even if tokens are zero — the request happened. func (h *CompletionHandler) logUsage( c *gin.Context, channelID, userID, configID, providerScope, modelID string, inputTokens, outputTokens, cacheCreation, cacheRead int, ) { if h.stores.Usage == nil { return } entry := &models.UsageEntry{ ChannelID: &channelID, UserID: userID, ProviderConfigID: &configID, ProviderScope: providerScope, ModelID: modelID, PromptTokens: inputTokens, CompletionTokens: outputTokens, CacheCreationTokens: cacheCreation, CacheReadTokens: cacheRead, } // Look up pricing for cost calculation if h.stores.Pricing != nil { pricing, err := h.stores.Pricing.GetForModel(c.Request.Context(), configID, modelID) if err == nil && pricing != nil { entry.CostInput = calcCost(inputTokens, pricing.InputPerM) entry.CostOutput = calcCost(outputTokens, pricing.OutputPerM) } } if err := h.stores.Usage.Log(c.Request.Context(), entry); err != nil { log.Printf("⚠ Failed to log usage: %v", err) } } // calcCost computes the cost for a given token count and price-per-million. func calcCost(tokens int, pricePerM *float64) *float64 { if pricePerM == nil || tokens == 0 { return nil } cost := float64(tokens) * *pricePerM / 1_000_000.0 return &cost }