Promote tracker patterns to org third parties via worker

Manual moves of a non-extension TrackerPattern lacking a ThirdPartyID
now request mapping, which the tracker-mapping worker resolves with a
four-stage pipeline: exact common_third_party_id link, heuristic
ranking, agent disambiguation, and finally CreateFromCommon. Existing
fuzzy-matched org rows are tagged with common_third_party_id so the
next promotion takes the O(1) exact-link path.

The matching primitives live in pkg/thirdparty (RankCandidates,
LinkToCommon, CreateFromCommon, ScoredCandidate, threshold constants)
so the disambiguation agent and the heuristic share one candidate
type. Cookiebanner orchestrates them; cookie-banner-specific concerns
(pattern -> common-pattern -> common-party navigation, the EXTENSION
gate, and structured logs) stay in the worker.

Signed-off-by: Émile Ré <emile@probo.com>
This commit is contained in:
Émile Ré
2026-05-28 17:20:36 +02:00
parent 33ddba9770
commit a99a4dde14
12 changed files with 1491 additions and 94 deletions

View File

@@ -42,17 +42,48 @@ pkg/cookiebanner/pattern_analysis_worker.go
## Agent files
When a package has a worker that uses an agent, the agent construction logic
goes in `<worker_prefix>_agent.go` alongside `<worker_prefix>_worker.go`. The
worker file stays focused on `Claim`/`Process` and handler methods; the agent
file owns agent construction, prompt building, constants, config, and the
`//go:embed` directive for prompt templates.
A file whose **sole purpose** is to construct and operate an agent uses the
`<name>_agent.go` suffix. The agent file owns agent construction, prompt
building, the `//go:embed` directive for prompt templates, the typed result,
the agent-specific config, and any agent-only constants (timeout, confidence
threshold). Callers (workers, services) hold a `*agent.Agent` field and import
the file's `Build…Agent` constructor.
This applies in two shapes:
1. **Paired with a worker** (most common). The worker file
`<worker_prefix>_worker.go` stays focused on `Claim`/`Process`, and the
agent it uses lives in `<worker_prefix>_agent.go` next to it.
```
pkg/cookiebanner/tracker_mapping_worker.go -- worker handler
pkg/cookiebanner/tracker_mapping_agent.go -- agent construction + prompts
```
2. **Standalone, called from elsewhere.** When the agent is consumed by a
different package (or by multiple packages — e.g. an agent that operates on
a domain entity, used by several feature workers), it lives in the package
that owns the domain, named `<purpose>_agent.go`.
```
pkg/thirdparty/disambiguation_agent.go -- catalog→org ThirdParty matcher
pkg/vetting/sub_agent.go -- generic vetting sub-agent
```
A file is NOT renamed to `_agent.go` when the agent is incidental to a service
that does substantially more than agent orchestration (e.g. CRUD, caching,
auth). In that case the file keeps its service name and the agent is built
inline:
```
pkg/cookiebanner/tracker_mapping_worker.go -- worker handler
pkg/cookiebanner/tracker_mapping_agent.go -- agent construction + prompts
pkg/evidencedescriber/evidencedescriber.go -- single-file describer service
pkg/vetting/assessment.go -- third-party assessment service
```
If the agent construction grows past a few dozen lines or sprouts its own
prompt embed / typed result / config struct, extract it into a sibling
`<purpose>_agent.go`.
## Tool files
Each agent tool lives in its own `<tool_name>_tool.go` file, named after the

View File

@@ -2628,6 +2628,20 @@ func (s *Service) MoveTrackerPatternToCategory(
return fmt.Errorf("cannot update tracker pattern: %w", err)
}
// A manual move is the user's signal that this is a
// real tracker. Enqueue the tracker-mapping worker so
// it can promote the pattern to an org ThirdParty (or
// link an existing one) — never EXTENSION-sourced
// patterns, and never patterns we already promoted.
// SetMappingRequested is idempotent: it short-circuits
// when mapping_requested_at is already non-NULL.
if pattern.ThirdPartyID == nil &&
(pattern.Source == nil || *pattern.Source != coredata.CookieSourceExtension) {
if err := pattern.SetMappingRequested(ctx, tx); err != nil {
return fmt.Errorf("cannot enqueue tracker mapping after move: %w", err)
}
}
var banner coredata.CookieBanner
if err := banner.LoadByID(ctx, tx, scope, pattern.CookieBannerID); err != nil {
return fmt.Errorf("cannot load cookie banner: %w", err)

View File

@@ -48,6 +48,9 @@ type TrackerMappingAgentResult struct {
Confidence float64 `json:"confidence" jsonschema:"Confidence level from 0.0 to 1.0. Set below 0.5 if unsure."`
}
// TrackerMappingConfig configures the tracker-mapping agent (catalog
// identification). The agent uses DB-backed search tools and may also
// use Firecrawl for web search when an API key is supplied.
type TrackerMappingConfig struct {
LLMClient *llm.Client
Model string

View File

@@ -28,18 +28,21 @@ import (
"go.probo.inc/probo/pkg/gid"
"go.probo.inc/probo/pkg/llm"
"go.probo.inc/probo/pkg/slug"
"go.probo.inc/probo/pkg/thirdparty"
)
type trackerMappingHandler struct {
pg *pg.Client
logger *log.Logger
agent *agent.Agent
pg *pg.Client
logger *log.Logger
mappingAgent *agent.Agent
disambiguationAgent *agent.Agent
}
func NewTrackerMappingWorker(
pgClient *pg.Client,
logger *log.Logger,
cfg TrackerMappingConfig,
mappingCfg TrackerMappingConfig,
disambiguationCfg thirdparty.DisambiguationConfig,
opts ...worker.Option,
) *worker.Worker[coredata.TrackerPattern] {
h := &trackerMappingHandler{
@@ -47,8 +50,12 @@ func NewTrackerMappingWorker(
logger: logger,
}
if cfg.LLMClient != nil {
h.agent = buildTrackerMappingAgent(cfg, pgClient, logger)
if mappingCfg.LLMClient != nil {
h.mappingAgent = buildTrackerMappingAgent(mappingCfg, pgClient, logger)
}
if disambiguationCfg.LLMClient != nil {
h.disambiguationAgent = thirdparty.BuildDisambiguationAgent(disambiguationCfg, logger)
}
return worker.New(
@@ -82,40 +89,62 @@ func (h *trackerMappingHandler) Claim(ctx context.Context) (coredata.TrackerPatt
return tp, nil
}
// Process resolves the catalog mapping (when missing) and then promotes
// the pattern to an org ThirdParty (when eligible). When a pattern is
// re-triggered by a manual move (it already carries a
// common_tracker_pattern_id), we MUST NOT re-resolve the catalog: the
// existing link is preserved and we jump straight to third-party
// promotion.
func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.TrackerPattern) error {
return h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
var (
commonPatternID *gid.GID
thirdPartyID *gid.GID
err error
)
commonPatternID, thirdPartyID, err = h.matchByPattern(ctx, tx, tp)
if err != nil {
return fmt.Errorf("cannot match by pattern: %w", err)
}
if commonPatternID == nil {
commonPatternID, thirdPartyID, err = h.matchByDomain(ctx, tx, tp)
if tp.CommonTrackerPatternID != nil {
commonPatternID = tp.CommonTrackerPatternID
} else {
commonPatternID, err = h.matchByPattern(ctx, tx, tp)
if err != nil {
return fmt.Errorf("cannot match by domain: %w", err)
return fmt.Errorf("cannot match by pattern: %w", err)
}
if commonPatternID == nil {
commonPatternID, err = h.matchByDomain(ctx, tx, tp)
if err != nil {
return fmt.Errorf("cannot match by domain: %w", err)
}
}
if commonPatternID == nil && h.mappingAgent != nil {
commonPatternID, err = h.identifyWithAgent(ctx, tx, tp)
if err != nil {
return fmt.Errorf("cannot identify with agent: %w", err)
}
}
if commonPatternID == nil {
commonPatternID, err = h.createUnmatchedPattern(ctx, tx, tp)
if err != nil {
return fmt.Errorf("cannot create unmatched pattern: %w", err)
}
}
}
if commonPatternID == nil && h.agent != nil {
commonPatternID, thirdPartyID, err = h.identifyWithAgent(ctx, tx, tp)
if err != nil {
return fmt.Errorf("cannot identify with agent: %w", err)
}
}
thirdPartyID := tp.ThirdPartyID
if commonPatternID == nil {
commonPatternID, err = h.createUnmatchedPattern(ctx, tx, tp)
if thirdPartyID == nil &&
commonPatternID != nil &&
(tp.Source == nil || *tp.Source != coredata.CookieSourceExtension) {
promoted, err := h.promoteThirdParty(ctx, tx, tp, *commonPatternID)
if err != nil {
return fmt.Errorf("cannot create unmatched pattern: %w", err)
return fmt.Errorf("cannot promote third party: %w", err)
}
thirdPartyID = promoted
}
if commonPatternID != nil || thirdPartyID != nil {
@@ -136,59 +165,55 @@ func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.Tracker
)
}
// matchByPattern looks for a catalog row with the same pattern. It now
// only returns the catalog ID; third-party resolution happens later in
// promoteThirdParty.
func (h *trackerMappingHandler) matchByPattern(
ctx context.Context,
conn pg.Querier,
tp coredata.TrackerPattern,
) (*gid.GID, *gid.GID, error) {
) (*gid.GID, error) {
var commonPattern coredata.CommonTrackerPattern
if err := commonPattern.LoadByPattern(ctx, conn, tp.TrackerType, tp.Pattern, tp.MaxAgeSeconds); err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
return nil, nil, nil
return nil, nil
}
return nil, nil, fmt.Errorf("cannot load common tracker pattern: %w", err)
return nil, fmt.Errorf("cannot load common tracker pattern: %w", err)
}
var thirdPartyID *gid.GID
if commonPattern.CommonThirdPartyID != nil {
var err error
thirdPartyID, err = h.resolveThirdParty(ctx, conn, tp, &commonPattern)
if err != nil {
return nil, nil, fmt.Errorf("cannot resolve third party from pattern match: %w", err)
}
}
return &commonPattern.ID, thirdPartyID, nil
return &commonPattern.ID, nil
}
// matchByDomain finds a CommonThirdParty whose registered domains
// overlap the pattern's observed initiator domains, and upserts a
// CommonTrackerPattern linking the two. As with matchByPattern,
// third-party resolution is deferred to promoteThirdParty.
func (h *trackerMappingHandler) matchByDomain(
ctx context.Context,
tx pg.Tx,
tp coredata.TrackerPattern,
) (*gid.GID, *gid.GID, error) {
) (*gid.GID, error) {
var trackers coredata.DetectedTrackers
domains, err := trackers.LoadInitiatorDomainsByTrackerPatternID(ctx, tx, tp.ID, 10)
if err != nil {
return nil, nil, fmt.Errorf("cannot load initiator domains: %w", err)
return nil, fmt.Errorf("cannot load initiator domains: %w", err)
}
if len(domains) == 0 {
return nil, nil, nil
return nil, nil
}
filter := coredata.NewCommonThirdPartyDomainFilter(domains)
var matchedDomains coredata.CommonThirdPartyDomains
if err := matchedDomains.Load(ctx, tx, 1, filter); err != nil {
return nil, nil, fmt.Errorf("cannot load common third party domain by domain match: %w", err)
return nil, fmt.Errorf("cannot load common third party domain by domain match: %w", err)
}
if len(matchedDomains) == 0 {
return nil, nil, nil
return nil, nil
}
commonThirdPartyID := matchedDomains[0].CommonThirdPartyID
@@ -208,22 +233,17 @@ func (h *trackerMappingHandler) matchByDomain(
}
if _, err := commonPattern.Upsert(ctx, tx); err != nil {
return nil, nil, fmt.Errorf("cannot upsert common tracker pattern from domain match: %w", err)
return nil, fmt.Errorf("cannot upsert common tracker pattern from domain match: %w", err)
}
thirdPartyID, err := h.resolveThirdParty(ctx, tx, tp, &commonPattern)
if err != nil {
return nil, nil, fmt.Errorf("cannot resolve third party from domain match: %w", err)
}
return &commonPattern.ID, thirdPartyID, nil
return &commonPattern.ID, nil
}
func (h *trackerMappingHandler) identifyWithAgent(
ctx context.Context,
tx pg.Tx,
tp coredata.TrackerPattern,
) (*gid.GID, *gid.GID, error) {
) (*gid.GID, error) {
var trackers coredata.DetectedTrackers
domains, err := trackers.LoadInitiatorDomainsByTrackerPatternID(ctx, tx, tp.ID, 5)
@@ -238,7 +258,7 @@ func (h *trackerMappingHandler) identifyWithAgent(
result, err := agent.RunTyped[TrackerMappingAgentResult](
agentCtx,
h.agent,
h.mappingAgent,
[]llm.Message{
{
Role: llm.RoleUser,
@@ -254,7 +274,7 @@ func (h *trackerMappingHandler) identifyWithAgent(
log.String("pattern", tp.Pattern),
)
return nil, nil, nil
return nil, nil
}
identification := result.Output
@@ -267,7 +287,7 @@ func (h *trackerMappingHandler) identifyWithAgent(
log.Float64("confidence", identification.Confidence),
)
return nil, nil, nil
return nil, nil
}
confidence := float32(identification.Confidence)
@@ -284,7 +304,7 @@ func (h *trackerMappingHandler) identifyWithAgent(
domains,
)
if err != nil {
return nil, nil, fmt.Errorf("cannot resolve or create common third party: %w", err)
return nil, fmt.Errorf("cannot resolve or create common third party: %w", err)
}
}
@@ -303,12 +323,7 @@ func (h *trackerMappingHandler) identifyWithAgent(
}
if _, err := commonPattern.Upsert(ctx, tx); err != nil {
return nil, nil, fmt.Errorf("cannot upsert common tracker pattern from agent: %w", err)
}
thirdPartyID, err := h.resolveThirdParty(ctx, tx, tp, &commonPattern)
if err != nil {
return nil, nil, fmt.Errorf("cannot resolve third party from agent match: %w", err)
return nil, fmt.Errorf("cannot upsert common tracker pattern from agent: %w", err)
}
h.logger.InfoCtx(
@@ -319,7 +334,7 @@ func (h *trackerMappingHandler) identifyWithAgent(
log.Float64("confidence", identification.Confidence),
)
return &commonPattern.ID, thirdPartyID, nil
return &commonPattern.ID, nil
}
func (h *trackerMappingHandler) resolveOrCreateCommonThirdParty(
@@ -408,32 +423,163 @@ func (h *trackerMappingHandler) createUnmatchedPattern(
return &commonPattern.ID, nil
}
func (h *trackerMappingHandler) resolveThirdParty(
// promoteThirdParty resolves an org ThirdParty for the given pattern
// once the catalog mapping is known. 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.
// 4. Fallback create from CommonThirdParty.
//
// A confident heuristic/agent match is auto-tagged with
// common_third_party_id so subsequent promotions hit the exact-link
// path in O(1). Returns (nil, nil) when the catalog row has no
// CommonThirdPartyID — there is nothing to promote to.
func (h *trackerMappingHandler) promoteThirdParty(
ctx context.Context,
conn pg.Querier,
tx pg.Tx,
tp coredata.TrackerPattern,
commonPattern *coredata.CommonTrackerPattern,
commonPatternID gid.GID,
) (*gid.GID, error) {
var commonPattern coredata.CommonTrackerPattern
if err := commonPattern.LoadByID(ctx, tx, commonPatternID); err != nil {
return nil, fmt.Errorf("cannot load common tracker pattern: %w", err)
}
if commonPattern.CommonThirdPartyID == nil {
return nil, nil
}
commonThirdPartyID := *commonPattern.CommonThirdPartyID
scope := coredata.NewScopeFromObjectID(tp.ID)
var t coredata.ThirdParty
if err := t.LoadByOrganizationIDAndCommonThirdPartyID(
var existing coredata.ThirdParty
err := existing.LoadByOrganizationIDAndCommonThirdPartyID(
ctx,
conn,
tx,
scope,
tp.OrganizationID,
*commonPattern.CommonThirdPartyID,
); err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
return nil, nil
}
return nil, fmt.Errorf("cannot resolve third party: %w", err)
commonThirdPartyID,
)
if err == nil {
return &existing.ID, nil
}
return &t.ID, nil
if !errors.Is(err, coredata.ErrResourceNotFound) {
return nil, fmt.Errorf("cannot load org third party by common id: %w", err)
}
var commonParty coredata.CommonThirdParty
if err := commonParty.LoadByID(ctx, tx, commonThirdPartyID); err != nil {
return nil, fmt.Errorf("cannot load common third party: %w", err)
}
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 {
matchedID, err := thirdparty.Disambiguate(
ctx,
h.disambiguationAgent,
h.logger,
commonParty,
commonDomains,
agentSet,
)
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 {
var picked *coredata.ThirdParty
for _, c := range agentSet {
if c.ThirdParty.ID == *matchedID {
picked = c.ThirdParty
break
}
}
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)
}
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()),
)
return &picked.ID, nil
}
}
}
created, err := thirdparty.CreateFromCommon(ctx, tx, scope, tp.OrganizationID, commonParty)
if err != nil {
return nil, 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()),
)
return &created.ID, nil
}

View File

@@ -0,0 +1,495 @@
// Copyright (c) 2026 Probo Inc <hello@getprobo.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 cookiebanner
import (
"context"
"io"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
)
//go:fix inline
func ptr[T any](v T) *T { return new(v) }
// promotionFixture extends workerFixture with a CommonThirdParty and a
// CommonTrackerPattern linking the catalog to the test pattern. It is
// the minimum scaffolding promoteThirdParty needs to run end-to-end.
type promotionFixture struct {
workerFixture
commonThirdParty coredata.CommonThirdParty
commonPatternID gid.GID
trackerPattern coredata.TrackerPattern
commonThirdPartyID gid.GID
}
func seedPromotionFixture(t *testing.T, ctx context.Context, client *pg.Client) promotionFixture {
t.Helper()
fx := seedWorkerFixture(t, ctx, client)
now := time.Now().UTC().Truncate(time.Microsecond)
commonThirdPartyID := gid.New(gid.NilTenant, coredata.CommonThirdPartyEntityType)
commonThirdParty := coredata.CommonThirdParty{
ID: commonThirdPartyID,
Name: "Google",
Slug: "google",
Category: coredata.ThirdPartyCategoryAnalytics,
WebsiteURL: new("https://google.com"),
Certifications: []string{},
CreatedAt: now,
UpdatedAt: now,
}
commonPattern := coredata.CommonTrackerPattern{
ID: gid.New(gid.NilTenant, coredata.CommonTrackerPatternEntityType),
CommonThirdPartyID: &commonThirdPartyID,
TrackerType: coredata.TrackerTypeCookie,
Pattern: "_ga",
MatchType: coredata.TrackerPatternMatchTypeExact,
Description: "",
Confidence: 0.9,
CreatedAt: now,
UpdatedAt: now,
}
pattern := coredata.TrackerPattern{
ID: gid.New(fx.scope.GetTenantID(), coredata.TrackerPatternEntityType),
OrganizationID: fx.organizationID,
CookieBannerID: fx.banner.ID,
CookieCategoryID: fx.normalCategoryID,
CommonTrackerPatternID: &commonPattern.ID,
TrackerType: coredata.TrackerTypeCookie,
Pattern: "_ga",
MatchType: coredata.TrackerPatternMatchTypeExact,
DisplayName: "_ga",
Description: "",
CreatedAt: now,
UpdatedAt: now,
}
require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error {
if err := commonThirdParty.Insert(ctx, tx); err != nil {
return err
}
if _, err := commonPattern.Upsert(ctx, tx); err != nil {
return err
}
return pattern.Insert(ctx, tx, fx.scope)
}))
t.Cleanup(func() {
_ = client.WithTx(context.Background(), func(ctx context.Context, tx pg.Tx) error {
if _, err := tx.Exec(ctx, `DELETE FROM common_third_party_domains WHERE common_third_party_id = $1`, commonThirdPartyID); err != nil {
return err
}
if _, err := tx.Exec(ctx, `DELETE FROM common_tracker_patterns WHERE id = $1`, commonPattern.ID); err != nil {
return err
}
if _, err := tx.Exec(ctx, `DELETE FROM common_third_parties WHERE id = $1`, commonThirdPartyID); err != nil {
return err
}
if _, err := tx.Exec(ctx, `DELETE FROM third_parties WHERE organization_id = $1`, fx.organizationID); err != nil {
return err
}
return nil
})
})
return promotionFixture{
workerFixture: fx,
commonThirdParty: commonThirdParty,
commonPatternID: commonPattern.ID,
commonThirdPartyID: commonThirdPartyID,
trackerPattern: pattern,
}
}
func newMappingHandler(client *pg.Client) *trackerMappingHandler {
return &trackerMappingHandler{
pg: client,
logger: log.NewLogger(log.WithOutput(io.Discard)),
}
}
// promote runs promoteThirdParty inside its own transaction so each
// test case starts from a clean state.
func promote(
t *testing.T,
ctx context.Context,
h *trackerMappingHandler,
client *pg.Client,
tp coredata.TrackerPattern,
commonPatternID gid.GID,
) *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.promoteThirdParty(ctx, tx, tp, commonPatternID)
return err
}))
return got
}
func TestPromoteThirdParty_ExactCommonLink(t *testing.T) {
t.Parallel()
client := newTestPgClient(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), client, fx.trackerPattern, fx.commonPatternID)
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 := newTestPgClient(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,
}
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), client, fx.trackerPattern, fx.commonPatternID)
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)
}
func TestPromoteThirdParty_FallbackCreate(t *testing.T) {
t.Parallel()
client := newTestPgClient(t)
ctx := context.Background()
fx := seedPromotionFixture(t, ctx, client)
got := promote(t, ctx, newMappingHandler(client), client, fx.trackerPattern, fx.commonPatternID)
require.NotNil(t, got, "fallback should create a new 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, *got)
}))
assert.Equal(t, fx.organizationID, reloaded.OrganizationID)
assert.Equal(t, "Google", reloaded.Name)
require.NotNil(t, reloaded.CommonThirdPartyID)
assert.Equal(t, fx.commonThirdPartyID, *reloaded.CommonThirdPartyID)
assert.Equal(t, coredata.ThirdPartyCategoryAnalytics, reloaded.Category)
assert.False(t, reloaded.FirstLevel)
assert.False(t, reloaded.ShowOnTrustCenter)
}
func TestPromoteThirdParty_NoCommonThirdPartyOnPattern(t *testing.T) {
t.Parallel()
client := newTestPgClient(t)
ctx := context.Background()
fx := seedWorkerFixture(t, ctx, client)
now := time.Now().UTC().Truncate(time.Microsecond)
commonPattern := coredata.CommonTrackerPattern{
ID: gid.New(gid.NilTenant, coredata.CommonTrackerPatternEntityType),
TrackerType: coredata.TrackerTypeCookie,
Pattern: "unknown_xyz",
MatchType: coredata.TrackerPatternMatchTypeExact,
Description: "",
Confidence: 0.5,
CreatedAt: now,
UpdatedAt: now,
}
pattern := coredata.TrackerPattern{
ID: gid.New(fx.scope.GetTenantID(), coredata.TrackerPatternEntityType),
OrganizationID: fx.organizationID,
CookieBannerID: fx.banner.ID,
CookieCategoryID: fx.normalCategoryID,
CommonTrackerPatternID: &commonPattern.ID,
TrackerType: coredata.TrackerTypeCookie,
Pattern: "unknown_xyz",
MatchType: coredata.TrackerPatternMatchTypeExact,
DisplayName: "unknown_xyz",
CreatedAt: now,
UpdatedAt: now,
}
require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error {
if _, err := commonPattern.Upsert(ctx, tx); err != nil {
return err
}
return pattern.Insert(ctx, tx, fx.scope)
}))
t.Cleanup(func() {
_ = client.WithTx(context.Background(), func(ctx context.Context, tx pg.Tx) error {
_, err := tx.Exec(ctx, `DELETE FROM common_tracker_patterns WHERE id = $1`, commonPattern.ID)
return err
})
})
got := promote(t, ctx, newMappingHandler(client), client, pattern, commonPattern.ID)
assert.Nil(t, got, "patterns whose catalog row has no CommonThirdPartyID should not be promoted")
}
// 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
// existing catalog link is preserved verbatim.
func TestProcess_PreservesCatalogMappingOnReTrigger(t *testing.T) {
t.Parallel()
client := newTestPgClient(t)
ctx := context.Background()
fx := seedPromotionFixture(t, ctx, client)
require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error {
return fx.trackerPattern.SetMappingRequested(ctx, tx)
}))
h := newMappingHandler(client)
require.NoError(t, h.Process(ctx, fx.trackerPattern))
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.CommonTrackerPatternID, "common tracker pattern link must be preserved")
assert.Equal(t, fx.commonPatternID, *reloaded.CommonTrackerPatternID)
require.NotNil(t, reloaded.ThirdPartyID, "the worker should have promoted to an org ThirdParty")
}
// TestProcess_ExtensionPatternIsNotPromoted asserts that even when a
// pattern has a catalog link, a Source=EXTENSION pattern stays
// un-promoted.
func TestProcess_ExtensionPatternIsNotPromoted(t *testing.T) {
t.Parallel()
client := newTestPgClient(t)
ctx := context.Background()
fx := seedPromotionFixture(t, ctx, client)
now := time.Now().UTC().Truncate(time.Microsecond)
source := coredata.CookieSourceExtension
require.NoError(t, client.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error {
_, err := tx.Exec(
ctx,
`UPDATE tracker_patterns
SET source = $1,
mapping_requested_at = $2
WHERE id = $3`,
source,
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)
}))
assert.Nil(t, reloaded.ThirdPartyID, "EXTENSION-sourced pattern must not be promoted")
}
// TestProcess_NoOpWhenAlreadyPromoted asserts that re-running the
// worker on a pattern that already has a third_party_id leaves the
// row alone (the guard in Process).
func TestProcess_NoOpWhenAlreadyPromoted(t *testing.T) {
t.Parallel()
client := newTestPgClient(t)
ctx := context.Background()
fx := seedPromotionFixture(t, ctx, client)
now := time.Now().UTC().Truncate(time.Microsecond)
preExisting := 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 := preExisting.Insert(ctx, tx, fx.scope); err != nil {
return err
}
fx.trackerPattern.ThirdPartyID = &preExisting.ID
_, err := tx.Exec(
ctx,
`UPDATE tracker_patterns
SET third_party_id = $1,
mapping_requested_at = $2
WHERE id = $3`,
preExisting.ID,
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)
assert.Equal(t, preExisting.ID, *reloaded.ThirdPartyID, "third_party_id must not be overwritten")
}
func TestPromoteThirdParty_ExactCommonLinkIgnoresSimilarUnlinked(t *testing.T) {
t.Parallel()
client := newTestPgClient(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), client, fx.trackerPattern, fx.commonPatternID)
require.NotNil(t, got)
assert.Equal(t, linked.ID, *got, "exact-link path must short-circuit before the heuristic fires")
}

View File

@@ -645,6 +645,7 @@ func (v *ThirdParty) Update(
q := `
UPDATE third_parties
SET
common_third_party_id = @common_third_party_id,
name = @name,
description = @description,
category = @category,
@@ -675,6 +676,7 @@ WHERE %s
args := pgx.StrictNamedArgs{
"third_party_id": v.ID,
"updated_at": time.Now(),
"common_third_party_id": v.CommonThirdPartyID,
"name": v.Name,
"description": v.Description,
"category": v.Category,

View File

@@ -313,7 +313,7 @@ func (impl *Implm) Run(
return err
}
trackerMappingCfg, err := impl.buildTrackerMappingConfig(l, tp, r)
trackerMappingCfg, thirdPartyDisambiguationCfg, err := impl.buildTrackerMappingConfig(l, tp, r)
if err != nil {
return err
}
@@ -725,7 +725,7 @@ func (impl *Implm) Run(
},
)
trackerMappingWorker := cookiebanner.NewTrackerMappingWorker(pgClient, l, trackerMappingCfg)
trackerMappingWorker := cookiebanner.NewTrackerMappingWorker(pgClient, l, trackerMappingCfg, thirdPartyDisambiguationCfg)
trackerMappingWorkerCtx, stopTrackerMappingWorker := context.WithCancel(context.Background())
wg.Go(

View File

@@ -21,18 +21,28 @@ import (
"go.gearno.de/kit/log"
"go.opentelemetry.io/otel/trace"
"go.probo.inc/probo/pkg/cookiebanner"
"go.probo.inc/probo/pkg/thirdparty"
)
// buildTrackerMappingConfig wires the tracker-mapping agent. It is opt-in:
// deployments that do not set `llm.tracker-mapping.provider` get a zero
// config (nil LLM client) so the worker runs without agent fallback.
// buildTrackerMappingConfig wires the tracker-mapping agent (catalog
// identification) and the third-party disambiguation agent that the
// tracker-mapping worker uses to promote patterns to org ThirdParties.
// Both are opt-in: deployments that do not set
// `llm.tracker-mapping.provider` get zero configs (nil LLM client) so
// the worker runs without agent fallback.
//
// Both agents are sourced from the same `tracker-mapping` config slot
// because they share the LLM client, model, and lifecycle. The
// disambiguation agent has no Firecrawl/DB tools, so its config
// surface is narrower and it lives in the cross-domain pkg/thirdparty
// package.
func (impl *Implm) buildTrackerMappingConfig(
l *log.Logger,
tp trace.TracerProvider,
r prometheus.Registerer,
) (cookiebanner.TrackerMappingConfig, error) {
) (cookiebanner.TrackerMappingConfig, thirdparty.DisambiguationConfig, error) {
if impl.cfg.Agents.TrackerMapping.Provider == "" {
return cookiebanner.TrackerMappingConfig{}, nil
return cookiebanner.TrackerMappingConfig{}, thirdparty.DisambiguationConfig{}, nil
}
agentCfg, llmClient, err := impl.resolveAgentClient(
@@ -43,12 +53,19 @@ func (impl *Implm) buildTrackerMappingConfig(
r,
)
if err != nil {
return cookiebanner.TrackerMappingConfig{}, fmt.Errorf("cannot resolve tracker mapping agent client: %w", err)
return cookiebanner.TrackerMappingConfig{}, thirdparty.DisambiguationConfig{}, fmt.Errorf("cannot resolve tracker mapping agent client: %w", err)
}
return cookiebanner.TrackerMappingConfig{
mappingCfg := cookiebanner.TrackerMappingConfig{
LLMClient: llmClient,
Model: agentCfg.ModelName,
FirecrawlAPIKey: impl.cfg.Agents.Tools.FirecrawlAPIKey,
}, nil
}
disambiguationCfg := thirdparty.DisambiguationConfig{
LLMClient: llmClient,
Model: agentCfg.ModelName,
}
return mappingCfg, disambiguationCfg, nil
}

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

@@ -0,0 +1,209 @@
// Copyright (c) 2026 Probo Inc <hello@getprobo.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
// 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
)
// 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.
type DisambiguationConfig struct {
LLMClient *llm.Client
Model string
}
// 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 DisambiguationConfig,
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))
}
return agent.New(
"third-party-disambiguation",
cfg.LLMClient,
agent.WithInstructions(disambiguationPrompt),
agent.WithModel(cfg.Model),
agent.WithOutputType(outputType),
agent.WithMaxTurns(1),
agent.WithLogger(logger),
)
}
// 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,
) (*gid.GID, error) {
if a == nil || len(candidates) == 0 {
return nil, nil
}
prompt := buildDisambiguationPrompt(commonParty, commonDomains, candidates)
agentCtx, cancel := context.WithTimeout(ctx, disambiguationTimeout)
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()
}

274
pkg/thirdparty/match.go vendored Normal file
View File

@@ -0,0 +1,274 @@
// Copyright (c) 2026 Probo Inc <hello@getprobo.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"
"time"
"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, fall back to creating, 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 commonID; otherwise writes the field via ThirdParty.Update and
// updates the receiver in place.
func LinkToCommon(
ctx context.Context,
tx pg.Tx,
scope coredata.Scoper,
orgThirdParty *coredata.ThirdParty,
commonID gid.GID,
) error {
if orgThirdParty.CommonThirdPartyID != nil && *orgThirdParty.CommonThirdPartyID == commonID {
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
}
// CreateFromCommon inserts a new org ThirdParty seeded from the catalog
// row (name, category, addresses, URLs, certifications, …). The new row
// has common_third_party_id pointed at commonParty, an empty Countries
// list, and ShowOnTrustCenter / FirstLevel both false — mirroring the
// front-end CreateThirdPartyDialog's "pick from catalog" seeding shape
// so the result is indistinguishable from a manual creation.
//
// Deliberately bypasses any service-level webhook emission: callers
// that need a webhook for the implicit creation should emit it
// themselves.
func CreateFromCommon(
ctx context.Context,
tx pg.Tx,
scope coredata.Scoper,
organizationID gid.GID,
commonParty coredata.CommonThirdParty,
) (*coredata.ThirdParty, error) {
commonID := commonParty.ID
now := time.Now()
tp := &coredata.ThirdParty{
ID: gid.New(scope.GetTenantID(), coredata.ThirdPartyEntityType),
OrganizationID: organizationID,
CommonThirdPartyID: &commonID,
Name: commonParty.Name,
Category: commonParty.Category,
HeadquarterAddress: commonParty.HeadquarterAddress,
LegalName: commonParty.LegalName,
WebsiteURL: commonParty.WebsiteURL,
PrivacyPolicyURL: commonParty.PrivacyPolicyURL,
ServiceLevelAgreementURL: commonParty.ServiceLevelAgreementURL,
DataProcessingAgreementURL: commonParty.DataProcessingAgreementURL,
BusinessAssociateAgreementURL: commonParty.BusinessAssociateAgreementURL,
SubprocessorsListURL: commonParty.SubprocessorsListURL,
Certifications: commonParty.Certifications,
Countries: coredata.CountryCodes{},
StatusPageURL: commonParty.StatusPageURL,
TermsOfServiceURL: commonParty.TermsOfServiceURL,
SecurityPageURL: commonParty.SecurityPageURL,
TrustPageURL: commonParty.TrustPageURL,
ShowOnTrustCenter: false,
FirstLevel: false,
CreatedAt: now,
UpdatedAt: now,
}
if tp.Certifications == nil {
tp.Certifications = []string{}
}
if err := tp.Insert(ctx, tx, scope); err != nil {
return nil, fmt.Errorf("cannot insert org third party: %w", err)
}
return tp, nil
}

181
pkg/thirdparty/match_test.go vendored Normal file
View File

@@ -0,0 +1,181 @@
// Copyright (c) 2026 Probo Inc <hello@getprobo.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"
)
//go:fix inline
func ptr[T any](v T) *T { return new(v) }
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)
})
}

View File

@@ -0,0 +1,25 @@
<role>
You are a third-party catalog matcher. The product groups web trackers under a global catalog of "common" third parties (Google Analytics, Stripe, Meta Pixel, …). Each customer organisation also maintains its own list of "third parties" — sometimes seeded from the catalog, sometimes typed manually. Your only job is to decide whether one of the organisation's existing third parties already represents a given catalog entry, so we don't create a duplicate.
</role>
<task>
You are given:
- A catalog third party: name, website, and known domains.
- A small list of candidate organisation third parties: each with a stable id, a name, and (optionally) a website.
Pick the candidate that best represents the catalog third party, or none.
Return a structured JSON response with:
- matched_id: the candidate id, or null if none of them is a confident match.
- confidence: 0.0 to 1.0; below 0.6 means "no confident match" (set matched_id to null in that case).
- reasoning: one short sentence describing the rationale.
</task>
<instructions>
1. The candidates have already been ranked by a heuristic; the list is small (usually a handful). Use that as a hint, but do not trust it blindly.
2. Treat corporate suffixes (LLC, Inc, Ltd, GmbH, SA, …) as noise. "Google" and "Google LLC" are the same company.
3. Treat brand/product names as the same when the parent company is obvious: "Google Analytics" maps to "Google" if the org only has the parent. Only do this when the catalog domains plainly match the parent's domains.
4. Website hostnames and known domains are the strongest signal. If the catalog domain (or its eTLD+1) matches a candidate's website host, they are almost certainly the same.
5. Be conservative. If two candidates look plausible and you cannot rule one out, return matched_id=null with confidence < 0.6 — we will create a fresh org third party from the catalog rather than risk a wrong link.
6. Do not invent ids. Return only ids that appear verbatim in the candidate list.
</instructions>