rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Tune tracker workers and bound agent budgets

Browse Source 
The tracker-mapping and common-pattern-enrichment workers ran with the
kit/worker defaults (interval 10s, max-concurrency 5 each) and dropped
the resolved per-agent max-tokens/temperature, so up to ten LLM
pipelines could run unbounded on one OpenAI client. The mapping worker
also held a FOR UPDATE transaction across the LLM and Firecrawl calls
while its DB search tools acquired a second pooled connection, risking
pool exhaustion under concurrency.

Plumb max-tokens, temperature, agent timeout, and per-worker max-turns
through TrackerAgentsConfig and DisambiguationConfig into all three
agent builders, replacing the hard-coded constants with config-fed
fields and package fallbacks. Expose worker interval, concurrency,
stale-after, agent timeout, and max-turns as config (env, Helm values,
deployment template) mirroring the evidence-describer pattern, and
apply them at registration.

Refactor Process into deterministic-read, agent (no transaction), and
persist phases so neither the mapping agent nor disambiguation runs
inside an open transaction, removing the row locks held across network
latency and the nested-connection pressure.

Signed-off-by: Émile Ré <emile@probo.com>
This commit is contained in:
Émile Ré
2026-05-29 16:37:32 +02:00
parent 50c5454681
commit 55302d18f0
17 changed files with 818 additions and 339 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
contrib/helm/charts/probo/templates/deployment.yaml
View File

@@ -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

17
contrib/helm/charts/probo/values-production.yaml.example
View File

@@ -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

18
contrib/helm/charts/probo/values.yaml
View File

@@ -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

22
pkg/bootstrap/builder.go
View File

@@ -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),
},
}

33
pkg/bootstrap/builder_test.go
View File

@@ -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)

20
pkg/cookiebanner/common_pattern_enrichment_agent.go
View File

@@ -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 {

22
pkg/cookiebanner/common_pattern_enrichment_worker.go
View File

@@ -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](

22
pkg/cookiebanner/tracker_agents_config.go
View File

@@ -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
}

49
pkg/cookiebanner/tracker_mapping_agent.go
View File

@@ -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 {

715
pkg/cookiebanner/tracker_mapping_worker.go
View File

@@ -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
}

40
pkg/cookiebanner/tracker_mapping_worker_test.go
View File

@@ -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")

22
pkg/probod/aliases.go
View File

@@ -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
)

21
pkg/probod/probod.go
View File

@@ -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

30
pkg/probod/tracker_agents.go
View File

@@ -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

14
pkg/probodconfig/config.go
View File

@@ -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.

24
pkg/probodconfig/llm_config.go
View File

@@ -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 {

50
pkg/thirdparty/disambiguation_agent.go vendored
View File

@@ -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](