diff --git a/contrib/helm/charts/probo/templates/deployment.yaml b/contrib/helm/charts/probo/templates/deployment.yaml index 1c813ba32..d62d64e1d 100644 --- a/contrib/helm/charts/probo/templates/deployment.yaml +++ b/contrib/helm/charts/probo/templates/deployment.yaml @@ -264,6 +264,44 @@ spec: - name: AGENT_TRACKER_MAPPING_MAX_TOKENS value: {{ .Values.probo.trackerMapping.maxTokens | quote }} {{- end }} + # Tracker Mapping Worker + {{- if .Values.probo.trackerMappingWorker.interval }} + - name: TRACKER_MAPPING_INTERVAL + value: {{ .Values.probo.trackerMappingWorker.interval | quote }} + {{- end }} + {{- if .Values.probo.trackerMappingWorker.maxConcurrency }} + - name: TRACKER_MAPPING_MAX_CONCURRENCY + value: {{ .Values.probo.trackerMappingWorker.maxConcurrency | quote }} + {{- end }} + {{- if .Values.probo.trackerMappingWorker.agentTimeout }} + - name: TRACKER_MAPPING_AGENT_TIMEOUT + value: {{ .Values.probo.trackerMappingWorker.agentTimeout | quote }} + {{- end }} + {{- if .Values.probo.trackerMappingWorker.agentMaxTurns }} + - name: TRACKER_MAPPING_AGENT_MAX_TURNS + value: {{ .Values.probo.trackerMappingWorker.agentMaxTurns | quote }} + {{- end }} + # Common Pattern Enrichment Worker + {{- if .Values.probo.commonPatternEnrichmentWorker.interval }} + - name: COMMON_PATTERN_ENRICHMENT_INTERVAL + value: {{ .Values.probo.commonPatternEnrichmentWorker.interval | quote }} + {{- end }} + {{- if .Values.probo.commonPatternEnrichmentWorker.maxConcurrency }} + - name: COMMON_PATTERN_ENRICHMENT_MAX_CONCURRENCY + value: {{ .Values.probo.commonPatternEnrichmentWorker.maxConcurrency | quote }} + {{- end }} + {{- if .Values.probo.commonPatternEnrichmentWorker.staleAfter }} + - name: COMMON_PATTERN_ENRICHMENT_STALE_AFTER + value: {{ .Values.probo.commonPatternEnrichmentWorker.staleAfter | quote }} + {{- end }} + {{- if .Values.probo.commonPatternEnrichmentWorker.agentTimeout }} + - name: COMMON_PATTERN_ENRICHMENT_AGENT_TIMEOUT + value: {{ .Values.probo.commonPatternEnrichmentWorker.agentTimeout | quote }} + {{- end }} + {{- if .Values.probo.commonPatternEnrichmentWorker.agentMaxTurns }} + - name: COMMON_PATTERN_ENRICHMENT_AGENT_MAX_TURNS + value: {{ .Values.probo.commonPatternEnrichmentWorker.agentMaxTurns | quote }} + {{- end }} # Custom Domains {{- if .Values.probo.customDomains.enabled }} - name: CUSTOM_DOMAINS_RENEWAL_INTERVAL diff --git a/contrib/helm/charts/probo/values-production.yaml.example b/contrib/helm/charts/probo/values-production.yaml.example index 5978d1b86..3b2128d3e 100644 --- a/contrib/helm/charts/probo/values-production.yaml.example +++ b/contrib/helm/charts/probo/values-production.yaml.example @@ -176,6 +176,23 @@ probo: # temperature: "0.1" # maxTokens: "1024" + # Tracker mapping worker tuning (optional; seconds for interval/agentTimeout). + # Keep concurrency modest to stay under OpenAI/Firecrawl limits and the DB pool. + # trackerMappingWorker: + # interval: 10 + # maxConcurrency: 3 + # agentTimeout: 45 + # agentMaxTurns: 4 + + # Common-pattern enrichment worker tuning (optional; seconds for + # interval/staleAfter/agentTimeout). + # commonPatternEnrichmentWorker: + # interval: 10 + # maxConcurrency: 2 + # staleAfter: 600 + # agentTimeout: 45 + # agentMaxTurns: 3 + # OpenTelemetry tracing (optional) tracing: enabled: true diff --git a/contrib/helm/charts/probo/values.yaml b/contrib/helm/charts/probo/values.yaml index 55f50489b..6f23205d4 100644 --- a/contrib/helm/charts/probo/values.yaml +++ b/contrib/helm/charts/probo/values.yaml @@ -276,6 +276,24 @@ probo: temperature: "" maxTokens: "" + # Tracker mapping background worker tuning (optional). interval and + # agentTimeout are in seconds. Keep concurrency modest to stay under + # OpenAI/Firecrawl rate limits and the database connection pool. + trackerMappingWorker: + interval: 10 + maxConcurrency: 3 + agentTimeout: 45 + agentMaxTurns: 4 + + # Common-pattern enrichment background worker tuning (optional). + # interval, staleAfter, and agentTimeout are in seconds. + commonPatternEnrichmentWorker: + interval: 10 + maxConcurrency: 2 + staleAfter: 600 + agentTimeout: 45 + agentMaxTurns: 3 + # Custom domains configuration (optional) customDomains: enabled: false diff --git a/pkg/bootstrap/builder.go b/pkg/bootstrap/builder.go index 91fbda44b..017a6eb3f 100644 --- a/pkg/bootstrap/builder.go +++ b/pkg/bootstrap/builder.go @@ -207,10 +207,13 @@ func (b *Builder) Build() (*probodconfig.FullConfig, error) { MaxTokens: b.getEnvIntPtr("AGENT_EVIDENCE_DESCRIBER_MAX_TOKENS"), }, TrackerMapping: probodconfig.LLMAgentConfig{ - Provider: b.getEnvOrDefault("AGENT_TRACKER_MAPPING_PROVIDER", ""), - ModelName: b.getEnvOrDefault("AGENT_TRACKER_MAPPING_MODEL_NAME", ""), + Provider: b.getEnvOrDefault("AGENT_TRACKER_MAPPING_PROVIDER", ""), + ModelName: b.getEnvOrDefault("AGENT_TRACKER_MAPPING_MODEL_NAME", ""), + // The tracker agents emit tiny structured JSON, so + // they default to a smaller token budget than the + // shared default rather than inheriting it. Temperature: b.getEnvFloatPtr("AGENT_TRACKER_MAPPING_TEMPERATURE"), - MaxTokens: b.getEnvIntPtr("AGENT_TRACKER_MAPPING_MAX_TOKENS"), + MaxTokens: new(b.getEnvIntOrDefault("AGENT_TRACKER_MAPPING_MAX_TOKENS", 1024)), }, Tools: probodconfig.AgentToolsConfig{ FirecrawlAPIKey: b.getEnv("FIRECRAWL_API_KEY"), @@ -242,6 +245,19 @@ func (b *Builder) Build() (*probodconfig.FullConfig, error) { StaleAfter: b.getEnvIntOrDefault("EVIDENCE_DESCRIBER_STALE_AFTER", 300), MaxConcurrency: b.getEnvIntOrDefault("EVIDENCE_DESCRIBER_MAX_CONCURRENCY", 10), }, + TrackerMappingWorker: probodconfig.TrackerMappingWorkerConfig{ + Interval: b.getEnvIntOrDefault("TRACKER_MAPPING_INTERVAL", 10), + MaxConcurrency: b.getEnvIntOrDefault("TRACKER_MAPPING_MAX_CONCURRENCY", 3), + AgentTimeout: b.getEnvIntOrDefault("TRACKER_MAPPING_AGENT_TIMEOUT", 45), + AgentMaxTurns: b.getEnvIntOrDefault("TRACKER_MAPPING_AGENT_MAX_TURNS", 4), + }, + CommonPatternEnrichmentWorker: probodconfig.CommonPatternEnrichmentWorkerConfig{ + Interval: b.getEnvIntOrDefault("COMMON_PATTERN_ENRICHMENT_INTERVAL", 10), + MaxConcurrency: b.getEnvIntOrDefault("COMMON_PATTERN_ENRICHMENT_MAX_CONCURRENCY", 2), + StaleAfter: b.getEnvIntOrDefault("COMMON_PATTERN_ENRICHMENT_STALE_AFTER", 600), + AgentTimeout: b.getEnvIntOrDefault("COMMON_PATTERN_ENRICHMENT_AGENT_TIMEOUT", 45), + AgentMaxTurns: b.getEnvIntOrDefault("COMMON_PATTERN_ENRICHMENT_AGENT_MAX_TURNS", 3), + }, Branding: b.getEnvBoolOrDefault("BRANDING", true), }, } diff --git a/pkg/bootstrap/builder_test.go b/pkg/bootstrap/builder_test.go index 9df51f351..a096ad292 100644 --- a/pkg/bootstrap/builder_test.go +++ b/pkg/bootstrap/builder_test.go @@ -218,7 +218,18 @@ func TestBuilder_Build_Defaults(t *testing.T) { assert.Empty(t, cfg.Probod.Agents.TrackerMapping.Provider) assert.Empty(t, cfg.Probod.Agents.TrackerMapping.ModelName) assert.Nil(t, cfg.Probod.Agents.TrackerMapping.Temperature) - assert.Nil(t, cfg.Probod.Agents.TrackerMapping.MaxTokens) + assert.Equal(t, new(1024), cfg.Probod.Agents.TrackerMapping.MaxTokens) + + // Tracker worker tuning — defaults + assert.Equal(t, 10, cfg.Probod.TrackerMappingWorker.Interval) + assert.Equal(t, 3, cfg.Probod.TrackerMappingWorker.MaxConcurrency) + assert.Equal(t, 45, cfg.Probod.TrackerMappingWorker.AgentTimeout) + assert.Equal(t, 4, cfg.Probod.TrackerMappingWorker.AgentMaxTurns) + assert.Equal(t, 10, cfg.Probod.CommonPatternEnrichmentWorker.Interval) + assert.Equal(t, 2, cfg.Probod.CommonPatternEnrichmentWorker.MaxConcurrency) + assert.Equal(t, 600, cfg.Probod.CommonPatternEnrichmentWorker.StaleAfter) + assert.Equal(t, 45, cfg.Probod.CommonPatternEnrichmentWorker.AgentTimeout) + assert.Equal(t, 3, cfg.Probod.CommonPatternEnrichmentWorker.AgentMaxTurns) // Custom domains config assert.Equal(t, 3600, cfg.Probod.CustomDomains.RenewalInterval) @@ -313,6 +324,16 @@ func TestBuilder_Build_CustomValues(t *testing.T) { env["AGENT_TRACKER_MAPPING_MODEL_NAME"] = "gpt-4o-mini" env["AGENT_TRACKER_MAPPING_TEMPERATURE"] = "0.1" env["AGENT_TRACKER_MAPPING_MAX_TOKENS"] = "1024" + // Tracker worker tuning override + env["TRACKER_MAPPING_INTERVAL"] = "20" + env["TRACKER_MAPPING_MAX_CONCURRENCY"] = "5" + env["TRACKER_MAPPING_AGENT_TIMEOUT"] = "30" + env["TRACKER_MAPPING_AGENT_MAX_TURNS"] = "6" + env["COMMON_PATTERN_ENRICHMENT_INTERVAL"] = "15" + env["COMMON_PATTERN_ENRICHMENT_MAX_CONCURRENCY"] = "4" + env["COMMON_PATTERN_ENRICHMENT_STALE_AFTER"] = "900" + env["COMMON_PATTERN_ENRICHMENT_AGENT_TIMEOUT"] = "50" + env["COMMON_PATTERN_ENRICHMENT_AGENT_MAX_TURNS"] = "5" // Custom domains env["CUSTOM_DOMAINS_RESOLVER_ADDR"] = "1.1.1.1:53" env["ACME_ACCOUNT_KEY"] = "-----BEGIN EC PRIVATE KEY-----\ntest\n-----END EC PRIVATE KEY-----" @@ -404,6 +425,16 @@ func TestBuilder_Build_CustomValues(t *testing.T) { assert.Equal(t, "gpt-4o-mini", cfg.Probod.Agents.TrackerMapping.ModelName) assert.Equal(t, new(0.1), cfg.Probod.Agents.TrackerMapping.Temperature) assert.Equal(t, new(1024), cfg.Probod.Agents.TrackerMapping.MaxTokens) + // Tracker worker tuning — overrides + assert.Equal(t, 20, cfg.Probod.TrackerMappingWorker.Interval) + assert.Equal(t, 5, cfg.Probod.TrackerMappingWorker.MaxConcurrency) + assert.Equal(t, 30, cfg.Probod.TrackerMappingWorker.AgentTimeout) + assert.Equal(t, 6, cfg.Probod.TrackerMappingWorker.AgentMaxTurns) + assert.Equal(t, 15, cfg.Probod.CommonPatternEnrichmentWorker.Interval) + assert.Equal(t, 4, cfg.Probod.CommonPatternEnrichmentWorker.MaxConcurrency) + assert.Equal(t, 900, cfg.Probod.CommonPatternEnrichmentWorker.StaleAfter) + assert.Equal(t, 50, cfg.Probod.CommonPatternEnrichmentWorker.AgentTimeout) + assert.Equal(t, 5, cfg.Probod.CommonPatternEnrichmentWorker.AgentMaxTurns) // Custom domains assert.Equal(t, "1.1.1.1:53", cfg.Probod.CustomDomains.ResolverAddr) assert.Equal(t, "-----BEGIN EC PRIVATE KEY-----\ntest\n-----END EC PRIVATE KEY-----", cfg.Probod.CustomDomains.ACME.AccountKey) diff --git a/pkg/cookiebanner/common_pattern_enrichment_agent.go b/pkg/cookiebanner/common_pattern_enrichment_agent.go index a1d216637..ab997723f 100644 --- a/pkg/cookiebanner/common_pattern_enrichment_agent.go +++ b/pkg/cookiebanner/common_pattern_enrichment_agent.go @@ -53,16 +53,26 @@ func buildCommonPatternEnrichmentAgent( panic(fmt.Sprintf("cookiebanner: cannot build tracker enrichment output type: %s", err)) } - return agent.New( - "common-pattern-enrichment", - cfg.LLMClient, + maxTurns := cfg.EnrichmentMaxTurns + if maxTurns < 1 { + maxTurns = defaultEnrichmentMaxTurns + } + + opts := []agent.Option{ agent.WithInstructions(trackerEnrichmentPrompt), agent.WithModel(cfg.Model), agent.WithTools(tools...), agent.WithOutputType(outputType), - agent.WithMaxTurns(agentMaxTurns), + agent.WithMaxTurns(maxTurns), + agent.WithMaxTokens(resolveAgentMaxTokens(cfg.MaxTokens)), agent.WithLogger(logger), - ) + } + + if cfg.Temperature != nil { + opts = append(opts, agent.WithTemperature(*cfg.Temperature)) + } + + return agent.New("common-pattern-enrichment", cfg.LLMClient, opts...) } func buildEnrichmentPrompt(cp coredata.CommonTrackerPattern, thirdPartyName string) string { diff --git a/pkg/cookiebanner/common_pattern_enrichment_worker.go b/pkg/cookiebanner/common_pattern_enrichment_worker.go index 317d69343..83dc1586d 100644 --- a/pkg/cookiebanner/common_pattern_enrichment_worker.go +++ b/pkg/cookiebanner/common_pattern_enrichment_worker.go @@ -29,13 +29,14 @@ import ( "go.probo.inc/probo/pkg/llm" ) -const enrichmentStaleAfter = 10 * time.Minute +const defaultEnrichmentStaleAfter = 10 * time.Minute type commonPatternEnrichmentHandler struct { pg *pg.Client logger *log.Logger enrichmentAgent *agent.Agent staleAfter time.Duration + agentTimeout time.Duration } // NewCommonPatternEnrichmentWorker builds the worker that fills @@ -49,12 +50,23 @@ func NewCommonPatternEnrichmentWorker( pgClient *pg.Client, logger *log.Logger, cfg TrackerAgentsConfig, + staleAfter time.Duration, opts ...worker.Option, ) *worker.Worker[coredata.CommonTrackerPattern] { + if staleAfter <= 0 { + staleAfter = defaultEnrichmentStaleAfter + } + + agentTimeout := cfg.AgentTimeout + if agentTimeout <= 0 { + agentTimeout = defaultAgentTimeout + } + h := &commonPatternEnrichmentHandler{ - pg: pgClient, - logger: logger, - staleAfter: enrichmentStaleAfter, + pg: pgClient, + logger: logger, + staleAfter: staleAfter, + agentTimeout: agentTimeout, } if cfg.LLMClient != nil { @@ -187,7 +199,7 @@ func (h *commonPatternEnrichmentHandler) research( ) (string, error) { prompt := buildEnrichmentPrompt(cp, thirdPartyName) - agentCtx, cancel := context.WithTimeout(ctx, agentTimeout) + agentCtx, cancel := context.WithTimeout(ctx, h.agentTimeout) defer cancel() result, err := agent.RunTyped[CommonPatternEnrichmentResult]( diff --git a/pkg/cookiebanner/tracker_agents_config.go b/pkg/cookiebanner/tracker_agents_config.go index ee1fbdff9..519700191 100644 --- a/pkg/cookiebanner/tracker_agents_config.go +++ b/pkg/cookiebanner/tracker_agents_config.go @@ -14,15 +14,29 @@ package cookiebanner -import "go.probo.inc/probo/pkg/llm" +import ( + "time" + + "go.probo.inc/probo/pkg/llm" +) // TrackerAgentsConfig configures the tracker agents that share one LLM // client, model, and tool surface: the tracker-mapping agent (catalog // identification) and the common-pattern enrichment agent (description // research). Both use DB-backed search tools and may also use Firecrawl // for web search when an API key is supplied. +// +// MaxTokens and Temperature bound and steer each LLM call (both +// outputs are tiny structured JSON). AgentTimeout caps a single agent +// run, and the per-worker max-turns bound the agent reasoning loop. +// Zero-valued tuning fields fall back to package defaults. type TrackerAgentsConfig struct { - LLMClient *llm.Client - Model string - FirecrawlAPIKey string + LLMClient *llm.Client + Model string + FirecrawlAPIKey string + MaxTokens *int + Temperature *float64 + AgentTimeout time.Duration + MappingMaxTurns int + EnrichmentMaxTurns int } diff --git a/pkg/cookiebanner/tracker_mapping_agent.go b/pkg/cookiebanner/tracker_mapping_agent.go index 18d9d0162..90b366d21 100644 --- a/pkg/cookiebanner/tracker_mapping_agent.go +++ b/pkg/cookiebanner/tracker_mapping_agent.go @@ -29,8 +29,22 @@ import ( ) const ( - agentTimeout = 60 * time.Second - agentMaxTurns = 5 + // defaultAgentTimeout caps a single mapping or enrichment agent run + // when the worker config does not supply one. It guards against a + // hung LLM provider or a slow web search. + defaultAgentTimeout = 45 * time.Second + + // defaultMappingMaxTurns and defaultEnrichmentMaxTurns bound the + // agent reasoning loop (LLM call + tool round-trips) when the worker + // config does not supply a value. + defaultMappingMaxTurns = 4 + defaultEnrichmentMaxTurns = 3 + + // defaultAgentMaxTokens caps the structured output of the mapping + // and enrichment agents when the agent config carries no max-tokens + // budget. Both outputs are tiny structured JSON. + defaultAgentMaxTokens = 1024 + agentThirdPartyConfidenceThreshold = 0.6 // agentSourceConfidence is the fixed confidence stored on catalog // rows the agent attributes to a third party. The agent's own @@ -70,16 +84,37 @@ func buildTrackerMappingAgent( panic(fmt.Sprintf("cookiebanner: cannot build tracker identification output type: %s", err)) } - return agent.New( - "tracker-mapping", - cfg.LLMClient, + maxTurns := cfg.MappingMaxTurns + if maxTurns < 1 { + maxTurns = defaultMappingMaxTurns + } + + opts := []agent.Option{ agent.WithInstructionsFunc(trackerMappingInstructions), agent.WithModel(cfg.Model), agent.WithTools(tools...), agent.WithOutputType(outputType), - agent.WithMaxTurns(agentMaxTurns), + agent.WithMaxTurns(maxTurns), + agent.WithMaxTokens(resolveAgentMaxTokens(cfg.MaxTokens)), agent.WithLogger(logger), - ) + } + + if cfg.Temperature != nil { + opts = append(opts, agent.WithTemperature(*cfg.Temperature)) + } + + return agent.New("tracker-mapping", cfg.LLMClient, opts...) +} + +// resolveAgentMaxTokens returns the configured max-tokens budget for the +// mapping and enrichment agents, falling back to defaultAgentMaxTokens +// when none is set. +func resolveAgentMaxTokens(configured *int) int { + if configured != nil && *configured > 0 { + return *configured + } + + return defaultAgentMaxTokens } func trackerMappingInstructions(_ context.Context, _ *agent.Agent) string { diff --git a/pkg/cookiebanner/tracker_mapping_worker.go b/pkg/cookiebanner/tracker_mapping_worker.go index ecc7895c8..11df9292e 100644 --- a/pkg/cookiebanner/tracker_mapping_worker.go +++ b/pkg/cookiebanner/tracker_mapping_worker.go @@ -33,10 +33,12 @@ import ( ) type trackerMappingHandler struct { - pg *pg.Client - logger *log.Logger - mappingAgent *agent.Agent - disambiguationAgent *agent.Agent + pg *pg.Client + logger *log.Logger + mappingAgent *agent.Agent + disambiguationAgent *agent.Agent + agentTimeout time.Duration + disambiguationTimeout time.Duration } func NewTrackerMappingWorker( @@ -46,9 +48,16 @@ func NewTrackerMappingWorker( disambiguationCfg thirdparty.DisambiguationConfig, opts ...worker.Option, ) *worker.Worker[coredata.TrackerPattern] { + agentTimeout := mappingCfg.AgentTimeout + if agentTimeout <= 0 { + agentTimeout = defaultAgentTimeout + } + h := &trackerMappingHandler{ - pg: pgClient, - logger: logger, + pg: pgClient, + logger: logger, + agentTimeout: agentTimeout, + disambiguationTimeout: disambiguationCfg.Timeout, } if mappingCfg.LLMClient != nil { @@ -121,110 +130,86 @@ type catalogMatch struct { // brand new org ThirdParty stays gated behind categorisation and a // non-extension source. func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.TrackerPattern) error { + scope := coredata.NewScopeFromObjectID(tp.ID) + + // Phase 1: deterministic catalog resolution in a short transaction. + // The existing-link, pattern, sibling, and domain signals (and their + // idempotent upserts) run here. No LLM or web-search call is made + // while the transaction — and its FOR UPDATE row lock — is held. + var det deterministicResult + + if err := h.pg.WithTx( + ctx, + func(ctx context.Context, tx pg.Tx) error { + var err error + + det, err = h.resolveDeterministic(ctx, tx, tp) + + return err + }, + ); err != nil { + return err + } + + commonPatternID := det.commonPatternID + commonThirdPartyID := det.commonThirdPartyID + directThirdPartyID := det.directThirdPartyID + + // Phase 2: tracker-mapping agent (no transaction). It runs only when + // the deterministic signals could not resolve a catalog third party. + // The LLM and web-search calls happen outside any transaction; the + // result is persisted in its own short transaction. + if commonThirdPartyID == nil && h.mappingAgent != nil { + ident, err := h.identifyWithAgent(ctx, tp, det.origin) + if err != nil { + return fmt.Errorf("cannot identify with agent: %w", err) + } + + if ident != nil { + if err := h.pg.WithTx( + ctx, + func(ctx context.Context, tx pg.Tx) error { + match, err := h.persistAgentIdentification(ctx, tx, tp, *ident) + if err != nil { + return err + } + + commonPatternID = firstNonNil(commonPatternID, match.commonPatternID) + commonThirdPartyID = match.commonThirdPartyID + + return nil + }, + ); err != nil { + return err + } + } + } + + // Phase 3: org ThirdParty resolution. The heuristic ranking and the + // disambiguation agent run without a transaction; only the final link + // or create touches the database (in a short transaction). + thirdPartyID := tp.ThirdPartyID + + if thirdPartyID == nil { + switch { + case directThirdPartyID != nil: + thirdPartyID = directThirdPartyID + case commonThirdPartyID != nil: + resolved, err := h.resolveOrgThirdParty(ctx, tp, *commonThirdPartyID) + if err != nil { + return fmt.Errorf("cannot resolve org third party: %w", err) + } + + thirdPartyID = resolved + } + } + + // Phase 4: persist the pattern mapping in a short transaction. The + // unmatched fallback keeps catalog coverage complete even when no + // vendor was resolved. return h.pg.WithTx( ctx, func(ctx context.Context, tx pg.Tx) error { - scope := coredata.NewScopeFromObjectID(tp.ID) - - var banner coredata.CookieBanner - if err := banner.LoadByID(ctx, tx, scope, tp.CookieBannerID); err != nil { - return fmt.Errorf("cannot load cookie banner for domain filtering: %w", err) - } - - var ( - commonPatternID *gid.GID - commonThirdPartyID *gid.GID - directThirdPartyID *gid.GID - ) - - if tp.CommonTrackerPatternID != nil { - commonPatternID = tp.CommonTrackerPatternID - - var commonPattern coredata.CommonTrackerPattern - if err := commonPattern.LoadByID(ctx, tx, *commonPatternID); err != nil { - return fmt.Errorf("cannot load linked common tracker pattern: %w", err) - } - - commonThirdPartyID = commonPattern.CommonThirdPartyID - } else { - match, err := h.matchByPattern(ctx, tx, tp) - if err != nil { - return fmt.Errorf("cannot match by pattern: %w", err) - } - - if match != nil { - commonPatternID = match.commonPatternID - commonThirdPartyID = match.commonThirdPartyID - } - } - - // Whether a catalog third party was already known before the - // signal pipeline ran this round. Re-enqueuing siblings is - // only useful when this run is the one that resolves the - // vendor; a pre-existing link adds no new signal and gating - // on it keeps cascades finite. - commonThirdPartyPreexisted := commonThirdPartyID != nil - - var domains []string - - if commonThirdPartyID == nil { - loaded, err := h.loadInitiatorDomains(ctx, tx, tp) - if err != nil { - return err - } - - domains = loaded - - // Sibling matching is an org-local co-occurrence signal: - // two patterns served from the same origin on the same - // banner are likely the same vendor, even when that origin - // is the site's own (first-party) host — a tracker proxied - // through first-party still co-occurs with its siblings. - // So it intentionally runs on the unfiltered domains; the - // ambiguity guard in resolveThirdPartyFromSiblings prevents - // grouping unrelated first-party scripts. - match, err := h.matchBySiblingOrigin(ctx, tx, tp, domains) - if err != nil { - return fmt.Errorf("cannot match by sibling origin: %w", err) - } - - if match != nil { - commonPatternID = firstNonNil(commonPatternID, match.commonPatternID) - commonThirdPartyID = match.commonThirdPartyID - directThirdPartyID = match.thirdPartyID - } - - if commonThirdPartyID == nil { - // Domain matching hits the global catalog, so - // first-party domains must be stripped: a tracker - // proxied through the site's own host would otherwise - // match the site owner's own CommonThirdParty entry. - catalogDomains := uri.FilterFirstPartyDomains(domains, banner.Origin) - - match, err := h.matchByDomain(ctx, tx, tp, catalogDomains) - if err != nil { - return fmt.Errorf("cannot match by domain: %w", err) - } - - if match != nil { - commonPatternID = firstNonNil(commonPatternID, match.commonPatternID) - commonThirdPartyID = match.commonThirdPartyID - } - } - - if commonThirdPartyID == nil && h.mappingAgent != nil { - match, err := h.identifyWithAgent(ctx, tx, tp, banner.Origin) - if err != nil { - return fmt.Errorf("cannot identify with agent: %w", err) - } - - if match != nil { - commonPatternID = firstNonNil(commonPatternID, match.commonPatternID) - commonThirdPartyID = match.commonThirdPartyID - } - } - } - if commonPatternID == nil { id, err := h.createUnmatchedPattern(ctx, tx, tp) if err != nil { @@ -234,64 +219,41 @@ func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.Tracker commonPatternID = id } - thirdPartyID := tp.ThirdPartyID + tp.CommonTrackerPatternID = commonPatternID + tp.ThirdPartyID = thirdPartyID + tp.UpdatedAt = time.Now() - if thirdPartyID == nil { - switch { - case directThirdPartyID != nil: - thirdPartyID = directThirdPartyID - case commonThirdPartyID != nil: - allowCreate, err := h.creationAllowed(ctx, tx, scope, tp) - if err != nil { - return err - } - - resolved, err := h.resolveOrgThirdParty(ctx, tx, tp, *commonThirdPartyID, allowCreate) - if err != nil { - return fmt.Errorf("cannot resolve org third party: %w", err) - } - - thirdPartyID = resolved + // Descriptions are owned by the common-pattern enrichment + // worker. Here we only propagate: if the linked catalog row + // is already enriched, copy its description onto this + // pattern. A pattern linked before enrichment is filled + // later by the enrichment worker's fan-out instead. + if commonPatternID != nil && tp.Description == "" { + var commonPattern coredata.CommonTrackerPattern + if err := commonPattern.LoadByID(ctx, tx, *commonPatternID); err == nil && commonPattern.Description != "" { + tp.Description = commonPattern.Description } } - if commonPatternID != nil || thirdPartyID != nil { - tp.CommonTrackerPatternID = commonPatternID - tp.ThirdPartyID = thirdPartyID - tp.UpdatedAt = time.Now() + if err := tp.UpdateMapping(ctx, tx, scope); err != nil { + return fmt.Errorf("cannot update tracker pattern mapping: %w", err) + } - // Descriptions are owned by the common-pattern enrichment - // worker. Here we only propagate: if the linked catalog row - // is already enriched, copy its description onto this - // pattern. A pattern linked before enrichment is filled - // later by the enrichment worker's fan-out instead. - if tp.Description == "" && commonPatternID != nil { - var commonPattern coredata.CommonTrackerPattern - if err := commonPattern.LoadByID(ctx, tx, *commonPatternID); err == nil && commonPattern.Description != "" { - tp.Description = commonPattern.Description - } - } + h.logger.InfoCtx( + ctx, + "mapped tracker pattern", + log.String("pattern", tp.Pattern), + log.String("tracker_pattern_id", tp.ID.String()), + ) - if err := tp.UpdateMapping(ctx, tx, scope); err != nil { - return fmt.Errorf("cannot update tracker pattern mapping: %w", err) - } - - h.logger.InfoCtx( - ctx, - "mapped tracker pattern", - log.String("pattern", tp.Pattern), - log.String("tracker_pattern_id", tp.ID.String()), - ) - - // This run newly resolved a catalog third party, so - // same-banner siblings that share an initiator domain but - // were processed earlier and left unmatched can now match - // against it. Re-arm their mapping so the worker revisits - // them; the guards keep already-mapped siblings untouched. - if commonThirdPartyID != nil && !commonThirdPartyPreexisted { - if err := h.reenqueueUnmappedSiblings(ctx, tx, tp, domains); err != nil { - return err - } + // This run newly resolved a catalog third party, so + // same-banner siblings that share an initiator domain but + // were processed earlier and left unmatched can now match + // against it. Re-arm their mapping so the worker revisits + // them; the guards keep already-mapped siblings untouched. + if commonThirdPartyID != nil && !det.commonThirdPartyPreexisted { + if err := h.reenqueueUnmappedSiblings(ctx, tx, tp, det.domains); err != nil { + return err } } @@ -300,6 +262,118 @@ func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.Tracker ) } +// deterministicResult carries the outcome of the pure-SQL catalog +// signals (existing link, pattern, sibling origin, domain overlap) from +// the read phase to the agent and persist phases. domains holds the +// unfiltered initiator domains observed for the pattern (used by the +// sibling re-enqueue cascade); commonThirdPartyPreexisted records +// whether a catalog third party was already known before this run, so +// the cascade only fires when this run is the one that resolved it. +type deterministicResult struct { + origin string + commonPatternID *gid.GID + commonThirdPartyID *gid.GID + directThirdPartyID *gid.GID + domains []string + commonThirdPartyPreexisted bool +} + +// resolveDeterministic runs the catalog signals that need no network +// call (existing link, pattern, sibling origin, domain overlap) inside a +// single short transaction and reports what they resolved. It never +// invokes the mapping agent; the caller runs that outside any +// transaction. +func (h *trackerMappingHandler) resolveDeterministic( + ctx context.Context, + tx pg.Tx, + tp coredata.TrackerPattern, +) (deterministicResult, error) { + scope := coredata.NewScopeFromObjectID(tp.ID) + + var res deterministicResult + + var banner coredata.CookieBanner + if err := banner.LoadByID(ctx, tx, scope, tp.CookieBannerID); err != nil { + return res, fmt.Errorf("cannot load cookie banner for domain filtering: %w", err) + } + + res.origin = banner.Origin + + if tp.CommonTrackerPatternID != nil { + res.commonPatternID = tp.CommonTrackerPatternID + + var commonPattern coredata.CommonTrackerPattern + if err := commonPattern.LoadByID(ctx, tx, *res.commonPatternID); err != nil { + return res, fmt.Errorf("cannot load linked common tracker pattern: %w", err) + } + + res.commonThirdPartyID = commonPattern.CommonThirdPartyID + } else { + match, err := h.matchByPattern(ctx, tx, tp) + if err != nil { + return res, fmt.Errorf("cannot match by pattern: %w", err) + } + + if match != nil { + res.commonPatternID = match.commonPatternID + res.commonThirdPartyID = match.commonThirdPartyID + } + } + + res.commonThirdPartyPreexisted = res.commonThirdPartyID != nil + + if res.commonThirdPartyID != nil { + return res, nil + } + + loaded, err := h.loadInitiatorDomains(ctx, tx, tp) + if err != nil { + return res, err + } + + res.domains = loaded + + // Sibling matching is an org-local co-occurrence signal: two + // patterns served from the same origin on the same banner are likely + // the same vendor, even when that origin is the site's own + // (first-party) host — a tracker proxied through first-party still + // co-occurs with its siblings. So it intentionally runs on the + // unfiltered domains; the ambiguity guard in + // resolveThirdPartyFromSiblings prevents grouping unrelated + // first-party scripts. + siblingMatch, err := h.matchBySiblingOrigin(ctx, tx, tp, res.domains) + if err != nil { + return res, fmt.Errorf("cannot match by sibling origin: %w", err) + } + + if siblingMatch != nil { + res.commonPatternID = firstNonNil(res.commonPatternID, siblingMatch.commonPatternID) + res.commonThirdPartyID = siblingMatch.commonThirdPartyID + res.directThirdPartyID = siblingMatch.thirdPartyID + } + + if res.commonThirdPartyID != nil { + return res, nil + } + + // Domain matching hits the global catalog, so first-party domains + // must be stripped: a tracker proxied through the site's own host + // would otherwise match the site owner's own CommonThirdParty entry. + catalogDomains := uri.FilterFirstPartyDomains(res.domains, banner.Origin) + + domainMatch, err := h.matchByDomain(ctx, tx, tp, catalogDomains) + if err != nil { + return res, fmt.Errorf("cannot match by domain: %w", err) + } + + if domainMatch != nil { + res.commonPatternID = firstNonNil(res.commonPatternID, domainMatch.commonPatternID) + res.commonThirdPartyID = domainMatch.commonThirdPartyID + } + + return res, nil +} + // reenqueueUnmappedSiblings re-arms mapping_requested_at on same-banner // siblings sharing an initiator domain with tp that are still unpromoted, // so the worker re-evaluates them now that tp resolved a vendor. @@ -472,16 +546,41 @@ func (h *trackerMappingHandler) matchByDomain( }, nil } +// agentIdentification carries a confident tracker-mapping agent result +// and the (first-party-filtered) domains it observed, from the no-tx +// agent phase to the short transaction that persists it. +type agentIdentification struct { + result TrackerMappingAgentResult + domains []string +} + +// identifyWithAgent runs the tracker-mapping agent outside any +// transaction. It loads the observed initiator domains with a +// short-lived connection, calls the LLM (and any web-search tool), and +// returns a confident identification or nil. It performs no writes; the +// caller persists the result via persistAgentIdentification. func (h *trackerMappingHandler) identifyWithAgent( ctx context.Context, - tx pg.Tx, tp coredata.TrackerPattern, siteOrigin string, -) (*catalogMatch, error) { - var trackers coredata.DetectedTrackers +) (*agentIdentification, error) { + var domains []string - domains, err := trackers.LoadInitiatorDomainsByTrackerPatternID(ctx, tx, tp.ID, 5) - if err != nil { + if err := h.pg.WithConn( + ctx, + func(ctx context.Context, conn pg.Querier) error { + var trackers coredata.DetectedTrackers + + loaded, err := trackers.LoadInitiatorDomainsByTrackerPatternID(ctx, conn, tp.ID, 5) + if err != nil { + return err + } + + domains = loaded + + return nil + }, + ); err != nil { h.logger.WarnCtx(ctx, "cannot load initiator domains for agent", log.Error(err)) } @@ -489,7 +588,7 @@ func (h *trackerMappingHandler) identifyWithAgent( prompt := buildAgentPrompt(tp, domains) - agentCtx, cancel := context.WithTimeout(ctx, agentTimeout) + agentCtx, cancel := context.WithTimeout(ctx, h.agentTimeout) defer cancel() result, err := agent.RunTyped[TrackerMappingAgentResult]( @@ -530,11 +629,27 @@ func (h *trackerMappingHandler) identifyWithAgent( return nil, nil } + return &agentIdentification{ + result: identification, + domains: domains, + }, nil +} + +// persistAgentIdentification writes a confident agent identification: +// it resolves or creates the catalog third party and upserts the +// catalog pattern row that links to it. It runs inside the caller's +// short transaction. +func (h *trackerMappingHandler) persistAgentIdentification( + ctx context.Context, + tx pg.Tx, + tp coredata.TrackerPattern, + ident agentIdentification, +) (*catalogMatch, error) { commonThirdPartyID, err := h.resolveOrCreateCommonThirdParty( ctx, tx, - identification, - domains, + ident.result, + ident.domains, ) if err != nil { return nil, fmt.Errorf("cannot resolve or create common third party: %w", err) @@ -561,8 +676,8 @@ func (h *trackerMappingHandler) identifyWithAgent( ctx, "agent identified tracker pattern", log.String("pattern", tp.Pattern), - log.String("third_party", identification.ThirdPartyName), - log.Float64("third_party_confidence", identification.ThirdPartyConfidence), + log.String("third_party", ident.result.ThirdPartyName), + log.Float64("third_party_confidence", ident.result.ThirdPartyConfidence), ) return &catalogMatch{ @@ -851,88 +966,46 @@ func (h *trackerMappingHandler) createUnmatchedPattern( // path in O(1). func (h *trackerMappingHandler) resolveOrgThirdParty( ctx context.Context, - tx pg.Tx, tp coredata.TrackerPattern, commonThirdPartyID gid.GID, - allowCreate bool, ) (*gid.GID, error) { scope := coredata.NewScopeFromObjectID(tp.ID) - var existing coredata.ThirdParty + // Read phase: exact link, candidate ranking, eligibility, and + // creation gating. No write or LLM call happens here. + var prep orgThirdPartyPrep - err := existing.LoadByOrganizationIDAndCommonThirdPartyID( + if err := h.pg.WithConn( ctx, - tx, - scope, - tp.OrganizationID, - commonThirdPartyID, - ) - if err == nil { - return &existing.ID, nil + func(ctx context.Context, conn pg.Querier) error { + var err error + + prep, err = h.prepareOrgThirdParty(ctx, conn, scope, tp, commonThirdPartyID) + + return err + }, + ); err != nil { + return nil, err } - if !errors.Is(err, coredata.ErrResourceNotFound) { - return nil, fmt.Errorf("cannot load org third party by common id: %w", err) + if prep.existingID != nil { + return prep.existingID, nil } - var commonParty coredata.CommonThirdParty - if err := commonParty.LoadByID(ctx, tx, commonThirdPartyID); err != nil { - return nil, fmt.Errorf("cannot load common third party: %w", err) - } + picked := prep.highConfidence + viaAgent := false - var commonDomains coredata.CommonThirdPartyDomains - if err := commonDomains.LoadByCommonThirdPartyID(ctx, tx, commonThirdPartyID); err != nil { - return nil, fmt.Errorf("cannot load common third party domains: %w", err) - } - - var orgThirdParties coredata.ThirdParties - if err := orgThirdParties.LoadAllByOrganizationID(ctx, tx, scope, tp.OrganizationID); err != nil { - return nil, fmt.Errorf("cannot load org third parties: %w", err) - } - - ranked := thirdparty.RankCandidates(commonParty, commonDomains, orgThirdParties) - - if len(ranked) > 0 && ranked[0].Score >= thirdparty.HighConfidenceScore { - picked := ranked[0].ThirdParty - - if err := thirdparty.LinkToCommon(ctx, tx, scope, picked, commonThirdPartyID); err != nil { - return nil, fmt.Errorf("cannot link fuzzy-matched third party to common: %w", err) - } - - h.logger.InfoCtx( - ctx, - "promoted tracker pattern via heuristic match", - log.String("tracker_pattern_id", tp.ID.String()), - log.String("third_party_id", picked.ID.String()), - log.Float64("score", ranked[0].Score), - ) - - return &picked.ID, nil - } - - agentSet := ranked - if len(agentSet) > thirdparty.MaxAgentCandidates { - agentSet = agentSet[:thirdparty.MaxAgentCandidates] - } - - eligibleForAgent := false - - for _, c := range agentSet { - if c.Score >= thirdparty.MinAgentScore { - eligibleForAgent = true - - break - } - } - - if eligibleForAgent && h.disambiguationAgent != nil { + // Agent phase (no transaction): disambiguate among the heuristic + // candidates when none scored high enough on its own. + if picked == nil && prep.eligibleForAgent && h.disambiguationAgent != nil { matchedID, err := thirdparty.Disambiguate( ctx, h.disambiguationAgent, h.logger, - commonParty, - commonDomains, - agentSet, + prep.commonParty, + prep.commonDomains, + prep.agentSet, + h.disambiguationTimeout, ) if err != nil { h.logger.WarnCtx( @@ -944,49 +1017,171 @@ func (h *trackerMappingHandler) resolveOrgThirdParty( } if matchedID != nil { - var picked *coredata.ThirdParty - - for _, c := range agentSet { + for _, c := range prep.agentSet { if c.ThirdParty.ID == *matchedID { picked = c.ThirdParty + viaAgent = true break } } + } + } + // Nothing to link and creation is not allowed: leave the pattern + // without an org third party. + if picked == nil && !prep.allowCreate { + return nil, nil + } + + // Write phase: link the picked candidate or create a new org third + // party from the catalog entry, in a short transaction. + var resolved *gid.GID + + if err := h.pg.WithTx( + ctx, + func(ctx context.Context, tx pg.Tx) error { if picked != nil { if err := thirdparty.LinkToCommon(ctx, tx, scope, picked, commonThirdPartyID); err != nil { - return nil, fmt.Errorf("cannot link agent-matched third party to common: %w", err) + return fmt.Errorf("cannot link third party to common: %w", err) } - h.logger.InfoCtx( - ctx, - "promoted tracker pattern via disambiguation agent", - log.String("tracker_pattern_id", tp.ID.String()), - log.String("third_party_id", picked.ID.String()), - ) + if viaAgent { + h.logger.InfoCtx( + ctx, + "promoted tracker pattern via disambiguation agent", + log.String("tracker_pattern_id", tp.ID.String()), + log.String("third_party_id", picked.ID.String()), + ) + } else { + h.logger.InfoCtx( + ctx, + "promoted tracker pattern via heuristic match", + log.String("tracker_pattern_id", tp.ID.String()), + log.String("third_party_id", picked.ID.String()), + log.Float64("score", prep.highScore), + ) + } - return &picked.ID, nil + resolved = &picked.ID + + return nil + } + + created, err := thirdparty.CreateFromCommon(ctx, tx, scope, tp.OrganizationID, prep.commonParty) + if err != nil { + return fmt.Errorf("cannot create third party from common: %w", err) + } + + h.logger.InfoCtx( + ctx, + "promoted tracker pattern by creating org third party from catalog", + log.String("tracker_pattern_id", tp.ID.String()), + log.String("third_party_id", created.ID.String()), + log.String("common_third_party_id", commonThirdPartyID.String()), + ) + + resolved = &created.ID + + return nil + }, + ); err != nil { + return nil, err + } + + return resolved, nil +} + +// orgThirdPartyPrep is the read-phase outcome for org ThirdParty +// resolution. existingID is set when an exact common-id link already +// exists (the other fields are then unused). Otherwise highConfidence +// holds a heuristic match at or above HighConfidenceScore (with +// highScore), or agentSet/eligibleForAgent describe the disambiguation +// candidates. allowCreate gates falling back to creating a new org +// ThirdParty from the catalog entry. +type orgThirdPartyPrep struct { + existingID *gid.GID + commonParty coredata.CommonThirdParty + commonDomains coredata.CommonThirdPartyDomains + agentSet []thirdparty.ScoredCandidate + highConfidence *coredata.ThirdParty + highScore float64 + eligibleForAgent bool + allowCreate bool +} + +// prepareOrgThirdParty performs the read-only work for org ThirdParty +// resolution: it checks for an exact common-id link, loads the catalog +// entry and the org's existing third parties, ranks the candidates, and +// computes creation eligibility. It makes no writes and no LLM call. +func (h *trackerMappingHandler) prepareOrgThirdParty( + ctx context.Context, + conn pg.Querier, + scope coredata.Scoper, + tp coredata.TrackerPattern, + commonThirdPartyID gid.GID, +) (orgThirdPartyPrep, error) { + var prep orgThirdPartyPrep + + var existing coredata.ThirdParty + + err := existing.LoadByOrganizationIDAndCommonThirdPartyID( + ctx, + conn, + scope, + tp.OrganizationID, + commonThirdPartyID, + ) + if err == nil { + id := existing.ID + prep.existingID = &id + + return prep, nil + } + + if !errors.Is(err, coredata.ErrResourceNotFound) { + return prep, fmt.Errorf("cannot load org third party by common id: %w", err) + } + + if err := prep.commonParty.LoadByID(ctx, conn, commonThirdPartyID); err != nil { + return prep, fmt.Errorf("cannot load common third party: %w", err) + } + + if err := prep.commonDomains.LoadByCommonThirdPartyID(ctx, conn, commonThirdPartyID); err != nil { + return prep, fmt.Errorf("cannot load common third party domains: %w", err) + } + + var orgThirdParties coredata.ThirdParties + if err := orgThirdParties.LoadAllByOrganizationID(ctx, conn, scope, tp.OrganizationID); err != nil { + return prep, fmt.Errorf("cannot load org third parties: %w", err) + } + + ranked := thirdparty.RankCandidates(prep.commonParty, prep.commonDomains, orgThirdParties) + + if len(ranked) > 0 && ranked[0].Score >= thirdparty.HighConfidenceScore { + prep.highConfidence = ranked[0].ThirdParty + prep.highScore = ranked[0].Score + } else { + prep.agentSet = ranked + if len(prep.agentSet) > thirdparty.MaxAgentCandidates { + prep.agentSet = prep.agentSet[:thirdparty.MaxAgentCandidates] + } + + for _, c := range prep.agentSet { + if c.Score >= thirdparty.MinAgentScore { + prep.eligibleForAgent = true + + break } } } - if !allowCreate { - return nil, nil - } - - created, err := thirdparty.CreateFromCommon(ctx, tx, scope, tp.OrganizationID, commonParty) + allowCreate, err := h.creationAllowed(ctx, conn, scope, tp) if err != nil { - return nil, fmt.Errorf("cannot create third party from common: %w", err) + return prep, err } - h.logger.InfoCtx( - ctx, - "promoted tracker pattern by creating org third party from catalog", - log.String("tracker_pattern_id", tp.ID.String()), - log.String("third_party_id", created.ID.String()), - log.String("common_third_party_id", commonThirdPartyID.String()), - ) + prep.allowCreate = allowCreate - return &created.ID, nil + return prep, nil } diff --git a/pkg/cookiebanner/tracker_mapping_worker_test.go b/pkg/cookiebanner/tracker_mapping_worker_test.go index 2eba9bd61..78709ecf8 100644 --- a/pkg/cookiebanner/tracker_mapping_worker_test.go +++ b/pkg/cookiebanner/tracker_mapping_worker_test.go @@ -143,28 +143,20 @@ func newMappingHandler(client *pg.Client) *trackerMappingHandler { } } -// promote runs resolveOrgThirdParty inside its own transaction so each -// test case starts from a clean state. +// promote runs resolveOrgThirdParty, which manages its own short +// transactions internally (creation gating is derived from the +// pattern's category, not passed in). func promote( t *testing.T, ctx context.Context, h *trackerMappingHandler, - client *pg.Client, tp coredata.TrackerPattern, commonThirdPartyID gid.GID, - allowCreate bool, ) *gid.GID { t.Helper() - var got *gid.GID - - require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { - var err error - - got, err = h.resolveOrgThirdParty(ctx, tx, tp, commonThirdPartyID, allowCreate) - - return err - })) + got, err := h.resolveOrgThirdParty(ctx, tp, commonThirdPartyID) + require.NoError(t, err) return got } @@ -193,7 +185,7 @@ func TestPromoteThirdParty_ExactCommonLink(t *testing.T) { return existing.Insert(ctx, tx, fx.scope) })) - got := promote(t, ctx, newMappingHandler(client), client, fx.trackerPattern, fx.commonThirdPartyID, true) + got := promote(t, ctx, newMappingHandler(client), fx.trackerPattern, fx.commonThirdPartyID) require.NotNil(t, got) assert.Equal(t, existing.ID, *got, "should return the existing org ThirdParty linked by common id") @@ -225,7 +217,7 @@ func TestPromoteThirdParty_HeuristicMatch(t *testing.T) { return manualEntry.Insert(ctx, tx, fx.scope) })) - got := promote(t, ctx, newMappingHandler(client), client, fx.trackerPattern, fx.commonThirdPartyID, true) + got := promote(t, ctx, newMappingHandler(client), fx.trackerPattern, fx.commonThirdPartyID) require.NotNil(t, got) assert.Equal(t, manualEntry.ID, *got, "heuristic match should return the manually-entered ThirdParty") @@ -247,7 +239,7 @@ func TestPromoteThirdParty_FallbackCreate(t *testing.T) { ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) - got := promote(t, ctx, newMappingHandler(client), client, fx.trackerPattern, fx.commonThirdPartyID, true) + got := promote(t, ctx, newMappingHandler(client), fx.trackerPattern, fx.commonThirdPartyID) require.NotNil(t, got, "fallback should create a new ThirdParty") @@ -268,7 +260,8 @@ func TestPromoteThirdParty_FallbackCreate(t *testing.T) { // TestResolveOrgThirdParty_CreationGated asserts that when no existing // org ThirdParty matches the catalog third party, creating a new one is -// suppressed unless allowCreate is true. +// suppressed for an uncategorised pattern (creation gating is derived +// from the pattern's category) and proceeds for a categorised one. func TestResolveOrgThirdParty_CreationGated(t *testing.T) { t.Parallel() @@ -276,11 +269,14 @@ func TestResolveOrgThirdParty_CreationGated(t *testing.T) { ctx := context.Background() fx := seedPromotionFixture(t, ctx, client) - gated := promote(t, ctx, newMappingHandler(client), client, fx.trackerPattern, fx.commonThirdPartyID, false) - assert.Nil(t, gated, "creation must be suppressed when allowCreate is false and nothing exists to link") + gatedPattern := fx.trackerPattern + gatedPattern.CookieCategoryID = fx.uncategorisedID - allowed := promote(t, ctx, newMappingHandler(client), client, fx.trackerPattern, fx.commonThirdPartyID, true) - require.NotNil(t, allowed, "creation must proceed when allowCreate is true") + gated := promote(t, ctx, newMappingHandler(client), gatedPattern, fx.commonThirdPartyID) + assert.Nil(t, gated, "creation must be suppressed for an uncategorised pattern with nothing to link") + + allowed := promote(t, ctx, newMappingHandler(client), fx.trackerPattern, fx.commonThirdPartyID) + require.NotNil(t, allowed, "creation must proceed for a categorised pattern") } // TestProcess_PreservesCatalogMappingOnReTrigger asserts that when @@ -1054,7 +1050,7 @@ func TestPromoteThirdParty_ExactCommonLinkIgnoresSimilarUnlinked(t *testing.T) { return linked.Insert(ctx, tx, fx.scope) })) - got := promote(t, ctx, newMappingHandler(client), client, fx.trackerPattern, fx.commonThirdPartyID, true) + got := promote(t, ctx, newMappingHandler(client), fx.trackerPattern, fx.commonThirdPartyID) require.NotNil(t, got) assert.Equal(t, linked.ID, *got, "exact-link path must short-circuit before the heuristic fires") diff --git a/pkg/probod/aliases.go b/pkg/probod/aliases.go index ff2c36570..9a81e4c2a 100644 --- a/pkg/probod/aliases.go +++ b/pkg/probod/aliases.go @@ -41,13 +41,17 @@ type ( LLMAgentConfig = probodconfig.LLMAgentConfig EvidenceDescriberConfig = probodconfig.EvidenceDescriberConfig AgentsConfig = probodconfig.AgentsConfig - MailerConfig = probodconfig.MailerConfig - SMTPConfig = probodconfig.SMTPConfig - NotificationsConfig = probodconfig.NotificationsConfig - WebhookConfig = probodconfig.WebhookConfig - OIDCProviderConfig = probodconfig.OIDCProviderConfig - PgConfig = probodconfig.PgConfig - SAMLConfig = probodconfig.SAMLConfig - SCIMBridgeConfig = probodconfig.SCIMBridgeConfig - SlackConfig = probodconfig.SlackConfig + + TrackerMappingWorkerConfig = probodconfig.TrackerMappingWorkerConfig + CommonPatternEnrichmentWorkerConfig = probodconfig.CommonPatternEnrichmentWorkerConfig + + MailerConfig = probodconfig.MailerConfig + SMTPConfig = probodconfig.SMTPConfig + NotificationsConfig = probodconfig.NotificationsConfig + WebhookConfig = probodconfig.WebhookConfig + OIDCProviderConfig = probodconfig.OIDCProviderConfig + PgConfig = probodconfig.PgConfig + SAMLConfig = probodconfig.SAMLConfig + SCIMBridgeConfig = probodconfig.SCIMBridgeConfig + SlackConfig = probodconfig.SlackConfig ) diff --git a/pkg/probod/probod.go b/pkg/probod/probod.go index 9686f0d81..aa6918dd8 100644 --- a/pkg/probod/probod.go +++ b/pkg/probod/probod.go @@ -725,7 +725,14 @@ func (impl *Implm) Run( }, ) - trackerMappingWorker := cookiebanner.NewTrackerMappingWorker(pgClient, l, trackerAgentsCfg, thirdPartyDisambiguationCfg) + trackerMappingWorker := cookiebanner.NewTrackerMappingWorker( + pgClient, + l, + trackerAgentsCfg, + thirdPartyDisambiguationCfg, + worker.WithInterval(time.Duration(impl.cfg.TrackerMappingWorker.Interval)*time.Second), + worker.WithMaxConcurrency(impl.cfg.TrackerMappingWorker.MaxConcurrency), + ) trackerMappingWorkerCtx, stopTrackerMappingWorker := context.WithCancel(context.Background()) wg.Go( @@ -742,7 +749,17 @@ func (impl *Implm) Run( stopCommonPatternEnrichmentWorker := func() {} if trackerAgentsCfg.LLMClient != nil { - commonPatternEnrichmentWorker := cookiebanner.NewCommonPatternEnrichmentWorker(pgClient, l, trackerAgentsCfg) + enrichmentCfg := trackerAgentsCfg + enrichmentCfg.AgentTimeout = time.Duration(impl.cfg.CommonPatternEnrichmentWorker.AgentTimeout) * time.Second + + commonPatternEnrichmentWorker := cookiebanner.NewCommonPatternEnrichmentWorker( + pgClient, + l, + enrichmentCfg, + time.Duration(impl.cfg.CommonPatternEnrichmentWorker.StaleAfter)*time.Second, + worker.WithInterval(time.Duration(impl.cfg.CommonPatternEnrichmentWorker.Interval)*time.Second), + worker.WithMaxConcurrency(impl.cfg.CommonPatternEnrichmentWorker.MaxConcurrency), + ) var commonPatternEnrichmentWorkerCtx context.Context diff --git a/pkg/probod/tracker_agents.go b/pkg/probod/tracker_agents.go index dd7048bf1..22921badb 100644 --- a/pkg/probod/tracker_agents.go +++ b/pkg/probod/tracker_agents.go @@ -16,6 +16,7 @@ package probod import ( "fmt" + "time" "github.com/prometheus/client_golang/prometheus" "go.gearno.de/kit/log" @@ -57,15 +58,34 @@ func (impl *Implm) buildTrackerAgentsConfig( return cookiebanner.TrackerAgentsConfig{}, thirdparty.DisambiguationConfig{}, fmt.Errorf("cannot resolve tracker mapping agent client: %w", err) } + mappingWorkerCfg := impl.cfg.TrackerMappingWorker + enrichmentWorkerCfg := impl.cfg.CommonPatternEnrichmentWorker + + // The mapping and enrichment agents share one config slot but run + // from separate workers with separate max-turns. AgentTimeout here + // carries the mapping worker's value (also reused by the + // disambiguation agent); the enrichment worker overrides it on its + // own copy at registration. trackerAgentsCfg := cookiebanner.TrackerAgentsConfig{ - LLMClient: llmClient, - Model: agentCfg.ModelName, - FirecrawlAPIKey: impl.cfg.Agents.Tools.FirecrawlAPIKey, + LLMClient: llmClient, + Model: agentCfg.ModelName, + FirecrawlAPIKey: impl.cfg.Agents.Tools.FirecrawlAPIKey, + MaxTokens: agentCfg.MaxTokens, + Temperature: agentCfg.Temperature, + AgentTimeout: time.Duration(mappingWorkerCfg.AgentTimeout) * time.Second, + MappingMaxTurns: mappingWorkerCfg.AgentMaxTurns, + EnrichmentMaxTurns: enrichmentWorkerCfg.AgentMaxTurns, } + // The disambiguation agent emits a single id plus a short rationale, + // so it keeps its own smaller token budget (left unset here) rather + // than inheriting the mapping agent's. It shares the mapping worker's + // timeout. disambiguationCfg := thirdparty.DisambiguationConfig{ - LLMClient: llmClient, - Model: agentCfg.ModelName, + LLMClient: llmClient, + Model: agentCfg.ModelName, + Temperature: agentCfg.Temperature, + Timeout: time.Duration(mappingWorkerCfg.AgentTimeout) * time.Second, } return trackerAgentsCfg, disambiguationCfg, nil diff --git a/pkg/probodconfig/config.go b/pkg/probodconfig/config.go index 9e789abaf..7eedc602f 100644 --- a/pkg/probodconfig/config.go +++ b/pkg/probodconfig/config.go @@ -60,11 +60,15 @@ type ( Connectors []ConnectorConfig `json:"connectors"` Agents AgentsConfig `json:"llm"` EvidenceDescriber EvidenceDescriberConfig `json:"evidence-describer"` - ChromeDPAddr string `json:"chrome-dp-addr"` - CustomDomains CustomDomainsConfig `json:"custom-domains"` - SCIMBridge SCIMBridgeConfig `json:"scim-bridge"` - ESign ESignConfig `json:"esign"` - Branding bool `json:"branding"` + + TrackerMappingWorker TrackerMappingWorkerConfig `json:"tracker-mapping-worker"` + CommonPatternEnrichmentWorker CommonPatternEnrichmentWorkerConfig `json:"common-pattern-enrichment-worker"` + + ChromeDPAddr string `json:"chrome-dp-addr"` + CustomDomains CustomDomainsConfig `json:"custom-domains"` + SCIMBridge SCIMBridgeConfig `json:"scim-bridge"` + ESign ESignConfig `json:"esign"` + Branding bool `json:"branding"` } // TrustCenterConfig contains trust center server configuration. diff --git a/pkg/probodconfig/llm_config.go b/pkg/probodconfig/llm_config.go index fe2aa00ce..70021d5af 100644 --- a/pkg/probodconfig/llm_config.go +++ b/pkg/probodconfig/llm_config.go @@ -40,6 +40,30 @@ type ( MaxConcurrency int `json:"max-concurrency"` } + // TrackerMappingWorkerConfig holds worker-side tuning for the + // tracker-mapping background worker. LLM parameters for the agents + // it runs live under AgentsConfig.TrackerMapping. AgentTimeout and + // AgentMaxTurns bound a single agent run (the identification and + // disambiguation agents). + TrackerMappingWorkerConfig struct { + Interval int `json:"interval"` // seconds between polls + MaxConcurrency int `json:"max-concurrency"` + AgentTimeout int `json:"agent-timeout"` // seconds, single agent run + AgentMaxTurns int `json:"agent-max-turns"` + } + + // CommonPatternEnrichmentWorkerConfig holds worker-side tuning for + // the common-pattern enrichment background worker. LLM parameters + // for the enrichment agent live under AgentsConfig.TrackerMapping + // (the agents share one config slot). + CommonPatternEnrichmentWorkerConfig struct { + Interval int `json:"interval"` // seconds between polls + MaxConcurrency int `json:"max-concurrency"` + StaleAfter int `json:"stale-after"` // seconds before a claim is recycled + AgentTimeout int `json:"agent-timeout"` // seconds, single agent run + AgentMaxTurns int `json:"agent-max-turns"` + } + // AgentToolsConfig holds API keys and settings for external tools // that agents can use (web search, scraping, etc.). AgentToolsConfig struct { diff --git a/pkg/thirdparty/disambiguation_agent.go b/pkg/thirdparty/disambiguation_agent.go index 93ecaec41..c995f65e6 100644 --- a/pkg/thirdparty/disambiguation_agent.go +++ b/pkg/thirdparty/disambiguation_agent.go @@ -38,19 +38,32 @@ const ( // described in the prompt. disambiguationConfidenceThreshold = 0.6 - // disambiguationTimeout caps a single disambiguation run. The - // agent has no tools and a single turn, so this is mostly a - // guard against a hung LLM provider, not a real budget. - disambiguationTimeout = 60 * time.Second + // defaultDisambiguationTimeout caps a single disambiguation run + // when the config supplies none. The agent has no tools and a + // single turn, so this is mostly a guard against a hung LLM + // provider, not a real budget. + defaultDisambiguationTimeout = 45 * time.Second + + // defaultDisambiguationMaxTokens caps the agent's structured + // output when the config carries no max-tokens budget. The output + // is a single id plus a one-sentence rationale. + defaultDisambiguationMaxTokens = 512 ) // DisambiguationConfig configures the third-party disambiguation // agent. The agent has no DB tools and no web-search tools: the // candidate list is supplied entirely in the prompt and the agent // only picks among it. +// +// MaxTokens and Temperature bound and steer the single LLM call, and +// Timeout caps a single run. Zero-valued fields fall back to package +// defaults. type DisambiguationConfig struct { - LLMClient *llm.Client - Model string + LLMClient *llm.Client + Model string + MaxTokens *int + Temperature *float64 + Timeout time.Duration } // DisambiguationResult is the structured output the disambiguation @@ -75,15 +88,25 @@ func BuildDisambiguationAgent( panic(fmt.Sprintf("thirdparty: cannot build disambiguation output type: %s", err)) } - return agent.New( - "third-party-disambiguation", - cfg.LLMClient, + maxTokens := defaultDisambiguationMaxTokens + if cfg.MaxTokens != nil && *cfg.MaxTokens > 0 { + maxTokens = *cfg.MaxTokens + } + + opts := []agent.Option{ agent.WithInstructions(disambiguationPrompt), agent.WithModel(cfg.Model), agent.WithOutputType(outputType), agent.WithMaxTurns(1), + agent.WithMaxTokens(maxTokens), agent.WithLogger(logger), - ) + } + + if cfg.Temperature != nil { + opts = append(opts, agent.WithTemperature(*cfg.Temperature)) + } + + return agent.New("third-party-disambiguation", cfg.LLMClient, opts...) } // Disambiguate runs the agent against the given catalog third party @@ -102,14 +125,19 @@ func Disambiguate( commonParty coredata.CommonThirdParty, commonDomains coredata.CommonThirdPartyDomains, candidates []ScoredCandidate, + timeout time.Duration, ) (*gid.GID, error) { if a == nil || len(candidates) == 0 { return nil, nil } + if timeout <= 0 { + timeout = defaultDisambiguationTimeout + } + prompt := buildDisambiguationPrompt(commonParty, commonDomains, candidates) - agentCtx, cancel := context.WithTimeout(ctx, disambiguationTimeout) + agentCtx, cancel := context.WithTimeout(ctx, timeout) defer cancel() result, err := agent.RunTyped[DisambiguationResult](