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

Restore tracker mapping linking, drop only create

Browse Source 
The tracker-mapping worker had been reduced to catalog resolution only,
which removed not just the auto-creation of an org ThirdParty but also
the auto-linking of an existing one. Only the creation needed to go: it
raced the load-then-create check and produced duplicate vendors.

Restore the full org ThirdParty resolution (exact common-id link,
sibling direct-link, high-confidence heuristic, and the disambiguation
agent) and remove only the CreateFromCommon branch and its
categorisation gate. When nothing matches, the worker now leaves
third_party_id unset rather than creating a vendor; creation happens
exclusively through the explicit ImportFromCommon action. Drop the
now-dead CreateFromCommon helper and rename match.go to common_match.go.

Fix a latent test bug surfaced by actually running the DB-backed suite
(skipped in CI without Postgres): the heuristic-match candidate lacked
Level 1, so the level-filtered candidate loader excluded it and the old
fallback create masked the miss.

Signed-off-by: Émile Ré <emile@probo.com>
This commit is contained in:
Émile Ré
2026-06-11 10:55:29 +02:00
parent 9a7bc26d49
commit 145aeaf402
17 changed files with 1323 additions and 238 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
contrib/helm/charts/probo/CHANGELOG.md
View File

@@ -4,10 +4,6 @@ All notable changes to the Probo Helm chart will be documented in this file.
## Unreleased
### Removed
- `thirdPartyDisambiguation` agent config slot and `trackerMappingWorker.disambiguationAgentTimeout`; the tracker-mapping worker no longer runs a third-party disambiguation agent
## [0.5.0] - 2026-06-09
### Added

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

@@ -281,6 +281,23 @@ spec:
- name: THIRD_PARTY_VETTING_STALE_AFTER
value: {{ .Values.probo.thirdPartyVettingWorker.staleAfter | quote }}
{{- end }}
# Third-party Disambiguation Agent
{{- if .Values.probo.thirdPartyDisambiguation.provider }}
- name: AGENT_THIRD_PARTY_DISAMBIGUATION_PROVIDER
value: {{ .Values.probo.thirdPartyDisambiguation.provider | quote }}
{{- end }}
{{- if .Values.probo.thirdPartyDisambiguation.modelName }}
- name: AGENT_THIRD_PARTY_DISAMBIGUATION_MODEL_NAME
value: {{ .Values.probo.thirdPartyDisambiguation.modelName | quote }}
{{- end }}
{{- if .Values.probo.thirdPartyDisambiguation.temperature }}
- name: AGENT_THIRD_PARTY_DISAMBIGUATION_TEMPERATURE
value: {{ .Values.probo.thirdPartyDisambiguation.temperature | quote }}
{{- end }}
{{- if .Values.probo.thirdPartyDisambiguation.maxTokens }}
- name: AGENT_THIRD_PARTY_DISAMBIGUATION_MAX_TOKENS
value: {{ .Values.probo.thirdPartyDisambiguation.maxTokens | quote }}
{{- end }}
# Tracker Mapping Agent
{{- if .Values.probo.trackerMapping.provider }}
- name: AGENT_TRACKER_MAPPING_PROVIDER
@@ -336,6 +353,10 @@ spec:
- name: TRACKER_MAPPING_AGENT_MAX_TURNS
value: {{ .Values.probo.trackerMappingWorker.agentMaxTurns | quote }}
{{- end }}
{{- if .Values.probo.trackerMappingWorker.disambiguationAgentTimeout }}
- name: TRACKER_MAPPING_DISAMBIGUATION_AGENT_TIMEOUT
value: {{ .Values.probo.trackerMappingWorker.disambiguationAgentTimeout | quote }}
{{- end }}
# Common Pattern Enrichment Worker
{{- if .Values.probo.commonPatternEnrichmentWorker.interval }}
- name: COMMON_PATTERN_ENRICHMENT_INTERVAL

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

@@ -185,6 +185,15 @@ probo:
# maxConcurrency: 1
# staleAfter: 1500
# Third-party disambiguation agent (optional; runs inside the tracker
# mapping worker to pick the best matching org third party). Falls back
# to trackerMapping when its provider is unset.
# thirdPartyDisambiguation:
# provider: "openai"
# modelName: "gpt-4o"
# temperature: "0.4"
# maxTokens: "4096"
# Tracker mapping agent (optional, auto-links tracker patterns to vendors)
# trackerMapping:
# provider: "openai"
@@ -201,7 +210,7 @@ probo:
# maxTokens: "4096"
# Tracker mapping worker tuning (optional; seconds for
# interval/staleAfter/agentTimeout).
# interval/staleAfter/agentTimeout/disambiguationAgentTimeout).
# Keep concurrency modest to stay under OpenAI/Firecrawl limits and the DB pool.
# trackerMappingWorker:
# interval: 10
@@ -209,6 +218,7 @@ probo:
# staleAfter: 600
# agentTimeout: 45
# agentMaxTurns: 10
# disambiguationAgentTimeout: 45
# Common-pattern enrichment worker tuning (optional; seconds for
# interval/staleAfter/agentTimeout).

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

@@ -286,6 +286,15 @@ probo:
maxConcurrency: 1
staleAfter: 1500
# Third-party disambiguation agent (optional; runs inside the tracker
# mapping worker to pick the best matching org third party). Falls back
# to trackerMapping when its provider is unset.
thirdPartyDisambiguation:
provider: ""
modelName: ""
temperature: ""
maxTokens: ""
# Tracker mapping agent (optional, requires openai.apiKey or anthropic key)
trackerMapping:
provider: ""
@@ -302,15 +311,16 @@ probo:
maxTokens: ""
# Tracker mapping background worker tuning (optional). interval,
# staleAfter, and agentTimeout are in seconds. Keep concurrency modest
# to stay under OpenAI/Firecrawl rate limits and the database
# connection pool.
# staleAfter, agentTimeout, and disambiguationAgentTimeout are in
# seconds. Keep concurrency modest to stay under OpenAI/Firecrawl rate
# limits and the database connection pool.
trackerMappingWorker:
interval: 10
maxConcurrency: 3
staleAfter: 600
agentTimeout: 45
agentMaxTurns: 10
disambiguationAgentTimeout: 45
# Common-pattern enrichment background worker tuning (optional).
# interval, staleAfter, and agentTimeout are in seconds.

22
pkg/bootstrap/builder.go
View File

@@ -213,6 +213,17 @@ func (b *Builder) Build() (*probodconfig.FullConfig, error) {
Temperature: b.getEnvFloatPtr("AGENT_THIRD_PARTY_VETTER_TEMPERATURE"),
MaxTokens: b.getEnvIntPtr("AGENT_THIRD_PARTY_VETTER_MAX_TOKENS"),
},
ThirdPartyDisambiguation: probodconfig.LLMAgentConfig{
Provider: b.getEnvOrDefault("AGENT_THIRD_PARTY_DISAMBIGUATION_PROVIDER", ""),
ModelName: b.getEnvOrDefault("AGENT_THIRD_PARTY_DISAMBIGUATION_MODEL_NAME", ""),
// The disambiguation agent emits a single id plus a
// short rationale, but the budget must leave headroom
// for reasoning models whose reasoning tokens count
// against max_tokens; too small a budget truncates the
// JSON.
Temperature: b.getEnvFloatPtr("AGENT_THIRD_PARTY_DISAMBIGUATION_TEMPERATURE"),
MaxTokens: new(b.getEnvIntOrDefault("AGENT_THIRD_PARTY_DISAMBIGUATION_MAX_TOKENS", 4096)),
},
TrackerMapping: probodconfig.LLMAgentConfig{
Provider: b.getEnvOrDefault("AGENT_TRACKER_MAPPING_PROVIDER", ""),
ModelName: b.getEnvOrDefault("AGENT_TRACKER_MAPPING_MODEL_NAME", ""),
@@ -267,11 +278,12 @@ func (b *Builder) Build() (*probodconfig.FullConfig, error) {
MaxConcurrency: b.getEnvIntOrDefault("THIRD_PARTY_VETTING_MAX_CONCURRENCY", 1),
},
TrackerMappingWorker: probodconfig.TrackerMappingWorkerConfig{
Interval: b.getEnvIntOrDefault("TRACKER_MAPPING_INTERVAL", 10),
MaxConcurrency: b.getEnvIntOrDefault("TRACKER_MAPPING_MAX_CONCURRENCY", 3),
StaleAfter: b.getEnvIntOrDefault("TRACKER_MAPPING_STALE_AFTER", 600),
AgentTimeout: b.getEnvIntOrDefault("TRACKER_MAPPING_AGENT_TIMEOUT", 45),
AgentMaxTurns: b.getEnvIntOrDefault("TRACKER_MAPPING_AGENT_MAX_TURNS", 10),
Interval: b.getEnvIntOrDefault("TRACKER_MAPPING_INTERVAL", 10),
MaxConcurrency: b.getEnvIntOrDefault("TRACKER_MAPPING_MAX_CONCURRENCY", 3),
StaleAfter: b.getEnvIntOrDefault("TRACKER_MAPPING_STALE_AFTER", 600),
AgentTimeout: b.getEnvIntOrDefault("TRACKER_MAPPING_AGENT_TIMEOUT", 45),
AgentMaxTurns: b.getEnvIntOrDefault("TRACKER_MAPPING_AGENT_MAX_TURNS", 10),
DisambiguationAgentTimeout: b.getEnvIntOrDefault("TRACKER_MAPPING_DISAMBIGUATION_AGENT_TIMEOUT", 45),
},
CommonPatternEnrichmentWorker: probodconfig.CommonPatternEnrichmentWorkerConfig{
Interval: b.getEnvIntOrDefault("COMMON_PATTERN_ENRICHMENT_INTERVAL", 10),

17
pkg/bootstrap/builder_test.go
View File

@@ -220,6 +220,10 @@ func TestBuilder_Build_Defaults(t *testing.T) {
assert.Empty(t, cfg.Probod.Agents.ThirdPartyVetter.ModelName)
assert.Nil(t, cfg.Probod.Agents.ThirdPartyVetter.Temperature)
assert.Nil(t, cfg.Probod.Agents.ThirdPartyVetter.MaxTokens)
assert.Empty(t, cfg.Probod.Agents.ThirdPartyDisambiguation.Provider)
assert.Empty(t, cfg.Probod.Agents.ThirdPartyDisambiguation.ModelName)
assert.Nil(t, cfg.Probod.Agents.ThirdPartyDisambiguation.Temperature)
assert.Equal(t, new(4096), cfg.Probod.Agents.ThirdPartyDisambiguation.MaxTokens)
assert.Empty(t, cfg.Probod.Agents.TrackerMapping.Provider)
assert.Empty(t, cfg.Probod.Agents.TrackerMapping.ModelName)
assert.Nil(t, cfg.Probod.Agents.TrackerMapping.Temperature)
@@ -235,6 +239,7 @@ func TestBuilder_Build_Defaults(t *testing.T) {
assert.Equal(t, 600, cfg.Probod.TrackerMappingWorker.StaleAfter)
assert.Equal(t, 45, cfg.Probod.TrackerMappingWorker.AgentTimeout)
assert.Equal(t, 10, cfg.Probod.TrackerMappingWorker.AgentMaxTurns)
assert.Equal(t, 45, cfg.Probod.TrackerMappingWorker.DisambiguationAgentTimeout)
assert.Equal(t, 10, cfg.Probod.CommonPatternEnrichmentWorker.Interval)
assert.Equal(t, 2, cfg.Probod.CommonPatternEnrichmentWorker.MaxConcurrency)
assert.Equal(t, 600, cfg.Probod.CommonPatternEnrichmentWorker.StaleAfter)
@@ -337,6 +342,11 @@ func TestBuilder_Build_CustomValues(t *testing.T) {
env["AGENT_THIRD_PARTY_VETTER_MODEL_NAME"] = "gpt-4o"
env["AGENT_THIRD_PARTY_VETTER_TEMPERATURE"] = "0.3"
env["AGENT_THIRD_PARTY_VETTER_MAX_TOKENS"] = "8192"
// Agents — third-party-disambiguation override
env["AGENT_THIRD_PARTY_DISAMBIGUATION_PROVIDER"] = "anthropic"
env["AGENT_THIRD_PARTY_DISAMBIGUATION_MODEL_NAME"] = "claude-sonnet-4-20250514"
env["AGENT_THIRD_PARTY_DISAMBIGUATION_TEMPERATURE"] = "0.4"
env["AGENT_THIRD_PARTY_DISAMBIGUATION_MAX_TOKENS"] = "2048"
// Agents — tracker-mapping override
env["AGENT_TRACKER_MAPPING_PROVIDER"] = "openai"
env["AGENT_TRACKER_MAPPING_MODEL_NAME"] = "gpt-4o-mini"
@@ -353,6 +363,7 @@ func TestBuilder_Build_CustomValues(t *testing.T) {
env["TRACKER_MAPPING_STALE_AFTER"] = "1200"
env["TRACKER_MAPPING_AGENT_TIMEOUT"] = "30"
env["TRACKER_MAPPING_AGENT_MAX_TURNS"] = "6"
env["TRACKER_MAPPING_DISAMBIGUATION_AGENT_TIMEOUT"] = "35"
env["COMMON_PATTERN_ENRICHMENT_INTERVAL"] = "15"
env["COMMON_PATTERN_ENRICHMENT_MAX_CONCURRENCY"] = "4"
env["COMMON_PATTERN_ENRICHMENT_STALE_AFTER"] = "900"
@@ -452,6 +463,11 @@ func TestBuilder_Build_CustomValues(t *testing.T) {
assert.Equal(t, "gpt-4o", cfg.Probod.Agents.ThirdPartyVetter.ModelName)
assert.Equal(t, new(0.3), cfg.Probod.Agents.ThirdPartyVetter.Temperature)
assert.Equal(t, new(8192), cfg.Probod.Agents.ThirdPartyVetter.MaxTokens)
// Agents — third-party-disambiguation overrides
assert.Equal(t, "anthropic", cfg.Probod.Agents.ThirdPartyDisambiguation.Provider)
assert.Equal(t, "claude-sonnet-4-20250514", cfg.Probod.Agents.ThirdPartyDisambiguation.ModelName)
assert.Equal(t, new(0.4), cfg.Probod.Agents.ThirdPartyDisambiguation.Temperature)
assert.Equal(t, new(2048), cfg.Probod.Agents.ThirdPartyDisambiguation.MaxTokens)
// Agents — tracker-mapping overrides
assert.Equal(t, "openai", cfg.Probod.Agents.TrackerMapping.Provider)
assert.Equal(t, "gpt-4o-mini", cfg.Probod.Agents.TrackerMapping.ModelName)
@@ -468,6 +484,7 @@ func TestBuilder_Build_CustomValues(t *testing.T) {
assert.Equal(t, 1200, cfg.Probod.TrackerMappingWorker.StaleAfter)
assert.Equal(t, 30, cfg.Probod.TrackerMappingWorker.AgentTimeout)
assert.Equal(t, 6, cfg.Probod.TrackerMappingWorker.AgentMaxTurns)
assert.Equal(t, 35, cfg.Probod.TrackerMappingWorker.DisambiguationAgentTimeout)
assert.Equal(t, 15, cfg.Probod.CommonPatternEnrichmentWorker.Interval)
assert.Equal(t, 4, cfg.Probod.CommonPatternEnrichmentWorker.MaxConcurrency)
assert.Equal(t, 900, cfg.Probod.CommonPatternEnrichmentWorker.StaleAfter)

13
pkg/cookiebanner/pattern_analysis_worker.go
View File

@@ -230,9 +230,10 @@ func (h *patternAnalysisHandler) Process(ctx context.Context, banner coredata.Co
return fmt.Errorf("cannot promote source on glob pattern %q: %w", key.template, err)
}
// A stronger source can unblock mapping (e.g.
// EXTENSION->SCRIPT lifts the creationAllowed
// gate), so re-arm mapping on the existing glob.
// A stronger source can unblock mapping (a fresh
// initiator domain lets matchByDomain/
// matchBySiblingOrigin resolve a vendor), so
// re-arm mapping on the existing glob.
if err := globPattern.SetMappingRequested(ctx, tx); err != nil {
return fmt.Errorf("cannot request mapping after source promotion on glob pattern %q: %w", key.template, err)
}
@@ -894,9 +895,9 @@ func (h *patternAnalysisHandler) adoptUncategorisedPatterns(
return false, fmt.Errorf("cannot promote source on glob pattern %q: %w", match.Pattern, err)
}
// A stronger source can unblock mapping (e.g.
// EXTENSION->SCRIPT lifts the creationAllowed gate), so
// re-arm mapping on the adopted glob.
// A stronger source can unblock mapping (a fresh
// initiator domain lets matchByDomain/matchBySiblingOrigin
// resolve a vendor), so re-arm mapping on the adopted glob.
if err := match.SetMappingRequested(ctx, tx); err != nil {
return false, fmt.Errorf("cannot request mapping after source promotion on glob pattern %q: %w", match.Pattern, err)
}

5
pkg/cookiebanner/service.go
View File

@@ -2284,9 +2284,8 @@ func (s *Service) reportDetectedTracker(
// A stronger source can unblock mapping: the detection
// upserted below carries a fresh initiator domain that
// matchByDomain/matchBySiblingOrigin can now use, and an
// EXTENSION->SCRIPT promotion lifts the creationAllowed
// gate. Re-arm mapping so the worker revisits the pattern.
// matchByDomain/matchBySiblingOrigin can now use. Re-arm
// mapping so the worker revisits the pattern.
if err := matchedPattern.SetMappingRequested(ctx, tx); err != nil {
return fmt.Errorf("cannot request mapping after source promotion on tracker pattern %q: %w", matchedPattern.Pattern, err)
}

341
pkg/cookiebanner/tracker_mapping_worker.go
View File

@@ -39,17 +39,20 @@ import (
const defaultMappingStaleAfter = 10 * time.Minute
type trackerMappingHandler struct {
pg *pg.Client
logger *log.Logger
mappingAgent *agent.Agent
agentTimeout time.Duration
staleAfter time.Duration
pg *pg.Client
logger *log.Logger
mappingAgent *agent.Agent
disambiguationAgent *agent.Agent
agentTimeout time.Duration
disambiguationTimeout time.Duration
staleAfter time.Duration
}
func NewTrackerMappingWorker(
pgClient *pg.Client,
logger *log.Logger,
mappingCfg TrackerMappingAgentConfig,
disambiguationCfg thirdparty.DisambiguationAgentConfig,
staleAfter time.Duration,
opts ...worker.Option,
) *worker.Worker[coredata.TrackerPattern] {
@@ -63,16 +66,21 @@ func NewTrackerMappingWorker(
}
h := &trackerMappingHandler{
pg: pgClient,
logger: logger,
agentTimeout: agentTimeout,
staleAfter: staleAfter,
pg: pgClient,
logger: logger,
agentTimeout: agentTimeout,
disambiguationTimeout: disambiguationCfg.Timeout,
staleAfter: staleAfter,
}
if mappingCfg.LLMClient != nil {
h.mappingAgent = buildTrackerMappingAgent(mappingCfg, pgClient, logger)
}
if disambiguationCfg.LLMClient != nil {
h.disambiguationAgent = thirdparty.BuildDisambiguationAgent(disambiguationCfg, logger)
}
return worker.New(
"tracker-mapping-worker",
h,
@@ -124,11 +132,14 @@ func (h *trackerMappingHandler) RecoverStale(ctx context.Context) error {
// catalogMatch is the result of a single catalog signal. commonPatternID
// is the catalog row the signal resolved (or backfilled); commonThirdPartyID
// is the catalog third party the signal discovered, when any. A nil
// *catalogMatch means the signal produced nothing.
// is the catalog third party the signal discovered, when any; thirdPartyID
// is an existing org ThirdParty the signal knows directly (e.g. a sibling
// pattern already promoted in the same organization). A nil *catalogMatch
// means the signal produced nothing.
type catalogMatch struct {
commonPatternID *gid.GID
commonThirdPartyID *gid.GID
thirdPartyID *gid.GID
}
// Process resolves the catalog mapping for a tracker pattern and links it
@@ -145,11 +156,10 @@ type catalogMatch struct {
// common_tracker_pattern_id but its catalog row has no common third
// party yet.
//
// The worker no longer materializes per-org ThirdParty rows: it resolves
// the shared catalog link only. An org ThirdParty is created exclusively
// through the explicit per-vendor import action, which also backfills
// tracker_patterns.third_party_id; an already-set third_party_id is
// preserved here untouched.
// Org ThirdParty resolution only links to an existing party (even for
// uncategorised or extension-sourced patterns); it never creates a brand
// new org ThirdParty. Creating an org ThirdParty from a catalog vendor is
// done exclusively through the explicit ImportFromCommon action.
func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.TrackerPattern) error {
scope := coredata.NewScopeFromObjectID(tp.ID)
@@ -174,6 +184,7 @@ func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.Tracker
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.
@@ -205,6 +216,25 @@ func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.Tracker
}
}
// Phase 3: org ThirdParty resolution. The heuristic ranking and the
// disambiguation agent run without a transaction; only the final link
// 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.
@@ -221,6 +251,7 @@ func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.Tracker
}
tp.CommonTrackerPatternID = commonPatternID
tp.ThirdPartyID = thirdPartyID
tp.UpdatedAt = time.Now()
// Descriptions are owned by the common-pattern enrichment
@@ -290,6 +321,7 @@ type deterministicResult struct {
origin string
commonPatternID *gid.GID
commonThirdPartyID *gid.GID
directThirdPartyID *gid.GID
domains []string
commonThirdPartyPreexisted bool
}
@@ -371,6 +403,7 @@ func (h *trackerMappingHandler) resolveDeterministic(
if siblingMatch != nil {
res.commonPatternID = firstNonNil(res.commonPatternID, siblingMatch.commonPatternID)
res.commonThirdPartyID = siblingMatch.commonThirdPartyID
res.directThirdPartyID = siblingMatch.thirdPartyID
}
if res.commonThirdPartyID != nil {
@@ -805,8 +838,10 @@ func (h *trackerMappingHandler) persistAgentIdentification(
// matchBySiblingOrigin finds other tracker patterns on the same banner
// that share initiator domains with the current pattern. Sharing an
// origin across multiple detected patterns is a strong indicator of the
// same third party, so the common third party the siblings resolve to is
// upserted onto the catalog row.
// same third party. When the siblings resolve to a single existing org
// ThirdParty, that id is returned directly so promotion can link to it
// without re-running heuristics; otherwise the resolved common third
// party is upserted onto the catalog row.
func (h *trackerMappingHandler) matchBySiblingOrigin(
ctx context.Context,
tx pg.Tx,
@@ -837,14 +872,19 @@ func (h *trackerMappingHandler) matchBySiblingOrigin(
scope := coredata.NewScopeFromObjectID(tp.ID)
commonThirdPartyID, err := h.resolveThirdPartyFromSiblings(ctx, tx, scope, siblingIDs)
commonThirdPartyID, thirdPartyID, err := h.resolveThirdPartyFromSiblings(ctx, tx, scope, siblingIDs)
if err != nil {
return nil, fmt.Errorf("cannot resolve third party from siblings: %w", err)
}
// No catalog third party to record: leave catalog creation to a later
// signal or the unmatched fallback.
// No catalog third party to record: surface a directly-known org
// third party (if any) so promotion can still link to it, and leave
// catalog creation to a later signal or the unmatched fallback.
if commonThirdPartyID == nil {
if thirdPartyID != nil {
return &catalogMatch{thirdPartyID: thirdPartyID}, nil
}
return nil, nil
}
@@ -876,25 +916,37 @@ func (h *trackerMappingHandler) matchBySiblingOrigin(
return &catalogMatch{
commonPatternID: &commonPattern.ID,
commonThirdPartyID: commonPattern.CommonThirdPartyID,
thirdPartyID: thirdPartyID,
}, nil
}
// resolveThirdPartyFromSiblings inspects sibling patterns to resolve a
// single unambiguous catalog third party for backfill. It is resolved
// first from the siblings' org ThirdParties, then, when those carry none,
// from siblings' common_tracker_pattern rows. It returns nil when the
// siblings carry no catalog third party or disagree on one.
// third party. It returns two independent signals: a direct org
// ThirdParty (set only when the siblings share a single one — the
// strongest, same-org signal), and a single unambiguous catalog third
// party for backfill. The catalog third party is resolved first from the
// siblings' org ThirdParties, then, when those carry none, from siblings'
// common_tracker_pattern rows. Either signal may be nil; siblings that
// disagree on the catalog third party resolve it to nothing.
func (h *trackerMappingHandler) resolveThirdPartyFromSiblings(
ctx context.Context,
conn pg.Querier,
scope coredata.Scoper,
siblingIDs []gid.GID,
) (commonThirdPartyID *gid.GID, err error) {
) (commonThirdPartyID *gid.GID, thirdPartyID *gid.GID, err error) {
var patterns coredata.TrackerPatterns
thirdPartyIDs, err := patterns.LoadDistinctThirdPartyIDsByIDs(ctx, conn, scope, siblingIDs)
if err != nil {
return nil, fmt.Errorf("cannot load distinct third party ids from siblings: %w", err)
return nil, nil, fmt.Errorf("cannot load distinct third party ids from siblings: %w", err)
}
// A single org third party shared across the siblings is the
// strongest, same-org signal: link to it directly. This is resolved
// independently from the catalog third party used for backfill.
if len(thirdPartyIDs) == 1 {
directID := thirdPartyIDs[0]
thirdPartyID = &directID
}
if len(thirdPartyIDs) > 0 {
@@ -913,26 +965,29 @@ func (h *trackerMappingHandler) resolveThirdPartyFromSiblings(
if len(commonIDs) == 1 {
for id := range commonIDs {
return &id, nil
return &id, thirdPartyID, nil
}
}
// Siblings are linked to several different catalog third
// parties: do not guess one.
// Siblings are promoted to several different catalog third
// parties: do not guess one. A single shared org third party (if
// any) is still a safe direct link.
if len(commonIDs) > 1 {
return nil, nil
return nil, thirdPartyID, nil
}
}
// Fall back to siblings carrying only a common_tracker_pattern_id, or
// whose org ThirdParty is not itself linked to the catalog.
// whose org ThirdParty is not itself linked to the catalog. This is
// reached when the org-third-party scan above found no catalog third
// party, so it must not be short-circuited by a direct match.
commonPatternIDs, err := patterns.LoadDistinctCommonTrackerPatternIDsByIDs(ctx, conn, scope, siblingIDs)
if err != nil {
return nil, fmt.Errorf("cannot load distinct common tracker pattern ids from siblings: %w", err)
return nil, nil, fmt.Errorf("cannot load distinct common tracker pattern ids from siblings: %w", err)
}
if len(commonPatternIDs) == 0 {
return nil, nil
return nil, thirdPartyID, nil
}
commonIDs := make(map[gid.GID]struct{})
@@ -950,11 +1005,11 @@ func (h *trackerMappingHandler) resolveThirdPartyFromSiblings(
if len(commonIDs) == 1 {
for id := range commonIDs {
return &id, nil
return &id, thirdPartyID, nil
}
}
return nil, nil
return nil, thirdPartyID, nil
}
func (h *trackerMappingHandler) createUnmatchedPattern(
@@ -980,3 +1035,217 @@ func (h *trackerMappingHandler) createUnmatchedPattern(
return &commonPattern.ID, nil
}
// resolveOrgThirdParty resolves an org ThirdParty for the given pattern
// from a known catalog third party by linking to an existing party. The
// resolution order is:
//
// 1. Exact link by common_third_party_id (O(1)).
// 2. Heuristic match against the org's existing ThirdParty rows
// (lowercased name, suffix-stripped name, slug, website host,
// CommonThirdPartyDomain overlap).
// 3. Agent disambiguation when the heuristic is ambiguous.
//
// When none of these resolve an existing party, the function returns
// (nil, nil); it never creates a brand new org ThirdParty (that happens
// only through the explicit ImportFromCommon action). A confident
// heuristic/agent match is auto-tagged with common_third_party_id so
// subsequent resolutions hit the exact-link path in O(1).
func (h *trackerMappingHandler) resolveOrgThirdParty(
ctx context.Context,
tp coredata.TrackerPattern,
commonThirdPartyID gid.GID,
) (*gid.GID, error) {
scope := coredata.NewScopeFromObjectID(tp.ID)
// Read phase: exact link, candidate ranking, and eligibility. No
// write or LLM call happens here.
var prep orgThirdPartyPrep
if err := h.pg.WithConn(
ctx,
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 prep.existingID != nil {
return prep.existingID, nil
}
picked := prep.highConfidence
viaAgent := false
// 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,
prep.commonParty,
prep.commonDomains,
prep.agentSet,
h.disambiguationTimeout,
)
if err != nil {
h.logger.WarnCtx(
ctx,
"third-party disambiguation agent failed",
log.Error(err),
log.String("tracker_pattern_id", tp.ID.String()),
)
}
if matchedID != nil {
for _, c := range prep.agentSet {
if c.ThirdParty.ID == *matchedID {
picked = c.ThirdParty
viaAgent = true
break
}
}
}
}
// Nothing to link: leave the pattern without an org third party. An
// org ThirdParty is created only through the explicit ImportFromCommon
// action, never here.
if picked == nil {
return nil, nil
}
// Write phase: link the picked candidate to the catalog entry in a
// short transaction.
if err := h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
if err := thirdparty.LinkToCommon(ctx, tx, scope, picked, commonThirdPartyID); err != nil {
return fmt.Errorf("cannot link third party to common: %w", err)
}
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 nil
},
); err != nil {
return nil, err
}
return &picked.ID, 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.
type orgThirdPartyPrep struct {
existingID *gid.GID
commonParty coredata.CommonThirdParty
commonDomains coredata.CommonThirdPartyDomains
agentSet []thirdparty.ScoredCandidate
highConfidence *coredata.ThirdParty
highScore float64
eligibleForAgent 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, and ranks the candidates.
// 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)
}
firstLevel := 1
var orgThirdParties coredata.ThirdParties
if err := orgThirdParties.LoadAllByOrganizationID(
ctx,
conn,
scope,
tp.OrganizationID,
coredata.NewThirdPartyFilter(nil, &firstLevel, nil),
); 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
}
}
}
return prep, nil
}

254
pkg/cookiebanner/tracker_mapping_worker_test.go
View File

@@ -31,10 +31,10 @@ import (
// promotionFixture extends workerFixture with a CommonThirdParty and a
// CommonTrackerPattern linking the catalog to the test pattern. It is
// the minimum scaffolding the catalog-resolution paths need to run
// end-to-end.
// the minimum scaffolding resolveOrgThirdParty needs to run end-to-end.
type promotionFixture struct {
workerFixture
commonThirdParty coredata.CommonThirdParty
commonPatternID gid.GID
trackerPattern coredata.TrackerPattern
commonThirdPartyID gid.GID
@@ -130,6 +130,7 @@ func seedPromotionFixture(t *testing.T, ctx context.Context, client *pg.Client)
return promotionFixture{
workerFixture: fx,
commonThirdParty: commonThirdParty,
commonPatternID: commonPattern.ID,
commonThirdPartyID: commonThirdPartyID,
trackerPattern: pattern,
@@ -143,6 +144,112 @@ func newMappingHandler(client *pg.Client) *trackerMappingHandler {
}
}
// promote runs resolveOrgThirdParty, which manages its own short
// transactions internally and only links to an existing org ThirdParty
// (it never creates one).
func promote(
t *testing.T,
ctx context.Context,
h *trackerMappingHandler,
tp coredata.TrackerPattern,
commonThirdPartyID gid.GID,
) *gid.GID {
t.Helper()
got, err := h.resolveOrgThirdParty(ctx, tp, commonThirdPartyID)
require.NoError(t, err)
return got
}
func TestPromoteThirdParty_ExactCommonLink(t *testing.T) {
t.Parallel()
client := test.PGClient(t)
ctx := context.Background()
fx := seedPromotionFixture(t, ctx, client)
now := time.Now().UTC().Truncate(time.Microsecond)
existing := coredata.ThirdParty{
ID: gid.New(fx.scope.GetTenantID(), coredata.ThirdPartyEntityType),
OrganizationID: fx.organizationID,
CommonThirdPartyID: &fx.commonThirdPartyID,
Name: "Google LLC",
Category: coredata.ThirdPartyCategoryAnalytics,
Certifications: []string{},
Countries: coredata.CountryCodes{},
CreatedAt: now,
UpdatedAt: now,
}
require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error {
return existing.Insert(ctx, tx, fx.scope)
}))
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")
}
func TestPromoteThirdParty_HeuristicMatch(t *testing.T) {
t.Parallel()
client := test.PGClient(t)
ctx := context.Background()
fx := seedPromotionFixture(t, ctx, client)
now := time.Now().UTC().Truncate(time.Microsecond)
// Append a corporate suffix to the catalog name so the heuristic
// matches on the suffix-stripped name (score 0.9) rather than an
// exact link.
manualEntry := coredata.ThirdParty{
ID: gid.New(fx.scope.GetTenantID(), coredata.ThirdPartyEntityType),
OrganizationID: fx.organizationID,
Name: fx.commonThirdParty.Name + " LLC",
Category: coredata.ThirdPartyCategoryAnalytics,
Certifications: []string{},
Countries: coredata.CountryCodes{},
Level: 1,
CreatedAt: now,
UpdatedAt: now,
}
require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error {
return manualEntry.Insert(ctx, tx, fx.scope)
}))
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")
var reloaded coredata.ThirdParty
require.NoError(t, client.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error {
return reloaded.LoadByID(ctx, conn, fx.scope, manualEntry.ID)
}))
require.NotNil(t, reloaded.CommonThirdPartyID, "matched row must be tagged with common_third_party_id")
assert.Equal(t, fx.commonThirdPartyID, *reloaded.CommonThirdPartyID)
}
// TestPromoteThirdParty_NoCreateWithoutMatch asserts that when no
// existing org ThirdParty matches the catalog third party, resolution
// returns nothing: the worker never creates a brand new org ThirdParty
// (that is done only through the explicit ImportFromCommon action).
func TestPromoteThirdParty_NoCreateWithoutMatch(t *testing.T) {
t.Parallel()
client := test.PGClient(t)
ctx := context.Background()
fx := seedPromotionFixture(t, ctx, client)
got := promote(t, ctx, newMappingHandler(client), fx.trackerPattern, fx.commonThirdPartyID)
assert.Nil(t, got, "resolution must not create a new org ThirdParty")
}
// TestProcess_PreservesCatalogMappingOnReTrigger asserts that when
// Process is called for a pattern that already carries a
// common_tracker_pattern_id, the catalog pipeline is skipped and the
@@ -169,7 +276,7 @@ func TestProcess_PreservesCatalogMappingOnReTrigger(t *testing.T) {
require.NotNil(t, reloaded.CommonTrackerPatternID, "common tracker pattern link must be preserved")
assert.Equal(t, fx.commonPatternID, *reloaded.CommonTrackerPatternID)
assert.Nil(t, reloaded.ThirdPartyID, "the worker must not auto-create or link an org ThirdParty")
assert.Nil(t, reloaded.ThirdPartyID, "no org ThirdParty exists to link, so third_party_id stays unset")
}
// TestProcess_UncategorisedPatternIsNotPromoted asserts that a pattern
@@ -442,6 +549,8 @@ func TestMatchBySiblingOrigin_SiblingWithThirdPartyID(t *testing.T) {
require.NotNil(t, got, "sibling origin match should return a catalog match")
require.NotNil(t, got.commonPatternID, "sibling origin match should return a common tracker pattern ID")
require.NotNil(t, got.thirdPartyID, "sibling origin match should surface the sibling's org third party directly")
assert.Equal(t, orgThirdParty.ID, *got.thirdPartyID)
var commonPattern coredata.CommonTrackerPattern
@@ -874,11 +983,56 @@ func TestMatchBySiblingOrigin_ConvergentSiblings(t *testing.T) {
assert.Equal(t, fx.commonThirdPartyID, *commonPattern.CommonThirdPartyID)
}
func TestPromoteThirdParty_ExactCommonLinkIgnoresSimilarUnlinked(t *testing.T) {
t.Parallel()
client := test.PGClient(t)
ctx := context.Background()
fx := seedPromotionFixture(t, ctx, client)
now := time.Now().UTC().Truncate(time.Microsecond)
manualEntry := coredata.ThirdParty{
ID: gid.New(fx.scope.GetTenantID(), coredata.ThirdPartyEntityType),
OrganizationID: fx.organizationID,
Name: "Google LLC",
Category: coredata.ThirdPartyCategoryAnalytics,
Certifications: []string{},
Countries: coredata.CountryCodes{},
CreatedAt: now,
UpdatedAt: now,
}
linked := coredata.ThirdParty{
ID: gid.New(fx.scope.GetTenantID(), coredata.ThirdPartyEntityType),
OrganizationID: fx.organizationID,
CommonThirdPartyID: &fx.commonThirdPartyID,
Name: "Google",
Category: coredata.ThirdPartyCategoryAnalytics,
Certifications: []string{},
Countries: coredata.CountryCodes{},
CreatedAt: now,
UpdatedAt: now,
}
require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error {
if err := manualEntry.Insert(ctx, tx, fx.scope); err != nil {
return err
}
return linked.Insert(ctx, tx, fx.scope)
}))
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")
}
// TestProcess_BackfillsCommonThirdPartyFromSibling asserts that a pattern
// linked to an unlinked catalog row (no common_third_party_id) gets its
// catalog row backfilled from a sibling signal. The worker resolves the
// catalog link only; it must not promote the pattern to an org
// ThirdParty.
// catalog row backfilled from a sibling signal, and is promoted directly
// to the sibling's existing org ThirdParty.
func TestProcess_BackfillsCommonThirdPartyFromSibling(t *testing.T) {
t.Parallel()
@@ -1022,18 +1176,81 @@ func TestProcess_BackfillsCommonThirdPartyFromSibling(t *testing.T) {
return reloadedTarget.LoadByID(ctx, conn, fx.scope, target.ID)
}))
assert.Nil(t, reloadedTarget.ThirdPartyID, "target must not be auto-promoted to an org third party")
require.NotNil(t, reloadedTarget.ThirdPartyID, "target must be promoted to the sibling's org third party")
assert.Equal(t, orgThirdParty.ID, *reloadedTarget.ThirdPartyID)
require.NotNil(t, reloadedTarget.CommonTrackerPatternID)
assert.Equal(t, unlinkedCommon.ID, *reloadedTarget.CommonTrackerPatternID, "the existing catalog link must be preserved")
}
// TestProcess_SiblingCatalogResolutionOnFirstPartyOrigin asserts that a
// pattern detected on the banner's own (first-party) origin is still
// grouped with its siblings sharing that origin for catalog resolution.
// Sibling matching is an org-local co-occurrence signal and must not be
// defeated by the first-party domain filter that only protects the global
// catalog (domain) match.
func TestProcess_SiblingCatalogResolutionOnFirstPartyOrigin(t *testing.T) {
// TestProcess_UncategorisedLinksExistingThirdParty asserts that an
// uncategorised pattern is still linked to an already-existing matching
// org ThirdParty (linking to an existing party is ungated); only the
// creation of a new party stays gated, as covered by
// TestProcess_UncategorisedPatternIsNotPromoted.
func TestProcess_UncategorisedLinksExistingThirdParty(t *testing.T) {
t.Parallel()
client := test.PGClient(t)
ctx := context.Background()
fx := seedPromotionFixture(t, ctx, client)
now := time.Now().UTC().Truncate(time.Microsecond)
existing := coredata.ThirdParty{
ID: gid.New(fx.scope.GetTenantID(), coredata.ThirdPartyEntityType),
OrganizationID: fx.organizationID,
CommonThirdPartyID: &fx.commonThirdPartyID,
Name: "Google LLC",
Category: coredata.ThirdPartyCategoryAnalytics,
Certifications: []string{},
Countries: coredata.CountryCodes{},
CreatedAt: now,
UpdatedAt: now,
}
require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error {
if err := existing.Insert(ctx, tx, fx.scope); err != nil {
return err
}
_, err := tx.Exec(
ctx,
`UPDATE tracker_patterns
SET cookie_category_id = $1,
mapping_requested_at = $2
WHERE id = $3`,
fx.uncategorisedID,
now,
fx.trackerPattern.ID,
)
return err
}))
var reloadedBefore coredata.TrackerPattern
require.NoError(t, client.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error {
return reloadedBefore.LoadByID(ctx, conn, fx.scope, fx.trackerPattern.ID)
}))
h := newMappingHandler(client)
require.NoError(t, h.Process(ctx, reloadedBefore))
var reloaded coredata.TrackerPattern
require.NoError(t, client.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error {
return reloaded.LoadByID(ctx, conn, fx.scope, fx.trackerPattern.ID)
}))
require.NotNil(t, reloaded.ThirdPartyID, "uncategorised pattern must still link to an existing org third party")
assert.Equal(t, existing.ID, *reloaded.ThirdPartyID)
}
// TestProcess_SiblingPromotionOnFirstPartyOrigin asserts that a pattern
// detected on the banner's own (first-party) origin is still grouped with
// its siblings sharing that origin. Sibling matching is an org-local
// co-occurrence signal and must not be defeated by the first-party domain
// filter that only protects the global catalog (domain) match.
func TestProcess_SiblingPromotionOnFirstPartyOrigin(t *testing.T) {
t.Parallel()
client := test.PGClient(t)
@@ -1177,7 +1394,8 @@ func TestProcess_SiblingCatalogResolutionOnFirstPartyOrigin(t *testing.T) {
return reloadedTarget.LoadByID(ctx, conn, fx.scope, target.ID)
}))
assert.Nil(t, reloadedTarget.ThirdPartyID, "target must not be auto-promoted to an org third party")
require.NotNil(t, reloadedTarget.ThirdPartyID, "target sharing a first-party origin must be promoted via its sibling")
assert.Equal(t, orgThirdParty.ID, *reloadedTarget.ThirdPartyID)
}
// TestProcess_ReenqueuesUnmappedSiblingOnResolve asserts that when a
@@ -1295,8 +1513,7 @@ func TestProcess_ReenqueuesUnmappedSiblingOnResolve(t *testing.T) {
return reloadedTarget.LoadByID(ctx, conn, fx.scope, target.ID)
}))
require.NotNil(t, reloadedTarget.CommonTrackerPatternID, "target must resolve a catalog link via its sibling")
assert.Nil(t, reloadedTarget.ThirdPartyID, "target must not be auto-promoted to an org third party")
require.NotNil(t, reloadedTarget.ThirdPartyID, "target must resolve via its promoted sibling")
var reloadedUnmapped coredata.TrackerPattern
@@ -1450,8 +1667,7 @@ func TestProcess_DoesNotReenqueuePromotedOrExtensionSiblings(t *testing.T) {
return p
}
require.NotNil(t, reload(target.ID).CommonTrackerPatternID, "target must resolve a catalog link via its sibling")
assert.Nil(t, reload(target.ID).ThirdPartyID, "target must not be auto-promoted to an org third party")
require.NotNil(t, reload(target.ID).ThirdPartyID, "target must resolve via its promoted sibling")
require.NotNil(t, reload(plainSibling.ID).MappingRequestedAt, "plain unmapped sibling must be re-enqueued")
assert.Nil(t, reload(mappedSibling.ID).MappingRequestedAt, "promoted sibling must not be re-enqueued")
assert.Nil(t, reload(extensionSibling.ID).MappingRequestedAt, "EXTENSION-sourced sibling must not be re-enqueued")

134
pkg/coredata/migrations/20260611T120000Z.sql
View File

@@ -1,134 +0,0 @@
-- Copyright (c) 2026 Probo Inc <hello@probo.com>.
--
-- Permission to use, copy, modify, and/or distribute this software for any
-- purpose with or without fee is hereby granted, provided that the above
-- copyright notice and this permission notice appear in all copies.
--
-- THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
-- REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
-- AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
-- INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
-- LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
-- OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
-- PERFORMANCE OF THIS SOFTWARE.
-- Enforce a single org ThirdParty per catalog vendor per organization.
--
-- The tracker-mapping worker used to materialize org third parties from
-- the catalog and could race the load-then-create check, leaving several
-- rows that share the same (organization_id, common_third_party_id).
-- Org third parties are now created only through the explicit import
-- action; this migration both cleans up the historical duplicates and
-- adds the partial unique index that makes the invariant enforceable.
--
-- DESTRUCTIVE: the DO block below merges duplicate org third parties onto
-- the earliest-created survivor by repointing every foreign key that
-- references third_parties(id) and then deleting the extras. It is driven
-- by catalog introspection (pg_constraint / pg_index) so it covers every
-- referencing table without hard-coding the list, and it removes link
-- rows that would collide on a referencing table's unique key before
-- repointing. Validate it against a production dump before deploying.
DO $$
DECLARE
fk RECORD;
uq RECORD;
has_dupes boolean;
BEGIN
-- Map each duplicate row to the survivor it should be merged into:
-- the earliest-created row for the (organization_id,
-- common_third_party_id) pair, ties broken by id.
CREATE TEMP TABLE _tp_dedupe ON COMMIT DROP AS
SELECT t.id AS dup_id, k.survivor_id
FROM third_parties t
JOIN (
SELECT
organization_id,
common_third_party_id,
(array_agg(id ORDER BY created_at ASC, id ASC))[1] AS survivor_id
FROM third_parties
WHERE common_third_party_id IS NOT NULL
GROUP BY organization_id, common_third_party_id
HAVING count(*) > 1
) k
ON t.organization_id = k.organization_id
AND t.common_third_party_id = k.common_third_party_id
WHERE t.id <> k.survivor_id;
SELECT EXISTS (SELECT 1 FROM _tp_dedupe) INTO has_dupes;
IF has_dupes THEN
-- For every single-column foreign key that references
-- third_parties(id) ...
FOR fk IN
SELECT c.conrelid::regclass::text AS tbl,
a.attname::text AS col
FROM pg_constraint c
JOIN pg_attribute a
ON a.attrelid = c.conrelid
AND a.attnum = c.conkey[1]
WHERE c.contype = 'f'
AND c.confrelid = 'third_parties'::regclass
AND array_length(c.conkey, 1) = 1
LOOP
-- ... drop duplicate-side rows that would collide with a
-- survivor-side row on any unique key that includes the FK
-- column (the survivor's row wins; this realizes the union of
-- the two vendors' links).
FOR uq IN
SELECT (
SELECT string_agg(
format('s.%I IS NOT DISTINCT FROM d.%I',
att.attname, att.attname),
' AND ')
FROM pg_attribute att
WHERE att.attrelid = i.indrelid
AND att.attnum = ANY (i.indkey)
AND att.attname <> fk.col
) AS match_pred
FROM pg_index i
WHERE i.indrelid = fk.tbl::regclass
AND (i.indisunique OR i.indisprimary)
AND EXISTS (
SELECT 1
FROM pg_attribute att
WHERE att.attrelid = i.indrelid
AND att.attnum = ANY (i.indkey)
AND att.attname = fk.col
)
LOOP
IF uq.match_pred IS NOT NULL AND length(uq.match_pred) > 0 THEN
EXECUTE format(
'DELETE FROM %1$s d
USING _tp_dedupe m
WHERE d.%2$I = m.dup_id
AND EXISTS (
SELECT 1 FROM %1$s s
WHERE s.%2$I = m.survivor_id
AND %3$s
)',
fk.tbl, fk.col, uq.match_pred
);
END IF;
END LOOP;
-- Repoint the remaining references onto the survivor.
EXECUTE format(
'UPDATE %1$s d
SET %2$I = m.survivor_id
FROM _tp_dedupe m
WHERE d.%2$I = m.dup_id',
fk.tbl, fk.col
);
END LOOP;
-- Drop the now-unreferenced duplicate rows.
DELETE FROM third_parties t
USING _tp_dedupe m
WHERE t.id = m.dup_id;
END IF;
END $$;
CREATE UNIQUE INDEX third_parties_org_common_key
ON third_parties (organization_id, common_third_party_id)
WHERE common_third_party_id IS NOT NULL;

3
pkg/probod/probod.go
View File

@@ -321,7 +321,7 @@ func (impl *Implm) Run(
return err
}
trackerMappingCfg, trackerEnrichmentCfg, err := impl.buildTrackerAgents(l, tp, r)
trackerMappingCfg, trackerEnrichmentCfg, thirdPartyDisambiguationCfg, err := impl.buildTrackerAgents(l, tp, r)
if err != nil {
return err
}
@@ -775,6 +775,7 @@ func (impl *Implm) Run(
pgClient,
l,
trackerMappingCfg,
thirdPartyDisambiguationCfg,
time.Duration(impl.cfg.TrackerMappingWorker.StaleAfter)*time.Second,
worker.WithInterval(time.Duration(impl.cfg.TrackerMappingWorker.Interval)*time.Second),
worker.WithMaxConcurrency(impl.cfg.TrackerMappingWorker.MaxConcurrency),

52
pkg/probod/tracker_agents.go
View File

@@ -22,25 +22,27 @@ import (
"go.gearno.de/kit/log"
"go.opentelemetry.io/otel/trace"
"go.probo.inc/probo/pkg/cookiebanner"
"go.probo.inc/probo/pkg/thirdparty"
)
// buildTrackerAgents wires the two tracker agents from the probod
// buildTrackerAgents wires the three tracker agents from the probod
// config, each with its own LLM client and tuning: the tracker-mapping
// agent (catalog identification) and the common-pattern enrichment agent
// (description research). Both are opt-in: when
// `llm.tracker-mapping.provider` is empty it returns zero configs (nil
// LLM clients) so callers run without agent fallback.
// agent (catalog identification), the common-pattern enrichment agent
// (description research), and the third-party disambiguation agent. All
// are opt-in: when `llm.tracker-mapping.provider` is empty it returns
// zero configs (nil LLM clients) so callers run without agent fallback.
//
// The enrichment agent falls back to the tracker-mapping config when its
// own provider slot is empty, so a deployment that configures only
// `tracker-mapping` keeps wiring both agents.
// The enrichment and disambiguation agents fall back to the
// tracker-mapping config when their own provider slot is empty, so a
// deployment that configures only `tracker-mapping` keeps wiring all
// three agents.
func (impl *Implm) buildTrackerAgents(
l *log.Logger,
tp trace.TracerProvider,
r prometheus.Registerer,
) (cookiebanner.TrackerMappingAgentConfig, cookiebanner.TrackerEnrichmentAgentConfig, error) {
) (cookiebanner.TrackerMappingAgentConfig, cookiebanner.TrackerEnrichmentAgentConfig, thirdparty.DisambiguationAgentConfig, error) {
if impl.cfg.Agents.TrackerMapping.Provider == "" {
return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, nil
return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, thirdparty.DisambiguationAgentConfig{}, nil
}
firecrawlAPIKey := impl.cfg.Agents.Tools.FirecrawlAPIKey
@@ -53,7 +55,7 @@ func (impl *Implm) buildTrackerAgents(
r,
)
if err != nil {
return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, fmt.Errorf("cannot resolve tracker mapping agent client: %w", err)
return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, thirdparty.DisambiguationAgentConfig{}, fmt.Errorf("cannot resolve tracker mapping agent client: %w", err)
}
mappingCfg := cookiebanner.TrackerMappingAgentConfig{
@@ -79,7 +81,7 @@ func (impl *Implm) buildTrackerAgents(
r,
)
if err != nil {
return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, fmt.Errorf("cannot resolve tracker enrichment agent client: %w", err)
return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, thirdparty.DisambiguationAgentConfig{}, fmt.Errorf("cannot resolve tracker enrichment agent client: %w", err)
}
enrichmentCfg := cookiebanner.TrackerEnrichmentAgentConfig{
@@ -92,5 +94,29 @@ func (impl *Implm) buildTrackerAgents(
MaxTurns: impl.cfg.CommonPatternEnrichmentWorker.AgentMaxTurns,
}
return mappingCfg, enrichmentCfg, nil
disambiguationSlot := impl.cfg.Agents.ThirdPartyDisambiguation
if disambiguationSlot.Provider == "" {
disambiguationSlot = impl.cfg.Agents.TrackerMapping
}
disambiguationAgentCfg, disambiguationClient, err := impl.resolveAgentClient(
"third-party-disambiguation",
disambiguationSlot,
l,
tp,
r,
)
if err != nil {
return cookiebanner.TrackerMappingAgentConfig{}, cookiebanner.TrackerEnrichmentAgentConfig{}, thirdparty.DisambiguationAgentConfig{}, fmt.Errorf("cannot resolve third party disambiguation agent client: %w", err)
}
disambiguationCfg := thirdparty.DisambiguationAgentConfig{
LLMClient: disambiguationClient,
Model: disambiguationAgentCfg.ModelName,
MaxTokens: disambiguationAgentCfg.MaxTokens,
Temperature: disambiguationAgentCfg.Temperature,
Timeout: time.Duration(impl.cfg.TrackerMappingWorker.DisambiguationAgentTimeout) * time.Second,
}
return mappingCfg, enrichmentCfg, disambiguationCfg, nil
}

31
pkg/probodconfig/llm_config.go
View File

@@ -53,12 +53,16 @@ type (
// tracker-mapping background worker. LLM parameters for the mapping
// agent it runs live under AgentsConfig.TrackerMapping. AgentTimeout
// and AgentMaxTurns bound a single mapping agent run.
// DisambiguationAgentTimeout caps a single third-party
// disambiguation agent run; that agent runs inside this worker but
// uses its own LLM parameters from AgentsConfig.ThirdPartyDisambiguation.
TrackerMappingWorkerConfig 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"`
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"`
DisambiguationAgentTimeout int `json:"disambiguation-agent-timeout"` // seconds, single disambiguation run
}
// CommonPatternEnrichmentWorkerConfig holds worker-side tuning for
@@ -82,14 +86,15 @@ type (
// settings. Default is used as a fallback when an agent-specific field
// is zero-valued.
AgentsConfig struct {
Providers map[string]LLMProviderConfig `json:"providers"`
Default LLMAgentConfig `json:"defaults"`
Probo LLMAgentConfig `json:"probo"`
EvidenceDescriber LLMAgentConfig `json:"evidence-describer"`
ThirdPartyVetter LLMAgentConfig `json:"third-party-vetter"`
TrackerMapping LLMAgentConfig `json:"tracker-mapping"`
TrackerEnrichment LLMAgentConfig `json:"tracker-enrichment"`
Tools AgentToolsConfig `json:"tools"`
Providers map[string]LLMProviderConfig `json:"providers"`
Default LLMAgentConfig `json:"defaults"`
Probo LLMAgentConfig `json:"probo"`
EvidenceDescriber LLMAgentConfig `json:"evidence-describer"`
ThirdPartyVetter LLMAgentConfig `json:"third-party-vetter"`
ThirdPartyDisambiguation LLMAgentConfig `json:"third-party-disambiguation"`
TrackerMapping LLMAgentConfig `json:"tracker-mapping"`
TrackerEnrichment LLMAgentConfig `json:"tracker-enrichment"`
Tools AgentToolsConfig `json:"tools"`
}
)

217
pkg/thirdparty/common_match.go vendored Normal file
View File

@@ -0,0 +1,217 @@
// Copyright (c) 2026 Probo Inc <hello@probo.com>.
//
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
// PERFORMANCE OF THIS SOFTWARE.
package thirdparty
import (
"context"
"fmt"
"sort"
"strings"
"go.gearno.de/kit/pg"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
"go.probo.inc/probo/pkg/slug"
"go.probo.inc/probo/pkg/uri"
)
// Heuristic thresholds for matching a CommonThirdParty to an existing
// org ThirdParty. Exported so callers can short-circuit explicitly
// (skip the agent, leave the pattern unlinked, etc.) instead of
// duplicating magic numbers.
const (
// HighConfidenceScore is the floor at which a heuristic match
// is treated as obvious (exact name, suffix-stripped name, slug
// equality). Callers typically link without consulting the
// agent at this score.
HighConfidenceScore = 0.85
// MinAgentScore is the floor below which a candidate is
// statistical noise and should not be shown to the
// disambiguation agent. Below this, callers typically prefer
// to create a fresh row over asking the model to pick among
// weak candidates.
MinAgentScore = 0.6
// MaxAgentCandidates caps the candidate list shown to the
// disambiguation agent. The list is heuristic-ranked, so the
// top few are the only ones worth the agent's tokens.
MaxAgentCandidates = 5
)
// ScoredCandidate is a scored heuristic-match candidate. It is the
// unified currency between the heuristic ranker (RankCandidates) and
// the disambiguation agent (Disambiguate): the agent renders the
// `ThirdParty` fields plus the score directly into its prompt, with
// no intermediate DTO.
type ScoredCandidate struct {
ThirdParty *coredata.ThirdParty
Score float64
}
// corporateSuffixes are the legal-form noise words stripped when
// comparing third-party names heuristically. The list is intentionally
// short and conservative: matching "Foo Inc" to "Foo" is safe, but
// stripping "Group" or "Services" would over-match unrelated entries.
//
// Order matters: stripCorporateSuffixes returns on the first match,
// so longer / comma-prefixed forms must come before their shorter
// siblings (", inc." before " inc.", which itself comes before " inc").
var corporateSuffixes = []string{
" incorporated",
" corporation",
", inc.",
", inc",
" l.l.c.",
" s.a.s.",
" inc.",
" inc",
" llc",
" ltd.",
" ltd",
" limited",
" gmbh",
" s.a.",
" sas",
" sa",
" ag",
" plc",
" corp.",
" corp",
" co.",
" co",
" b.v.",
" bv",
}
// RankCandidates ranks org ThirdParty rows by how likely each is to
// represent the given CommonThirdParty. Returned slice is sorted by
// descending score; only candidates with score > 0 are kept. Pure
// function: no I/O, deterministic on its inputs.
//
// Scoring (highest match wins; website-host overlap can lift a name
// miss to 0.8):
//
// - exact lowercase name = 1.0
// - lowercase name with corporate suffix stripped, equal = 0.9
// - slug equality (slug.Make on the org's name) = 0.85
// - website host (eTLD+1) overlap with the catalog domain set = 0.8
func RankCandidates(
commonParty coredata.CommonThirdParty,
commonDomains coredata.CommonThirdPartyDomains,
candidates coredata.ThirdParties,
) []ScoredCandidate {
commonName := strings.ToLower(strings.TrimSpace(commonParty.Name))
commonStripped := stripCorporateSuffixes(commonName)
commonSlug := commonParty.Slug
commonHost := ""
if commonParty.WebsiteURL != nil {
commonHost = uri.ExtractDomain(*commonParty.WebsiteURL)
}
commonDomainSet := make(map[string]struct{}, len(commonDomains))
for _, d := range commonDomains {
commonDomainSet[strings.ToLower(d.Domain)] = struct{}{}
}
if commonHost != "" {
commonDomainSet[commonHost] = struct{}{}
}
scored := make([]ScoredCandidate, 0, len(candidates))
for _, tp := range candidates {
score := 0.0
orgName := strings.ToLower(strings.TrimSpace(tp.Name))
orgStripped := stripCorporateSuffixes(orgName)
switch {
case orgName != "" && orgName == commonName:
score = 1.0
case orgStripped != "" && orgStripped == commonStripped:
score = 0.9
case commonSlug != "" && slug.Make(tp.Name) == commonSlug:
score = 0.85
}
if tp.WebsiteURL != nil {
orgHost := uri.ExtractDomain(*tp.WebsiteURL)
if orgHost != "" {
if _, hit := commonDomainSet[orgHost]; hit {
if score < 0.8 {
score = 0.8
}
}
}
}
if score == 0 {
continue
}
scored = append(scored, ScoredCandidate{
ThirdParty: tp,
Score: score,
})
}
sort.SliceStable(scored, func(i, j int) bool {
return scored[i].Score > scored[j].Score
})
return scored
}
// stripCorporateSuffixes removes a single trailing legal-form suffix
// from a lowercased name. Only one suffix is stripped to avoid
// mangling names that happen to end in two stop-words (e.g. "Foo Inc
// LLC" → "Foo Inc", not "Foo").
func stripCorporateSuffixes(lowerName string) string {
for _, s := range corporateSuffixes {
if before, ok := strings.CutSuffix(lowerName, s); ok {
return strings.TrimSpace(before)
}
}
return lowerName
}
// LinkToCommon writes common_third_party_id onto an org ThirdParty so
// future matches against the same CommonThirdParty can short-circuit
// to the exact-link path in O(1). No-op when the field is already set
// (to any value) — we never overwrite an existing catalog link because
// a heuristic or agent false-positive must not corrupt a previous,
// possibly more accurate, association.
func LinkToCommon(
ctx context.Context,
tx pg.Tx,
scope coredata.Scoper,
orgThirdParty *coredata.ThirdParty,
commonID gid.GID,
) error {
if orgThirdParty.CommonThirdPartyID != nil {
return nil
}
orgThirdParty.CommonThirdPartyID = &commonID
if err := orgThirdParty.Update(ctx, tx, scope); err != nil {
return fmt.Errorf("cannot update third party with common id: %w", err)
}
return nil
}

178
pkg/thirdparty/common_match_test.go vendored Normal file
View File

@@ -0,0 +1,178 @@
// Copyright (c) 2026 Probo Inc <hello@probo.com>.
//
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
// PERFORMANCE OF THIS SOFTWARE.
package thirdparty
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
)
func TestStripCorporateSuffixes(t *testing.T) {
t.Parallel()
tests := []struct {
name string
in string
want string
}{
{name: "llc suffix", in: "google llc", want: "google"},
{name: "comma inc", in: "stripe, inc", want: "stripe"},
{name: "inc dot", in: "meta inc.", want: "meta"},
{name: "ltd", in: "deepmind ltd", want: "deepmind"},
{name: "gmbh", in: "n8n gmbh", want: "n8n"},
{name: "no suffix", in: "cloudflare", want: "cloudflare"},
{name: "trailing space", in: "github inc", want: "github"},
{name: "only one suffix stripped", in: "foo inc llc", want: "foo inc"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
assert.Equal(t, tt.want, stripCorporateSuffixes(tt.in))
})
}
}
func TestRankCandidates(t *testing.T) {
t.Parallel()
tenantID := gid.NewTenantID()
mkTP := func(name, website string) *coredata.ThirdParty {
tp := &coredata.ThirdParty{
ID: gid.New(tenantID, coredata.ThirdPartyEntityType),
Name: name,
}
if website != "" {
tp.WebsiteURL = new(website)
}
return tp
}
t.Run("exact name match scores 1.0", func(t *testing.T) {
t.Parallel()
common := coredata.CommonThirdParty{Name: "Google", Slug: "google"}
got := RankCandidates(common, nil, coredata.ThirdParties{
mkTP("Google", ""),
mkTP("Stripe", ""),
})
require.Len(t, got, 1)
assert.Equal(t, 1.0, got[0].Score)
assert.Equal(t, "Google", got[0].ThirdParty.Name)
})
t.Run("suffix-stripped name scores 0.9", func(t *testing.T) {
t.Parallel()
common := coredata.CommonThirdParty{Name: "Google", Slug: "google"}
got := RankCandidates(common, nil, coredata.ThirdParties{
mkTP("Google LLC", ""),
})
require.Len(t, got, 1)
assert.Equal(t, 0.9, got[0].Score)
})
t.Run("slug equality scores 0.85", func(t *testing.T) {
t.Parallel()
common := coredata.CommonThirdParty{Name: "Google", Slug: "google"}
got := RankCandidates(common, nil, coredata.ThirdParties{
mkTP("google!", ""),
})
require.Len(t, got, 1)
assert.Equal(t, 0.85, got[0].Score)
})
t.Run("website host overlap scores 0.8 when name does not match", func(t *testing.T) {
t.Parallel()
common := coredata.CommonThirdParty{
Name: "Google Analytics",
Slug: "google-analytics",
WebsiteURL: new("https://google.com"),
}
got := RankCandidates(common, nil, coredata.ThirdParties{
mkTP("Sundar's Search Co", "https://www.google.com/about"),
})
require.Len(t, got, 1)
assert.Equal(t, 0.8, got[0].Score)
})
t.Run("domain set overlap scores 0.8", func(t *testing.T) {
t.Parallel()
common := coredata.CommonThirdParty{Name: "Stripe", Slug: "stripe"}
domains := coredata.CommonThirdPartyDomains{
{Domain: "stripe.com"},
{Domain: "stripe.network"},
}
got := RankCandidates(common, domains, coredata.ThirdParties{
mkTP("Payment Processor", "https://api.stripe.com/v1"),
})
require.Len(t, got, 1)
assert.Equal(t, 0.8, got[0].Score)
})
t.Run("no match returns empty", func(t *testing.T) {
t.Parallel()
common := coredata.CommonThirdParty{Name: "Stripe", Slug: "stripe"}
got := RankCandidates(common, nil, coredata.ThirdParties{
mkTP("Acme", "https://acme.example"),
mkTP("Widgets Inc", "https://widgets.example"),
})
assert.Empty(t, got)
})
t.Run("ranks descending by score", func(t *testing.T) {
t.Parallel()
common := coredata.CommonThirdParty{
Name: "Google",
Slug: "google",
WebsiteURL: new("https://google.com"),
}
got := RankCandidates(common, nil, coredata.ThirdParties{
mkTP("Random", "https://google.com"),
mkTP("Google", ""),
mkTP("Google LLC", ""),
})
require.Len(t, got, 3)
assert.Equal(t, "Google", got[0].ThirdParty.Name)
assert.Equal(t, 1.0, got[0].Score)
assert.Equal(t, "Google LLC", got[1].ThirdParty.Name)
assert.Equal(t, 0.9, got[1].Score)
assert.Equal(t, "Random", got[2].ThirdParty.Name)
assert.Equal(t, 0.8, got[2].Score)
})
}

241
pkg/thirdparty/disambiguation_agent.go vendored Normal file
View File

@@ -0,0 +1,241 @@
// Copyright (c) 2026 Probo Inc <hello@probo.com>.
//
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
// PERFORMANCE OF THIS SOFTWARE.
package thirdparty
import (
"context"
_ "embed"
"fmt"
"strings"
"time"
"go.gearno.de/kit/log"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
"go.probo.inc/probo/pkg/llm"
)
//go:embed prompts/disambiguation.txt.tmpl
var disambiguationPrompt string
const (
// disambiguationConfidenceThreshold is the floor below which we
// treat the agent's pick as "no confident match" even when it
// returned a non-nil matched_id. Mirrors the conservative bias
// described in the prompt.
disambiguationConfidenceThreshold = 0.6
// 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 output when the
// config carries no max-tokens budget. The final output is tiny (a
// single id plus a one-sentence rationale), but the budget must
// leave ample headroom for reasoning models (e.g. the GPT-5
// family): their reasoning tokens count against max_tokens, so too
// small a budget gets consumed by reasoning and truncates the JSON,
// surfacing as "unexpected end of JSON input".
defaultDisambiguationMaxTokens = 4096
)
// DisambiguationAgentConfig 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 DisambiguationAgentConfig struct {
LLMClient *llm.Client
Model string
MaxTokens *int
Temperature *float64
Timeout time.Duration
}
// DisambiguationResult is the structured output the disambiguation
// agent returns when picking the best existing org ThirdParty for a
// catalog entry.
type DisambiguationResult struct {
MatchedID *string `json:"matched_id" jsonschema:"GID of the org third party that best matches, or null if none of the candidates is a confident match."`
Confidence float64 `json:"confidence" jsonschema:"Confidence level from 0.0 to 1.0. Below 0.6 means 'no confident match' and matched_id MUST be null."`
Reasoning string `json:"reasoning" jsonschema:"One short sentence describing the rationale."`
}
// BuildDisambiguationAgent wires the agent that picks the best
// existing org ThirdParty for a catalog entry. It deliberately has
// no tools: the candidate list is supplied in the prompt and the
// agent must only choose among it.
func BuildDisambiguationAgent(
cfg DisambiguationAgentConfig,
logger *log.Logger,
) *agent.Agent {
outputType, err := agent.NewOutputType[DisambiguationResult]("third_party_disambiguation")
if err != nil {
panic(fmt.Sprintf("thirdparty: cannot build disambiguation output type: %s", err))
}
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
// and candidate list, and returns the matched candidate's ID — or
// nil when the agent picks "none", returns a confidence below the
// threshold, or fails. Errors from the agent itself are returned;
// "no confident match" is not an error.
//
// The matched candidate is identified by string equality against the
// IDs supplied in `candidates`; we never invent IDs from the agent's
// output, so a model that hallucinates an ID is treated as "none".
func Disambiguate(
ctx context.Context,
a *agent.Agent,
logger *log.Logger,
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, timeout)
defer cancel()
result, err := agent.RunTyped[DisambiguationResult](
agentCtx,
a,
[]llm.Message{
{
Role: llm.RoleUser,
Parts: []llm.Part{llm.TextPart{Text: prompt}},
},
},
)
if err != nil {
return nil, fmt.Errorf("cannot run disambiguation agent: %w", err)
}
out := result.Output
if out.MatchedID == nil || *out.MatchedID == "" {
return nil, nil
}
if out.Confidence < disambiguationConfidenceThreshold {
logger.InfoCtx(
ctx,
"disambiguation agent below confidence threshold",
log.String("matched_id", *out.MatchedID),
log.Float64("confidence", out.Confidence),
)
return nil, nil
}
for _, c := range candidates {
if c.ThirdParty.ID.String() == *out.MatchedID {
id := c.ThirdParty.ID
return &id, nil
}
}
logger.WarnCtx(
ctx,
"disambiguation agent returned id not in candidate list",
log.String("matched_id", *out.MatchedID),
)
return nil, nil
}
// buildDisambiguationPrompt formats the catalog third party and the
// heuristic-ranked candidate list into the user message for the
// disambiguation agent. The prompt is intentionally compact: the
// agent only needs ids, names, websites, and the heuristic score to
// decide.
func buildDisambiguationPrompt(
commonParty coredata.CommonThirdParty,
commonDomains coredata.CommonThirdPartyDomains,
candidates []ScoredCandidate,
) string {
var b strings.Builder
b.WriteString("Catalog third party:\n")
fmt.Fprintf(&b, " name: %s\n", commonParty.Name)
if commonParty.WebsiteURL != nil && *commonParty.WebsiteURL != "" {
fmt.Fprintf(&b, " website: %s\n", *commonParty.WebsiteURL)
}
if len(commonDomains) > 0 {
domains := make([]string, len(commonDomains))
for i, d := range commonDomains {
domains[i] = d.Domain
}
fmt.Fprintf(&b, " domains: %s\n", strings.Join(domains, ", "))
}
b.WriteString("\nCandidate organisation third parties (heuristic-ranked):\n")
for i, c := range candidates {
fmt.Fprintf(&b, "- id: %s\n", c.ThirdParty.ID.String())
fmt.Fprintf(&b, " name: %s\n", c.ThirdParty.Name)
if c.ThirdParty.WebsiteURL != nil && *c.ThirdParty.WebsiteURL != "" {
fmt.Fprintf(&b, " website: %s\n", *c.ThirdParty.WebsiteURL)
}
fmt.Fprintf(&b, " heuristic_score: %.2f\n", c.Score)
if i < len(candidates)-1 {
b.WriteString("\n")
}
}
return b.String()
}