From b488ef44d66a4b901ec90e9383a039fc371c69c2 Mon Sep 17 00:00:00 2001 From: Luan Vu Date: Thu, 12 Mar 2026 09:35:04 +0700 Subject: [PATCH] fix: media tag enrichment, Gemini file polling, credential merge (#158) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. Media tag enrichment (audio/video/document): - Add enrichVideoIDs() — video media_id was never injected into tags, causing LLM to hallucinate UUIDs - Fix all enrich functions to replace the LAST bare tag instead of the first. When group history prepends older media tags, the first occurrence belongs to history — injecting the current turn's ID there causes the LLM to reference the wrong file 2. Gemini File API polling: - Upload response returns fileURI immediately but file may still be in PROCESSING state. Check state field; only skip polling when file is already ACTIVE. Fixes "not in an ACTIVE state" errors 3. Channel instance credential merge: - Partial credential updates (e.g. updating just token) now merge with existing credentials instead of wiping other fields - Loads, decrypts, merges, re-encrypts in a single Update() call Co-authored-by: Luvu182 <208665161+Luvu182@users.noreply.github.com> --- internal/agent/loop.go | 2 + internal/agent/media.go | 68 ++++++++++++++++++++------ internal/store/pg/channel_instances.go | 61 +++++++++++++++++++---- internal/tools/gemini_file_api.go | 11 +++-- 4 files changed, 116 insertions(+), 26 deletions(-) diff --git a/internal/agent/loop.go b/internal/agent/loop.go index 044e28c5..e4b4e9aa 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -279,6 +279,8 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) } if len(videoRefs) > 0 { ctx = tools.WithMediaVideoRefs(ctx, videoRefs) + // Embed media IDs into tags so LLM can reference them. + l.enrichVideoIDs(messages, mediaRefs) } // 2e. Cross-session recovery: notify team leads about orphaned pending tasks diff --git a/internal/agent/media.go b/internal/agent/media.go index 086f48c7..e2c280c3 100644 --- a/internal/agent/media.go +++ b/internal/agent/media.go @@ -139,24 +139,25 @@ func (l *Loop) enrichDocumentPaths(messages []providers.Message, refs []provider pathAttr := fmt.Sprintf(" path=%q", p) old1 := "" new1 := "" - if strings.Contains(content, old1) { - content = strings.Replace(content, old1, new1, 1) + // Replace the LAST bare tag (current message, not group history). + if idx := strings.LastIndex(content, old1); idx >= 0 { + content = content[:idx] + new1 + content[idx+len(old1):] continue } - // For named variant, inject path attribute - if idx := strings.Index(content, "") if closeIdx >= 0 { tag := content[idx : idx+closeIdx] - content = strings.Replace(content, tag+">", tag+pathAttr+">", 1) + content = content[:idx] + tag + pathAttr + ">" + content[idx+closeIdx+1:] } } - // For Slack variant with file= attribute - if idx := strings.Index(content, "") if closeIdx >= 0 { tag := content[idx : idx+closeIdx] - content = strings.Replace(content, tag+">", tag+pathAttr+">", 1) + content = content[:idx] + tag + pathAttr + ">" + content[idx+closeIdx+1:] } } } @@ -166,6 +167,8 @@ func (l *Loop) enrichDocumentPaths(messages []providers.Message, refs []provider // enrichAudioIDs updates the last user message to embed persisted media IDs // in and tags so the LLM can reference them. // Without this, the LLM sees plain and cannot pass a valid media_id. +// Replaces the LAST bare tag (current message) rather than the first (which may be +// in group history context), so the current turn's media gets the correct ID. func (l *Loop) enrichAudioIDs(messages []providers.Message, refs []providers.MediaRef) { if len(messages) == 0 { return @@ -188,16 +191,53 @@ func (l *Loop) enrichAudioIDs(messages []providers.Message, refs []providers.Med } idAttr := fmt.Sprintf(" id=%q", ref.ID) - // Replace bare with + // Replace the LAST bare with bare := "" - if strings.Contains(content, bare) { - content = strings.Replace(content, bare, "", 1) + if idx := strings.LastIndex(content, bare); idx >= 0 { + content = content[:idx] + "" + content[idx+len(bare):] continue } - // Replace bare with + // Replace the LAST bare with bareVoice := "" - if strings.Contains(content, bareVoice) { - content = strings.Replace(content, bareVoice, "", 1) + if idx := strings.LastIndex(content, bareVoice); idx >= 0 { + content = content[:idx] + "" + content[idx+len(bareVoice):] + continue + } + } + messages[lastIdx].Content = content +} + +// enrichVideoIDs updates the last user message to embed persisted media IDs +// in tags so the LLM can reference them via read_video tool. +// Without this, the LLM sees plain and hallucinates a media_id. +// Replaces the LAST bare tag (current message) rather than the first (which may be +// in group history context), so the current turn's media gets the correct ID. +func (l *Loop) enrichVideoIDs(messages []providers.Message, refs []providers.MediaRef) { + if len(messages) == 0 { + return + } + lastIdx := -1 + for i := len(messages) - 1; i >= 0; i-- { + if messages[i].Role == "user" { + lastIdx = i + break + } + } + if lastIdx < 0 { + return + } + + content := messages[lastIdx].Content + for _, ref := range refs { + if ref.Kind != "video" { + continue + } + idAttr := fmt.Sprintf(" id=%q", ref.ID) + + // Replace the LAST bare with + bare := "" + if idx := strings.LastIndex(content, bare); idx >= 0 { + content = content[:idx] + "" + content[idx+len(bare):] continue } } diff --git a/internal/store/pg/channel_instances.go b/internal/store/pg/channel_instances.go index 94b8abaf..0c4c829b 100644 --- a/internal/store/pg/channel_instances.go +++ b/internal/store/pg/channel_instances.go @@ -145,19 +145,43 @@ func (s *PGChannelInstanceStore) scanInstances(rows *sql.Rows) ([]store.ChannelI } func (s *PGChannelInstanceStore) Update(ctx context.Context, id uuid.UUID, updates map[string]any) error { - // Encrypt credentials if present + // Merge and encrypt credentials if present if credsVal, ok := updates["credentials"]; ok && credsVal != nil { - var credsBytes []byte + var newCreds map[string]any switch v := credsVal.(type) { - case []byte: - credsBytes = v - case string: - credsBytes = []byte(v) + case map[string]any: + newCreds = v default: - // Object/map from JSON — marshal to []byte - if b, err := json.Marshal(v); err == nil { - credsBytes = b + var raw []byte + switch vv := v.(type) { + case []byte: + raw = vv + case string: + raw = []byte(vv) + default: + if b, err := json.Marshal(v); err == nil { + raw = b + } } + if len(raw) > 0 { + if err := json.Unmarshal(raw, &newCreds); err != nil { + newCreds = nil + } + } + } + + // Merge with existing credentials so partial updates don't wipe other fields + if len(newCreds) > 0 { + existing, _ := s.loadExistingCreds(ctx, id) + for k, v := range newCreds { + existing[k] = v + } + newCreds = existing + } + + var credsBytes []byte + if len(newCreds) > 0 { + credsBytes, _ = json.Marshal(newCreds) } if len(credsBytes) > 0 && s.encKey != "" { encrypted, err := crypto.Encrypt(string(credsBytes), s.encKey) @@ -172,6 +196,25 @@ func (s *PGChannelInstanceStore) Update(ctx context.Context, id uuid.UUID, updat return execMapUpdate(ctx, s.db, "channel_instances", id, updates) } +// loadExistingCreds reads and decrypts the current credentials for merging. +func (s *PGChannelInstanceStore) loadExistingCreds(ctx context.Context, id uuid.UUID) (map[string]any, error) { + var raw []byte + err := s.db.QueryRowContext(ctx, "SELECT credentials FROM channel_instances WHERE id = $1", id).Scan(&raw) + if err != nil || len(raw) == 0 { + return make(map[string]any), err + } + if s.encKey != "" { + if dec, err := crypto.Decrypt(string(raw), s.encKey); err == nil { + raw = []byte(dec) + } + } + var m map[string]any + if err := json.Unmarshal(raw, &m); err != nil { + return make(map[string]any), nil + } + return m, nil +} + func (s *PGChannelInstanceStore) Delete(ctx context.Context, id uuid.UUID) error { _, err := s.db.ExecContext(ctx, "DELETE FROM channel_instances WHERE id = $1", id) return err diff --git a/internal/tools/gemini_file_api.go b/internal/tools/gemini_file_api.go index 351e516c..fe96b941 100644 --- a/internal/tools/gemini_file_api.go +++ b/internal/tools/gemini_file_api.go @@ -81,15 +81,20 @@ func geminiFileUpload(ctx context.Context, apiKey, displayName string, data []by var uploadResult struct { File struct { - Name string `json:"name"` - URI string `json:"uri"` + Name string `json:"name"` + URI string `json:"uri"` + State string `json:"state"` } `json:"file"` } if err := json.Unmarshal(respBody, &uploadResult); err != nil { return "", "", fmt.Errorf("parse upload response: %w", err) } - return uploadResult.File.Name, uploadResult.File.URI, nil + // Only return URI if file is already ACTIVE; otherwise caller must poll. + if uploadResult.File.State == "ACTIVE" { + return uploadResult.File.Name, uploadResult.File.URI, nil + } + return uploadResult.File.Name, "", nil } // geminiFilePoll polls the Gemini File API until the file reaches ACTIVE state.