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

Add vendor assessment agent

Browse Source 
Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com>
This commit is contained in:
Aurélien Sibiril
2026-04-22 22:36:14 +02:00
parent 25c590ffe6
commit 509d0c88b1
108 changed files with 9445 additions and 645 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
GNUmakefile
View File

@@ -38,11 +38,24 @@ GO_TOOL= $(GO_BASE) tool
TEST_FLAGS?= -race -cover -coverprofile=coverage.out
E2E_CONFIG ?= $(CURDIR)/e2e/console/testdata/config.yaml
E2E_COVER_DIR ?= $(CURDIR)/coverage/e2e
DOCKER_IMAGE_NAME= ghcr.io/getprobo/probo
DOCKER_TAG_NAME?= latest
PROBOD_BIN_DEPS= pkg/server/api/connect/v1/schema/schema.go \
pkg/server/api/connect/v1/types/types.go \
pkg/server/api/console/v1/schema/schema.go \
pkg/server/api/console/v1/types/types.go \
pkg/server/api/trust/v1/schema/schema.go \
pkg/server/api/trust/v1/types/types.go \
pkg/server/api/mcp/v1/server/server.go \
pkg/server/api/mcp/v1/types/types.go \
apps/console/dist/index.html \
apps/trust/dist/index.html \
@probo/emails
PROBOD_BIN_EXTRA_DEPS=
PROBOD_BIN= bin/probod
PROBOD_SRC= cmd/probod/main.go
@@ -126,8 +139,9 @@ test-bench: test ## Run benchmark tests
.PHONY: test-e2e
test-e2e: CGO_ENABLED=1
test-e2e: bin/probod ## Run console e2e tests
PROBO_E2E_BINARY=$(CURDIR)/bin/probod \
test-e2e: $(PROBOD_BIN) ## Run console e2e tests
PROBO_E2E_BINARY=$(CURDIR)/$(PROBOD_BIN) \
PROBO_E2E_CONFIG=$(E2E_CONFIG) \
GOTESTSUM_FORMAT=testname $(GO_TEST) -count=1 ./e2e/console/...
bin/probod-coverage:
@@ -138,6 +152,7 @@ test-e2e-coverage: bin/probod-coverage ## Run e2e tests with coverage
@$(RM) -rf $(E2E_COVER_DIR) && $(MKDIR) -p $(E2E_COVER_DIR)
PROBO_E2E_BINARY=$(CURDIR)/bin/probod-coverage \
PROBO_E2E_COVERDIR=$(E2E_COVER_DIR) \
PROBO_E2E_CONFIG=$(E2E_CONFIG) \
CGO_ENABLED=1 $(GO) test -count=1 -v ./e2e/console/...
$(GO) tool covdata textfmt -i=$(E2E_COVER_DIR) -o=coverage-e2e.out
$(GO) tool cover -html=coverage-e2e.out -o=coverage-e2e.html
@@ -149,7 +164,7 @@ coverage-combined: coverage-report test-e2e-coverage ## Generate combined covera
$(GO) tool cover -html=coverage-combined.out -o=coverage-combined.html
.PHONY: build
build: bin/probod bin/prb bin/probod-bootstrap
build: $(PROBOD_BIN) bin/prb bin/probod-bootstrap
CFG_DEV_OAUTH2_KEY = cfg/.dev-oauth2-signing-key.pem
DEV_ENV = .env
@@ -217,19 +232,8 @@ scan-license: ## Check dependencies licenses compliance
docker-build:
$(DOCKER_BUILD) --tag $(DOCKER_IMAGE_NAME):$(DOCKER_TAG_NAME) --file Dockerfile .
.PHONY: bin/probod
bin/probod: pkg/server/api/connect/v1/schema/schema.go \
pkg/server/api/connect/v1/types/types.go \
pkg/server/api/console/v1/schema/schema.go \
pkg/server/api/console/v1/types/types.go \
pkg/server/api/trust/v1/schema/schema.go \
pkg/server/api/trust/v1/types/types.go \
pkg/server/api/mcp/v1/server/server.go \
pkg/server/api/mcp/v1/types/types.go \
apps/console/dist/index.html \
apps/trust/dist/index.html \
$(PROBOD_BIN_EXTRA_DEPS) \
@probo/emails
.PHONY: $(PROBOD_BIN)
$(PROBOD_BIN): $(PROBOD_BIN_DEPS) $(PROBOD_BIN_EXTRA_DEPS)
$(GO_BUILD) -o $(PROBOD_BIN) $(PROBOD_SRC)
.PHONY: bin/prb

112
e2e/console/vendor_test.go
View File

@@ -986,6 +986,118 @@ func TestVendor_OmittableWebsiteUrl(t *testing.T) {
})
}
// TestVendor_Assess exercises the assessVendor mutation through authorization
// and tenant-isolation paths without running the real LLM/browser pipeline.
// The e2e config deliberately omits `llm.vendor-assessor.provider`, so an
// authorized call reaches DisabledVendorAssessor and surfaces a stable
// UNAVAILABLE error. Happy-path payload shape is covered by unit tests in
// pkg/probo.
func TestVendor_Assess(t *testing.T) {
t.Parallel()
const query = `
mutation AssessVendor($input: AssessVendorInput!) {
assessVendor(input: $input) {
vendor {
id
}
}
}
`
type resultShape struct {
AssessVendor struct {
Vendor struct {
ID string `json:"id"`
} `json:"vendor"`
} `json:"assessVendor"`
}
t.Run("owner call surfaces the disabled error", func(t *testing.T) {
t.Parallel()
owner := testutil.NewClient(t, testutil.RoleOwner)
vendorID := factory.NewVendor(owner).WithName("Unconfigured assess").Create()
var result resultShape
err := owner.Execute(query, map[string]any{
"input": map[string]any{
"id": vendorID,
"websiteUrl": "https://vendor.example.com",
},
}, &result)
testutil.RequireErrorCode(t, err, "UNAVAILABLE")
})
t.Run("admin call surfaces the disabled error", func(t *testing.T) {
t.Parallel()
owner := testutil.NewClient(t, testutil.RoleOwner)
admin := testutil.NewClientInOrg(t, testutil.RoleAdmin, owner)
vendorID := factory.NewVendor(owner).WithName("Admin-assessed vendor").Create()
var result resultShape
err := admin.Execute(query, map[string]any{
"input": map[string]any{
"id": vendorID,
"websiteUrl": "https://admin.example.com",
},
}, &result)
testutil.RequireErrorCode(t, err, "UNAVAILABLE")
})
t.Run("viewer cannot assess a vendor", func(t *testing.T) {
t.Parallel()
owner := testutil.NewClient(t, testutil.RoleOwner)
viewer := testutil.NewClientInOrg(t, testutil.RoleViewer, owner)
vendorID := factory.NewVendor(owner).WithName("Viewer attempt").Create()
var result resultShape
err := viewer.Execute(query, map[string]any{
"input": map[string]any{
"id": vendorID,
"websiteUrl": "https://viewer.example.com",
},
}, &result)
testutil.RequireForbiddenError(t, err)
})
t.Run("cannot assess vendor from another organization", func(t *testing.T) {
t.Parallel()
org1Owner := testutil.NewClient(t, testutil.RoleOwner)
org2Owner := testutil.NewClient(t, testutil.RoleOwner)
vendorID := factory.NewVendor(org1Owner).WithName("Org1 vendor").Create()
var result resultShape
err := org2Owner.Execute(query, map[string]any{
"input": map[string]any{
"id": vendorID,
"websiteUrl": "https://cross-tenant.example.com",
},
}, &result)
require.Error(t, err, "vendor assess must not cross tenant boundaries")
})
t.Run("procedure is accepted on the input", func(t *testing.T) {
t.Parallel()
owner := testutil.NewClient(t, testutil.RoleOwner)
vendorID := factory.NewVendor(owner).WithName("Procedure test").Create()
var result resultShape
err := owner.Execute(query, map[string]any{
"input": map[string]any{
"id": vendorID,
"websiteUrl": "https://procedure.example.com",
"procedure": "Focus on SOC 2 controls and data residency",
},
}, &result)
testutil.RequireErrorCode(t, err, "UNAVAILABLE")
})
}
func TestVendor_TenantIsolation(t *testing.T) {
t.Parallel()

90
pkg/agent/agent.go
View File

@@ -23,48 +23,53 @@ import (
"go.probo.inc/probo/pkg/llm"
)
const DefaultMaxTurns = 10
const (
DefaultMaxTurns = 10
DefaultMaxEmptyOutputRetries = 2
)
type (
Option func(*Agent)
Agent struct {
name string
handoffDescription string
instructions string
instructionsFunc func(ctx context.Context, a *Agent) string
model string
modelSettings ModelSettings
tools []Tool
handoffs []*Handoff
mcpServers []*MCPServer
maxTurns int
maxToolDepth int
client *llm.Client
logger *log.Logger
hooks []RunHooks
agentHooks AgentHooks
inputGuardrails []InputGuardrail
outputGuardrails []OutputGuardrail
session Session
sessionID string
outputType *OutputType
toolUseBehavior ToolUseBehavior
resetToolChoice bool
responseFormat *llm.ResponseFormat
approval *ApprovalConfig
name string
handoffDescription string
instructions string
instructionsFunc func(ctx context.Context, a *Agent) string
model string
modelSettings ModelSettings
tools []Tool
handoffs []*Handoff
mcpServers []*MCPServer
maxTurns int
maxEmptyOutputRetries int
maxToolDepth int
client *llm.Client
logger *log.Logger
hooks []RunHooks
agentHooks AgentHooks
inputGuardrails []InputGuardrail
outputGuardrails []OutputGuardrail
session Session
sessionID string
outputType *OutputType
toolUseBehavior ToolUseBehavior
resetToolChoice bool
responseFormat *llm.ResponseFormat
approval *ApprovalConfig
}
)
func New(name string, client *llm.Client, opts ...Option) *Agent {
a := &Agent{
name: name,
client: client,
maxTurns: DefaultMaxTurns,
maxToolDepth: DefaultMaxToolDepth,
toolUseBehavior: RunLLMAgain(),
resetToolChoice: true,
logger: log.NewLogger(log.WithOutput(io.Discard)),
name: name,
client: client,
maxTurns: DefaultMaxTurns,
maxEmptyOutputRetries: DefaultMaxEmptyOutputRetries,
maxToolDepth: DefaultMaxToolDepth,
toolUseBehavior: RunLLMAgain(),
resetToolChoice: true,
logger: log.NewLogger(log.WithOutput(io.Discard)),
}
for _, opt := range opts {
@@ -204,6 +209,18 @@ func WithMaxTurns(n int) Option {
}
}
// WithMaxEmptyOutputRetries bounds the number of times the core loop
// will re-ask the model to produce a structured output after it
// returned a thinking-only empty response on a synthesis turn.
func WithMaxEmptyOutputRetries(n int) Option {
return func(a *Agent) {
if n < 0 {
n = 0
}
a.maxEmptyOutputRetries = n
}
}
func WithMaxToolDepth(n int) Option {
return func(a *Agent) {
if n < 1 {
@@ -255,6 +272,15 @@ func WithParallelToolCalls(enabled bool) Option {
}
}
func WithThinking(budgetTokens int) Option {
return func(a *Agent) {
a.modelSettings.Thinking = &llm.ThinkingConfig{
Enabled: true,
BudgetTokens: budgetTokens,
}
}
}
func WithLogger(l *log.Logger) Option {
return func(a *Agent) {
a.logger = l

113
pkg/agent/agent_test.go
View File

@@ -351,14 +351,13 @@ func TestRun(t *testing.T) {
City string `json:"city"`
}
weatherTool, err := agent.FunctionTool[Params](
weatherTool := agent.FunctionTool[Params](
"get_weather",
"Get weather for a city",
func(_ context.Context, p Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "Sunny, 22°C in " + p.City}, nil
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -411,14 +410,13 @@ func TestRun(t *testing.T) {
}
type Params struct{}
noopTool, err := agent.FunctionTool[Params](
noopTool := agent.FunctionTool[Params](
"noop",
"No-op",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil
},
)
require.NoError(t, err)
ag := agent.New(
"assistant",
@@ -428,7 +426,7 @@ func TestRun(t *testing.T) {
agent.WithMaxTurns(2),
)
_, err = ag.Run(
_, err := ag.Run(
context.Background(),
[]llm.Message{userMessage("loop")},
)
@@ -447,14 +445,13 @@ func TestRun(t *testing.T) {
type Params struct{}
makeTool := func(name string) agent.Tool {
tool, err := agent.FunctionTool[Params](
tool := agent.FunctionTool[Params](
name,
"desc",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil
},
)
require.NoError(t, err)
return tool
}
@@ -553,22 +550,20 @@ func TestRun(t *testing.T) {
type Params struct{}
tool1, err := agent.FunctionTool[Params](
tool1 := agent.FunctionTool[Params](
"first",
"First tool",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "result_1"}, nil
},
)
require.NoError(t, err)
tool2, err := agent.FunctionTool[Params](
tool2 := agent.FunctionTool[Params](
"second",
"Second tool",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "result_2"}, nil
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -622,22 +617,20 @@ func TestRun(t *testing.T) {
type Params struct{}
successTool, err := agent.FunctionTool[Params](
successTool := agent.FunctionTool[Params](
"succeed",
"Always succeeds",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "success_result"}, nil
},
)
require.NoError(t, err)
failTool, err := agent.FunctionTool[Params](
failTool := agent.FunctionTool[Params](
"fail",
"Always fails",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, errors.New("tool exploded")
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -696,7 +689,7 @@ func TestRun(t *testing.T) {
var capturedTenantID string
type Params struct{}
tool, err := agent.FunctionTool[Params](
tool := agent.FunctionTool[Params](
"check_tenant",
"Check current tenant",
func(ctx context.Context, _ Params) (agent.ToolResult, error) {
@@ -705,7 +698,6 @@ func TestRun(t *testing.T) {
return agent.ToolResult{Content: "tenant: " + rc.TenantID}, nil
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -1069,14 +1061,13 @@ func TestRun_Hooks(t *testing.T) {
t.Parallel()
type Params struct{}
noopTool, err := agent.FunctionTool[Params](
noopTool := agent.FunctionTool[Params](
"noop",
"No-op",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -1354,14 +1345,13 @@ func TestRun_ToolUseBehavior(t *testing.T) {
t.Parallel()
type Params struct{}
tool, err := agent.FunctionTool[Params](
tool := agent.FunctionTool[Params](
"compute",
"Compute something",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "computed_result"}, nil
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -1399,22 +1389,20 @@ func TestRun_ToolUseBehavior(t *testing.T) {
type Params struct{}
tool1, err := agent.FunctionTool[Params](
tool1 := agent.FunctionTool[Params](
"search",
"Search",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "search_result"}, nil
},
)
require.NoError(t, err)
tool2, err := agent.FunctionTool[Params](
tool2 := agent.FunctionTool[Params](
"submit",
"Submit",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "submitted"}, nil
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -1449,14 +1437,13 @@ func TestRun_ToolUseBehavior(t *testing.T) {
t.Parallel()
type Params struct{}
tool, err := agent.FunctionTool[Params](
tool := agent.FunctionTool[Params](
"noop",
"No-op",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -1492,14 +1479,13 @@ func TestRun_ToolUseBehavior(t *testing.T) {
t.Parallel()
type Params struct{}
tool, err := agent.FunctionTool[Params](
tool := agent.FunctionTool[Params](
"compute",
"Compute something",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "result"}, nil
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -1520,7 +1506,7 @@ func TestRun_ToolUseBehavior(t *testing.T) {
})),
)
_, err = ag.Run(
_, err := ag.Run(
context.Background(),
[]llm.Message{userMessage("compute")},
)
@@ -1581,14 +1567,13 @@ func TestRun_Approval(t *testing.T) {
},
}
deleteTool, err := agent.FunctionTool[struct{}](
deleteTool := agent.FunctionTool[struct{}](
"delete_account",
"Deletes the user account",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
return agent.ToolResult{Content: "deleted"}, nil
},
)
require.NoError(t, err)
ag := agent.New(
"assistant",
@@ -1600,7 +1585,7 @@ func TestRun_Approval(t *testing.T) {
}),
)
_, err = ag.Run(
_, err := ag.Run(
context.Background(),
[]llm.Message{userMessage("Delete my account")},
)
@@ -1623,7 +1608,7 @@ func TestRun_Approval(t *testing.T) {
var toolExecuted bool
deleteTool, err := agent.FunctionTool[struct{}](
deleteTool := agent.FunctionTool[struct{}](
"delete_account",
"Deletes the user account",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
@@ -1631,7 +1616,6 @@ func TestRun_Approval(t *testing.T) {
return agent.ToolResult{Content: "account deleted"}, nil
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -1653,7 +1637,7 @@ func TestRun_Approval(t *testing.T) {
}),
)
_, err = ag.Run(
_, err := ag.Run(
context.Background(),
[]llm.Message{userMessage("Delete my account")},
)
@@ -1683,7 +1667,7 @@ func TestRun_Approval(t *testing.T) {
func(t *testing.T) {
t.Parallel()
deleteTool, err := agent.FunctionTool[struct{}](
deleteTool := agent.FunctionTool[struct{}](
"delete_account",
"Deletes the user account",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
@@ -1691,7 +1675,6 @@ func TestRun_Approval(t *testing.T) {
return agent.ToolResult{}, nil
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -1713,7 +1696,7 @@ func TestRun_Approval(t *testing.T) {
}),
)
_, err = ag.Run(
_, err := ag.Run(
context.Background(),
[]llm.Message{userMessage("Delete my account")},
)
@@ -1751,14 +1734,13 @@ func TestRun_Approval(t *testing.T) {
},
}
safeTool, err := agent.FunctionTool[struct{}](
safeTool := agent.FunctionTool[struct{}](
"safe_tool",
"A safe tool",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
return agent.ToolResult{Content: "safe result"}, nil
},
)
require.NoError(t, err)
ag := agent.New(
"assistant",
@@ -1789,7 +1771,7 @@ func TestRun_Approval(t *testing.T) {
var safeExecuted, dangerExecuted bool
safeTool, err := agent.FunctionTool[struct{}](
safeTool := agent.FunctionTool[struct{}](
"safe_action",
"A safe action",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
@@ -1797,9 +1779,8 @@ func TestRun_Approval(t *testing.T) {
return agent.ToolResult{Content: "safe done"}, nil
},
)
require.NoError(t, err)
dangerTool, err := agent.FunctionTool[struct{}](
dangerTool := agent.FunctionTool[struct{}](
"danger_action",
"A dangerous action",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
@@ -1807,7 +1788,6 @@ func TestRun_Approval(t *testing.T) {
return agent.ToolResult{Content: "danger done"}, nil
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -1835,7 +1815,7 @@ func TestRun_Approval(t *testing.T) {
}),
)
_, err = ag.Run(
_, err := ag.Run(
context.Background(),
[]llm.Message{userMessage("Do both")},
)
@@ -1943,14 +1923,13 @@ func TestResume(t *testing.T) {
func(t *testing.T) {
t.Parallel()
deleteTool, err := agent.FunctionTool[struct{}](
deleteTool := agent.FunctionTool[struct{}](
"delete_account",
"Deletes the user account",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
return agent.ToolResult{Content: "deleted"}, nil
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -1972,7 +1951,7 @@ func TestResume(t *testing.T) {
}),
)
_, err = ag.Run(
_, err := ag.Run(
context.Background(),
[]llm.Message{userMessage("Delete my account")},
)
@@ -2152,14 +2131,13 @@ func TestRunStreamed(t *testing.T) {
t.Parallel()
type Params struct{}
tool, err := agent.FunctionTool[Params](
tool := agent.FunctionTool[Params](
"noop",
"No-op",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil
},
)
require.NoError(t, err)
stream1 := &mockChatStream{
events: []llm.ChatCompletionStreamEvent{
@@ -2378,22 +2356,20 @@ func TestClone(t *testing.T) {
t.Parallel()
type Params struct{}
tool1, err := agent.FunctionTool[Params](
tool1 := agent.FunctionTool[Params](
"t1",
"desc",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil
},
)
require.NoError(t, err)
tool2, err := agent.FunctionTool[Params](
tool2 := agent.FunctionTool[Params](
"t2",
"desc",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -2500,14 +2476,13 @@ func TestGenerateSchema_EmbeddedStruct(t *testing.T) {
Name string `json:"name"`
}
tool, err := agent.FunctionTool[Params](
tool := agent.FunctionTool[Params](
"create",
"Create item",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil
},
)
require.NoError(t, err)
var schema map[string]any
require.NoError(t, json.Unmarshal(tool.Definition().Parameters, &schema))
@@ -2675,14 +2650,13 @@ func TestRun_UnknownToolCall(t *testing.T) {
}
type Params struct{}
tool, err := agent.FunctionTool[Params](
tool := agent.FunctionTool[Params](
"real_tool",
"A real tool",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil
},
)
require.NoError(t, err)
ag := agent.New(
"assistant",
@@ -2691,7 +2665,7 @@ func TestRun_UnknownToolCall(t *testing.T) {
agent.WithTools(tool),
)
_, err = ag.Run(
_, err := ag.Run(
context.Background(),
[]llm.Message{userMessage("test")},
)
@@ -2743,14 +2717,13 @@ func TestClone_WithApprovalConfig(t *testing.T) {
},
}
deleteTool, err := agent.FunctionTool[struct{}](
deleteTool := agent.FunctionTool[struct{}](
"delete",
"Delete something",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
return agent.ToolResult{Content: "deleted"}, nil
},
)
require.NoError(t, err)
original := agent.New(
"assistant",
@@ -2764,7 +2737,7 @@ func TestClone_WithApprovalConfig(t *testing.T) {
cloned := original.Clone()
_, err = cloned.Run(
_, err := cloned.Run(
context.Background(),
[]llm.Message{userMessage("delete it")},
)
@@ -2795,7 +2768,7 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) {
var executionOrder []string
type Params struct{}
tool1, err := agent.FunctionTool[Params](
tool1 := agent.FunctionTool[Params](
"prepare",
"Prepare data",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
@@ -2803,7 +2776,6 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) {
return agent.ToolResult{Content: "prepared"}, nil
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -2855,15 +2827,14 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) {
t.Parallel()
type Params struct{}
tool1, err := agent.FunctionTool[Params](
tool1 := agent.FunctionTool[Params](
"prepare",
"Prepare data",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "prepared"}, nil
},
)
require.NoError(t, err)
tool2, err := agent.FunctionTool[Params](
tool2 := agent.FunctionTool[Params](
"finalize",
"Finalize data",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
@@ -2871,7 +2842,6 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) {
return agent.ToolResult{}, nil
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -2941,14 +2911,13 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) {
t.Parallel()
type Params struct{}
failingTool, err := agent.FunctionTool[Params](
failingTool := agent.FunctionTool[Params](
"prepare",
"Prepare data",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, errors.New("preparation failed")
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -2981,7 +2950,7 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) {
agent.WithHandoffs(specialist),
)
_, err = router.Run(
_, err := router.Run(
context.Background(),
[]llm.Message{userMessage("prepare and transfer")},
)

17
pkg/agent/agent_tool.go
View File

@@ -116,5 +116,20 @@ func (t *agentTool) Execute(ctx context.Context, arguments string) (ToolResult,
return ToolResult{}, err
}
return ToolResult{Content: result.FinalMessage().Text()}, nil
text := result.FinalMessage().Text()
if t.agent.outputType != nil {
if !json.Valid([]byte(text)) {
preview := text
if len(preview) > 500 {
preview = preview[:500] + "... (truncated)"
}
return ToolResult{
Content: fmt.Sprintf("Sub-agent %q returned invalid JSON. Raw output:\n%s", t.agent.name, preview),
IsError: true,
}, nil
}
}
return ToolResult{Content: text}, nil
}

19
pkg/agent/agent_tool_test.go
View File

@@ -273,7 +273,7 @@ func TestAgentTool_Execute(t *testing.T) {
var captured string
type Params struct{}
tenantTool, err := agent.FunctionTool[Params](
tenantTool := agent.FunctionTool[Params](
"get_tenant",
"Get tenant",
func(ctx context.Context, _ Params) (agent.ToolResult, error) {
@@ -282,7 +282,6 @@ func TestAgentTool_Execute(t *testing.T) {
return agent.ToolResult{Content: rc.TenantID}, nil
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -325,14 +324,13 @@ func TestAgentTool_Execute(t *testing.T) {
Expr string `json:"expr"`
}
calcTool, err := agent.FunctionTool[Params](
calcTool := agent.FunctionTool[Params](
"calc",
"Calculate expression",
func(_ context.Context, p Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "42"}, nil
},
)
require.NoError(t, err)
provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -413,12 +411,11 @@ func newNestedApprovalFixture(
) nestedApprovalFixture {
t.Helper()
deleteTool, err := agent.FunctionTool[struct{}](
deleteTool := agent.FunctionTool[struct{}](
"delete_file",
"Delete a file",
deleteFunc,
)
require.NoError(t, err)
innerProvider := &mockProvider{responses: innerResponses}
@@ -608,7 +605,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
var siblingCalled bool
type Params struct{}
siblingTool, err := agent.FunctionTool[Params](
siblingTool := agent.FunctionTool[Params](
"list_files",
"List files",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
@@ -616,7 +613,6 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
return agent.ToolResult{Content: "file1.txt, file2.txt"}, nil
},
)
require.NoError(t, err)
f := newNestedApprovalFixture(
t,
@@ -644,7 +640,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
agent.WithTools(siblingTool),
)
_, err = f.outerAgent.Run(
_, err := f.outerAgent.Run(
context.Background(),
[]llm.Message{userMessage("List and delete files")},
)
@@ -675,7 +671,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
var toolExecuted bool
dangerTool, err := agent.FunctionTool[struct{}](
dangerTool := agent.FunctionTool[struct{}](
"danger",
"Dangerous operation",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
@@ -683,7 +679,6 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
return agent.ToolResult{Content: "danger executed"}, nil
},
)
require.NoError(t, err)
cProvider := &mockProvider{
responses: []*llm.ChatCompletionResponse{
@@ -739,7 +734,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
agent.WithTools(agentB.AsTool("call_b", "Call agent B")),
)
_, err = agentA.Run(
_, err := agentA.Run(
context.Background(),
[]llm.Message{userMessage("start")},
)

1
pkg/agent/model_settings.go
View File

@@ -24,4 +24,5 @@ type ModelSettings struct {
MaxTokens *int
ToolChoice *llm.ToolChoice
ParallelToolCalls *bool
Thinking *llm.ThinkingConfig
}

36
pkg/agent/progress.go Normal file
View File

@@ -0,0 +1,36 @@
// 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 agent
import "context"
type (
ProgressEventType string
ProgressEvent struct {
Type ProgressEventType `json:"type"`
Step string `json:"step"`
ParentStep string `json:"parent_step,omitempty"`
Message string `json:"message"`
}
ProgressReporter func(ctx context.Context, event ProgressEvent)
)
const (
ProgressEventStepStarted ProgressEventType = "step_started"
ProgressEventStepCompleted ProgressEventType = "step_completed"
ProgressEventStepFailed ProgressEventType = "step_failed"
)

166
pkg/agent/run.go
View File

@@ -28,7 +28,14 @@ import (
"go.probo.inc/probo/pkg/llm"
)
const tracerName = "go.probo.inc/probo/pkg/agent"
const (
tracerName = "go.probo.inc/probo/pkg/agent"
// synthesisNudge is the static user message appended after tool
// exploration completes, asking the model to produce the final
// structured output on the next (synthesis) turn.
synthesisNudge = "Based on everything you have gathered, produce the final structured output now."
)
type (
CallLLMFunc func(ctx context.Context, agent *Agent, req *llm.ChatCompletionRequest) (*llm.ChatCompletionResponse, error)
@@ -68,7 +75,32 @@ type (
func noopEvent(_ context.Context, _ StreamEvent) {}
func blockingCallLLM(ctx context.Context, agent *Agent, req *llm.ChatCompletionRequest) (*llm.ChatCompletionResponse, error) {
return agent.client.ChatCompletion(ctx, req)
resp, err := agent.client.ChatCompletion(ctx, req)
if err == nil {
return resp, nil
}
// Some providers (e.g. Anthropic) require streaming for large
// max_tokens or when thinking is enabled. Fall back to streaming
// transparently when the blocking call returns ErrStreamingRequired.
var streamRequired *llm.ErrStreamingRequired
if !errors.As(err, &streamRequired) {
return nil, err
}
stream, sErr := agent.client.ChatCompletionStream(ctx, req)
if sErr != nil {
return nil, err // return the original error
}
defer stream.Close()
acc := llm.NewStreamAccumulator(stream)
for acc.Next() {
}
if sErr := acc.Err(); sErr != nil {
return nil, sErr
}
return acc.Response(), nil
}
func (a *Agent) Run(ctx context.Context, messages []llm.Message) (*Result, error) {
@@ -273,6 +305,24 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
log.Int("tool_count", len(s.toolDefs)),
)
emptyOutputRetries := 0
structuredFormat := resolveStructuredFormat(s.agent)
// When the agent has both tools and a structured output request,
// we delay structured output enforcement until a dedicated
// synthesis turn. Enforcing the schema during tool exploration
// causes models with extended thinking to stuff planning prose
// into the first text field of the schema as a scratchpad,
// burning the entire max_tokens budget on thinking-inside-JSON
// before ever producing a valid object. Instead, we let the
// model freely call tools without a schema, then force one final
// synthesis turn with ToolChoice=none + schema enforced once the
// model signals it has enough information (finish_reason=stop).
// Agents without tools or without a structured output request
// do not need this dance and enforce the schema immediately.
exploring := structuredFormat != nil && len(s.toolDefs) > 0
for {
if err := ctx.Err(); err != nil {
return s.finishRun(ctx, nil, fmt.Errorf("cannot complete: %w", err))
@@ -284,15 +334,21 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
fullMessages := buildFullMessages(s.systemPrompt, s.messages)
responseFormat := s.agent.responseFormat
if responseFormat == nil && s.agent.outputType != nil {
responseFormat = s.agent.outputType.responseFormat()
var responseFormat *llm.ResponseFormat
if !exploring {
responseFormat = structuredFormat
}
toolChoice := s.agent.modelSettings.ToolChoice
if s.toolUsedInRun && s.agent.resetToolChoice && toolChoice != nil {
toolChoice = nil
}
if !exploring && structuredFormat != nil && len(s.toolDefs) > 0 {
// On the synthesis turn, forbid further tool calls so the
// model is forced to convert what it has into JSON.
none := llm.ToolChoice{Type: llm.ToolChoiceNone}
toolChoice = &none
}
req := &llm.ChatCompletionRequest{
Model: s.agent.model,
@@ -306,6 +362,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
ToolChoice: toolChoice,
ParallelToolCalls: s.agent.modelSettings.ParallelToolCalls,
ResponseFormat: responseFormat,
Thinking: s.agent.modelSettings.Thinking,
}
s.logger.InfoCtx(
@@ -336,6 +393,62 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
switch resp.FinishReason {
case llm.FinishReasonStop, llm.FinishReasonLength:
// Model signalled it has nothing more to do with tools.
// If we have a structured output request but haven't
// enforced the schema yet, promote this turn to the
// synthesis turn: the next iteration runs with
// ToolChoice=none and the schema enforced, so the model
// converts what it has gathered into JSON in one shot.
//
// Anthropic requires the last message in the conversation
// to be a user message, so we cannot simply continue after
// an assistant stop turn. Drop empty (thinking-only) turns
// from history and append a user nudge that asks for the
// final structured output. Non-empty assistant turns stay
// in history so the model can reference its own
// conclusions during synthesis.
if exploring && s.turns < s.agent.maxTurns {
exploring = false
if resp.Message.Text() == "" {
s.messages = s.messages[:len(s.messages)-1]
}
s.messages = append(
s.messages,
llm.Message{
Role: llm.RoleUser,
Parts: []llm.Part{llm.TextPart{Text: synthesisNudge}},
},
)
s.logger.WarnCtx(
ctx,
"entering synthesis turn: forcing structured output with tool_choice=none",
log.Int("turn", s.turns),
log.Int("output_tokens", resp.Usage.OutputTokens),
)
continue
}
// Anthropic extended-thinking models can return a synthesis turn
// that contains only thinking blocks and no text part, leaving us
// with no structured output to validate. Retry the same turn a
// bounded number of times so the model gets another chance to
// emit the required JSON output. The empty assistant turn must be
// dropped from history because Anthropic rejects requests where
// the last message is a thinking-only assistant turn.
if structuredFormat != nil && resp.Message.Text() == "" && emptyOutputRetries < s.agent.maxEmptyOutputRetries && s.turns < s.agent.maxTurns {
emptyOutputRetries++
s.messages = s.messages[:len(s.messages)-1]
s.logger.WarnCtx(
ctx,
"retrying turn: structured output expected but got empty text",
log.Int("turn", s.turns),
log.Int("retry", emptyOutputRetries),
log.Int("output_tokens", resp.Usage.OutputTokens),
)
continue
}
emptyOutputRetries = 0
if err := runOutputGuardrails(ctx, s.agent, resp.Message); err != nil {
return s.finishRun(ctx, nil, err)
}
@@ -354,6 +467,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
case llm.FinishReasonToolCalls:
s.toolUsedInRun = true
emptyOutputRetries = 0
s.logger.InfoCtx(
ctx,
@@ -442,7 +556,8 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
if isFinal {
s.messages = append(
s.messages, llm.Message{
s.messages,
llm.Message{
Role: llm.RoleAssistant,
Parts: []llm.Part{llm.TextPart{Text: finalOutput}},
},
@@ -852,12 +967,24 @@ func executeSingleTool(
emitHook(agent, func(h RunHooks) { h.OnToolEnd(ctx, agent, tool, result, nil) })
emitAgentHook(agent, func(h AgentHooks) { h.OnToolEnd(ctx, agent, tool, result) })
logger.InfoCtx(
ctx,
"tool execution completed",
log.String("tool", tool.Name()),
log.Bool("is_error", result.IsError),
)
if result.IsError {
content := result.Content
if len(content) > 200 {
content = content[:200] + "... (truncated)"
}
logger.WarnCtx(
ctx,
"tool returned error",
log.String("tool", tool.Name()),
log.String("content", content),
)
} else {
logger.InfoCtx(
ctx,
"tool execution completed",
log.String("tool", tool.Name()),
)
}
return result, nil
}
@@ -1178,3 +1305,18 @@ func emitAgentHook(agent *Agent, fn func(AgentHooks)) {
fn(agent.agentHooks)
}
}
// resolveStructuredFormat returns the structured output request the
// agent wants enforced on its final turn, or nil if none. An agent can
// declare structured output through either WithOutputType (typed
// sub-agents) or a directly-set responseFormat (the RunTyped
// convenience wrapper).
func resolveStructuredFormat(a *Agent) *llm.ResponseFormat {
if a.responseFormat != nil {
return a.responseFormat
}
if a.outputType != nil {
return a.outputType.responseFormat()
}
return nil
}

43
pkg/agent/tool.go
View File

@@ -29,6 +29,7 @@ type (
IsError bool
}
// ToolDescriptor describes a tool's name and LLM definition.
ToolDescriptor interface {
Name() string
Definition() llm.Tool
@@ -38,7 +39,31 @@ type (
ToolDescriptor
Execute(ctx context.Context, arguments string) (ToolResult, error)
}
)
// ResultJSON marshals v to JSON and returns a successful ToolResult.
func ResultJSON(v any) ToolResult {
data, err := json.Marshal(v)
if err != nil {
return ToolResult{
Content: fmt.Sprintf("cannot marshal tool result: %s", err),
IsError: true,
}
}
return ToolResult{Content: string(data)}
}
// ResultError returns an error ToolResult with the given message.
func ResultError(msg string) ToolResult {
return ToolResult{Content: msg, IsError: true}
}
// ResultErrorf returns an error ToolResult with a formatted message.
func ResultErrorf(format string, args ...any) ToolResult {
return ToolResult{Content: fmt.Sprintf(format, args...), IsError: true}
}
type (
functionTool[P any] struct {
name string
description string
@@ -48,20 +73,30 @@ type (
}
)
// FunctionTool creates a tool whose parameters are typed by P. The JSON
// schema advertised to the LLM is generated from P at construction time.
//
// Schema generation is derived from a compile-time Go type: a failure
// here is a programmer error (bad struct tag, unsupported type), not a
// runtime condition, so we panic rather than returning an error. The
// same applies to the required-fields metadata parsed back out of the
// generated schema.
func FunctionTool[P any](
name string,
description string,
fn func(ctx context.Context, params P) (ToolResult, error),
) (Tool, error) {
) Tool {
schema, err := jsonSchemaFor[P]()
if err != nil {
return nil, fmt.Errorf("cannot create tool %q: %w", name, err)
panic(fmt.Sprintf("agent: cannot generate JSON schema for tool %q: %s", name, err))
}
var parsed struct {
Required []string `json:"required"`
}
_ = json.Unmarshal(schema, &parsed)
if err := json.Unmarshal(schema, &parsed); err != nil {
panic(fmt.Sprintf("agent: cannot parse generated schema for tool %q: %s", name, err))
}
return &functionTool[P]{
name: name,
@@ -69,7 +104,7 @@ func FunctionTool[P any](
fn: fn,
schema: schema,
requiredFields: parsed.Required,
}, nil
}
}
func (t *functionTool[P]) Name() string { return t.name }

58
pkg/agent/tool_test.go
View File

@@ -30,14 +30,13 @@ func TestFunctionTool_Name(t *testing.T) {
type Params struct{}
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"my_tool",
"does things",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, nil
},
)
require.NoError(t, err)
assert.Equal(t, "my_tool", tool.Name())
}
@@ -54,14 +53,13 @@ func TestFunctionTool_Definition(t *testing.T) {
Query string `json:"query"`
}
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"search",
"Search for items",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, nil
},
)
require.NoError(t, err)
def := tool.Definition()
assert.Equal(t, "search", def.Name)
@@ -79,14 +77,13 @@ func TestFunctionTool_Definition(t *testing.T) {
Count int `json:"count"`
}
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"create",
"Create items",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, nil
},
)
require.NoError(t, err)
def := tool.Definition()
require.NotNil(t, def.Parameters)
@@ -116,14 +113,13 @@ func TestFunctionTool_Definition(t *testing.T) {
type Params struct{}
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"noop",
"No-op",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, nil
},
)
require.NoError(t, err)
var schema map[string]any
require.NoError(t, json.Unmarshal(tool.Definition().Parameters, &schema))
@@ -140,14 +136,13 @@ func TestFunctionTool_Definition(t *testing.T) {
Title *string `json:"title,omitempty"`
}
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"update",
"Update",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, nil
},
)
require.NoError(t, err)
var schema map[string]any
require.NoError(t, json.Unmarshal(tool.Definition().Parameters, &schema))
@@ -173,14 +168,13 @@ func TestFunctionTool_Execute(t *testing.T) {
Y int `json:"y"`
}
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"add",
"Add two numbers",
func(_ context.Context, p Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "42"}, nil
},
)
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{"x": 1, "y": 2}`)
require.NoError(t, err)
@@ -199,7 +193,7 @@ func TestFunctionTool_Execute(t *testing.T) {
}
var received string
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"weather",
"Get weather",
func(_ context.Context, p Params) (agent.ToolResult, error) {
@@ -207,9 +201,8 @@ func TestFunctionTool_Execute(t *testing.T) {
return agent.ToolResult{Content: "sunny"}, nil
},
)
require.NoError(t, err)
_, err = tool.Execute(context.Background(), `{"city":"Paris"}`)
_, err := tool.Execute(context.Background(), `{"city":"Paris"}`)
require.NoError(t, err)
assert.Equal(t, "Paris", received)
},
@@ -222,14 +215,13 @@ func TestFunctionTool_Execute(t *testing.T) {
type Params struct{}
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"noop",
"No-op",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil
},
)
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{invalid`)
require.NoError(t, err)
@@ -245,16 +237,15 @@ func TestFunctionTool_Execute(t *testing.T) {
type Params struct{}
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"fail",
"Always fails",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, errors.New("db down")
},
)
require.NoError(t, err)
_, err = tool.Execute(context.Background(), `{}`)
_, err := tool.Execute(context.Background(), `{}`)
require.Error(t, err)
assert.Contains(t, err.Error(), "db down")
},
@@ -268,7 +259,7 @@ func TestFunctionTool_Execute(t *testing.T) {
type ctxKey struct{}
type Params struct{}
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"ctx_check",
"Check context",
func(ctx context.Context, _ Params) (agent.ToolResult, error) {
@@ -276,7 +267,6 @@ func TestFunctionTool_Execute(t *testing.T) {
return agent.ToolResult{Content: val}, nil
},
)
require.NoError(t, err)
ctx := context.WithValue(context.Background(), ctxKey{}, "hello")
result, err := tool.Execute(ctx, `{}`)
@@ -294,7 +284,7 @@ func TestFunctionTool_Execute(t *testing.T) {
City string `json:"city"`
}
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"weather",
"Get weather",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
@@ -302,7 +292,6 @@ func TestFunctionTool_Execute(t *testing.T) {
return agent.ToolResult{}, nil
},
)
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{}`)
require.NoError(t, err)
@@ -321,7 +310,7 @@ func TestFunctionTool_Execute(t *testing.T) {
Country string `json:"country"`
}
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"weather",
"Get weather",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
@@ -329,7 +318,6 @@ func TestFunctionTool_Execute(t *testing.T) {
return agent.ToolResult{}, nil
},
)
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{}`)
require.NoError(t, err)
@@ -349,7 +337,7 @@ func TestFunctionTool_Execute(t *testing.T) {
Country string `json:"country"`
}
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"weather",
"Get weather",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
@@ -357,7 +345,6 @@ func TestFunctionTool_Execute(t *testing.T) {
return agent.ToolResult{}, nil
},
)
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{"city":"Paris"}`)
require.NoError(t, err)
@@ -376,14 +363,13 @@ func TestFunctionTool_Execute(t *testing.T) {
Units *string `json:"units,omitempty"`
}
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"weather",
"Get weather",
func(_ context.Context, p Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "sunny in " + p.City}, nil
},
)
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{"city":"Paris"}`)
require.NoError(t, err)
@@ -401,14 +387,13 @@ func TestFunctionTool_Execute(t *testing.T) {
Name string `json:"name"`
}
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"greet",
"Greet",
func(_ context.Context, p Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "hi " + p.Name}, nil
},
)
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{"name":"Alice","extra":"ignored"}`)
require.NoError(t, err)
@@ -424,14 +409,13 @@ func TestFunctionTool_Execute(t *testing.T) {
type Params struct{}
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"ping",
"Ping",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "pong"}, nil
},
)
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{}`)
require.NoError(t, err)
@@ -446,14 +430,13 @@ func TestFunctionTool_Execute(t *testing.T) {
type Params struct{}
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"validate",
"Validate input",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "validation failed", IsError: true}, nil
},
)
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{}`)
require.NoError(t, err)
@@ -468,14 +451,13 @@ func TestFunctionTool_InterfaceSatisfaction(t *testing.T) {
type Params struct{}
tool, err := agent.FunctionTool(
tool := agent.FunctionTool(
"test",
"test tool",
func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, nil
},
)
require.NoError(t, err)
assert.Implements(t, (*agent.Tool)(nil), tool)
assert.Implements(t, (*agent.ToolDescriptor)(nil), tool)

170
pkg/agent/tools/browser/browser.go Normal file
View File

@@ -0,0 +1,170 @@
// 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 browser
import (
"context"
"errors"
"fmt"
"net/url"
"strings"
"time"
"github.com/chromedp/chromedp"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/agent/tools/internal/netcheck"
)
const (
defaultToolTimeout = 60 * time.Second
)
type Browser struct {
addr string
allocCtx context.Context
cancel context.CancelFunc
allowedDomains []string
}
func NewBrowser(ctx context.Context, addr string) *Browser {
if !strings.HasPrefix(addr, "ws://") && !strings.HasPrefix(addr, "wss://") {
addr = "ws://" + addr
}
allocCtx, cancel := chromedp.NewRemoteAllocator(ctx, addr)
return &Browser{
addr: addr,
allocCtx: allocCtx,
cancel: cancel,
}
}
// SetAllowedDomain restricts navigation to URLs under the given domain and
// its subdomains. For example, setting "getprobo.com" allows navigation to
// getprobo.com, www.getprobo.com, and compliance.getprobo.com.
// This replaces any previously set domains.
func (b *Browser) SetAllowedDomain(domain string) {
domain = strings.ToLower(strings.TrimSpace(domain))
// Strip "www." prefix so that setting either "www.example.com" or
// "example.com" allows navigation to *.example.com.
domain = strings.TrimPrefix(domain, "www.")
b.allowedDomains = []string{domain}
}
// checkURL validates that the URL is allowed. It returns an error tool result
// if the URL uses a disallowed scheme, resolves to a non-public IP, or is
// outside the allowed domains.
func (b *Browser) checkURL(rawURL string) *agent.ToolResult {
u, err := url.Parse(rawURL)
if err != nil {
return &agent.ToolResult{
Content: fmt.Sprintf("invalid URL: %s", err),
IsError: true,
}
}
if u.Scheme != "http" && u.Scheme != "https" {
return &agent.ToolResult{
Content: fmt.Sprintf("cannot navigate to URL with scheme %q: only http and https are allowed", u.Scheme),
IsError: true,
}
}
// Always reject URLs that resolve to non-public IPs, even when no
// allowed-domain list is set. This closes the SSRF path on browsers
// used for open-ended external research (e.g. the research browser
// in vendor assessments).
if err := netcheck.ValidatePublicURL(rawURL); err != nil {
return &agent.ToolResult{
Content: fmt.Sprintf("navigation blocked: %s", err),
IsError: true,
}
}
if len(b.allowedDomains) == 0 {
return nil
}
host := strings.ToLower(u.Hostname())
for _, allowed := range b.allowedDomains {
if host == allowed || strings.HasSuffix(host, "."+allowed) {
return nil
}
}
return &agent.ToolResult{
Content: fmt.Sprintf("navigation blocked: %s is outside the allowed domains", host),
IsError: true,
}
}
// checkAlive returns a tool error result if the browser connection has been
// lost. Call this at the start of every tool to fail fast with a clear
// message instead of waiting for the tool timeout.
func (b *Browser) checkAlive() *agent.ToolResult {
if err := b.allocCtx.Err(); err != nil {
return &agent.ToolResult{
Content: "browser connection lost: the remote Chrome instance is no longer reachable",
IsError: true,
}
}
return nil
}
// classifyError inspects the caller's timeout context and the browser's
// allocator context to produce a human-readable error message. Without this,
// both a tool timeout and a dropped Chrome connection appear as the opaque
// "context canceled".
func (b *Browser) classifyError(timeoutCtx context.Context, rawURL string, err error) string {
if b.allocCtx.Err() != nil {
return fmt.Sprintf(
"browser connection lost while loading %s: the remote Chrome instance is no longer reachable",
rawURL,
)
}
if errors.Is(timeoutCtx.Err(), context.DeadlineExceeded) {
return fmt.Sprintf(
"page load timed out after %s for %s: the page may be too slow or unresponsive",
defaultToolTimeout,
rawURL,
)
}
return fmt.Sprintf("cannot load %s: %s", rawURL, err)
}
func (b *Browser) NewTab(ctx context.Context) (context.Context, context.CancelFunc) {
tabCtx, tabCancel := chromedp.NewContext(b.allocCtx)
// Propagate the caller's cancellation to the Chrome tab so that
// tool-level timeouts and context deadlines actually stop the browser.
go func() {
select {
case <-ctx.Done():
tabCancel()
case <-tabCtx.Done():
}
}()
return tabCtx, tabCancel
}
func (b *Browser) Close() {
b.cancel()
}

88
pkg/agent/tools/browser/click.go Normal file
View File

@@ -0,0 +1,88 @@
// 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 browser
import (
"context"
"github.com/chromedp/chromedp"
"go.probo.inc/probo/pkg/agent"
)
type (
clickParams struct {
URL string `json:"url" jsonschema:"The URL to navigate to before clicking"`
Selector string `json:"selector" jsonschema:"CSS selector of the element to click (e.g. button.next, a[href*=page])"`
}
)
func ClickElementTool(b *Browser) agent.Tool {
return agent.FunctionTool(
"click_element",
"Navigate to a URL, click an element matching a CSS selector, and return the page text after the click. Useful for pagination buttons, 'show all' links, tabs, and other interactive elements.",
func(ctx context.Context, p clickParams) (agent.ToolResult, error) {
if r := b.checkAlive(); r != nil {
return *r, nil
}
if r := b.checkURL(p.URL); r != nil {
return *r, nil
}
ctx, timeoutCancel := withToolTimeout(ctx)
defer timeoutCancel()
tabCtx, cancel := b.NewTab(ctx)
defer cancel()
var (
text string
postClickURL string
)
err := chromedp.Run(
tabCtx,
chromedp.Navigate(p.URL),
waitForPage(),
chromedp.WaitVisible(p.Selector),
chromedp.Click(p.Selector),
waitForPage(),
chromedp.Location(&postClickURL),
chromedp.Evaluate(`document.body.innerText`, &text),
)
if err != nil {
return agent.ResultError(b.classifyError(ctx, p.URL, err)), nil
}
// Revalidate the post-click URL: a click may navigate
// the page to a different host (redirect, JS navigation,
// <a href>), bypassing the initial checkURL. Reject the
// result if the new URL is outside the allowed scope or
// resolves to a non-public IP.
if postClickURL != "" && postClickURL != p.URL {
if r := b.checkURL(postClickURL); r != nil {
return *r, nil
}
}
runes := []rune(text)
if len(runes) > maxTextLength {
text = string(runes[:maxTextLength])
}
return agent.ToolResult{Content: text}, nil
},
)
}

157
pkg/agent/tools/browser/download_pdf.go Normal file
View File

@@ -0,0 +1,157 @@
// 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 browser
import (
"bytes"
"context"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"strings"
"time"
"github.com/pdfcpu/pdfcpu/pkg/api"
"github.com/pdfcpu/pdfcpu/pkg/pdfcpu/model"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/agent/tools/internal/netcheck"
)
type (
downloadPDFParams struct {
URL string `json:"url" jsonschema:"The URL of the PDF document to download and extract text from"`
}
downloadPDFResult struct {
Text string `json:"text"`
PageCount int `json:"page_count"`
ErrorDetail string `json:"error_detail,omitempty"`
}
)
func DownloadPDFTool() agent.Tool {
client := &http.Client{
Timeout: 30 * time.Second,
Transport: netcheck.NewPinnedTransport(),
}
return agent.FunctionTool(
"download_pdf",
"Download a PDF document from a URL and extract its text content. Use this for DPAs, SOC 2 reports, privacy policies, and other documents hosted as PDFs.",
func(ctx context.Context, p downloadPDFParams) (agent.ToolResult, error) {
if err := validatePublicURL(p.URL); err != nil {
return agent.ResultJSON(downloadPDFResult{
ErrorDetail: fmt.Sprintf("URL not allowed: %s", err),
}), nil
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, p.URL, nil)
if err != nil {
return agent.ResultJSON(downloadPDFResult{
ErrorDetail: fmt.Sprintf("cannot create request: %s", err),
}), nil
}
resp, err := client.Do(req)
if err != nil {
return agent.ResultJSON(downloadPDFResult{
ErrorDetail: fmt.Sprintf("cannot download PDF: %s", err),
}), nil
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode != http.StatusOK {
return agent.ResultJSON(downloadPDFResult{
ErrorDetail: fmt.Sprintf("PDF download returned status %d", resp.StatusCode),
}), nil
}
// Read PDF into memory (max 20MB).
body, err := io.ReadAll(io.LimitReader(resp.Body, 20*1024*1024))
if err != nil {
return agent.ResultJSON(downloadPDFResult{
ErrorDetail: fmt.Sprintf("cannot read PDF body: %s", err),
}), nil
}
// Write to temp file for pdfcpu.
tmpDir, err := os.MkdirTemp("", "pdf-extract-*")
if err != nil {
return agent.ResultJSON(downloadPDFResult{
ErrorDetail: fmt.Sprintf("cannot create temp dir: %s", err),
}), nil
}
defer os.RemoveAll(tmpDir)
tmpFile := filepath.Join(tmpDir, "input.pdf")
if err := os.WriteFile(tmpFile, body, 0o600); err != nil {
return agent.ResultJSON(downloadPDFResult{
ErrorDetail: fmt.Sprintf("cannot write temp file: %s", err),
}), nil
}
// Get page count.
conf := model.NewDefaultConfiguration()
pageCount, err := api.PageCountFile(tmpFile)
if err != nil {
return agent.ResultJSON(downloadPDFResult{
ErrorDetail: fmt.Sprintf("cannot read PDF: %s", err),
}), nil
}
// Extract content to output dir.
outDir := filepath.Join(tmpDir, "out")
if err := os.MkdirAll(outDir, 0o700); err != nil {
return agent.ResultJSON(downloadPDFResult{
ErrorDetail: fmt.Sprintf("cannot create output dir: %s", err),
}), nil
}
reader := bytes.NewReader(body)
if err := api.ExtractContent(reader, outDir, "content", nil, conf); err != nil {
return agent.ResultJSON(downloadPDFResult{
ErrorDetail: fmt.Sprintf("cannot extract PDF content: %s", err),
}), nil
}
// Read all extracted content files.
var sb strings.Builder
entries, _ := os.ReadDir(outDir)
for _, entry := range entries {
if entry.IsDir() {
continue
}
content, err := os.ReadFile(filepath.Join(outDir, entry.Name()))
if err != nil {
continue
}
sb.Write(content)
sb.WriteString("\n")
}
text := sb.String()
if len(text) > maxTextLength {
text = text[:maxTextLength] + "\n[... truncated]"
}
return agent.ResultJSON(downloadPDFResult{
Text: text,
PageCount: pageCount,
}), nil
},
)
}

81
pkg/agent/tools/browser/extract_links.go Normal file
View File

@@ -0,0 +1,81 @@
// 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 browser
import (
"context"
"net/url"
"github.com/chromedp/chromedp"
"go.probo.inc/probo/pkg/agent"
)
type (
extractLinksParams struct {
URL string `json:"url" jsonschema:"The URL to extract links from"`
}
link struct {
Href string `json:"href"`
Text string `json:"text"`
}
)
func ExtractLinksTool(b *Browser) agent.Tool {
return agent.FunctionTool(
"extract_links",
"Navigate to a URL and extract all links (<a> elements) with their href and text.",
func(ctx context.Context, p extractLinksParams) (agent.ToolResult, error) {
if r := b.checkAlive(); r != nil {
return *r, nil
}
u, err := url.Parse(p.URL)
if err != nil || (u.Scheme != "http" && u.Scheme != "https") {
return agent.ResultError("invalid URL scheme: only http and https are allowed"), nil
}
if r := b.checkURL(p.URL); r != nil {
return *r, nil
}
ctx, timeoutCancel := withToolTimeout(ctx)
defer timeoutCancel()
tabCtx, cancel := b.NewTab(ctx)
defer cancel()
var links []link
err = chromedp.Run(
tabCtx,
chromedp.Navigate(p.URL),
waitForPage(),
chromedp.Evaluate(
`Array.from(document.querySelectorAll("a[href]")).map(a => ({
href: a.href,
text: a.innerText.trim().substring(0, 200)
}))`,
&links,
),
)
if err != nil {
return agent.ResultError(b.classifyError(ctx, p.URL, err)), nil
}
return agent.ResultJSON(links), nil
},
)
}

95
pkg/agent/tools/browser/extract_text.go Normal file
View File

@@ -0,0 +1,95 @@
// 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 browser
import (
"context"
"fmt"
"time"
"github.com/chromedp/chromedp"
"go.probo.inc/probo/pkg/agent"
)
const (
maxTextLength = 32000
)
type (
extractTextParams struct {
URL string `json:"url" jsonschema:"The URL to extract text from"`
}
)
func ExtractPageTextTool(b *Browser) agent.Tool {
return agent.FunctionTool(
"extract_page_text",
"Navigate to a URL and extract the visible text content of the page, truncated to 32000 characters.",
func(ctx context.Context, p extractTextParams) (agent.ToolResult, error) {
if r := b.checkAlive(); r != nil {
return *r, nil
}
if r := b.checkURL(p.URL); r != nil {
return *r, nil
}
if r := checkPDF(p.URL); r != nil {
return *r, nil
}
ctx, timeoutCancel := withToolTimeout(ctx)
defer timeoutCancel()
tabCtx, cancel := b.NewTab(ctx)
defer cancel()
var text string
// Cap the JS-side slice at 4 code units per rune so the
// DevTools transfer stays bounded even for huge pages;
// the Go-side rune truncation below then produces the
// final exact-length output.
jsMaxLen := maxTextLength * 4
extractJS := fmt.Sprintf(
`String(document.body?.innerText ?? '').slice(0, %d)`,
jsMaxLen,
)
err := chromedp.Run(
tabCtx,
chromedp.Navigate(p.URL),
waitForPage(),
// Scroll to bottom to trigger lazy-loaded content,
// then back to top and wait briefly for rendering.
chromedp.Evaluate(`window.scrollTo(0, document.body.scrollHeight)`, nil),
chromedp.Sleep(500*time.Millisecond),
chromedp.Evaluate(`window.scrollTo(0, 0)`, nil),
chromedp.Sleep(200*time.Millisecond),
chromedp.Evaluate(extractJS, &text),
)
if err != nil {
return agent.ResultError(b.classifyError(ctx, p.URL, err)), nil
}
runes := []rune(text)
if len(runes) > maxTextLength {
text = string(runes[:maxTextLength])
}
return agent.ToolResult{Content: text}, nil
},
)
}

107
pkg/agent/tools/browser/fetch_robots.go Normal file
View File

@@ -0,0 +1,107 @@
// 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 browser
import (
"bufio"
"context"
"fmt"
"net/http"
"strings"
"time"
"go.probo.inc/probo/pkg/agent"
)
type (
robotsParams struct {
Domain string `json:"domain" jsonschema:"The domain to fetch robots.txt from (e.g. example.com)"`
}
robotsResult struct {
Found bool `json:"found"`
Sitemaps []string `json:"sitemaps,omitempty"`
Disallowed []string `json:"disallowed_paths,omitempty"`
ErrorDetail string `json:"error_detail,omitempty"`
}
)
func FetchRobotsTxtTool() agent.Tool {
client := &http.Client{Timeout: 10 * time.Second}
return agent.FunctionTool(
"fetch_robots_txt",
"Fetch and parse the robots.txt file for a domain. Returns sitemap URLs and disallowed paths, which can reveal hidden pages the crawler might miss.",
func(ctx context.Context, p robotsParams) (agent.ToolResult, error) {
if err := validatePublicDomain(p.Domain); err != nil {
return agent.ResultJSON(robotsResult{
Found: false,
ErrorDetail: fmt.Sprintf("domain not allowed: %s", err),
}), nil
}
u := "https://" + p.Domain + "/robots.txt"
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
if err != nil {
return agent.ResultJSON(robotsResult{
Found: false,
ErrorDetail: fmt.Sprintf("cannot create request: %s", err),
}), nil
}
resp, err := client.Do(req)
if err != nil {
return agent.ResultJSON(robotsResult{
Found: false,
ErrorDetail: fmt.Sprintf("cannot fetch robots.txt: %s", err),
}), nil
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode != http.StatusOK {
return agent.ResultJSON(robotsResult{
Found: false,
ErrorDetail: fmt.Sprintf("robots.txt returned status %d", resp.StatusCode),
}), nil
}
var result robotsResult
result.Found = true
scanner := bufio.NewScanner(resp.Body)
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
// Directive names are case-insensitive but values
// (URLs, paths) are case-sensitive, so extract the
// original-case suffix from the raw line rather than
// reading it off the lowercased copy.
if after, ok := strings.CutPrefix(strings.ToLower(line), "sitemap:"); ok {
result.Sitemaps = append(result.Sitemaps, strings.TrimSpace(line[len(line)-len(after):]))
}
if after, ok := strings.CutPrefix(strings.ToLower(line), "disallow:"); ok {
path := strings.TrimSpace(line[len(line)-len(after):])
if path != "" && len(result.Disallowed) < 50 {
result.Disallowed = append(result.Disallowed, path)
}
}
}
return agent.ResultJSON(result), nil
},
)
}

151
pkg/agent/tools/browser/fetch_sitemap.go Normal file
View File

@@ -0,0 +1,151 @@
// 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 browser
import (
"compress/gzip"
"context"
"encoding/xml"
"fmt"
"io"
"net/http"
"strings"
"time"
"go.probo.inc/probo/pkg/agent"
)
type (
sitemapParams struct {
URL string `json:"url" jsonschema:"The full URL of the sitemap to fetch (e.g. https://example.com/sitemap.xml)"`
}
sitemapResult struct {
Found bool `json:"found"`
URLs []string `json:"urls,omitempty"`
URLCount int `json:"url_count"`
ErrorDetail string `json:"error_detail,omitempty"`
}
)
const (
maxSitemapURLs = 200
)
func FetchSitemapTool() agent.Tool {
client := &http.Client{Timeout: 15 * time.Second}
return agent.FunctionTool(
"fetch_sitemap",
"Fetch and parse a sitemap XML file. Returns discovered URLs which can reveal pages not linked from the main navigation (trust centers, legal docs, status pages).",
func(ctx context.Context, p sitemapParams) (agent.ToolResult, error) {
if err := validatePublicURL(p.URL); err != nil {
return agent.ResultJSON(sitemapResult{
Found: false,
ErrorDetail: fmt.Sprintf("URL not allowed: %s", err),
}), nil
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, p.URL, nil)
if err != nil {
return agent.ResultJSON(sitemapResult{
Found: false,
ErrorDetail: fmt.Sprintf("cannot create request: %s", err),
}), nil
}
resp, err := client.Do(req)
if err != nil {
return agent.ResultJSON(sitemapResult{
Found: false,
ErrorDetail: fmt.Sprintf("cannot fetch sitemap: %s", err),
}), nil
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode != http.StatusOK {
return agent.ResultJSON(sitemapResult{
Found: false,
ErrorDetail: fmt.Sprintf("sitemap returned status %d", resp.StatusCode),
}), nil
}
var reader io.Reader = resp.Body
if strings.HasSuffix(strings.ToLower(p.URL), ".gz") ||
resp.Header.Get("Content-Encoding") == "gzip" {
gz, err := gzip.NewReader(resp.Body)
if err != nil {
return agent.ResultJSON(sitemapResult{
Found: false,
ErrorDetail: fmt.Sprintf("cannot decompress gzipped sitemap: %s", err),
}), nil
}
defer gz.Close()
reader = gz
}
// Limit read to 5MB.
reader = io.LimitReader(reader, 5*1024*1024)
urls, err := parseSitemapXML(reader)
if err != nil {
return agent.ResultJSON(sitemapResult{
Found: false,
ErrorDetail: fmt.Sprintf("cannot parse sitemap XML: %s", err),
}), nil
}
result := sitemapResult{
Found: true,
URLCount: len(urls),
}
if len(urls) > maxSitemapURLs {
result.URLs = urls[:maxSitemapURLs]
} else {
result.URLs = urls
}
return agent.ResultJSON(result), nil
},
)
}
func parseSitemapXML(r io.Reader) ([]string, error) {
var urls []string
decoder := xml.NewDecoder(r)
for {
tok, err := decoder.Token()
if err == io.EOF {
break
}
if err != nil {
return urls, err
}
if se, ok := tok.(xml.StartElement); ok && se.Name.Local == "loc" {
var loc string
if err := decoder.DecodeElement(&loc, &se); err == nil {
loc = strings.TrimSpace(loc)
if loc != "" {
urls = append(urls, loc)
}
}
}
}
return urls, nil
}

97
pkg/agent/tools/browser/find_links.go Normal file
View File

@@ -0,0 +1,97 @@
// 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 browser
import (
"context"
"encoding/json"
"fmt"
"github.com/chromedp/chromedp"
"go.probo.inc/probo/pkg/agent"
)
type (
findLinksParams struct {
URL string `json:"url" jsonschema:"The URL to search for links"`
Pattern string `json:"pattern" jsonschema:"Keyword to filter links by (case-insensitive match on href or text)"`
}
)
func FindLinksMatchingTool(b *Browser) agent.Tool {
return agent.FunctionTool(
"find_links_matching",
"Navigate to a URL and extract links whose href or text matches a keyword (case-insensitive).",
func(ctx context.Context, p findLinksParams) (agent.ToolResult, error) {
if r := b.checkAlive(); r != nil {
return *r, nil
}
if r := b.checkURL(p.URL); r != nil {
return *r, nil
}
if p.Pattern == "" {
return agent.ResultError("pattern must not be empty"), nil
}
ctx, timeoutCancel := withToolTimeout(ctx)
defer timeoutCancel()
tabCtx, cancel := b.NewTab(ctx)
defer cancel()
var links []link
patternJSON, err := json.Marshal(p.Pattern)
if err != nil {
return agent.ResultErrorf("cannot encode pattern: %s", err), nil
}
js := fmt.Sprintf(
`(() => {
const pattern = JSON.parse(%s).toLowerCase();
const normalize = s => s.replace(/[-_\s]+/g, "");
const normalizedPattern = normalize(pattern);
return Array.from(document.querySelectorAll("a[href]"))
.filter(a => {
const href = a.href.toLowerCase();
const text = a.innerText.toLowerCase();
return href.includes(pattern) || text.includes(pattern)
|| normalize(href).includes(normalizedPattern)
|| normalize(text).includes(normalizedPattern);
})
.map(a => ({
href: a.href,
text: a.innerText.trim().substring(0, 200)
}));
})()`,
string(patternJSON),
)
err = chromedp.Run(
tabCtx,
chromedp.Navigate(p.URL),
waitForPage(),
chromedp.Evaluate(js, &links),
)
if err != nil {
return agent.ResultError(b.classifyError(ctx, p.URL, err)), nil
}
return agent.ResultJSON(links), nil
},
)
}

118
pkg/agent/tools/browser/helpers.go Normal file
View File

@@ -0,0 +1,118 @@
// 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 browser
import (
"context"
"fmt"
"strings"
"github.com/chromedp/chromedp"
"go.probo.inc/probo/pkg/agent"
)
// waitForPage returns chromedp actions that wait for the page to fully load,
// including SPA content rendered by JavaScript. It first waits for the body to
// be ready, then polls until the page content stabilizes (innerText stops
// changing) with a short debounce. After stabilization, it attempts to dismiss
// common cookie consent banners so they don't interfere with content
// extraction.
func waitForPage() chromedp.Action {
return chromedp.ActionFunc(func(ctx context.Context) error {
if err := chromedp.WaitReady("body").Do(ctx); err != nil {
return err
}
// Wait for SPA content to stabilize by checking if innerText
// length stops changing over a 500ms window. Gives up after 5s.
// EvaluateAsDevTools is required to await the Promise.
if err := chromedp.EvaluateAsDevTools(`
new Promise((resolve) => {
let lastLen = -1;
let stableCount = 0;
const interval = setInterval(() => {
const curLen = document.body.innerText.length;
if (curLen === lastLen && curLen > 0) {
stableCount++;
} else {
stableCount = 0;
}
lastLen = curLen;
if (stableCount >= 2) {
clearInterval(interval);
resolve(true);
}
}, 250);
setTimeout(() => {
clearInterval(interval);
resolve(true);
}, 5000);
})
`, nil).Do(ctx); err != nil {
return err
}
// Dismiss common cookie consent banners. This is best-effort;
// failures are silently ignored because not every page has a
// banner and the selectors may not match.
return chromedp.Evaluate(`
(() => {
const selectors = [
"#onetrust-accept-btn-handler",
"#CybotCookiebotDialogBodyLevelButtonLevelOptinAllowAll",
"#CybotCookiebotDialogBodyButtonAccept",
".cky-btn-accept",
"[data-testid='cookie-policy-dialog-accept-button']",
"button.accept-cookies",
"#cookie-accept",
"#accept-cookies",
".cc-accept",
".cc-btn.cc-dismiss",
];
for (const sel of selectors) {
const btn = document.querySelector(sel);
if (btn) { btn.click(); return; }
}
const buttons = document.querySelectorAll(
"button, a[role='button'], [role='button']"
);
const patterns = /^(accept all|accept|agree|i agree|allow all|allow|got it|ok|okay|consent)$/i;
for (const btn of buttons) {
if (patterns.test(btn.innerText.trim())) {
btn.click();
return;
}
}
})()
`, nil).Do(ctx)
})
}
// checkPDF returns an error tool result if the URL points to a PDF file,
// which cannot be rendered by the headless browser.
func checkPDF(rawURL string) *agent.ToolResult {
if strings.HasSuffix(strings.ToLower(rawURL), ".pdf") {
return &agent.ToolResult{
Content: fmt.Sprintf("cannot load %s: PDF files are not supported by the browser", rawURL),
IsError: true,
}
}
return nil
}
func withToolTimeout(ctx context.Context) (context.Context, context.CancelFunc) {
return context.WithTimeout(ctx, defaultToolTimeout)
}

92
pkg/agent/tools/browser/helpers_test.go Normal file
View File

@@ -0,0 +1,92 @@
// 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 browser
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestCheckPDF(t *testing.T) {
t.Parallel()
tests := []struct {
name string
url string
wantError bool
}{
{
name: "lowercase .pdf returns error",
url: "https://example.com/document.pdf",
wantError: true,
},
{
name: "uppercase .PDF returns error",
url: "https://example.com/document.PDF",
wantError: true,
},
{
name: "mixed case .Pdf returns error",
url: "https://example.com/document.Pdf",
wantError: true,
},
{
name: "normal URL returns nil",
url: "https://example.com/page",
wantError: false,
},
{
name: "URL with .pdf in path but not at end returns nil",
url: "https://example.com/pdf-viewer/document",
wantError: false,
},
{
name: "URL with .pdf in query but not at end returns nil",
url: "https://example.com/view?file=report.pdf&page=1",
wantError: false,
},
{
name: "html URL returns nil",
url: "https://example.com/page.html",
wantError: false,
},
{
name: "URL ending with .pdf and path segments",
url: "https://example.com/files/reports/annual.pdf",
wantError: true,
},
}
for _, tt := range tests {
t.Run(
tt.name,
func(t *testing.T) {
t.Parallel()
result := checkPDF(tt.url)
if tt.wantError {
require.NotNil(t, result)
assert.True(t, result.IsError)
assert.Contains(t, result.Content, "PDF files are not supported")
} else {
assert.Nil(t, result)
}
},
)
}
}

90
pkg/agent/tools/browser/navigate.go Normal file
View File

@@ -0,0 +1,90 @@
// 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 browser
import (
"context"
"github.com/chromedp/chromedp"
"go.probo.inc/probo/pkg/agent"
)
type (
navigateParams struct {
URL string `json:"url" jsonschema:"The URL to navigate to"`
}
navigateResult struct {
Title string `json:"title"`
Description string `json:"description"`
FinalURL string `json:"final_url"`
}
)
func NavigateToURLTool(b *Browser) agent.Tool {
return agent.FunctionTool(
"navigate_to_url",
"Navigate to a URL and return the page title, meta description, and final URL after redirects.",
func(ctx context.Context, p navigateParams) (agent.ToolResult, error) {
if r := b.checkAlive(); r != nil {
return *r, nil
}
if r := b.checkURL(p.URL); r != nil {
return *r, nil
}
if r := checkPDF(p.URL); r != nil {
return *r, nil
}
ctx, timeoutCancel := withToolTimeout(ctx)
defer timeoutCancel()
tabCtx, cancel := b.NewTab(ctx)
defer cancel()
var (
title string
description string
finalURL string
)
err := chromedp.Run(
tabCtx,
chromedp.Navigate(p.URL),
waitForPage(),
chromedp.Title(&title),
chromedp.Evaluate(
`(() => {
const meta = document.querySelector('meta[name="description"]');
return meta ? meta.getAttribute("content") : "";
})()`,
&description,
),
chromedp.Location(&finalURL),
)
if err != nil {
return agent.ResultError(b.classifyError(ctx, p.URL, err)), nil
}
return agent.ResultJSON(navigateResult{
Title: title,
Description: description,
FinalURL: finalURL,
}), nil
},
)
}

82
pkg/agent/tools/browser/select.go Normal file
View File

@@ -0,0 +1,82 @@
// 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 browser
import (
"context"
"fmt"
"github.com/chromedp/chromedp"
"go.probo.inc/probo/pkg/agent"
)
type (
selectParams struct {
URL string `json:"url" jsonschema:"The URL to navigate to before selecting"`
Selector string `json:"selector" jsonschema:"CSS selector of the select element"`
Value string `json:"value" jsonschema:"The option value to select"`
}
)
func SelectOptionTool(b *Browser) agent.Tool {
return agent.FunctionTool(
"select_option",
"Navigate to a URL, select an option from a <select> dropdown, and return the page text after selection. Useful for changing page size dropdowns (e.g. 'show 100 per page').",
func(ctx context.Context, p selectParams) (agent.ToolResult, error) {
if r := b.checkAlive(); r != nil {
return *r, nil
}
if r := b.checkURL(p.URL); r != nil {
return *r, nil
}
ctx, timeoutCancel := withToolTimeout(ctx)
defer timeoutCancel()
tabCtx, cancel := b.NewTab(ctx)
defer cancel()
var text string
err := chromedp.Run(
tabCtx,
chromedp.Navigate(p.URL),
waitForPage(),
chromedp.WaitVisible(p.Selector),
chromedp.SetValue(p.Selector, p.Value),
chromedp.Evaluate(
fmt.Sprintf(
`document.querySelector(%q).dispatchEvent(new Event('change', {bubbles: true}))`,
p.Selector,
),
nil,
),
waitForPage(),
chromedp.Evaluate(`document.body.innerText`, &text),
)
if err != nil {
return agent.ResultError(b.classifyError(ctx, p.URL, err)), nil
}
runes := []rune(text)
if len(runes) > maxTextLength {
text = string(runes[:maxTextLength])
}
return agent.ToolResult{Content: text}, nil
},
)
}

191
pkg/agent/tools/browser/sitemap_test.go Normal file
View File

@@ -0,0 +1,191 @@
// 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 browser
import (
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestParseSitemapXML(t *testing.T) {
t.Parallel()
t.Run(
"valid urlset with multiple URLs",
func(t *testing.T) {
t.Parallel()
xml := `<?xml version="1.0" encoding="UTF-8"?>
<urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">
<url><loc>https://example.com/</loc></url>
<url><loc>https://example.com/about</loc></url>
<url><loc>https://example.com/contact</loc></url>
</urlset>`
urls, err := parseSitemapXML(strings.NewReader(xml))
require.NoError(t, err)
require.Len(t, urls, 3)
assert.Equal(t, "https://example.com/", urls[0])
assert.Equal(t, "https://example.com/about", urls[1])
assert.Equal(t, "https://example.com/contact", urls[2])
},
)
t.Run(
"valid sitemapindex with sitemap locations",
func(t *testing.T) {
t.Parallel()
xml := `<?xml version="1.0" encoding="UTF-8"?>
<sitemapindex xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">
<sitemap><loc>https://example.com/sitemap-pages.xml</loc></sitemap>
<sitemap><loc>https://example.com/sitemap-posts.xml</loc></sitemap>
</sitemapindex>`
urls, err := parseSitemapXML(strings.NewReader(xml))
require.NoError(t, err)
require.Len(t, urls, 2)
assert.Equal(t, "https://example.com/sitemap-pages.xml", urls[0])
assert.Equal(t, "https://example.com/sitemap-posts.xml", urls[1])
},
)
t.Run(
"empty urlset returns empty slice",
func(t *testing.T) {
t.Parallel()
xml := `<?xml version="1.0" encoding="UTF-8"?>
<urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">
</urlset>`
urls, err := parseSitemapXML(strings.NewReader(xml))
require.NoError(t, err)
assert.Empty(t, urls)
},
)
t.Run(
"malformed XML returns error",
func(t *testing.T) {
t.Parallel()
xml := `<urlset><url><loc>https://example.com/</loc></url`
_, err := parseSitemapXML(strings.NewReader(xml))
assert.Error(t, err)
},
)
t.Run(
"urlset without namespace",
func(t *testing.T) {
t.Parallel()
xml := `<?xml version="1.0" encoding="UTF-8"?>
<urlset>
<url><loc>https://example.com/page1</loc></url>
<url><loc>https://example.com/page2</loc></url>
</urlset>`
urls, err := parseSitemapXML(strings.NewReader(xml))
require.NoError(t, err)
require.Len(t, urls, 2)
assert.Equal(t, "https://example.com/page1", urls[0])
assert.Equal(t, "https://example.com/page2", urls[1])
},
)
t.Run(
"trims whitespace in loc elements",
func(t *testing.T) {
t.Parallel()
xml := `<?xml version="1.0" encoding="UTF-8"?>
<urlset>
<url><loc> https://example.com/padded </loc></url>
</urlset>`
urls, err := parseSitemapXML(strings.NewReader(xml))
require.NoError(t, err)
require.Len(t, urls, 1)
assert.Equal(t, "https://example.com/padded", urls[0])
},
)
t.Run(
"skips empty loc elements",
func(t *testing.T) {
t.Parallel()
xml := `<?xml version="1.0" encoding="UTF-8"?>
<urlset>
<url><loc></loc></url>
<url><loc>https://example.com/valid</loc></url>
<url><loc> </loc></url>
</urlset>`
urls, err := parseSitemapXML(strings.NewReader(xml))
require.NoError(t, err)
require.Len(t, urls, 1)
assert.Equal(t, "https://example.com/valid", urls[0])
},
)
t.Run(
"empty reader returns empty slice",
func(t *testing.T) {
t.Parallel()
urls, err := parseSitemapXML(strings.NewReader(""))
require.NoError(t, err)
assert.Empty(t, urls)
},
)
t.Run(
"urlset with additional elements besides loc",
func(t *testing.T) {
t.Parallel()
xml := `<?xml version="1.0" encoding="UTF-8"?>
<urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">
<url>
<loc>https://example.com/page</loc>
<lastmod>2024-01-01</lastmod>
<changefreq>weekly</changefreq>
<priority>0.8</priority>
</url>
</urlset>`
urls, err := parseSitemapXML(strings.NewReader(xml))
require.NoError(t, err)
require.Len(t, urls, 1)
assert.Equal(t, "https://example.com/page", urls[0])
},
)
}

65
pkg/agent/tools/browser/toolset.go Normal file
View File

@@ -0,0 +1,65 @@
// 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 browser
import (
"go.probo.inc/probo/pkg/agent"
)
// ReadOnlyToolset provides browser tools that only read page content.
type ReadOnlyToolset struct {
browser *Browser
}
// NewReadOnlyToolset creates a read-only browser toolset.
func NewReadOnlyToolset(b *Browser) *ReadOnlyToolset {
return &ReadOnlyToolset{browser: b}
}
func (t *ReadOnlyToolset) Tools() []agent.Tool {
return []agent.Tool{
NavigateToURLTool(t.browser),
ExtractPageTextTool(t.browser),
ExtractLinksTool(t.browser),
FindLinksMatchingTool(t.browser),
FetchRobotsTxtTool(),
FetchSitemapTool(),
DownloadPDFTool(),
}
}
// InteractiveToolset provides all browser tools including click and select.
type InteractiveToolset struct {
browser *Browser
}
// NewInteractiveToolset creates an interactive browser toolset.
func NewInteractiveToolset(b *Browser) *InteractiveToolset {
return &InteractiveToolset{browser: b}
}
func (t *InteractiveToolset) Tools() []agent.Tool {
return []agent.Tool{
NavigateToURLTool(t.browser),
ExtractPageTextTool(t.browser),
ExtractLinksTool(t.browser),
FindLinksMatchingTool(t.browser),
ClickElementTool(t.browser),
SelectOptionTool(t.browser),
FetchRobotsTxtTool(),
FetchSitemapTool(),
DownloadPDFTool(),
}
}

33
pkg/agent/tools/browser/url_check.go Normal file
View File

@@ -0,0 +1,33 @@
// 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 browser
import (
"go.probo.inc/probo/pkg/agent/tools/internal/netcheck"
)
// validatePublicURL checks that a URL uses an http(s) scheme and that its
// host does not resolve to a private, loopback, or link-local IP address.
// This prevents SSRF attacks where the LLM could be tricked into requesting
// internal network endpoints.
func validatePublicURL(rawURL string) error {
return netcheck.ValidatePublicURL(rawURL)
}
// validatePublicDomain checks that a domain does not resolve to a private,
// loopback, or link-local IP address.
func validatePublicDomain(domain string) error {
return netcheck.ValidatePublicDomain(domain)
}

126
pkg/agent/tools/internal/netcheck/netcheck.go Normal file
View File

@@ -0,0 +1,126 @@
// 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 netcheck provides shared network validation functions to prevent
// SSRF attacks and DNS rebinding across agent tool packages.
package netcheck
import (
"context"
"fmt"
"net"
"net/http"
"net/url"
)
// IsPublicIP reports whether ip is a publicly routable address. It returns
// false for loopback, private, link-local, multicast (any range), and
// unspecified addresses.
func IsPublicIP(ip net.IP) bool {
if ip.IsLoopback() ||
ip.IsPrivate() ||
ip.IsLinkLocalUnicast() ||
ip.IsMulticast() ||
ip.IsUnspecified() {
return false
}
return true
}
// ValidatePublicURL checks that rawURL uses an http or https scheme and that
// its host does not resolve to a private, loopback, or link-local IP address.
// This prevents SSRF attacks where the LLM could be tricked into requesting
// internal network endpoints.
func ValidatePublicURL(rawURL string) error {
u, err := url.Parse(rawURL)
if err != nil {
return fmt.Errorf("cannot parse URL: %w", err)
}
if u.Scheme != "http" && u.Scheme != "https" {
return fmt.Errorf("unsupported URL scheme %q: only http and https are allowed", u.Scheme)
}
host := u.Hostname()
if host == "" {
return fmt.Errorf("URL has no host")
}
ips, err := net.LookupIP(host)
if err != nil {
return fmt.Errorf("cannot resolve host %q: %w", host, err)
}
for _, ip := range ips {
if !IsPublicIP(ip) {
return fmt.Errorf("host %q resolves to non-public IP %s", host, ip)
}
}
return nil
}
// ValidatePublicDomain checks that a domain does not resolve to a private,
// loopback, or link-local IP address.
func ValidatePublicDomain(domain string) error {
ips, err := net.LookupIP(domain)
if err != nil {
return fmt.Errorf("cannot resolve host %q: %w", domain, err)
}
for _, ip := range ips {
if !IsPublicIP(ip) {
return fmt.Errorf("host %q resolves to non-public IP %s", domain, ip)
}
}
return nil
}
// NewPinnedTransport returns an *http.Transport with a custom DialContext that
// resolves the target host once, validates all resolved IPs with IsPublicIP,
// and dials the validated IP directly. This prevents DNS rebinding attacks
// where the first lookup returns a public IP but a subsequent lookup (at
// connection time) returns a private IP.
func NewPinnedTransport() *http.Transport {
return &http.Transport{
DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
host, port, err := net.SplitHostPort(addr)
if err != nil {
return nil, fmt.Errorf("cannot parse address: %w", err)
}
ips, err := net.DefaultResolver.LookupIPAddr(ctx, host)
if err != nil {
return nil, fmt.Errorf("cannot resolve host: %w", err)
}
if len(ips) == 0 {
return nil, fmt.Errorf("cannot resolve host: no addresses found")
}
for _, ip := range ips {
if !IsPublicIP(ip.IP) {
return nil, fmt.Errorf("cannot connect to non-public IP %s", ip.IP)
}
}
// Dial the first validated IP directly to prevent DNS rebinding.
pinnedAddr := net.JoinHostPort(ips[0].IP.String(), port)
var d net.Dialer
return d.DialContext(ctx, network, pinnedAddr)
},
}
}

156
pkg/agent/tools/search/diff_documents.go Normal file
View File

@@ -0,0 +1,156 @@
// 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 search
import (
"context"
"fmt"
"strings"
"go.probo.inc/probo/pkg/agent"
)
type (
diffParams struct {
TextA string `json:"text_a" jsonschema:"The first document text to compare"`
TextB string `json:"text_b" jsonschema:"The second document text to compare"`
LabelA string `json:"label_a" jsonschema:"Label for the first document (e.g. 'current version')"`
LabelB string `json:"label_b" jsonschema:"Label for the second document (e.g. 'archived version')"`
}
diffResult struct {
HasDifferences bool `json:"has_differences"`
UnifiedDiff string `json:"unified_diff,omitempty"`
AddedLines int `json:"added_lines"`
RemovedLines int `json:"removed_lines"`
ErrorDetail string `json:"error_detail,omitempty"`
}
)
const (
maxDiffOutput = 16000
)
func DiffDocumentsTool() agent.Tool {
return agent.FunctionTool(
"diff_documents",
"Compare two document texts and return a unified diff showing the differences. Useful for comparing current vs. archived versions of privacy policies, terms of service, or other legal documents.",
func(ctx context.Context, p diffParams) (agent.ToolResult, error) {
labelA := p.LabelA
if labelA == "" {
labelA = "document_a"
}
labelB := p.LabelB
if labelB == "" {
labelB = "document_b"
}
linesA := strings.Split(p.TextA, "\n")
linesB := strings.Split(p.TextB, "\n")
diff := computeDiff(linesA, linesB, labelA, labelB)
if diff.tooLarge {
return agent.ResultJSON(diffResult{
HasDifferences: true,
ErrorDetail: diff.output,
}), nil
}
result := diffResult{
HasDifferences: diff.added > 0 || diff.removed > 0,
AddedLines: diff.added,
RemovedLines: diff.removed,
}
if result.HasDifferences {
output := diff.output
if len(output) > maxDiffOutput {
output = output[:maxDiffOutput] + "\n[... diff truncated]"
}
result.UnifiedDiff = output
}
return agent.ResultJSON(result), nil
},
)
}
type (
diffOutput struct {
output string
added int
removed int
tooLarge bool
}
)
func computeDiff(linesA, linesB []string, labelA, labelB string) diffOutput {
// Simple line-by-line LCS-based diff.
m, n := len(linesA), len(linesB)
// Build LCS table (bounded to prevent excessive memory for very large docs).
if m > 5000 || n > 5000 {
return diffOutput{
output: "documents too large for detailed diff (limit 5000 lines per side)",
tooLarge: true,
}
}
// LCS length table.
dp := make([][]int, m+1)
for i := range dp {
dp[i] = make([]int, n+1)
}
for i := m - 1; i >= 0; i-- {
for j := n - 1; j >= 0; j-- {
if linesA[i] == linesB[j] {
dp[i][j] = dp[i+1][j+1] + 1
} else if dp[i+1][j] >= dp[i][j+1] {
dp[i][j] = dp[i+1][j]
} else {
dp[i][j] = dp[i][j+1]
}
}
}
// Walk the LCS table to produce diff hunks.
var sb strings.Builder
fmt.Fprintf(&sb, "--- %s\n+++ %s\n", labelA, labelB)
var added, removed int
i, j := 0, 0
for i < m || j < n {
if i < m && j < n && linesA[i] == linesB[j] {
// Context line — only emit near changes.
i++
j++
} else if j < n && (i >= m || dp[i][j+1] >= dp[i+1][j]) {
sb.WriteString("+ " + linesB[j] + "\n")
added++
j++
} else if i < m {
sb.WriteString("- " + linesA[i] + "\n")
removed++
i++
}
}
return diffOutput{
output: sb.String(),
added: added,
removed: removed,
}
}

203
pkg/agent/tools/search/diff_test.go Normal file
View File

@@ -0,0 +1,203 @@
// 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 search
import (
"strings"
"testing"
"github.com/stretchr/testify/assert"
)
func TestComputeDiff(t *testing.T) {
t.Parallel()
t.Run(
"identical documents have no changes",
func(t *testing.T) {
t.Parallel()
lines := []string{"line one", "line two", "line three"}
diff := computeDiff(lines, lines, "a", "b")
assert.Equal(t, 0, diff.added)
assert.Equal(t, 0, diff.removed)
},
)
t.Run(
"completely different documents",
func(t *testing.T) {
t.Parallel()
linesA := []string{"alpha", "beta"}
linesB := []string{"gamma", "delta"}
diff := computeDiff(linesA, linesB, "a", "b")
assert.Equal(t, 2, diff.added)
assert.Equal(t, 2, diff.removed)
assert.Contains(t, diff.output, "- alpha")
assert.Contains(t, diff.output, "- beta")
assert.Contains(t, diff.output, "+ gamma")
assert.Contains(t, diff.output, "+ delta")
},
)
t.Run(
"added lines only",
func(t *testing.T) {
t.Parallel()
linesA := []string{"line one"}
linesB := []string{"line one", "line two", "line three"}
diff := computeDiff(linesA, linesB, "a", "b")
assert.Equal(t, 2, diff.added)
assert.Equal(t, 0, diff.removed)
assert.Contains(t, diff.output, "+ line two")
assert.Contains(t, diff.output, "+ line three")
},
)
t.Run(
"removed lines only",
func(t *testing.T) {
t.Parallel()
linesA := []string{"line one", "line two", "line three"}
linesB := []string{"line one"}
diff := computeDiff(linesA, linesB, "a", "b")
assert.Equal(t, 0, diff.added)
assert.Equal(t, 2, diff.removed)
assert.Contains(t, diff.output, "- line two")
assert.Contains(t, diff.output, "- line three")
},
)
t.Run(
"mixed changes",
func(t *testing.T) {
t.Parallel()
linesA := []string{"keep", "remove me", "also keep"}
linesB := []string{"keep", "add me", "also keep"}
diff := computeDiff(linesA, linesB, "a", "b")
assert.Equal(t, 1, diff.added)
assert.Equal(t, 1, diff.removed)
assert.Contains(t, diff.output, "- remove me")
assert.Contains(t, diff.output, "+ add me")
},
)
t.Run(
"both inputs empty",
func(t *testing.T) {
t.Parallel()
diff := computeDiff([]string{}, []string{}, "a", "b")
assert.Equal(t, 0, diff.added)
assert.Equal(t, 0, diff.removed)
},
)
t.Run(
"first input empty",
func(t *testing.T) {
t.Parallel()
linesB := []string{"new line"}
diff := computeDiff([]string{}, linesB, "a", "b")
assert.Equal(t, 1, diff.added)
assert.Equal(t, 0, diff.removed)
assert.Contains(t, diff.output, "+ new line")
},
)
t.Run(
"second input empty",
func(t *testing.T) {
t.Parallel()
linesA := []string{"old line"}
diff := computeDiff(linesA, []string{}, "a", "b")
assert.Equal(t, 0, diff.added)
assert.Equal(t, 1, diff.removed)
assert.Contains(t, diff.output, "- old line")
},
)
t.Run(
"single line documents identical",
func(t *testing.T) {
t.Parallel()
diff := computeDiff([]string{"same"}, []string{"same"}, "a", "b")
assert.Equal(t, 0, diff.added)
assert.Equal(t, 0, diff.removed)
},
)
t.Run(
"single line documents different",
func(t *testing.T) {
t.Parallel()
diff := computeDiff([]string{"old"}, []string{"new"}, "a", "b")
assert.Equal(t, 1, diff.added)
assert.Equal(t, 1, diff.removed)
},
)
t.Run(
"output contains labels",
func(t *testing.T) {
t.Parallel()
diff := computeDiff(
[]string{"a"},
[]string{"b"},
"current version",
"archived version",
)
assert.True(t, strings.HasPrefix(diff.output, "--- current version\n+++ archived version\n"))
},
)
t.Run(
"documents too large returns bounded message",
func(t *testing.T) {
t.Parallel()
large := make([]string, 5001)
for i := range large {
large[i] = "line"
}
diff := computeDiff(large, []string{"small"}, "a", "b")
assert.Equal(t, 0, diff.added)
assert.Equal(t, 0, diff.removed)
assert.Contains(t, diff.output, "too large")
},
)
}

164
pkg/agent/tools/search/government_db.go Normal file
View File

@@ -0,0 +1,164 @@
// 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 search
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"time"
"go.probo.inc/probo/pkg/agent"
)
type (
govDBParams struct {
CompanyName string `json:"company_name" jsonschema:"The company name to search for in government databases"`
Domain string `json:"domain" jsonschema:"The company domain for additional search context (optional)"`
}
govDBEntry struct {
Source string `json:"source"`
Title string `json:"title"`
URL string `json:"url"`
Snippet string `json:"snippet,omitempty"`
}
govDBResult struct {
SECFilings []govDBEntry `json:"sec_filings,omitempty"`
FTCActions []govDBEntry `json:"ftc_actions,omitempty"`
GDPRFines []govDBEntry `json:"gdpr_fines,omitempty"`
OtherActions []govDBEntry `json:"other_regulatory_actions,omitempty"`
ErrorDetail string `json:"error_detail,omitempty"`
}
)
func CheckGovernmentDBTool(searchEndpoint string) agent.Tool {
client := &http.Client{Timeout: 15 * time.Second}
return agent.FunctionTool(
"check_government_databases",
"Search government and regulatory databases for enforcement actions, SEC filings, FTC actions, and GDPR fines related to a company.",
func(ctx context.Context, p govDBParams) (agent.ToolResult, error) {
var result govDBResult
name := p.CompanyName
if p.Domain != "" {
name = name + " " + p.Domain
}
type searchSpec struct {
query string
source string
target *[]govDBEntry
}
searches := []searchSpec{
{
query: fmt.Sprintf(`site:sec.gov "%s"`, p.CompanyName),
source: "SEC",
target: &result.SECFilings,
},
{
query: fmt.Sprintf(`site:ftc.gov "%s"`, p.CompanyName),
source: "FTC",
target: &result.FTCActions,
},
{
query: fmt.Sprintf(`site:enforcementtracker.com "%s"`, p.CompanyName),
source: "GDPR Enforcement Tracker",
target: &result.GDPRFines,
},
{
query: fmt.Sprintf(`"%s" regulatory action OR enforcement OR fine OR penalty OR sanction`, name),
source: "General",
target: &result.OtherActions,
},
}
for _, s := range searches {
entries, err := searxngSearch(ctx, client, searchEndpoint, s.query, 3)
if err != nil {
continue
}
for _, e := range entries {
*s.target = append(*s.target, govDBEntry{
Source: s.source,
Title: e.Title,
URL: e.URL,
Snippet: e.Snippet,
})
}
}
return agent.ResultJSON(result), nil
},
)
}
func searxngSearch(ctx context.Context, client *http.Client, endpoint, query string, maxResults int) ([]searchResult, error) {
u, err := url.Parse(endpoint + "/search")
if err != nil {
return nil, err
}
q := u.Query()
q.Set("q", query)
q.Set("format", "json")
q.Set("categories", "general")
u.RawQuery = q.Encode()
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
if err != nil {
return nil, err
}
resp, err := client.Do(req)
if err != nil {
return nil, err
}
defer func() { _ = resp.Body.Close() }()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("search returned status %d", resp.StatusCode)
}
var searxResp searxngResponse
if err := json.Unmarshal(body, &searxResp); err != nil {
return nil, err
}
results := make([]searchResult, 0, maxResults)
for i, r := range searxResp.Results {
if i >= maxResults {
break
}
results = append(results, searchResult{
Title: r.Title,
URL: r.URL,
Snippet: r.Content,
})
}
return results, nil
}

30
pkg/agents/agents.go → pkg/agent/tools/search/search.go
View File

@@ -1,4 +1,4 @@
// Copyright (c) 2025-2026 Probo Inc <hello@getprobo.com>.
// 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
@@ -12,21 +12,27 @@
// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
// PERFORMANCE OF THIS SOFTWARE.
package agents
package search
import (
"go.gearno.de/kit/log"
"go.probo.inc/probo/pkg/llm"
"go.probo.inc/probo/pkg/agent"
)
type Agent struct {
l *log.Logger
client *llm.Client
model string
temp float64
maxTokens int
// Toolset provides web search tools.
type Toolset struct {
endpoint string
}
func NewAgent(l *log.Logger, client *llm.Client, model string, temp float64, maxTokens int) *Agent {
return &Agent{l: l, client: client, model: model, temp: temp, maxTokens: maxTokens}
// NewToolset creates a search toolset with the given SearXNG endpoint.
func NewToolset(endpoint string) *Toolset {
return &Toolset{endpoint: endpoint}
}
func (t *Toolset) Tools() []agent.Tool {
return []agent.Tool{
WebSearchTool(t.endpoint),
CheckGovernmentDBTool(t.endpoint),
CheckWaybackTool(),
DiffDocumentsTool(),
}
}

147
pkg/agent/tools/search/wayback.go Normal file
View File

@@ -0,0 +1,147 @@
// 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 search
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"time"
"go.probo.inc/probo/pkg/agent"
)
type (
waybackParams struct {
URL string `json:"url" jsonschema:"The URL to check in the Wayback Machine (e.g. https://example.com/privacy)"`
}
waybackSnapshot struct {
Timestamp string `json:"timestamp"`
URL string `json:"url"`
}
waybackResult struct {
Available bool `json:"available"`
OldestSnapshot *waybackSnapshot `json:"oldest_snapshot,omitempty"`
NewestSnapshot *waybackSnapshot `json:"newest_snapshot,omitempty"`
ErrorDetail string `json:"error_detail,omitempty"`
}
waybackAvailabilityResponse struct {
ArchivedSnapshots struct {
Closest struct {
Available bool `json:"available"`
URL string `json:"url"`
Timestamp string `json:"timestamp"`
} `json:"closest"`
} `json:"archived_snapshots"`
}
waybackCDXResponse = [][]string
)
func CheckWaybackTool() agent.Tool {
client := &http.Client{Timeout: 15 * time.Second}
return agent.FunctionTool(
"check_wayback",
"Check the Internet Archive Wayback Machine for archived versions of a URL. Useful for detecting changes in privacy policies, trust pages, or terms of service over time.",
func(ctx context.Context, p waybackParams) (agent.ToolResult, error) {
var result waybackResult
// Check availability.
availURL := "https://archive.org/wayback/available?url=" + url.QueryEscape(p.URL)
body, err := httpGet(ctx, client, availURL)
if err != nil {
result.ErrorDetail = fmt.Sprintf("cannot check Wayback Machine availability: %s", err)
return agent.ResultJSON(result), nil
}
var avail waybackAvailabilityResponse
if err := json.Unmarshal(body, &avail); err == nil {
result.Available = avail.ArchivedSnapshots.Closest.Available
}
if !result.Available {
return agent.ResultJSON(result), nil
}
// Get oldest snapshot.
oldestURL := fmt.Sprintf(
"https://web.archive.org/cdx/search/cdx?url=%s&output=json&fl=timestamp,original&limit=1",
url.QueryEscape(p.URL),
)
if body, err := httpGet(ctx, client, oldestURL); err == nil {
if snap := parseCDXSnapshot(body); snap != nil {
result.OldestSnapshot = snap
}
}
// Get newest snapshot.
newestURL := fmt.Sprintf(
"https://web.archive.org/cdx/search/cdx?url=%s&output=json&fl=timestamp,original&limit=1&sort=reverse",
url.QueryEscape(p.URL),
)
if body, err := httpGet(ctx, client, newestURL); err == nil {
if snap := parseCDXSnapshot(body); snap != nil {
result.NewestSnapshot = snap
}
}
return agent.ResultJSON(result), nil
},
)
}
func httpGet(ctx context.Context, client *http.Client, rawURL string) ([]byte, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
if err != nil {
return nil, err
}
resp, err := client.Do(req)
if err != nil {
return nil, err
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("status %d", resp.StatusCode)
}
return io.ReadAll(io.LimitReader(resp.Body, 1*1024*1024))
}
func parseCDXSnapshot(body []byte) *waybackSnapshot {
var rows waybackCDXResponse
if err := json.Unmarshal(body, &rows); err != nil || len(rows) < 2 {
return nil
}
// First row is headers ["timestamp", "original"], data starts at row 1.
row := rows[1]
if len(row) < 2 {
return nil
}
return &waybackSnapshot{
Timestamp: row[0],
URL: row[1],
}
}

109
pkg/agent/tools/search/wayback_test.go Normal file
View File

@@ -0,0 +1,109 @@
// 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 search
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestParseCDXSnapshot(t *testing.T) {
t.Parallel()
t.Run(
"valid JSON array response",
func(t *testing.T) {
t.Parallel()
body := []byte(`[["timestamp","original"],["20200115120000","https://example.com/privacy"]]`)
snap := parseCDXSnapshot(body)
require.NotNil(t, snap)
assert.Equal(t, "20200115120000", snap.Timestamp)
assert.Equal(t, "https://example.com/privacy", snap.URL)
},
)
t.Run(
"empty array returns nil",
func(t *testing.T) {
t.Parallel()
body := []byte(`[]`)
assert.Nil(t, parseCDXSnapshot(body))
},
)
t.Run(
"single row header only returns nil",
func(t *testing.T) {
t.Parallel()
body := []byte(`[["timestamp","original"]]`)
assert.Nil(t, parseCDXSnapshot(body))
},
)
t.Run(
"malformed JSON returns nil",
func(t *testing.T) {
t.Parallel()
body := []byte(`not valid json`)
assert.Nil(t, parseCDXSnapshot(body))
},
)
t.Run(
"data row with insufficient fields returns nil",
func(t *testing.T) {
t.Parallel()
body := []byte(`[["timestamp","original"],["20200115120000"]]`)
assert.Nil(t, parseCDXSnapshot(body))
},
)
t.Run(
"empty body returns nil",
func(t *testing.T) {
t.Parallel()
assert.Nil(t, parseCDXSnapshot([]byte{}))
},
)
t.Run(
"response with extra fields uses first two",
func(t *testing.T) {
t.Parallel()
body := []byte(`[["timestamp","original","extra"],["20210601000000","https://example.com/tos","200"]]`)
snap := parseCDXSnapshot(body)
require.NotNil(t, snap)
assert.Equal(t, "20210601000000", snap.Timestamp)
assert.Equal(t, "https://example.com/tos", snap.URL)
},
)
}

74
pkg/agent/tools/search/web_search.go Normal file
View File

@@ -0,0 +1,74 @@
// 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 search
import (
"context"
"net/http"
"time"
"go.probo.inc/probo/pkg/agent"
)
type (
searchParams struct {
Query string `json:"query" jsonschema:"The search query to execute"`
MaxResults int `json:"max_results" jsonschema:"Maximum number of results to return (default 5, max 10)"`
}
searchResult struct {
Title string `json:"title"`
URL string `json:"url"`
Snippet string `json:"snippet"`
}
searxngResponse struct {
Results []searxngResult `json:"results"`
}
searxngResult struct {
Title string `json:"title"`
URL string `json:"url"`
Content string `json:"content"`
}
)
// WebSearchTool creates a tool that searches the web using a SearXNG instance.
// The endpoint should be the base URL of the SearXNG instance (e.g.
// "http://localhost:8888").
func WebSearchTool(endpoint string) agent.Tool {
client := &http.Client{Timeout: 15 * time.Second}
return agent.FunctionTool(
"web_search",
"Search the web for information about a topic. Returns a list of results with title, URL, and snippet. Use this to find news, reviews, breach reports, regulatory actions, and other external information about a vendor.",
func(ctx context.Context, p searchParams) (agent.ToolResult, error) {
maxResults := p.MaxResults
if maxResults <= 0 {
maxResults = 5
}
if maxResults > 10 {
maxResults = 10
}
results, err := searxngSearch(ctx, client, endpoint, p.Query, maxResults)
if err != nil {
return agent.ResultErrorf("search request failed: %s", err), nil
}
return agent.ResultJSON(results), nil
},
)
}

122
pkg/agent/tools/security/cors.go Normal file
View File

@@ -0,0 +1,122 @@
// 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 security
import (
"context"
"fmt"
"net/http"
"strings"
"time"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/agent/tools/internal/netcheck"
)
type (
corsParams struct {
URL string `json:"url" jsonschema:"The URL to check CORS headers for"`
Origin string `json:"origin" jsonschema:"The Origin header value to send in the preflight request (e.g. https://evil.com)"`
}
corsResult struct {
AllowOrigin string `json:"access_control_allow_origin,omitempty"`
AllowMethods []string `json:"access_control_allow_methods,omitempty"`
AllowHeaders []string `json:"access_control_allow_headers,omitempty"`
AllowCredentials bool `json:"access_control_allow_credentials"`
ExposeHeaders []string `json:"access_control_expose_headers,omitempty"`
MaxAge string `json:"access_control_max_age,omitempty"`
WildcardOrigin bool `json:"wildcard_origin"`
ReflectsOrigin bool `json:"reflects_origin"`
ErrorDetail string `json:"error_detail,omitempty"`
}
)
func splitTrimmed(s, sep string) []string {
if s == "" {
return nil
}
parts := strings.Split(s, sep)
out := make([]string, 0, len(parts))
for _, p := range parts {
p = strings.TrimSpace(p)
if p != "" {
out = append(out, p)
}
}
return out
}
func CheckCORSTool() agent.Tool {
return agent.FunctionTool(
"check_cors",
"Send a CORS preflight (OPTIONS) request to a URL with a given Origin and analyze the Access-Control-* response headers, flagging wildcard origins and origin reflection.",
func(ctx context.Context, p corsParams) (agent.ToolResult, error) {
if err := netcheck.ValidatePublicURL(p.URL); err != nil {
return agent.ResultJSON(corsResult{
ErrorDetail: fmt.Sprintf("URL not allowed: %s", err),
}), nil
}
client := &http.Client{
Timeout: 10 * time.Second,
CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
return http.ErrUseLastResponse
},
}
req, err := http.NewRequestWithContext(
ctx,
http.MethodOptions,
p.URL,
nil,
)
if err != nil {
return agent.ResultJSON(corsResult{
ErrorDetail: fmt.Sprintf("cannot build request: %s", err),
}), nil
}
req.Header.Set("Origin", p.Origin)
req.Header.Set("Access-Control-Request-Method", "GET")
resp, err := client.Do(req)
if err != nil {
return agent.ResultJSON(corsResult{
ErrorDetail: fmt.Sprintf("cannot fetch %s: %s", p.URL, err),
}), nil
}
defer func() { _ = resp.Body.Close() }()
allowOrigin := resp.Header.Get("Access-Control-Allow-Origin")
result := corsResult{
AllowOrigin: allowOrigin,
AllowMethods: splitTrimmed(resp.Header.Get("Access-Control-Allow-Methods"), ","),
AllowHeaders: splitTrimmed(resp.Header.Get("Access-Control-Allow-Headers"), ","),
AllowCredentials: strings.EqualFold(resp.Header.Get("Access-Control-Allow-Credentials"), "true"),
ExposeHeaders: splitTrimmed(resp.Header.Get("Access-Control-Expose-Headers"), ","),
MaxAge: resp.Header.Get("Access-Control-Max-Age"),
WildcardOrigin: allowOrigin == "*",
ReflectsOrigin: p.Origin != "" && allowOrigin == p.Origin,
}
return agent.ResultJSON(result), nil
},
)
}

71
pkg/agent/tools/security/cors_test.go Normal file
View File

@@ -0,0 +1,71 @@
// 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 security
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestSplitTrimmed(t *testing.T) {
t.Parallel()
t.Run(
"splits and trims values",
func(t *testing.T) {
t.Parallel()
result := splitTrimmed("GET, POST, PUT", ",")
require.Len(t, result, 3)
assert.Equal(t, "GET", result[0])
assert.Equal(t, "POST", result[1])
assert.Equal(t, "PUT", result[2])
},
)
t.Run(
"returns nil for empty string",
func(t *testing.T) {
t.Parallel()
assert.Nil(t, splitTrimmed("", ","))
},
)
t.Run(
"skips empty parts",
func(t *testing.T) {
t.Parallel()
result := splitTrimmed("GET,,POST", ",")
require.Len(t, result, 2)
assert.Equal(t, "GET", result[0])
assert.Equal(t, "POST", result[1])
},
)
t.Run(
"single value",
func(t *testing.T) {
t.Parallel()
result := splitTrimmed("GET", ",")
require.Len(t, result, 1)
assert.Equal(t, "GET", result[0])
},
)
}

140
pkg/agent/tools/security/csp.go Normal file
View File

@@ -0,0 +1,140 @@
// 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 security
import (
"context"
"fmt"
"net/http"
"strings"
"time"
"go.probo.inc/probo/pkg/agent"
)
type (
cspParams struct {
URL string `json:"url" jsonschema:"The URL to analyze the Content-Security-Policy header for"`
}
cspDirective struct {
Name string `json:"name"`
Values []string `json:"values"`
}
cspResult struct {
Present bool `json:"present"`
ReportOnly bool `json:"report_only"`
RawHeader string `json:"raw_header,omitempty"`
Directives []cspDirective `json:"directives,omitempty"`
HasUnsafeEval bool `json:"has_unsafe_eval"`
HasUnsafeInline bool `json:"has_unsafe_inline"`
HasWildcard bool `json:"has_wildcard"`
ErrorDetail string `json:"error_detail,omitempty"`
}
)
func parseCSPDirectives(raw string) []cspDirective {
var directives []cspDirective
for part := range strings.SplitSeq(raw, ";") {
part = strings.TrimSpace(part)
if part == "" {
continue
}
tokens := strings.Fields(part)
if len(tokens) == 0 {
continue
}
directives = append(
directives,
cspDirective{
Name: tokens[0],
Values: tokens[1:],
},
)
}
return directives
}
func AnalyzeCSPTool() agent.Tool {
return agent.FunctionTool(
"analyze_csp",
"Analyze the Content-Security-Policy header for a URL, parsing directives and flagging unsafe patterns like unsafe-eval, unsafe-inline, and wildcard sources.",
func(ctx context.Context, p cspParams) (agent.ToolResult, error) {
client := &http.Client{Timeout: 10 * time.Second}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, p.URL, nil)
if err != nil {
return agent.ResultJSON(cspResult{
ErrorDetail: fmt.Sprintf("cannot create request for %s: %s", p.URL, err),
}), nil
}
resp, err := client.Do(req)
if err != nil {
return agent.ResultJSON(cspResult{
ErrorDetail: fmt.Sprintf("cannot fetch %s: %s", p.URL, err),
}), nil
}
defer func() { _ = resp.Body.Close() }()
raw := resp.Header.Get("Content-Security-Policy")
reportOnly := false
if raw == "" {
raw = resp.Header.Get("Content-Security-Policy-Report-Only")
if raw != "" {
reportOnly = true
}
}
if raw == "" {
return agent.ResultJSON(cspResult{Present: false}), nil
}
directives := parseCSPDirectives(raw)
var hasUnsafeEval, hasUnsafeInline, hasWildcard bool
for _, d := range directives {
for _, v := range d.Values {
switch v {
case "'unsafe-eval'":
hasUnsafeEval = true
case "'unsafe-inline'":
hasUnsafeInline = true
case "*":
hasWildcard = true
}
}
}
result := cspResult{
Present: true,
ReportOnly: reportOnly,
RawHeader: raw,
Directives: directives,
HasUnsafeEval: hasUnsafeEval,
HasUnsafeInline: hasUnsafeInline,
HasWildcard: hasWildcard,
}
return agent.ResultJSON(result), nil
},
)
}

81
pkg/agent/tools/security/csp_test.go Normal file
View File

@@ -0,0 +1,81 @@
// 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 security
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestParseCSPDirectives(t *testing.T) {
t.Parallel()
t.Run(
"parses multiple directives",
func(t *testing.T) {
t.Parallel()
raw := "default-src 'self'; script-src 'self' https://cdn.example.com; style-src 'unsafe-inline'"
directives := parseCSPDirectives(raw)
require.Len(t, directives, 3)
assert.Equal(t, "default-src", directives[0].Name)
assert.Equal(t, []string{"'self'"}, directives[0].Values)
assert.Equal(t, "script-src", directives[1].Name)
assert.Equal(t, []string{"'self'", "https://cdn.example.com"}, directives[1].Values)
assert.Equal(t, "style-src", directives[2].Name)
assert.Equal(t, []string{"'unsafe-inline'"}, directives[2].Values)
},
)
t.Run(
"handles empty string",
func(t *testing.T) {
t.Parallel()
directives := parseCSPDirectives("")
assert.Empty(t, directives)
},
)
t.Run(
"handles directive without values",
func(t *testing.T) {
t.Parallel()
raw := "upgrade-insecure-requests"
directives := parseCSPDirectives(raw)
require.Len(t, directives, 1)
assert.Equal(t, "upgrade-insecure-requests", directives[0].Name)
assert.Empty(t, directives[0].Values)
},
)
t.Run(
"ignores trailing semicolons",
func(t *testing.T) {
t.Parallel()
raw := "default-src 'self';"
directives := parseCSPDirectives(raw)
require.Len(t, directives, 1)
assert.Equal(t, "default-src", directives[0].Name)
},
)
}

106
pkg/agent/tools/security/dmarc.go Normal file
View File

@@ -0,0 +1,106 @@
// 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 security
import (
"context"
"fmt"
"strings"
"codeberg.org/miekg/dns"
"go.probo.inc/probo/pkg/agent"
)
type (
dmarcParams struct {
Domain string `json:"domain" jsonschema:"The domain to check DMARC record for (e.g. example.com)"`
}
dmarcResult struct {
Found bool `json:"found"`
RawRecord string `json:"raw_record,omitempty"`
Policy string `json:"policy,omitempty"`
Percentage string `json:"pct,omitempty"`
RUA string `json:"rua,omitempty"`
RUF string `json:"ruf,omitempty"`
ErrorDetail string `json:"error_detail,omitempty"`
}
)
func parseDMARCTag(record, tag string) string {
for part := range strings.SplitSeq(record, ";") {
part = strings.TrimSpace(part)
if after, ok := strings.CutPrefix(part, tag+"="); ok {
return after
}
}
return ""
}
func CheckDMARCTool() agent.Tool {
return agent.FunctionTool(
"check_dmarc",
"Check the DMARC DNS record for a domain, returning the policy, percentage, and reporting addresses.",
func(ctx context.Context, p dmarcParams) (agent.ToolResult, error) {
fqdn := "_dmarc." + p.Domain
if !strings.HasSuffix(fqdn, ".") {
fqdn = fqdn + "."
}
client := dns.NewClient()
answers, err := queryDNS(
ctx,
client,
&dns.TXT{
Hdr: dns.Header{
Name: fqdn,
Class: dns.ClassINET,
},
},
)
if err != nil {
return agent.ResultJSON(dmarcResult{
Found: false,
ErrorDetail: fmt.Sprintf("cannot lookup DMARC record: %s", err),
}), nil
}
for _, answer := range answers {
txt, ok := answer.(*dns.TXT)
if !ok {
continue
}
record := strings.Join(txt.Txt, "")
if !strings.HasPrefix(record, "v=DMARC1") {
continue
}
result := dmarcResult{
Found: true,
RawRecord: record,
Policy: parseDMARCTag(record, "p"),
Percentage: parseDMARCTag(record, "pct"),
RUA: parseDMARCTag(record, "rua"),
RUF: parseDMARCTag(record, "ruf"),
}
return agent.ResultJSON(result), nil
}
return agent.ResultJSON(dmarcResult{Found: false}), nil
},
)
}

65
pkg/agent/tools/security/dmarc_test.go Normal file
View File

@@ -0,0 +1,65 @@
// 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 security
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestParseDMARCTag(t *testing.T) {
t.Parallel()
t.Run(
"extracts policy tag",
func(t *testing.T) {
t.Parallel()
record := "v=DMARC1; p=reject; rua=mailto:dmarc@example.com"
assert.Equal(t, "reject", parseDMARCTag(record, "p"))
},
)
t.Run(
"extracts rua tag",
func(t *testing.T) {
t.Parallel()
record := "v=DMARC1; p=none; rua=mailto:reports@example.com"
assert.Equal(t, "mailto:reports@example.com", parseDMARCTag(record, "rua"))
},
)
t.Run(
"returns empty string for missing tag",
func(t *testing.T) {
t.Parallel()
record := "v=DMARC1; p=quarantine"
assert.Equal(t, "", parseDMARCTag(record, "ruf"))
},
)
t.Run(
"extracts pct tag",
func(t *testing.T) {
t.Parallel()
record := "v=DMARC1; p=reject; pct=50; rua=mailto:d@example.com"
assert.Equal(t, "50", parseDMARCTag(record, "pct"))
},
)
}

166
pkg/agent/tools/security/dns_records.go Normal file
View File

@@ -0,0 +1,166 @@
// 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 security
import (
"context"
"fmt"
"strings"
"codeberg.org/miekg/dns"
"go.probo.inc/probo/pkg/agent"
)
type (
dnsRecordsParams struct {
Domain string `json:"domain" jsonschema:"The domain to query DNS records for (e.g. example.com)"`
}
dnsRecordsResult struct {
A []string `json:"a_records,omitempty"`
AAAA []string `json:"aaaa_records,omitempty"`
MX []string `json:"mx_records,omitempty"`
CNAME []string `json:"cname_records,omitempty"`
TXT []string `json:"txt_records,omitempty"`
NS []string `json:"ns_records,omitempty"`
ErrorDetail string `json:"error_detail,omitempty"`
}
queryOption func(*dns.MsgHeader)
)
func CheckDNSRecordsTool() agent.Tool {
return agent.FunctionTool(
"check_dns_records",
"Query DNS records for a domain (A, AAAA, MX, CNAME, TXT, NS). Reveals hosting provider, email provider, and additional security signals.",
func(ctx context.Context, p dnsRecordsParams) (agent.ToolResult, error) {
fqdn := p.Domain
if !strings.HasSuffix(fqdn, ".") {
fqdn = fqdn + "."
}
hdr := dns.Header{Name: fqdn, Class: dns.ClassINET}
client := dns.NewClient()
var result dnsRecordsResult
var errs []string
// A records.
if answers, err := queryDNS(ctx, client, &dns.A{Hdr: hdr}); err != nil {
errs = append(errs, fmt.Sprintf("A query failed: %s", err))
} else {
for _, rr := range answers {
if a, ok := rr.(*dns.A); ok {
result.A = append(result.A, a.A.String())
}
}
}
// AAAA records.
if answers, err := queryDNS(ctx, client, &dns.AAAA{Hdr: hdr}); err != nil {
errs = append(errs, fmt.Sprintf("AAAA query failed: %s", err))
} else {
for _, rr := range answers {
if aaaa, ok := rr.(*dns.AAAA); ok {
result.AAAA = append(result.AAAA, aaaa.AAAA.String())
}
}
}
// MX records.
if answers, err := queryDNS(ctx, client, &dns.MX{Hdr: hdr}); err != nil {
errs = append(errs, fmt.Sprintf("MX query failed: %s", err))
} else {
for _, rr := range answers {
if mx, ok := rr.(*dns.MX); ok {
result.MX = append(result.MX, fmt.Sprintf("%d %s", mx.Preference, strings.TrimSuffix(mx.Mx, ".")))
}
}
}
// CNAME records.
if answers, err := queryDNS(ctx, client, &dns.CNAME{Hdr: hdr}); err != nil {
errs = append(errs, fmt.Sprintf("CNAME query failed: %s", err))
} else {
for _, rr := range answers {
if cname, ok := rr.(*dns.CNAME); ok {
result.CNAME = append(result.CNAME, strings.TrimSuffix(cname.Target, "."))
}
}
}
// TXT records.
if answers, err := queryDNS(ctx, client, &dns.TXT{Hdr: hdr}); err != nil {
errs = append(errs, fmt.Sprintf("TXT query failed: %s", err))
} else {
for _, rr := range answers {
if txt, ok := rr.(*dns.TXT); ok {
result.TXT = append(result.TXT, strings.Join(txt.Txt, ""))
}
}
}
// NS records.
if answers, err := queryDNS(ctx, client, &dns.NS{Hdr: hdr}); err != nil {
errs = append(errs, fmt.Sprintf("NS query failed: %s", err))
} else {
for _, rr := range answers {
if ns, ok := rr.(*dns.NS); ok {
result.NS = append(result.NS, strings.TrimSuffix(ns.Ns, "."))
}
}
}
if len(errs) > 0 {
result.ErrorDetail = strings.Join(errs, "; ")
}
return agent.ResultJSON(result), nil
},
)
}
func withDNSSEC() queryOption {
return func(h *dns.MsgHeader) {
h.UDPSize = 4096
h.Security = true
}
}
func queryDNS(ctx context.Context, client *dns.Client, question dns.RR, opts ...queryOption) ([]dns.RR, error) {
msg := &dns.Msg{
MsgHeader: dns.MsgHeader{
ID: dns.ID(),
RecursionDesired: true,
},
}
for _, opt := range opts {
opt(&msg.MsgHeader)
}
msg.Question = []dns.RR{question}
resp, _, err := client.Exchange(ctx, msg, "udp", defaultResolverAddr)
if err == nil && resp.Truncated {
resp, _, err = client.Exchange(ctx, msg, "tcp", defaultResolverAddr)
}
if err != nil {
return nil, err
}
if resp.Rcode != dns.RcodeSuccess {
return nil, fmt.Errorf("cannot execute DNS query: %s", dns.RcodeToString[resp.Rcode])
}
return resp.Answer, nil
}

101
pkg/agent/tools/security/dnssec.go Normal file
View File

@@ -0,0 +1,101 @@
// 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 security
import (
"context"
"fmt"
"strings"
"codeberg.org/miekg/dns"
"go.probo.inc/probo/pkg/agent"
)
type (
dnssecParams struct {
Domain string `json:"domain" jsonschema:"The domain to check DNSSEC for (e.g. example.com)"`
}
dnssecResult struct {
Enabled bool `json:"enabled"`
HasDNSKEY bool `json:"has_dnskey"`
KeyCount int `json:"key_count,omitempty"`
Details string `json:"details,omitempty"`
ErrorDetail string `json:"error_detail,omitempty"`
}
)
func CheckDNSSECTool() agent.Tool {
return agent.FunctionTool(
"check_dnssec",
"Check if DNSSEC is enabled for a domain by looking up DNSKEY records.",
func(ctx context.Context, p dnssecParams) (agent.ToolResult, error) {
fqdn := p.Domain
if !strings.HasSuffix(fqdn, ".") {
fqdn = fqdn + "."
}
client := dns.NewClient()
answers, err := queryDNS(
ctx,
client,
&dns.DNSKEY{
Hdr: dns.Header{
Name: fqdn,
Class: dns.ClassINET,
},
},
withDNSSEC(),
)
if err != nil {
return agent.ResultJSON(dnssecResult{
Enabled: false,
ErrorDetail: fmt.Sprintf("cannot query DNSKEY records: %s", err),
}), nil
}
var keyCount int
var keyDetails []string
for _, answer := range answers {
if key, ok := answer.(*dns.DNSKEY); ok {
keyCount++
flags := "ZSK"
// SEP (Secure Entry Point) flag is bit 15 (value 1)
if key.Flags&0x0001 != 0 {
flags = "KSK"
}
keyDetails = append(
keyDetails,
fmt.Sprintf("%s (algorithm=%d, flags=%d)", flags, key.Algorithm, key.Flags),
)
}
}
hasDNSKEY := keyCount > 0
result := dnssecResult{
Enabled: hasDNSKEY,
HasDNSKEY: hasDNSKEY,
KeyCount: keyCount,
Details: strings.Join(keyDetails, "; "),
}
if !hasDNSKEY {
result.Details = "no DNSKEY records found"
}
return agent.ResultJSON(result), nil
},
)
}

141
pkg/agent/tools/security/headers.go Normal file
View File

@@ -0,0 +1,141 @@
// 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 security
import (
"context"
"fmt"
"net/http"
"strings"
"time"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/agent/tools/internal/netcheck"
)
type (
headersParams struct {
URL string `json:"url" jsonschema:"The URL to check security headers for (e.g. https://example.com)"`
}
headerCheck struct {
Present bool `json:"present"`
Value string `json:"value,omitempty"`
}
headersResult struct {
HSTS headerCheck `json:"strict_transport_security"`
CSP headerCheck `json:"content_security_policy"`
XFrameOptions headerCheck `json:"x_frame_options"`
XContentTypeOptions headerCheck `json:"x_content_type_options"`
ReferrerPolicy headerCheck `json:"referrer_policy"`
PermissionsPolicy headerCheck `json:"permissions_policy"`
CrossOriginOpenerPolicy headerCheck `json:"cross_origin_opener_policy"`
CrossOriginEmbedderPolicy headerCheck `json:"cross_origin_embedder_policy"`
CrossOriginResourcePolicy headerCheck `json:"cross_origin_resource_policy"`
RedirectsToHTTPS bool `json:"redirects_to_https"`
ErrorDetail string `json:"error_detail,omitempty"`
}
)
func checkHeader(h http.Header, name string) headerCheck {
v := h.Get(name)
return headerCheck{
Present: v != "",
Value: v,
}
}
func headersFromResponse(resp *http.Response) headersResult {
return headersResult{
HSTS: checkHeader(resp.Header, "Strict-Transport-Security"),
CSP: checkHeader(resp.Header, "Content-Security-Policy"),
XFrameOptions: checkHeader(resp.Header, "X-Frame-Options"),
XContentTypeOptions: checkHeader(resp.Header, "X-Content-Type-Options"),
ReferrerPolicy: checkHeader(resp.Header, "Referrer-Policy"),
PermissionsPolicy: checkHeader(resp.Header, "Permissions-Policy"),
CrossOriginOpenerPolicy: checkHeader(resp.Header, "Cross-Origin-Opener-Policy"),
CrossOriginEmbedderPolicy: checkHeader(resp.Header, "Cross-Origin-Embedder-Policy"),
CrossOriginResourcePolicy: checkHeader(resp.Header, "Cross-Origin-Resource-Policy"),
}
}
func CheckSecurityHeadersTool() agent.Tool {
return agent.FunctionTool(
"check_security_headers",
"Check security-related HTTP headers for a URL (HSTS, CSP, X-Frame-Options, X-Content-Type-Options, Referrer-Policy, Permissions-Policy, Cross-Origin-*-Policy). Also checks if HTTP redirects to HTTPS.",
func(ctx context.Context, p headersParams) (agent.ToolResult, error) {
if err := netcheck.ValidatePublicURL(p.URL); err != nil {
return agent.ResultJSON(headersResult{
ErrorDetail: fmt.Sprintf("URL not allowed: %s", err),
}), nil
}
client := &http.Client{
Timeout: 10 * time.Second,
CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
return http.ErrUseLastResponse
},
}
// First check the HTTP version to detect HTTP→HTTPS redirect.
redirectsToHTTPS := false
httpURL := p.URL
if after, ok := strings.CutPrefix(httpURL, "https://"); ok {
httpURL = "http://" + after
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, httpURL, nil)
if err == nil {
httpResp, err := client.Do(httpReq)
if err == nil {
_ = httpResp.Body.Close()
if httpResp.StatusCode >= 300 && httpResp.StatusCode < 400 {
loc := httpResp.Header.Get("Location")
if strings.HasPrefix(loc, "https://") {
redirectsToHTTPS = true
}
}
}
}
// Now check the HTTPS version for the actual security headers.
httpsURL := p.URL
if after, ok := strings.CutPrefix(httpsURL, "http://"); ok {
httpsURL = "https://" + after
}
followClient := &http.Client{Timeout: 10 * time.Second}
httpsReq, err := http.NewRequestWithContext(ctx, http.MethodGet, httpsURL, nil)
if err != nil {
return agent.ResultJSON(headersResult{
ErrorDetail: fmt.Sprintf("cannot create request for %s: %s", httpsURL, err),
}), nil
}
resp, err := followClient.Do(httpsReq)
if err != nil {
return agent.ResultJSON(headersResult{
ErrorDetail: fmt.Sprintf("cannot fetch %s: %s", httpsURL, err),
}), nil
}
defer func() { _ = resp.Body.Close() }()
result := headersFromResponse(resp)
result.RedirectsToHTTPS = redirectsToHTTPS
return agent.ResultJSON(result), nil
},
)
}

197
pkg/agent/tools/security/headers_test.go Normal file
View File

@@ -0,0 +1,197 @@
// 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 security
import (
"net/http"
"testing"
"github.com/stretchr/testify/assert"
)
func TestCheckHeader(t *testing.T) {
t.Parallel()
t.Run(
"present header returns present true and value",
func(t *testing.T) {
t.Parallel()
h := http.Header{}
h.Set("X-Frame-Options", "DENY")
result := checkHeader(h, "X-Frame-Options")
assert.True(t, result.Present)
assert.Equal(t, "DENY", result.Value)
},
)
t.Run(
"missing header returns present false",
func(t *testing.T) {
t.Parallel()
h := http.Header{}
result := checkHeader(h, "X-Frame-Options")
assert.False(t, result.Present)
assert.Equal(t, "", result.Value)
},
)
t.Run(
"empty header map returns present false",
func(t *testing.T) {
t.Parallel()
result := checkHeader(http.Header{}, "Strict-Transport-Security")
assert.False(t, result.Present)
assert.Equal(t, "", result.Value)
},
)
t.Run(
"header lookup is case insensitive",
func(t *testing.T) {
t.Parallel()
h := http.Header{}
h.Set("content-security-policy", "default-src 'self'")
result := checkHeader(h, "Content-Security-Policy")
assert.True(t, result.Present)
assert.Equal(t, "default-src 'self'", result.Value)
},
)
}
func TestHeadersFromResponse(t *testing.T) {
t.Parallel()
t.Run(
"all security headers present",
func(t *testing.T) {
t.Parallel()
resp := &http.Response{
Header: http.Header{
"Strict-Transport-Security": {"max-age=31536000; includeSubDomains"},
"Content-Security-Policy": {"default-src 'self'"},
"X-Frame-Options": {"DENY"},
"X-Content-Type-Options": {"nosniff"},
"Referrer-Policy": {"strict-origin-when-cross-origin"},
"Permissions-Policy": {"camera=(), microphone=()"},
"Cross-Origin-Opener-Policy": {"same-origin"},
"Cross-Origin-Embedder-Policy": {"require-corp"},
"Cross-Origin-Resource-Policy": {"same-origin"},
},
}
result := headersFromResponse(resp)
assert.True(t, result.HSTS.Present)
assert.Equal(t, "max-age=31536000; includeSubDomains", result.HSTS.Value)
assert.True(t, result.CSP.Present)
assert.Equal(t, "default-src 'self'", result.CSP.Value)
assert.True(t, result.XFrameOptions.Present)
assert.Equal(t, "DENY", result.XFrameOptions.Value)
assert.True(t, result.XContentTypeOptions.Present)
assert.Equal(t, "nosniff", result.XContentTypeOptions.Value)
assert.True(t, result.ReferrerPolicy.Present)
assert.Equal(t, "strict-origin-when-cross-origin", result.ReferrerPolicy.Value)
assert.True(t, result.PermissionsPolicy.Present)
assert.Equal(t, "camera=(), microphone=()", result.PermissionsPolicy.Value)
assert.True(t, result.CrossOriginOpenerPolicy.Present)
assert.Equal(t, "same-origin", result.CrossOriginOpenerPolicy.Value)
assert.True(t, result.CrossOriginEmbedderPolicy.Present)
assert.Equal(t, "require-corp", result.CrossOriginEmbedderPolicy.Value)
assert.True(t, result.CrossOriginResourcePolicy.Present)
assert.Equal(t, "same-origin", result.CrossOriginResourcePolicy.Value)
},
)
t.Run(
"no security headers present",
func(t *testing.T) {
t.Parallel()
resp := &http.Response{
Header: http.Header{},
}
result := headersFromResponse(resp)
assert.False(t, result.HSTS.Present)
assert.False(t, result.CSP.Present)
assert.False(t, result.XFrameOptions.Present)
assert.False(t, result.XContentTypeOptions.Present)
assert.False(t, result.ReferrerPolicy.Present)
assert.False(t, result.PermissionsPolicy.Present)
assert.False(t, result.CrossOriginOpenerPolicy.Present)
assert.False(t, result.CrossOriginEmbedderPolicy.Present)
assert.False(t, result.CrossOriginResourcePolicy.Present)
assert.False(t, result.RedirectsToHTTPS)
},
)
t.Run(
"partial headers present",
func(t *testing.T) {
t.Parallel()
resp := &http.Response{
Header: http.Header{
"Strict-Transport-Security": {"max-age=86400"},
"X-Content-Type-Options": {"nosniff"},
},
}
result := headersFromResponse(resp)
assert.True(t, result.HSTS.Present)
assert.Equal(t, "max-age=86400", result.HSTS.Value)
assert.False(t, result.CSP.Present)
assert.False(t, result.XFrameOptions.Present)
assert.True(t, result.XContentTypeOptions.Present)
assert.Equal(t, "nosniff", result.XContentTypeOptions.Value)
assert.False(t, result.ReferrerPolicy.Present)
assert.False(t, result.PermissionsPolicy.Present)
assert.False(t, result.CrossOriginOpenerPolicy.Present)
assert.False(t, result.CrossOriginEmbedderPolicy.Present)
assert.False(t, result.CrossOriginResourcePolicy.Present)
},
)
t.Run(
"does not set redirects to https",
func(t *testing.T) {
t.Parallel()
resp := &http.Response{
Header: http.Header{
"Strict-Transport-Security": {"max-age=31536000"},
},
}
result := headersFromResponse(resp)
assert.False(t, result.RedirectsToHTTPS)
},
)
}

117
pkg/agent/tools/security/hibp.go Normal file
View File

@@ -0,0 +1,117 @@
// 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 security
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"time"
"go.probo.inc/probo/pkg/agent"
)
type (
hibpParams struct {
Domain string `json:"domain" jsonschema:"The domain to check for known data breaches (e.g. example.com)"`
}
breach struct {
Name string `json:"Name"`
BreachDate string `json:"BreachDate"`
PwnCount int `json:"PwnCount"`
DataClasses []string `json:"DataClasses"`
Description string `json:"Description"`
IsVerified bool `json:"IsVerified"`
IsSensitive bool `json:"IsSensitive"`
IsRetired bool `json:"IsRetired"`
IsSpamList bool `json:"IsSpamList"`
IsMalware bool `json:"IsMalware"`
IsSubscFree bool `json:"IsSubscriptionFree"`
IsFabricated bool `json:"IsFabricated"`
}
hibpResult struct {
Found bool `json:"found"`
Count int `json:"count"`
Breaches []breach `json:"breaches,omitempty"`
ErrorDetail string `json:"error_detail,omitempty"`
}
)
func CheckBreachesTool() agent.Tool {
return agent.FunctionTool(
"check_breaches",
"Check if a domain has been involved in known data breaches using the Have I Been Pwned API.",
func(ctx context.Context, p hibpParams) (agent.ToolResult, error) {
client := &http.Client{Timeout: 10 * time.Second}
req, err := http.NewRequestWithContext(
ctx,
http.MethodGet,
"https://haveibeenpwned.com/api/v3/breaches?domain="+url.QueryEscape(p.Domain),
nil,
)
if err != nil {
return agent.ResultJSON(hibpResult{
ErrorDetail: fmt.Sprintf("cannot create request: %s", err),
}), nil
}
req.Header.Set("User-Agent", "Probo-Vendor-Assessment")
resp, err := client.Do(req)
if err != nil {
return agent.ResultJSON(hibpResult{
ErrorDetail: fmt.Sprintf("cannot fetch breaches: %s", err),
}), nil
}
defer func() { _ = resp.Body.Close() }()
body, err := io.ReadAll(resp.Body)
if err != nil {
return agent.ResultJSON(hibpResult{
ErrorDetail: fmt.Sprintf("cannot read response: %s", err),
}), nil
}
if resp.StatusCode == http.StatusNotFound {
return agent.ResultJSON(hibpResult{Found: false, Count: 0}), nil
}
if resp.StatusCode != http.StatusOK {
return agent.ResultJSON(hibpResult{
ErrorDetail: fmt.Sprintf("HIBP API returned status %d", resp.StatusCode),
}), nil
}
var breaches []breach
if err := json.Unmarshal(body, &breaches); err != nil {
return agent.ResultJSON(hibpResult{
ErrorDetail: fmt.Sprintf("cannot parse response: %s", err),
}), nil
}
return agent.ResultJSON(hibpResult{
Found: len(breaches) > 0,
Count: len(breaches),
Breaches: breaches,
}), nil
},
)
}

51
pkg/agent/tools/security/security.go Normal file
View File

@@ -0,0 +1,51 @@
// 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 security
import (
"os"
"go.probo.inc/probo/pkg/agent"
)
var defaultResolverAddr = resolverAddr()
func resolverAddr() string {
if addr := os.Getenv("DNS_RESOLVER_ADDR"); addr != "" {
return addr
}
return "8.8.8.8:53"
}
// Toolset provides all security assessment tools.
type Toolset struct{}
// NewToolset creates a security toolset.
func NewToolset() *Toolset { return &Toolset{} }
func (t *Toolset) Tools() []agent.Tool {
return []agent.Tool{
CheckSSLCertificateTool(),
CheckSecurityHeadersTool(),
CheckDMARCTool(),
CheckSPFTool(),
CheckBreachesTool(),
CheckDNSSECTool(),
AnalyzeCSPTool(),
CheckCORSTool(),
CheckWhoisTool(),
CheckDNSRecordsTool(),
}
}

120
pkg/agent/tools/security/spf.go Normal file
View File

@@ -0,0 +1,120 @@
// 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 security
import (
"context"
"fmt"
"strings"
"codeberg.org/miekg/dns"
"go.probo.inc/probo/pkg/agent"
)
type (
spfParams struct {
Domain string `json:"domain" jsonschema:"The domain to check SPF record for (e.g. example.com)"`
}
spfResult struct {
Found bool `json:"found"`
RawRecord string `json:"raw_record,omitempty"`
Policy string `json:"policy,omitempty"`
Mechanisms string `json:"mechanisms,omitempty"`
ErrorDetail string `json:"error_detail,omitempty"`
}
)
func parseSPFPolicy(record string) string {
for part := range strings.FieldsSeq(strings.ToLower(record)) {
switch part {
case "-all":
return "fail"
case "~all":
return "softfail"
case "?all":
return "neutral"
case "+all":
return "pass"
}
}
return ""
}
func CheckSPFTool() agent.Tool {
return agent.FunctionTool(
"check_spf",
"Check the SPF (Sender Policy Framework) DNS record for a domain, returning the raw record and its policy qualifier.",
func(ctx context.Context, p spfParams) (agent.ToolResult, error) {
fqdn := p.Domain
if !strings.HasSuffix(fqdn, ".") {
fqdn = fqdn + "."
}
client := dns.NewClient()
answers, err := queryDNS(
ctx,
client,
&dns.TXT{
Hdr: dns.Header{
Name: fqdn,
Class: dns.ClassINET,
},
},
)
if err != nil {
return agent.ResultJSON(spfResult{
Found: false,
ErrorDetail: fmt.Sprintf("cannot lookup SPF record: %s", err),
}), nil
}
var spfRecords []string
for _, answer := range answers {
txt, ok := answer.(*dns.TXT)
if !ok {
continue
}
record := strings.Join(txt.Txt, "")
if !strings.HasPrefix(strings.ToLower(record), "v=spf1") {
continue
}
spfRecords = append(spfRecords, record)
}
if len(spfRecords) > 1 {
return agent.ResultJSON(spfResult{
Found: true,
ErrorDetail: fmt.Sprintf("multiple SPF records found (%d); this is an invalid configuration per RFC 7208", len(spfRecords)),
}), nil
}
if len(spfRecords) == 1 {
record := spfRecords[0]
return agent.ResultJSON(spfResult{
Found: true,
RawRecord: record,
Policy: parseSPFPolicy(record),
Mechanisms: record,
}), nil
}
return agent.ResultJSON(spfResult{Found: false}), nil
},
)
}

70
pkg/agent/tools/security/spf_test.go Normal file
View File

@@ -0,0 +1,70 @@
// 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 security
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestParseSPFPolicy(t *testing.T) {
t.Parallel()
t.Run(
"detects hard fail",
func(t *testing.T) {
t.Parallel()
assert.Equal(t, "fail", parseSPFPolicy("v=spf1 include:_spf.google.com -all"))
},
)
t.Run(
"detects soft fail",
func(t *testing.T) {
t.Parallel()
assert.Equal(t, "softfail", parseSPFPolicy("v=spf1 include:spf.example.com ~all"))
},
)
t.Run(
"detects neutral",
func(t *testing.T) {
t.Parallel()
assert.Equal(t, "neutral", parseSPFPolicy("v=spf1 ?all"))
},
)
t.Run(
"detects pass all",
func(t *testing.T) {
t.Parallel()
assert.Equal(t, "pass", parseSPFPolicy("v=spf1 +all"))
},
)
t.Run(
"returns empty for no all qualifier",
func(t *testing.T) {
t.Parallel()
assert.Equal(t, "", parseSPFPolicy("v=spf1 include:_spf.google.com"))
},
)
}

147
pkg/agent/tools/security/ssl.go Normal file
View File

@@ -0,0 +1,147 @@
// 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 security
import (
"context"
"crypto/tls"
"crypto/x509"
"fmt"
"net"
"time"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/agent/tools/internal/netcheck"
)
type (
sslParams struct {
Domain string `json:"domain" jsonschema:"The domain to check the SSL certificate for (e.g. example.com)"`
}
sslResult struct {
Valid bool `json:"valid"`
Issuer string `json:"issuer"`
Subject string `json:"subject"`
NotBefore string `json:"not_before"`
NotAfter string `json:"not_after"`
DaysLeft int `json:"days_left"`
Protocol string `json:"protocol"`
DNSNames []string `json:"dns_names"`
IsExpired bool `json:"is_expired"`
ErrorDetail string `json:"error_detail,omitempty"`
}
)
func protocolName(version uint16) string {
switch version {
case tls.VersionTLS10:
return "TLS 1.0"
case tls.VersionTLS11:
return "TLS 1.1"
case tls.VersionTLS12:
return "TLS 1.2"
case tls.VersionTLS13:
return "TLS 1.3"
default:
return fmt.Sprintf("unknown (0x%04x)", version)
}
}
func CheckSSLCertificateTool() agent.Tool {
return agent.FunctionTool(
"check_ssl_certificate",
"Check the SSL/TLS certificate for a domain, returning issuer, expiry, protocol version, and validity.",
func(ctx context.Context, p sslParams) (agent.ToolResult, error) {
if err := netcheck.ValidatePublicDomain(p.Domain); err != nil {
return agent.ResultJSON(sslResult{
Valid: false,
ErrorDetail: fmt.Sprintf("domain not allowed: %s", err),
}), nil
}
// This is a certificate inspection tool: we intentionally
// connect to servers whose certificates may be expired,
// self-signed, or otherwise invalid, because the whole
// point is to report back on the certificate state.
// InsecureSkipVerify disables the handshake's built-in
// verification; we then perform the verification manually
// below (x509.Verify) and surface the result in Valid.
// This pattern is safe here because we never send any
// credentials or confidential data over the connection.
dialer := &tls.Dialer{
NetDialer: &net.Dialer{Timeout: 10 * time.Second},
Config: &tls.Config{
InsecureSkipVerify: true, //nolint:gosec // cert inspector; verification happens manually below
ServerName: p.Domain,
},
}
netConn, err := dialer.DialContext(ctx, "tcp", p.Domain+":443")
var conn *tls.Conn
if netConn != nil {
conn = netConn.(*tls.Conn)
}
if err != nil {
return agent.ResultJSON(sslResult{
Valid: false,
ErrorDetail: err.Error(),
}), nil
}
defer func() { _ = conn.Close() }()
state := conn.ConnectionState()
if len(state.PeerCertificates) == 0 {
return agent.ResultJSON(sslResult{
Valid: false,
ErrorDetail: "no peer certificates",
}), nil
}
cert := state.PeerCertificates[0]
now := time.Now()
// Manually verify the certificate since we connected
// with InsecureSkipVerify to retrieve cert details
// even for expired/invalid certificates.
valid := now.Before(cert.NotAfter) && now.After(cert.NotBefore)
if valid {
opts := x509.VerifyOptions{
DNSName: p.Domain,
Intermediates: x509.NewCertPool(),
}
for _, ic := range state.PeerCertificates[1:] {
opts.Intermediates.AddCert(ic)
}
if _, err := cert.Verify(opts); err != nil {
valid = false
}
}
result := sslResult{
Valid: valid,
Issuer: cert.Issuer.String(),
Subject: cert.Subject.String(),
NotBefore: cert.NotBefore.Format(time.RFC3339),
NotAfter: cert.NotAfter.Format(time.RFC3339),
DaysLeft: int(time.Until(cert.NotAfter).Hours() / 24),
Protocol: protocolName(state.Version),
DNSNames: cert.DNSNames,
IsExpired: now.After(cert.NotAfter),
}
return agent.ResultJSON(result), nil
},
)
}

48
pkg/agent/tools/security/ssl_test.go Normal file
View File

@@ -0,0 +1,48 @@
// 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 security
import (
"crypto/tls"
"testing"
"github.com/stretchr/testify/assert"
)
func TestProtocolName(t *testing.T) {
t.Parallel()
t.Run(
"known protocols",
func(t *testing.T) {
t.Parallel()
assert.Equal(t, "TLS 1.0", protocolName(tls.VersionTLS10))
assert.Equal(t, "TLS 1.1", protocolName(tls.VersionTLS11))
assert.Equal(t, "TLS 1.2", protocolName(tls.VersionTLS12))
assert.Equal(t, "TLS 1.3", protocolName(tls.VersionTLS13))
},
)
t.Run(
"unknown protocol",
func(t *testing.T) {
t.Parallel()
result := protocolName(0x9999)
assert.Contains(t, result, "unknown")
},
)
}

253
pkg/agent/tools/security/whois.go Normal file
View File

@@ -0,0 +1,253 @@
// 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 security
import (
"bufio"
"context"
"fmt"
"net"
"strings"
"time"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/agent/tools/internal/netcheck"
)
type (
whoisParams struct {
Domain string `json:"domain" jsonschema:"The domain to perform a WHOIS lookup on (e.g. example.com)"`
}
whoisResult struct {
Registrar string `json:"registrar,omitempty"`
CreationDate string `json:"creation_date,omitempty"`
ExpiryDate string `json:"expiry_date,omitempty"`
UpdatedDate string `json:"updated_date,omitempty"`
RegistrantOrg string `json:"registrant_org,omitempty"`
RegistrantCC string `json:"registrant_country,omitempty"`
NameServers []string `json:"name_servers,omitempty"`
DomainAge string `json:"domain_age,omitempty"`
ErrorDetail string `json:"error_detail,omitempty"`
}
)
func CheckWhoisTool() agent.Tool {
return agent.FunctionTool(
"check_whois",
"Perform a WHOIS lookup on a domain to retrieve registration details including registrar, creation date, expiry date, registrant organization, and name servers.",
func(ctx context.Context, p whoisParams) (agent.ToolResult, error) {
if err := netcheck.ValidatePublicDomain(p.Domain); err != nil {
return agent.ResultJSON(whoisResult{
ErrorDetail: fmt.Sprintf("domain not allowed: %s", err),
}), nil
}
// Step 1: query IANA to find the referral WHOIS server.
referral, err := queryWhois(ctx, "whois.iana.org:43", p.Domain)
if err != nil {
return agent.ResultJSON(whoisResult{
ErrorDetail: fmt.Sprintf("cannot query IANA WHOIS: %s", err),
}), nil
}
whoisServer := parseWhoisField(referral, "refer")
if whoisServer == "" {
whoisServer = parseWhoisField(referral, "whois")
}
if whoisServer == "" {
// Try common TLD WHOIS servers as fallback.
parts := strings.Split(p.Domain, ".")
tld := parts[len(parts)-1]
whoisServer = "whois." + tld + ".com"
}
if !strings.Contains(whoisServer, ":") {
whoisServer = whoisServer + ":43"
}
// Validate the referral WHOIS server resolves to a public IP
// to prevent SSRF via crafted IANA responses.
whoisHost, _, _ := net.SplitHostPort(whoisServer)
if whoisHost == "" {
whoisHost = whoisServer
}
if err := netcheck.ValidatePublicDomain(whoisHost); err != nil {
return agent.ResultJSON(whoisResult{
ErrorDetail: fmt.Sprintf("WHOIS referral server not allowed: %s", err),
}), nil
}
// Step 2: query the registrar's WHOIS server.
raw, err := queryWhois(ctx, whoisServer, p.Domain)
if err != nil {
return agent.ResultJSON(whoisResult{
ErrorDetail: fmt.Sprintf("cannot query WHOIS server %s: %s", whoisServer, err),
}), nil
}
result := parseWhoisResponse(raw)
// Compute domain age from creation date.
if result.CreationDate != "" {
for _, layout := range []string{
"2006-01-02T15:04:05Z",
"2006-01-02",
"02-Jan-2006",
"2006-01-02 15:04:05",
time.RFC3339,
} {
if t, err := time.Parse(layout, result.CreationDate); err == nil {
age := time.Since(t)
years := int(age.Hours() / 24 / 365)
months := int(age.Hours()/24/30) % 12
result.DomainAge = fmt.Sprintf("%d years, %d months", years, months)
break
}
}
}
return agent.ResultJSON(result), nil
},
)
}
func queryWhois(ctx context.Context, server, domain string) (string, error) {
dialer := net.Dialer{Timeout: 10 * time.Second}
conn, err := dialer.DialContext(ctx, "tcp", server)
if err != nil {
return "", fmt.Errorf("cannot connect to %s: %w", server, err)
}
defer conn.Close()
_ = conn.SetDeadline(time.Now().Add(10 * time.Second))
_, err = fmt.Fprintf(conn, "%s\r\n", domain)
if err != nil {
return "", fmt.Errorf("cannot write to %s: %w", server, err)
}
var sb strings.Builder
scanner := bufio.NewScanner(conn)
for scanner.Scan() {
sb.WriteString(scanner.Text())
sb.WriteString("\n")
}
if err := scanner.Err(); err != nil {
return "", fmt.Errorf("cannot read from %s: %w", server, err)
}
return sb.String(), nil
}
func parseWhoisField(raw, field string) string {
field = strings.ToLower(field)
for line := range strings.SplitSeq(raw, "\n") {
line = strings.TrimSpace(line)
if line == "" || strings.HasPrefix(line, "%") || strings.HasPrefix(line, "#") {
continue
}
k, v, ok := strings.Cut(line, ":")
if !ok {
continue
}
if strings.ToLower(strings.TrimSpace(k)) == field {
return strings.TrimSpace(v)
}
}
return ""
}
var (
whoisFieldMap = map[string]string{
"registrar": "registrar",
"registrar name": "registrar",
"sponsoring registrar": "registrar",
"creation date": "creation_date",
"created": "creation_date",
"created on": "creation_date",
"registration date": "creation_date",
"domain name commencement date": "creation_date",
"registry expiry date": "expiry_date",
"registrar registration expiration date": "expiry_date",
"expiry date": "expiry_date",
"paid-till": "expiry_date",
"updated date": "updated_date",
"last updated": "updated_date",
"last modified": "updated_date",
"registrant organization": "registrant_org",
"registrant organisation": "registrant_org",
"org": "registrant_org",
"registrant country": "registrant_cc",
"registrant country/economy": "registrant_cc",
"name server": "name_server",
"nserver": "name_server",
}
)
func parseWhoisResponse(raw string) whoisResult {
var result whoisResult
for line := range strings.SplitSeq(raw, "\n") {
line = strings.TrimSpace(line)
if line == "" || strings.HasPrefix(line, "%") || strings.HasPrefix(line, "#") {
continue
}
k, v, ok := strings.Cut(line, ":")
if !ok {
continue
}
key := strings.ToLower(strings.TrimSpace(k))
val := strings.TrimSpace(v)
if val == "" {
continue
}
field, ok := whoisFieldMap[key]
if !ok {
continue
}
switch field {
case "registrar":
if result.Registrar == "" {
result.Registrar = val
}
case "creation_date":
if result.CreationDate == "" {
result.CreationDate = val
}
case "expiry_date":
if result.ExpiryDate == "" {
result.ExpiryDate = val
}
case "updated_date":
if result.UpdatedDate == "" {
result.UpdatedDate = val
}
case "registrant_org":
if result.RegistrantOrg == "" {
result.RegistrantOrg = val
}
case "registrant_cc":
if result.RegistrantCC == "" {
result.RegistrantCC = val
}
case "name_server":
result.NameServers = append(result.NameServers, strings.ToLower(val))
}
}
return result
}

271
pkg/agent/tools/security/whois_test.go Normal file
View File

@@ -0,0 +1,271 @@
// 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 security
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestParseWhoisField(t *testing.T) {
t.Parallel()
t.Run(
"extracts known field",
func(t *testing.T) {
t.Parallel()
raw := "refer: whois.verisign-grs.com\nstatus: ACTIVE\n"
assert.Equal(t, "whois.verisign-grs.com", parseWhoisField(raw, "refer"))
},
)
t.Run(
"returns first match",
func(t *testing.T) {
t.Parallel()
raw := "refer: first.example.com\nrefer: second.example.com\n"
assert.Equal(t, "first.example.com", parseWhoisField(raw, "refer"))
},
)
t.Run(
"handles missing field",
func(t *testing.T) {
t.Parallel()
raw := "status: ACTIVE\ncreated: 2020-01-01\n"
assert.Equal(t, "", parseWhoisField(raw, "refer"))
},
)
t.Run(
"handles empty input",
func(t *testing.T) {
t.Parallel()
assert.Equal(t, "", parseWhoisField("", "refer"))
},
)
t.Run(
"case insensitive field matching",
func(t *testing.T) {
t.Parallel()
raw := "Refer: whois.example.com\n"
assert.Equal(t, "whois.example.com", parseWhoisField(raw, "refer"))
},
)
t.Run(
"case insensitive field name argument",
func(t *testing.T) {
t.Parallel()
raw := "refer: whois.example.com\n"
assert.Equal(t, "whois.example.com", parseWhoisField(raw, "REFER"))
},
)
t.Run(
"skips comment lines",
func(t *testing.T) {
t.Parallel()
raw := "% This is a comment\n# Another comment\nrefer: whois.example.com\n"
assert.Equal(t, "whois.example.com", parseWhoisField(raw, "refer"))
},
)
t.Run(
"skips lines without colon",
func(t *testing.T) {
t.Parallel()
raw := "no colon here\nrefer: whois.example.com\n"
assert.Equal(t, "whois.example.com", parseWhoisField(raw, "refer"))
},
)
t.Run(
"trims whitespace around key and value",
func(t *testing.T) {
t.Parallel()
raw := " refer : whois.example.com \n"
assert.Equal(t, "whois.example.com", parseWhoisField(raw, "refer"))
},
)
}
func TestParseWhoisResponse(t *testing.T) {
t.Parallel()
t.Run(
"parses full realistic response",
func(t *testing.T) {
t.Parallel()
raw := `Domain Name: EXAMPLE.COM
Registrar: Example Registrar, Inc.
Sponsoring Registrar: Another Registrar
Creation Date: 2005-03-15T00:00:00Z
Registry Expiry Date: 2030-03-15T00:00:00Z
Updated Date: 2024-01-10T12:00:00Z
Registrant Organization: Example Corp
Registrant Country: US
Name Server: ns1.example.com
Name Server: ns2.example.com
`
result := parseWhoisResponse(raw)
assert.Equal(t, "Example Registrar, Inc.", result.Registrar)
assert.Equal(t, "2005-03-15T00:00:00Z", result.CreationDate)
assert.Equal(t, "2030-03-15T00:00:00Z", result.ExpiryDate)
assert.Equal(t, "2024-01-10T12:00:00Z", result.UpdatedDate)
assert.Equal(t, "Example Corp", result.RegistrantOrg)
assert.Equal(t, "US", result.RegistrantCC)
require.Len(t, result.NameServers, 2)
assert.Equal(t, "ns1.example.com", result.NameServers[0])
assert.Equal(t, "ns2.example.com", result.NameServers[1])
},
)
t.Run(
"uses first value for duplicate fields",
func(t *testing.T) {
t.Parallel()
raw := `Registrar: First Registrar
Registrar: Second Registrar
Creation Date: 2005-01-01
Creation Date: 2010-01-01
`
result := parseWhoisResponse(raw)
assert.Equal(t, "First Registrar", result.Registrar)
assert.Equal(t, "2005-01-01", result.CreationDate)
},
)
t.Run(
"accumulates all name servers",
func(t *testing.T) {
t.Parallel()
raw := `Name Server: NS1.EXAMPLE.COM
Name Server: NS2.EXAMPLE.COM
Name Server: NS3.EXAMPLE.COM
`
result := parseWhoisResponse(raw)
require.Len(t, result.NameServers, 3)
assert.Equal(t, "ns1.example.com", result.NameServers[0])
assert.Equal(t, "ns2.example.com", result.NameServers[1])
assert.Equal(t, "ns3.example.com", result.NameServers[2])
},
)
t.Run(
"maps alternative field names",
func(t *testing.T) {
t.Parallel()
raw := `Registrar Name: Alt Registrar
Created: 2010-06-01
Paid-Till: 2030-06-01
Last Modified: 2024-06-01
Registrant Organisation: Alt Org
nserver: ns1.alt.com
`
result := parseWhoisResponse(raw)
assert.Equal(t, "Alt Registrar", result.Registrar)
assert.Equal(t, "2010-06-01", result.CreationDate)
assert.Equal(t, "2030-06-01", result.ExpiryDate)
assert.Equal(t, "2024-06-01", result.UpdatedDate)
assert.Equal(t, "Alt Org", result.RegistrantOrg)
require.Len(t, result.NameServers, 1)
assert.Equal(t, "ns1.alt.com", result.NameServers[0])
},
)
t.Run(
"empty input returns zero value",
func(t *testing.T) {
t.Parallel()
result := parseWhoisResponse("")
assert.Equal(t, "", result.Registrar)
assert.Equal(t, "", result.CreationDate)
assert.Equal(t, "", result.ExpiryDate)
assert.Equal(t, "", result.UpdatedDate)
assert.Equal(t, "", result.RegistrantOrg)
assert.Equal(t, "", result.RegistrantCC)
assert.Nil(t, result.NameServers)
},
)
t.Run(
"skips comment and blank lines",
func(t *testing.T) {
t.Parallel()
raw := `% WHOIS server comment
# Another comment
Registrar: Good Registrar
Creation Date: 2020-01-01
`
result := parseWhoisResponse(raw)
assert.Equal(t, "Good Registrar", result.Registrar)
assert.Equal(t, "2020-01-01", result.CreationDate)
},
)
t.Run(
"skips lines with empty values",
func(t *testing.T) {
t.Parallel()
raw := `Registrar:
Registrar: Actual Registrar
`
result := parseWhoisResponse(raw)
assert.Equal(t, "Actual Registrar", result.Registrar)
},
)
t.Run(
"handles extra whitespace around keys and values",
func(t *testing.T) {
t.Parallel()
raw := " Registrar : Spaced Registrar \n Creation Date : 2023-05-01 \n"
result := parseWhoisResponse(raw)
assert.Equal(t, "Spaced Registrar", result.Registrar)
assert.Equal(t, "2023-05-01", result.CreationDate)
},
)
}

10
pkg/agent/typed_test.go
View File

@@ -422,15 +422,18 @@ func TestRunTyped(t *testing.T) {
City string `json:"city"`
}
weatherTool, err := FunctionTool[Params](
weatherTool := FunctionTool[Params](
"get_weather",
"Get weather for a city",
func(_ context.Context, p Params) (ToolResult, error) {
return ToolResult{Content: "Sunny, 22°C in " + p.City}, nil
},
)
require.NoError(t, err)
// Three responses: (1) tool call, (2) free-text summary
// that triggers promotion to the synthesis turn, (3) the
// forced structured output produced on the synthesis turn
// with ToolChoice=none + schema enforced.
provider := &typedMockProvider{
responses: []*llm.ChatCompletionResponse{
{
@@ -448,6 +451,7 @@ func TestRunTyped(t *testing.T) {
Usage: llm.Usage{InputTokens: 10, OutputTokens: 5},
FinishReason: llm.FinishReasonToolCalls,
},
typedStopResponse("Got the weather, ready to respond."),
typedStopResponse(`{"city":"Paris","weather":"Sunny, 22°C"}`),
},
}
@@ -471,7 +475,7 @@ func TestRunTyped(t *testing.T) {
require.NoError(t, err)
assert.Equal(t, "Paris", result.Output.City)
assert.Equal(t, "Sunny, 22°C", result.Output.Weather)
assert.Equal(t, 2, result.Turns)
assert.Equal(t, 3, result.Turns)
},
)
}

79
pkg/agents/changelog_generator.go
View File

@@ -1,79 +0,0 @@
// Copyright (c) 2025-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 agents
import (
"context"
"fmt"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/llm"
)
const (
changelogGeneratorSystemPrompt = `
# Role:You are an assistant that creates clear and concise changelogs.
# Objective
Given two versions of a document — the "old version" and the "new version" — identify and summarize all meaningful changes between them.
Focus on additions, deletions, modifications, and restructuring.
# Response Format
Respond with ONE simple phrase that describe the changes.
# Change types
If possible use the following words with additional context to describe the change types:
"Added", "Removed", "Updated", "Reworded", "Reorganized", "Fixed", etc.
# SOP
- Be objective and neutral in tone.
- Do not comment on the quality of the change.
- Use the language of the document.
**Example output format:**
Respond ONLY with the phrase that describes the changes. No explanation, no markdown, no preamble. Like this:
Added clauses about sharing personal information with trusted partners
`
)
func (a *Agent) GenerateChangelog(ctx context.Context, oldContent string, newContent string) (*string, error) {
ag := agent.New(
"changelog_generator",
a.client,
agent.WithInstructions(changelogGeneratorSystemPrompt),
agent.WithModel(a.model),
agent.WithTemperature(a.temp),
agent.WithMaxTokens(a.maxTokens),
)
result, err := ag.Run(
ctx,
[]llm.Message{
{
Role: llm.RoleUser,
Parts: []llm.Part{
llm.TextPart{Text: fmt.Sprintf("Old content: %s", oldContent)},
llm.TextPart{Text: fmt.Sprintf("New content: %s", newContent)},
},
},
},
)
if err != nil {
return nil, fmt.Errorf("cannot generate changelog: %w", err)
}
text := result.FinalMessage().Text()
return &text, nil
}

149
pkg/agents/vendor_assessment.go
View File

@@ -1,149 +0,0 @@
// Copyright (c) 2025-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 agents
import (
"context"
"fmt"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/llm"
)
type (
vendorInfo struct {
Name string `json:"name"`
Description string `json:"description"`
Category string `json:"category"`
HeadquarterAddress string `json:"headquarter_address"`
LegalName string `json:"legal_name"`
PrivacyPolicyURL string `json:"privacy_policy_url"`
ServiceLevelAgreementURL string `json:"service_level_agreement_url"`
DataProcessingAgreementURL string `json:"data_processing_agreement_url"`
BusinessAssociateAgreementURL string `json:"business_associate_agreement_url"`
SubprocessorsListURL string `json:"subprocessors_list_url"`
SecurityPageURL string `json:"security_page_url"`
TrustPageURL string `json:"trust_page_url"`
TermsOfServiceURL string `json:"terms_of_service_url"`
StatusPageURL string `json:"status_page_url"`
Certifications []string `json:"certifications"`
}
)
const (
assessVendorSystemPrompt = `
# Role: You are a compliance assistant.
# Objective
Your task is to fetch the provided company URL and to return comprehensive company information.
# For the company url, return the following fields in structured JSON format:
- name: The company's commonly used name
- description: One-sentence summary of the company's core offering
- headquarter_address: Company's main headquarter full address
- legal_name: Official registered company name
- privacy_policy_url: URL to privacy policy page
- service_level_agreement_url: URL to SLA page
- data_processing_agreement_url: URL to DPA page
- business_associate_agreement_url: URL to BAA page
- subprocessors_list_url: URL to subprocessors/subcontractors list page
- security_page_url: URL to security information page
- trust_page_url: URL to trust/compliance page
- terms_of_service_url: URL to terms of service page
- status_page_url: URL to system status page
- certifications: Array of security/compliance certifications (e.g., ["SOC2", "ISO27001"])
- category: One of the following enum values:
- "ANALYTICS"
- "CLOUD_MONITORING"
- "CLOUD_PROVIDER"
- "COLLABORATION"
- "CUSTOMER_SUPPORT"
- "DATA_STORAGE_AND_PROCESSING"
- "DOCUMENT_MANAGEMENT"
- "EMPLOYEE_MANAGEMENT"
- "ENGINEERING"
- "FINANCE"
- "IDENTITY_PROVIDER"
- "IT"
- "MARKETING"
- "OFFICE_OPERATIONS"
- "OTHER"
- "PASSWORD_MANAGEMENT"
- "PRODUCT_AND_DESIGN"
- "PROFESSIONAL_SERVICES"
- "RECRUITING"
- "SALES"
- "SECURITY"
- "VERSION_CONTROL"
# SOP
- Please ensure the output is clean, standardized JSON.
- Use web search to gather info, if you cannot find what you are looking for, just return an empty string instead
- For URLs, return the full URL if found, otherwise an empty string
- For certifications, return an empty array if none found
- For category, if you cannot determine the category, use "OTHER"
# **Example output format:**
Respond ONLY with a JSON object. No explanation, no markdown, no preamble. Like this:
{
"name": "Stripe",
"description": "Online payment processing platform that enables businesses to accept and manage digital payments, supporting various payment methods and currencies with integrated fraud protection and compliance features",
"headquarter_address": "San Francisco, CA",
"legal_name": "Stripe, Inc.",
"privacy_policy_url": "https://stripe.com/privacy",
"service_level_agreement_url": "https://stripe.com/sla",
"data_processing_agreement_url": "https://stripe.com/dpa",
"business_associate_agreement_url": "https://stripe.com/baa",
"subprocessors_list_url": "https://stripe.com/subprocessors",
"security_page_url": "https://stripe.com/security",
"trust_page_url": "https://stripe.com/trust",
"terms_of_service_url": "https://stripe.com/terms",
"status_page_url": "https://status.stripe.com",
"business_associate_agreement_url": "https://stripe.com/baa",
"subprocessors_list_url": "https://stripe.com/subprocessors",
"certifications": ["SOC1", "SOC2", "PCI DSS Level 1", "ISO 27001"]
"category": "FINANCE"
}
### Company url:
`
)
func (a *Agent) AssessVendor(ctx context.Context, websiteURL string) (*vendorInfo, error) {
ag := agent.New(
"vendor_assessor",
a.client,
agent.WithInstructions(assessVendorSystemPrompt),
agent.WithModel(a.model),
agent.WithTemperature(a.temp),
agent.WithMaxTokens(a.maxTokens),
)
typedResult, err := agent.RunTyped[vendorInfo](
ctx,
ag,
[]llm.Message{
{
Role: llm.RoleUser,
Parts: []llm.Part{llm.TextPart{Text: websiteURL}},
},
},
)
if err != nil {
return nil, fmt.Errorf("cannot assess vendor: %w", err)
}
return &typedResult.Output, nil
}

39
pkg/bootstrap/builder.go
View File

@@ -174,7 +174,7 @@ func (b *Builder) Build() (*probod.FullConfig, error) {
CacheTTL: b.getEnvIntOrDefault("WEBHOOK_CACHE_TTL", 86400),
},
},
LLM: probod.LLMSettings{
Agents: probod.AgentsConfig{
Providers: map[string]probod.LLMProviderConfig{
"openai": {
Type: "openai",
@@ -185,27 +185,24 @@ func (b *Builder) Build() (*probod.FullConfig, error) {
APIKey: b.getEnv("ANTHROPIC_API_KEY"),
},
},
Defaults: probod.LLMConfig{
Provider: b.getEnvOrDefault("LLM_DEFAULT_PROVIDER", "openai"),
ModelName: b.getEnvOrDefault("LLM_DEFAULT_MODEL_NAME", "gpt-4o"),
Temperature: new(b.getEnvFloatOrDefault("LLM_DEFAULT_TEMPERATURE", 0.1)),
MaxTokens: new(b.getEnvIntOrDefault("LLM_DEFAULT_MAX_TOKENS", 4096)),
Default: probod.LLMAgentConfig{
Provider: b.getEnvOrDefault("AGENT_DEFAULT_PROVIDER", "openai"),
ModelName: b.getEnvOrDefault("AGENT_DEFAULT_MODEL_NAME", "gpt-4o"),
Temperature: new(b.getEnvFloatOrDefault("AGENT_DEFAULT_TEMPERATURE", 0.1)),
MaxTokens: new(b.getEnvIntOrDefault("AGENT_DEFAULT_MAX_TOKENS", 4096)),
},
Probo: probod.LLMAgentConfig{
Provider: b.getEnvOrDefault("AGENT_PROBO_PROVIDER", ""),
ModelName: b.getEnvOrDefault("AGENT_PROBO_MODEL_NAME", ""),
Temperature: b.getEnvFloatPtr("AGENT_PROBO_TEMPERATURE"),
MaxTokens: b.getEnvIntPtr("AGENT_PROBO_MAX_TOKENS"),
},
EvidenceDescriber: probod.LLMAgentConfig{
Provider: b.getEnvOrDefault("AGENT_EVIDENCE_DESCRIBER_PROVIDER", ""),
ModelName: b.getEnvOrDefault("AGENT_EVIDENCE_DESCRIBER_MODEL_NAME", ""),
Temperature: b.getEnvFloatPtr("AGENT_EVIDENCE_DESCRIBER_TEMPERATURE"),
MaxTokens: b.getEnvIntPtr("AGENT_EVIDENCE_DESCRIBER_MAX_TOKENS"),
},
},
ProboAgent: probod.LLMConfig{
Provider: b.getEnvOrDefault("PROBO_AGENT_PROVIDER", ""),
ModelName: b.getEnvOrDefault("PROBO_AGENT_MODEL_NAME", ""),
Temperature: b.getEnvFloatPtr("PROBO_AGENT_TEMPERATURE"),
MaxTokens: b.getEnvIntPtr("PROBO_AGENT_MAX_TOKENS"),
},
EvidenceDescriber: probod.EvidenceDescriberConfig{
Interval: b.getEnvIntOrDefault("EVIDENCE_DESCRIBER_INTERVAL", 10),
StaleAfter: b.getEnvIntOrDefault("EVIDENCE_DESCRIBER_STALE_AFTER", 300),
MaxConcurrency: b.getEnvIntOrDefault("EVIDENCE_DESCRIBER_MAX_CONCURRENCY", 10),
Provider: b.getEnvOrDefault("EVIDENCE_DESCRIBER_PROVIDER", ""),
ModelName: b.getEnvOrDefault("EVIDENCE_DESCRIBER_MODEL_NAME", ""),
Temperature: b.getEnvFloatPtr("EVIDENCE_DESCRIBER_TEMPERATURE"),
MaxTokens: b.getEnvIntPtr("EVIDENCE_DESCRIBER_MAX_TOKENS"),
},
CustomDomains: probod.CustomDomainsConfig{
RenewalInterval: b.getEnvIntOrDefault("CUSTOM_DOMAINS_RENEWAL_INTERVAL", 3600),

98
pkg/bootstrap/builder_test.go
View File

@@ -174,25 +174,20 @@ func TestBuilder_Build_Defaults(t *testing.T) {
assert.Equal(t, 5, cfg.Probod.Notifications.Webhook.SenderInterval)
assert.Equal(t, 86400, cfg.Probod.Notifications.Webhook.CacheTTL)
// LLM config — defaults
assert.Equal(t, "openai", cfg.Probod.LLM.Defaults.Provider)
assert.Equal(t, "gpt-4o", cfg.Probod.LLM.Defaults.ModelName)
assert.Equal(t, new(0.1), cfg.Probod.LLM.Defaults.Temperature)
assert.Equal(t, new(4096), cfg.Probod.LLM.Defaults.MaxTokens)
// Probo agent — empty (inherits from defaults)
assert.Empty(t, cfg.Probod.ProboAgent.Provider)
assert.Empty(t, cfg.Probod.ProboAgent.ModelName)
assert.Nil(t, cfg.Probod.ProboAgent.Temperature)
assert.Nil(t, cfg.Probod.ProboAgent.MaxTokens)
// Evidence describer — LLM fields empty (inherits from defaults)
assert.Empty(t, cfg.Probod.EvidenceDescriber.Provider)
assert.Empty(t, cfg.Probod.EvidenceDescriber.ModelName)
assert.Nil(t, cfg.Probod.EvidenceDescriber.Temperature)
assert.Nil(t, cfg.Probod.EvidenceDescriber.MaxTokens)
// Evidence describer — worker defaults
assert.Equal(t, 10, cfg.Probod.EvidenceDescriber.Interval)
assert.Equal(t, 300, cfg.Probod.EvidenceDescriber.StaleAfter)
assert.Equal(t, 10, cfg.Probod.EvidenceDescriber.MaxConcurrency)
// Agents config — default
assert.Equal(t, "openai", cfg.Probod.Agents.Default.Provider)
assert.Equal(t, "gpt-4o", cfg.Probod.Agents.Default.ModelName)
assert.Equal(t, new(0.1), cfg.Probod.Agents.Default.Temperature)
assert.Equal(t, new(4096), cfg.Probod.Agents.Default.MaxTokens)
// Agents config — per-agent overrides are empty (inherit from default)
assert.Empty(t, cfg.Probod.Agents.Probo.Provider)
assert.Empty(t, cfg.Probod.Agents.Probo.ModelName)
assert.Nil(t, cfg.Probod.Agents.Probo.Temperature)
assert.Nil(t, cfg.Probod.Agents.Probo.MaxTokens)
assert.Empty(t, cfg.Probod.Agents.EvidenceDescriber.Provider)
assert.Empty(t, cfg.Probod.Agents.EvidenceDescriber.ModelName)
assert.Nil(t, cfg.Probod.Agents.EvidenceDescriber.Temperature)
assert.Nil(t, cfg.Probod.Agents.EvidenceDescriber.MaxTokens)
// Custom domains config
assert.Equal(t, 3600, cfg.Probod.CustomDomains.RenewalInterval)
@@ -265,22 +260,19 @@ func TestBuilder_Build_CustomValues(t *testing.T) {
env["WEBHOOK_SENDER_INTERVAL"] = "10"
env["WEBHOOK_CACHE_TTL"] = "3600"
env["CONNECTOR_SLACK_SIGNING_SECRET"] = "slack-signing-secret"
// LLM — providers
// Agents — providers
env["OPENAI_API_KEY"] = "sk-test-key"
env["ANTHROPIC_API_KEY"] = "sk-ant-test-key"
// LLM — defaults
env["LLM_DEFAULT_PROVIDER"] = "openai"
env["LLM_DEFAULT_MODEL_NAME"] = "gpt-4-turbo"
env["LLM_DEFAULT_TEMPERATURE"] = "0.5"
env["LLM_DEFAULT_MAX_TOKENS"] = "8192"
// Evidence describer
env["EVIDENCE_DESCRIBER_PROVIDER"] = "anthropic"
env["EVIDENCE_DESCRIBER_MODEL_NAME"] = "claude-sonnet-4-20250514"
env["EVIDENCE_DESCRIBER_TEMPERATURE"] = "0.2"
env["EVIDENCE_DESCRIBER_MAX_TOKENS"] = "4096"
env["EVIDENCE_DESCRIBER_INTERVAL"] = "15"
env["EVIDENCE_DESCRIBER_STALE_AFTER"] = "600"
env["EVIDENCE_DESCRIBER_MAX_CONCURRENCY"] = "20"
// Agents — default
env["AGENT_DEFAULT_PROVIDER"] = "openai"
env["AGENT_DEFAULT_MODEL_NAME"] = "gpt-4-turbo"
env["AGENT_DEFAULT_TEMPERATURE"] = "0.5"
env["AGENT_DEFAULT_MAX_TOKENS"] = "8192"
// Agents — evidence-describer override
env["AGENT_EVIDENCE_DESCRIBER_PROVIDER"] = "anthropic"
env["AGENT_EVIDENCE_DESCRIBER_MODEL_NAME"] = "claude-sonnet-4-20250514"
env["AGENT_EVIDENCE_DESCRIBER_TEMPERATURE"] = "0.2"
env["AGENT_EVIDENCE_DESCRIBER_MAX_TOKENS"] = "4096"
// Custom domains
env["CUSTOM_DOMAINS_RESOLVER_ADDR"] = "1.1.1.1:53"
env["ACME_ACCOUNT_KEY"] = "-----BEGIN EC PRIVATE KEY-----\ntest\n-----END EC PRIVATE KEY-----"
@@ -345,28 +337,24 @@ func TestBuilder_Build_CustomValues(t *testing.T) {
assert.Equal(t, "slack-signing-secret", cfg.Probod.Notifications.Slack.SigningSecret)
assert.Equal(t, 10, cfg.Probod.Notifications.Webhook.SenderInterval)
assert.Equal(t, 3600, cfg.Probod.Notifications.Webhook.CacheTTL)
// LLM — providers
assert.Equal(t, "openai", cfg.Probod.LLM.Providers["openai"].Type)
assert.Equal(t, "sk-test-key", cfg.Probod.LLM.Providers["openai"].APIKey)
assert.Equal(t, "anthropic", cfg.Probod.LLM.Providers["anthropic"].Type)
assert.Equal(t, "sk-ant-test-key", cfg.Probod.LLM.Providers["anthropic"].APIKey)
// LLM — defaults
assert.Equal(t, "openai", cfg.Probod.LLM.Defaults.Provider)
assert.Equal(t, "gpt-4-turbo", cfg.Probod.LLM.Defaults.ModelName)
assert.Equal(t, new(0.5), cfg.Probod.LLM.Defaults.Temperature)
assert.Equal(t, new(8192), cfg.Probod.LLM.Defaults.MaxTokens)
// Probo agent — inherits defaults (no overrides set)
assert.Empty(t, cfg.Probod.ProboAgent.Provider)
assert.Empty(t, cfg.Probod.ProboAgent.ModelName)
// Evidence describer — LLM overrides
assert.Equal(t, "anthropic", cfg.Probod.EvidenceDescriber.Provider)
assert.Equal(t, "claude-sonnet-4-20250514", cfg.Probod.EvidenceDescriber.ModelName)
assert.Equal(t, new(0.2), cfg.Probod.EvidenceDescriber.Temperature)
assert.Equal(t, new(4096), cfg.Probod.EvidenceDescriber.MaxTokens)
// Evidence describer — worker config
assert.Equal(t, 15, cfg.Probod.EvidenceDescriber.Interval)
assert.Equal(t, 600, cfg.Probod.EvidenceDescriber.StaleAfter)
assert.Equal(t, 20, cfg.Probod.EvidenceDescriber.MaxConcurrency)
// Agents — providers
assert.Equal(t, "openai", cfg.Probod.Agents.Providers["openai"].Type)
assert.Equal(t, "sk-test-key", cfg.Probod.Agents.Providers["openai"].APIKey)
assert.Equal(t, "anthropic", cfg.Probod.Agents.Providers["anthropic"].Type)
assert.Equal(t, "sk-ant-test-key", cfg.Probod.Agents.Providers["anthropic"].APIKey)
// Agents — default
assert.Equal(t, "openai", cfg.Probod.Agents.Default.Provider)
assert.Equal(t, "gpt-4-turbo", cfg.Probod.Agents.Default.ModelName)
assert.Equal(t, new(0.5), cfg.Probod.Agents.Default.Temperature)
assert.Equal(t, new(8192), cfg.Probod.Agents.Default.MaxTokens)
// Agents — probo inherits default (no overrides set)
assert.Empty(t, cfg.Probod.Agents.Probo.Provider)
assert.Empty(t, cfg.Probod.Agents.Probo.ModelName)
// Agents — evidence-describer overrides
assert.Equal(t, "anthropic", cfg.Probod.Agents.EvidenceDescriber.Provider)
assert.Equal(t, "claude-sonnet-4-20250514", cfg.Probod.Agents.EvidenceDescriber.ModelName)
assert.Equal(t, new(0.2), cfg.Probod.Agents.EvidenceDescriber.Temperature)
assert.Equal(t, new(4096), cfg.Probod.Agents.EvidenceDescriber.MaxTokens)
// Custom domains
assert.Equal(t, "1.1.1.1:53", cfg.Probod.CustomDomains.ResolverAddr)
assert.Equal(t, "-----BEGIN EC PRIVATE KEY-----\ntest\n-----END EC PRIVATE KEY-----", cfg.Probod.CustomDomains.ACME.AccountKey)

151
pkg/cmd/vendormgmt/assess/assess.go Normal file
View File

@@ -0,0 +1,151 @@
// 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 assess
import (
"encoding/json"
"fmt"
"os"
"time"
"github.com/spf13/cobra"
"go.probo.inc/probo/pkg/cli/api"
"go.probo.inc/probo/pkg/cmd/cmdutil"
)
const assessMutation = `
mutation($input: AssessVendorInput!) {
assessVendor(input: $input) {
report
subprocessors {
name
country
purpose
}
vendor {
id
name
}
}
}
`
type assessResponse struct {
AssessVendor struct {
Report string `json:"report"`
Subprocessors []struct {
Name string `json:"name"`
Country string `json:"country"`
Purpose string `json:"purpose"`
} `json:"subprocessors"`
Vendor struct {
ID string `json:"id"`
Name string `json:"name"`
} `json:"vendor"`
} `json:"assessVendor"`
}
func NewCmdAssess(f *cmdutil.Factory) *cobra.Command {
var (
flagOutput *string
)
cmd := &cobra.Command{
Use: "assess <vendor-id> --url <website-url>",
Short: "Run AI assessment on a vendor from its website",
Long: "Analyze a vendor's website using AI agents to extract security, compliance, and business information.",
Example: ` # Assess a vendor by website URL
prb vendor assess VND_123 --url https://example.com
# Assess with a custom procedure file
prb vendor assess VND_123 --url https://example.com --procedure-file ./my-procedure.txt
# Output as JSON
prb vendor assess VND_123 --url https://example.com -o json`,
Args: cobra.ExactArgs(1),
RunE: func(cmd *cobra.Command, args []string) error {
if err := cmdutil.ValidateOutputFlag(flagOutput); err != nil {
return err
}
cfg, err := f.Config()
if err != nil {
return err
}
host, hc, err := cfg.DefaultHost()
if err != nil {
return err
}
flagURL, _ := cmd.Flags().GetString("url")
flagProcedureFile, _ := cmd.Flags().GetString("procedure-file")
input := map[string]any{
"id": args[0],
"websiteUrl": flagURL,
}
if flagProcedureFile != "" {
data, err := os.ReadFile(flagProcedureFile)
if err != nil {
return fmt.Errorf("cannot read procedure file: %w", err)
}
input["procedure"] = string(data)
}
// The CLI timeout must outlast the server-side assessment
// timeout (vetting.AssessmentTimeout = 20m) plus HTTP overhead.
client := api.NewClient(
host,
hc.Token,
"/api/console/v1/graphql",
22*time.Minute,
)
_, _ = fmt.Fprintf(f.IOStreams.ErrOut, "Assessing vendor from %s (this may take a few minutes)...\n", flagURL)
data, err := client.Do(
assessMutation,
map[string]any{
"input": input,
},
)
if err != nil {
return err
}
var resp assessResponse
if err := json.Unmarshal(data, &resp); err != nil {
return fmt.Errorf("cannot parse response: %w", err)
}
if *flagOutput == cmdutil.OutputJSON {
return cmdutil.PrintJSON(f.IOStreams.Out, resp.AssessVendor)
}
_, _ = fmt.Fprintln(f.IOStreams.Out, resp.AssessVendor.Report)
return nil
},
}
cmd.Flags().String("url", "", "Vendor website URL to assess (required)")
_ = cmd.MarkFlagRequired("url")
cmd.Flags().String("procedure-file", "", "Path to a custom assessment procedure file")
flagOutput = cmdutil.AddOutputFlag(cmd)
return cmd
}

2
pkg/cmd/vendormgmt/vendormgmt.go
View File

@@ -17,6 +17,7 @@ package vendormgmt
import (
"github.com/spf13/cobra"
"go.probo.inc/probo/pkg/cmd/cmdutil"
"go.probo.inc/probo/pkg/cmd/vendormgmt/assess"
"go.probo.inc/probo/pkg/cmd/vendormgmt/create"
"go.probo.inc/probo/pkg/cmd/vendormgmt/delete"
"go.probo.inc/probo/pkg/cmd/vendormgmt/list"
@@ -35,6 +36,7 @@ func NewCmdVendor(f *cmdutil.Factory) *cobra.Command {
cmd.AddCommand(view.NewCmdView(f))
cmd.AddCommand(update.NewCmdUpdate(f))
cmd.AddCommand(delete.NewCmdDelete(f))
cmd.AddCommand(assess.NewCmdAssess(f))
return cmd
}

74
pkg/llm/anthropic/provider.go
View File

@@ -152,6 +152,28 @@ func buildParams(req *llm.ChatCompletionRequest) (anthropic.MessageNewParams, er
if req.ToolChoice != nil {
params.ToolChoice = buildToolChoice(req.ToolChoice)
}
if req.Thinking != nil && req.Thinking.Enabled {
params.Thinking = anthropic.ThinkingConfigParamOfEnabled(int64(req.Thinking.BudgetTokens))
}
if req.ResponseFormat != nil {
switch req.ResponseFormat.Type {
case llm.ResponseFormatJSONSchema:
if req.ResponseFormat.JSONSchema == nil {
return anthropic.MessageNewParams{}, fmt.Errorf("cannot apply JSON schema output format: schema is nil")
}
var schema map[string]any
if err := json.Unmarshal(req.ResponseFormat.JSONSchema.Schema, &schema); err != nil {
return anthropic.MessageNewParams{}, fmt.Errorf("cannot unmarshal JSON schema for output format: %w", err)
}
params.OutputConfig = anthropic.OutputConfigParam{
Format: anthropic.JSONOutputFormatParam{Schema: schema},
}
case llm.ResponseFormatJSONObject:
return anthropic.MessageNewParams{}, fmt.Errorf("anthropic does not support json_object response format without a schema; use json_schema instead")
case llm.ResponseFormatText:
// default behaviour, nothing to set
}
}
return params, nil
}
@@ -194,12 +216,21 @@ func buildMessages(messages []llm.Message) []anthropic.MessageParam {
out = append(out, anthropic.NewUserMessage(blocks...))
case llm.RoleAssistant:
var blocks []anthropic.ContentBlockParamUnion
if text := msg.Text(); text != "" {
blocks = append(blocks, anthropic.NewTextBlock(text))
for _, p := range msg.Parts {
switch part := p.(type) {
case llm.ThinkingPart:
blocks = append(blocks, anthropic.NewThinkingBlock(part.Signature, part.Text))
case llm.TextPart:
if part.Text != "" {
blocks = append(blocks, anthropic.NewTextBlock(part.Text))
}
}
}
for _, tc := range msg.ToolCalls {
var input any
_ = json.Unmarshal([]byte(tc.Function.Arguments), &input)
if err := json.Unmarshal([]byte(tc.Function.Arguments), &input); err != nil || input == nil {
input = map[string]any{}
}
blocks = append(blocks, anthropic.NewToolUseBlock(tc.ID, input, tc.Function.Name))
}
out = append(out, anthropic.NewAssistantMessage(blocks...))
@@ -295,6 +326,12 @@ func mapResponse(msg *anthropic.Message) *llm.ChatCompletionResponse {
for _, block := range msg.Content {
switch block.Type {
case "thinking":
tb := block.AsThinking()
resp.Message.Parts = append(resp.Message.Parts, llm.ThinkingPart{
Text: tb.Thinking,
Signature: tb.Signature,
})
case "text":
resp.Message.Parts = append(resp.Message.Parts, llm.TextPart{Text: block.Text})
case "tool_use":
@@ -326,6 +363,15 @@ func mapStopReason(reason anthropic.StopReason) llm.FinishReason {
}
func mapError(err error) error {
// The Anthropic SDK refuses non-streaming requests client-side when
// the expected response time exceeds 10 minutes (large max_tokens or
// model-specific non-streaming token limits). It returns a plain
// fmt.Errorf, not an *anthropic.Error, so we must match on the
// message before attempting the type assertion.
if err != nil && strings.Contains(err.Error(), "streaming is required") {
return &llm.ErrStreamingRequired{Err: err}
}
var apiErr *anthropic.Error
if !errors.As(err, &apiErr) {
return err
@@ -361,7 +407,9 @@ type anthropicStream struct {
stream *ssestream.Stream[anthropic.MessageStreamEventUnion]
current llm.ChatCompletionStreamEvent
// Track tool call indices for mapping content_block_start events.
toolCallIndex int
toolCallIndex int
inToolUse bool
thinkingSignature string
}
func (s *anthropicStream) Next() bool {
@@ -396,7 +444,9 @@ func (s *anthropicStream) mapStreamEvent(event *anthropic.MessageStreamEventUnio
switch event.Type {
case "content_block_start":
cb := event.ContentBlock
if cb.Type == "tool_use" {
switch cb.Type {
case "tool_use":
s.inToolUse = true
tu := cb.AsToolUse()
return llm.ChatCompletionStreamEvent{
Delta: llm.MessageDelta{
@@ -407,6 +457,8 @@ func (s *anthropicStream) mapStreamEvent(event *anthropic.MessageStreamEventUnio
}},
},
}, true
case "thinking":
return llm.ChatCompletionStreamEvent{}, false
}
return llm.ChatCompletionStreamEvent{}, false
@@ -417,6 +469,15 @@ func (s *anthropicStream) mapStreamEvent(event *anthropic.MessageStreamEventUnio
return llm.ChatCompletionStreamEvent{
Delta: llm.MessageDelta{Content: delta.Text},
}, true
case "thinking_delta":
return llm.ChatCompletionStreamEvent{
Delta: llm.MessageDelta{Thinking: delta.Thinking},
}, true
case "signature_delta":
s.thinkingSignature = delta.Signature
return llm.ChatCompletionStreamEvent{
Delta: llm.MessageDelta{ThinkingSignature: delta.Signature},
}, true
case "input_json_delta":
return llm.ChatCompletionStreamEvent{
Delta: llm.MessageDelta{
@@ -430,8 +491,9 @@ func (s *anthropicStream) mapStreamEvent(event *anthropic.MessageStreamEventUnio
return llm.ChatCompletionStreamEvent{}, false
case "content_block_stop":
if event.ContentBlock.Type == "tool_use" {
if s.inToolUse {
s.toolCallIndex++
s.inToolUse = false
}
return llm.ChatCompletionStreamEvent{}, false

43
pkg/llm/chat.go
View File

@@ -33,6 +33,12 @@ type (
ToolChoice *ToolChoice
ParallelToolCalls *bool
ResponseFormat *ResponseFormat
Thinking *ThinkingConfig
}
ThinkingConfig struct {
Enabled bool
BudgetTokens int
}
ToolChoiceType string
@@ -97,8 +103,10 @@ type (
}
MessageDelta struct {
Content string
ToolCalls []ToolCallDelta
Content string
Thinking string
ThinkingSignature string
ToolCalls []ToolCallDelta
}
ToolCallDelta struct {
@@ -144,13 +152,15 @@ func (u Usage) Add(other Usage) Usage {
// After the stream is exhausted (Next returns false), call Response
// to get the fully assembled ChatCompletionResponse.
type StreamAccumulator struct {
stream ChatCompletionStream
current ChatCompletionStreamEvent
content strings.Builder
toolCalls map[int]*ToolCall
usage Usage
finishReason FinishReason
model string
stream ChatCompletionStream
current ChatCompletionStreamEvent
content strings.Builder
thinking strings.Builder
thinkingSignature string
toolCalls map[int]*ToolCall
usage Usage
finishReason FinishReason
model string
}
func NewStreamAccumulator(stream ChatCompletionStream) *StreamAccumulator {
@@ -194,11 +204,20 @@ func (a *StreamAccumulator) Response() *ChatCompletionResponse {
}
}
var parts []Part
if thinking := a.thinking.String(); thinking != "" {
parts = append(parts, ThinkingPart{
Text: thinking,
Signature: a.thinkingSignature,
})
}
parts = append(parts, TextPart{Text: a.content.String()})
return &ChatCompletionResponse{
Model: a.model,
Message: Message{
Role: RoleAssistant,
Parts: []Part{TextPart{Text: a.content.String()}},
Parts: parts,
ToolCalls: toolCalls,
},
Usage: a.usage,
@@ -212,6 +231,10 @@ func (a *StreamAccumulator) accumulate(event ChatCompletionStreamEvent) {
}
a.content.WriteString(event.Delta.Content)
a.thinking.WriteString(event.Delta.Thinking)
if event.Delta.ThinkingSignature != "" {
a.thinkingSignature = event.Delta.ThinkingSignature
}
for _, tcd := range event.Delta.ToolCalls {
tc, ok := a.toolCalls[tcd.Index]

13
pkg/llm/errors.go
View File

@@ -37,6 +37,13 @@ type (
ErrAuthentication struct {
Err error
}
// ErrStreamingRequired is returned by a provider when a non-streaming
// request must be retried with the streaming endpoint (e.g. Anthropic
// requires streaming for responses that may take longer than 10 minutes).
ErrStreamingRequired struct {
Err error
}
)
func (e *ErrRateLimit) Error() string {
@@ -68,3 +75,9 @@ func (e *ErrAuthentication) Error() string {
}
func (e *ErrAuthentication) Unwrap() error { return e.Err }
func (e *ErrStreamingRequired) Error() string {
return fmt.Sprintf("streaming is required: %v", e.Err)
}
func (e *ErrStreamingRequired) Unwrap() error { return e.Err }

10
pkg/llm/message.go
View File

@@ -52,3 +52,13 @@ func (m Message) Text() string {
}
return s.String()
}
func (m Message) Thinking() string {
var s strings.Builder
for _, p := range m.Parts {
if tp, ok := p.(ThinkingPart); ok {
s.WriteString(tp.Text)
}
}
return s.String()
}

21
pkg/llm/openai/provider.go
View File

@@ -166,6 +166,16 @@ func buildParams(req *llm.ChatCompletionRequest) openai.ChatCompletionNewParams
if req.ResponseFormat != nil {
params.ResponseFormat = buildResponseFormat(req.ResponseFormat)
}
if req.Thinking != nil && req.Thinking.Enabled && isReasoningModel(req.Model) {
switch {
case req.Thinking.BudgetTokens <= 1024:
params.ReasoningEffort = shared.ReasoningEffortLow
case req.Thinking.BudgetTokens <= 8192:
params.ReasoningEffort = shared.ReasoningEffortMedium
default:
params.ReasoningEffort = shared.ReasoningEffortHigh
}
}
return params
}
@@ -456,6 +466,17 @@ func mapChunkToEvent(chunk *openai.ChatCompletionChunk) llm.ChatCompletionStream
return event
}
// isReasoningModel returns true for OpenAI models that support
// reasoning_effort (o1, o3-mini, o3, and their dated variants).
func isReasoningModel(model string) bool {
for _, prefix := range []string{"o1", "o3"} {
if model == prefix || strings.HasPrefix(model, prefix+"-") {
return true
}
}
return false
}
func buildFilePart(p llm.FilePart) openai.ChatCompletionContentPartUnionParam {
switch {
case strings.HasPrefix(p.MimeType, "image/"):

12
pkg/llm/part.go
View File

@@ -32,8 +32,14 @@ type (
MimeType string // e.g. "application/pdf", "text/csv", "image/png"
Filename string
}
ThinkingPart struct {
Text string
Signature string // Anthropic thinking signature for multi-turn continuity
}
)
func (TextPart) part() {}
func (ImagePart) part() {}
func (FilePart) part() {}
func (TextPart) part() {}
func (ImagePart) part() {}
func (FilePart) part() {}
func (ThinkingPart) part() {}

41
pkg/probo/document_service.go
View File

@@ -18,6 +18,7 @@ import (
"archive/zip"
"bytes"
"context"
_ "embed"
"encoding/json"
"errors"
"fmt"
@@ -33,12 +34,14 @@ import (
"go.gearno.de/crypto/uuid"
"go.gearno.de/kit/pg"
"go.probo.inc/probo/packages/emails"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/baseurl"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/docgen"
"go.probo.inc/probo/pkg/gid"
"go.probo.inc/probo/pkg/html2pdf"
"go.probo.inc/probo/pkg/iam"
"go.probo.inc/probo/pkg/llm"
"go.probo.inc/probo/pkg/mail"
"go.probo.inc/probo/pkg/page"
"go.probo.inc/probo/pkg/pdfutils"
@@ -458,7 +461,7 @@ func (s DocumentService) GenerateChangelog(
}
if changelog == nil {
changelog, err = s.svc.agent.GenerateChangelog(ctx, publishedVersion.Content, draftVersion.Content)
changelog, err = s.generateChangelog(ctx, publishedVersion.Content, draftVersion.Content)
if err != nil {
return nil, fmt.Errorf("cannot generate changelog: %w", err)
}
@@ -467,6 +470,42 @@ func (s DocumentService) GenerateChangelog(
return changelog, nil
}
//go:embed prompts/changelog_generator.txt
var changelogGeneratorSystemPrompt string
func (s DocumentService) generateChangelog(
ctx context.Context,
oldContent, newContent string,
) (*string, error) {
ag := agent.New(
"changelog_generator",
s.svc.llmClient,
agent.WithInstructions(changelogGeneratorSystemPrompt),
agent.WithModel(s.svc.llmModel),
agent.WithTemperature(s.svc.llmTemperature),
agent.WithMaxTokens(s.svc.llmMaxTokens),
)
result, err := ag.Run(
ctx,
[]llm.Message{
{
Role: llm.RoleUser,
Parts: []llm.Part{
llm.TextPart{Text: fmt.Sprintf("Old content: %s", oldContent)},
llm.TextPart{Text: fmt.Sprintf("New content: %s", newContent)},
},
},
},
)
if err != nil {
return nil, fmt.Errorf("cannot generate changelog: %w", err)
}
text := result.FinalMessage().Text()
return &text, nil
}
func (s *DocumentService) BulkPublishMinorVersions(
ctx context.Context,
req BulkPublishVersionsRequest,

21
pkg/probo/prompts/changelog_generator.txt Normal file
View File

@@ -0,0 +1,21 @@
# Role:You are an assistant that creates clear and concise changelogs.
# Objective
Given two versions of a document — the "old version" and the "new version" — identify and summarize all meaningful changes between them.
Focus on additions, deletions, modifications, and restructuring.
# Response Format
Respond with ONE simple phrase that describe the changes.
# Change types
If possible use the following words with additional context to describe the change types:
"Added", "Removed", "Updated", "Reworded", "Reorganized", "Fixed", etc.
# SOP
- Be objective and neutral in tone.
- Do not comment on the quality of the change.
- Use the language of the document.
**Example output format:**
Respond ONLY with the phrase that describes the changes. No explanation, no markdown, no preamble. Like this:
Added clauses about sharing personal information with trusted partners

34
pkg/probo/service.go
View File

@@ -22,7 +22,6 @@ import (
"github.com/aws/aws-sdk-go-v2/service/s3"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.probo.inc/probo/pkg/agents"
"go.probo.inc/probo/pkg/certmanager"
"go.probo.inc/probo/pkg/connector"
"go.probo.inc/probo/pkg/coredata"
@@ -69,6 +68,7 @@ type (
esign *esign.Service
connectorRegistry *connector.ConnectorRegistry
invitationTokenValidity time.Duration
vendorAssessor VendorAssessor
}
TenantService struct {
@@ -79,7 +79,11 @@ type (
scope coredata.Scoper
baseURL string
tokenSecret string
agent *agents.Agent
llmClient *llm.Client
llmModel string
llmTemperature float64
llmMaxTokens int
vendorAssessor VendorAssessor
fileManager *filemanager.Service
esign *esign.Service
Frameworks *FrameworkService
@@ -145,6 +149,7 @@ func NewService(
esignService *esign.Service,
connectorRegistry *connector.ConnectorRegistry,
invitationTokenValidity time.Duration,
vendorAssessor VendorAssessor,
) (*Service, error) {
if bucket == "" {
return nil, fmt.Errorf("bucket is required")
@@ -171,6 +176,7 @@ func NewService(
esign: esignService,
connectorRegistry: connectorRegistry,
invitationTokenValidity: invitationTokenValidity,
vendorAssessor: vendorAssessor,
}
return svc, nil
@@ -178,16 +184,20 @@ func NewService(
func (s *Service) WithTenant(tenantID gid.TenantID) *TenantService {
tenantService := &TenantService{
pg: s.pg,
s3: s.s3,
bucket: s.bucket,
encryptionKey: s.encryptionKey,
baseURL: s.baseURL,
scope: coredata.NewScope(tenantID),
tokenSecret: s.tokenSecret,
agent: agents.NewAgent(nil, s.llmClient, s.llmModel, s.llmTemperature, s.llmMaxTokens),
fileManager: s.fileManager,
esign: s.esign,
pg: s.pg,
s3: s.s3,
bucket: s.bucket,
encryptionKey: s.encryptionKey,
baseURL: s.baseURL,
scope: coredata.NewScope(tenantID),
tokenSecret: s.tokenSecret,
llmClient: s.llmClient,
llmModel: s.llmModel,
llmTemperature: s.llmTemperature,
llmMaxTokens: s.llmMaxTokens,
vendorAssessor: s.vendorAssessor,
fileManager: s.fileManager,
esign: s.esign,
}
tenantService.Frameworks = &FrameworkService{

200
pkg/probo/vendor_service.go
View File

@@ -16,18 +16,56 @@ package probo
import (
"context"
"errors"
"fmt"
"time"
"go.gearno.de/kit/pg"
"go.gearno.de/x/ref"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
"go.probo.inc/probo/pkg/page"
"go.probo.inc/probo/pkg/validator"
"go.probo.inc/probo/pkg/vetting"
"go.probo.inc/probo/pkg/webhook"
webhooktypes "go.probo.inc/probo/pkg/webhook/types"
)
// ErrVendorAssessmentDisabled is returned by VendorAssessor.Assess when the
// deployment has not configured an LLM provider for vendor assessment.
var ErrVendorAssessmentDisabled = errors.New("vendor assessment is not configured on this deployment")
// VendorAssessor produces a vendor assessment report from a website URL and
// an optional procedure description. Implementations that cannot perform
// assessment (missing LLM credentials, misconfigured provider) must return
// ErrVendorAssessmentDisabled from Assess so callers can surface a stable
// "feature unavailable" error instead of a generic internal error.
type VendorAssessor interface {
Assess(
ctx context.Context,
websiteURL string,
procedure string,
reporter agent.ProgressReporter,
) (*vetting.Result, error)
}
// DisabledVendorAssessor is the VendorAssessor implementation used when no
// LLM provider is configured for the vendor-assessor agent. Its Assess
// method always returns ErrVendorAssessmentDisabled.
type DisabledVendorAssessor struct{}
var _ VendorAssessor = DisabledVendorAssessor{}
func (DisabledVendorAssessor) Assess(
_ context.Context,
_ string,
_ string,
_ agent.ProgressReporter,
) (*vetting.Result, error) {
return nil, ErrVendorAssessmentDisabled
}
type (
VendorService struct {
svc *TenantService
@@ -83,6 +121,19 @@ type (
AssessVendorRequest struct {
ID gid.GID
WebsiteURL string
Procedure *string
}
AssessVendorResult struct {
Vendor *coredata.Vendor
Report string
Subprocessors []Subprocessor
}
Subprocessor struct {
Name string
Country string
Purpose string
}
CreateVendorRiskAssessmentRequest struct {
@@ -394,7 +445,14 @@ func (s VendorService) Update(
return fmt.Errorf("cannot update vendor: %w", err)
}
if err := webhook.InsertData(ctx, conn, s.svc.scope, vendor.OrganizationID, coredata.WebhookEventTypeVendorUpdated, webhooktypes.NewVendor(vendor)); err != nil {
if err := webhook.InsertData(
ctx,
conn,
s.svc.scope,
vendor.OrganizationID,
coredata.WebhookEventTypeVendorUpdated,
webhooktypes.NewVendor(vendor),
); err != nil {
return fmt.Errorf("cannot insert webhook event: %w", err)
}
@@ -470,7 +528,14 @@ func (s VendorService) Delete(
return fmt.Errorf("cannot load vendor: %w", err)
}
if err := webhook.InsertData(ctx, conn, s.svc.scope, vendor.OrganizationID, coredata.WebhookEventTypeVendorDeleted, webhooktypes.NewVendor(vendor)); err != nil {
if err := webhook.InsertData(
ctx,
conn,
s.svc.scope,
vendor.OrganizationID,
coredata.WebhookEventTypeVendorDeleted,
webhooktypes.NewVendor(vendor),
); err != nil {
return fmt.Errorf("cannot insert webhook event: %w", err)
}
@@ -547,7 +612,14 @@ func (s VendorService) Create(
return fmt.Errorf("cannot insert vendor: %w", err)
}
if err := webhook.InsertData(ctx, conn, s.svc.scope, organization.ID, coredata.WebhookEventTypeVendorCreated, webhooktypes.NewVendor(vendor)); err != nil {
if err := webhook.InsertData(
ctx,
conn,
s.svc.scope,
organization.ID,
coredata.WebhookEventTypeVendorCreated,
webhooktypes.NewVendor(vendor),
); err != nil {
return fmt.Errorf("cannot insert webhook event: %w", err)
}
@@ -763,32 +835,108 @@ func (s VendorService) GetByRiskAssessmentID(
func (s VendorService) Assess(
ctx context.Context,
req AssessVendorRequest,
) (*coredata.Vendor, error) {
vendorInfo, err := s.svc.agent.AssessVendor(ctx, req.WebsiteURL)
) (*AssessVendorResult, error) {
result, err := s.svc.vendorAssessor.Assess(ctx, req.WebsiteURL, ref.UnrefOrZero(req.Procedure), nil)
if err != nil {
return nil, fmt.Errorf("cannot assess vendor info: %w", err)
return nil, fmt.Errorf("cannot assess vendor: %w", err)
}
vendor := &coredata.Vendor{
ID: req.ID,
Name: vendorInfo.Name,
WebsiteURL: &req.WebsiteURL,
Description: &vendorInfo.Description,
Category: coredata.VendorCategory(vendorInfo.Category),
HeadquarterAddress: &vendorInfo.HeadquarterAddress,
LegalName: &vendorInfo.LegalName,
PrivacyPolicyURL: &vendorInfo.PrivacyPolicyURL,
ServiceLevelAgreementURL: &vendorInfo.ServiceLevelAgreementURL,
DataProcessingAgreementURL: &vendorInfo.DataProcessingAgreementURL,
BusinessAssociateAgreementURL: &vendorInfo.BusinessAssociateAgreementURL,
SubprocessorsListURL: &vendorInfo.SubprocessorsListURL,
SecurityPageURL: &vendorInfo.SecurityPageURL,
TrustPageURL: &vendorInfo.TrustPageURL,
TermsOfServiceURL: &vendorInfo.TermsOfServiceURL,
StatusPageURL: &vendorInfo.StatusPageURL,
Certifications: vendorInfo.Certifications,
UpdatedAt: time.Now(),
vendor := &coredata.Vendor{}
err = s.svc.pg.WithTx(
ctx,
func(ctx context.Context, conn pg.Tx) error {
if err := vendor.LoadByID(ctx, conn, s.svc.scope, req.ID); err != nil {
return fmt.Errorf("cannot load vendor %q: %w", req.ID, err)
}
info := result.Info
if info.Name != "" {
vendor.Name = info.Name
}
vendor.WebsiteURL = &req.WebsiteURL
if info.Category != "" {
vendor.Category = coredata.VendorCategory(info.Category)
}
vendor.UpdatedAt = time.Now()
if info.Description != "" {
vendor.Description = &info.Description
}
if info.HeadquarterAddress != "" {
vendor.HeadquarterAddress = &info.HeadquarterAddress
}
if info.LegalName != "" {
vendor.LegalName = &info.LegalName
}
if info.PrivacyPolicyURL != "" {
vendor.PrivacyPolicyURL = &info.PrivacyPolicyURL
}
if info.ServiceLevelAgreementURL != "" {
vendor.ServiceLevelAgreementURL = &info.ServiceLevelAgreementURL
}
if info.DataProcessingAgreementURL != "" {
vendor.DataProcessingAgreementURL = &info.DataProcessingAgreementURL
}
if info.BusinessAssociateAgreementURL != "" {
vendor.BusinessAssociateAgreementURL = &info.BusinessAssociateAgreementURL
}
if info.SubprocessorsListURL != "" {
vendor.SubprocessorsListURL = &info.SubprocessorsListURL
}
if info.SecurityPageURL != "" {
vendor.SecurityPageURL = &info.SecurityPageURL
}
if info.TrustPageURL != "" {
vendor.TrustPageURL = &info.TrustPageURL
}
if info.TermsOfServiceURL != "" {
vendor.TermsOfServiceURL = &info.TermsOfServiceURL
}
if info.StatusPageURL != "" {
vendor.StatusPageURL = &info.StatusPageURL
}
if len(info.Certifications) > 0 {
vendor.Certifications = info.Certifications
}
if err := vendor.Update(ctx, conn, s.svc.scope); err != nil {
return fmt.Errorf("cannot update vendor: %w", err)
}
if err := webhook.InsertData(
ctx,
conn,
s.svc.scope,
vendor.OrganizationID,
coredata.WebhookEventTypeVendorUpdated,
webhooktypes.NewVendor(vendor),
); err != nil {
return fmt.Errorf("cannot insert webhook event: %w", err)
}
return nil
},
)
if err != nil {
return nil, err
}
return vendor, nil
subprocessors := make([]Subprocessor, len(result.Info.Subprocessors))
for i, sp := range result.Info.Subprocessors {
subprocessors[i] = Subprocessor{
Name: sp.Name,
Country: sp.Country,
Purpose: sp.Purpose,
}
}
return &AssessVendorResult{
Vendor: vendor,
Report: result.Document,
Subprocessors: subprocessors,
}, nil
}

38
pkg/probod/evidence_describer_config.go
View File

@@ -1,38 +0,0 @@
// Copyright (c) 2025-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 probod
// EvidenceDescriberConfig holds both the worker settings and LLM overrides
// for the evidence description worker.
type EvidenceDescriberConfig struct {
Interval int `json:"interval"` // seconds
StaleAfter int `json:"stale-after"` // seconds
MaxConcurrency int `json:"max-concurrency"`
Provider string `json:"provider"`
ModelName string `json:"model-name"`
Temperature *float64 `json:"temperature"`
MaxTokens *int `json:"max-tokens"`
}
// LLMConfig extracts the LLM-specific fields as an LLMConfig.
func (c *EvidenceDescriberConfig) LLMConfig() LLMConfig {
return LLMConfig{
Provider: c.Provider,
ModelName: c.ModelName,
Temperature: c.Temperature,
MaxTokens: c.MaxTokens,
}
}

22
pkg/probod/llm.go
View File

@@ -26,6 +26,28 @@ import (
llmopenai "go.probo.inc/probo/pkg/llm/openai"
)
// resolveAgentClient resolves the agent's effective config from defaults and
// builds an LLM client for it. The name parameter is used in the logger and
// in error messages.
func (impl *Implm) resolveAgentClient(
name string,
agent LLMAgentConfig,
l *log.Logger,
tp trace.TracerProvider,
r prometheus.Registerer,
) (LLMAgentConfig, *llm.Client, error) {
resolved := impl.cfg.Agents.ResolveAgent(agent)
providerCfg, ok := impl.cfg.Agents.Providers[resolved.Provider]
if !ok {
return LLMAgentConfig{}, nil, fmt.Errorf("unknown LLM provider %q for %s agent", resolved.Provider, name)
}
client, err := buildLLMClient(providerCfg, l.Named("llm."+name), tp, r)
if err != nil {
return LLMAgentConfig{}, nil, fmt.Errorf("cannot create %s LLM client: %w", name, err)
}
return resolved, client, nil
}
func buildLLMClient(cfg LLMProviderConfig, l *log.Logger, tp trace.TracerProvider, r prometheus.Registerer) (*llm.Client, error) {
providerType := cfg.Type
if providerType == "" {

56
pkg/probod/llm_config.go
View File

@@ -22,38 +22,50 @@ type (
APIKey string `json:"api-key"` // for OpenAI and Anthropic
}
// LLMConfig holds model parameters for a single LLM consumer. Provider
// references one of the keys in LLMSettings.Providers.
LLMConfig struct {
Provider string `json:"provider"` // key into LLMSettings.Providers
// LLMAgentConfig holds model parameters for a single agent. Provider
// references one of the keys in AgentsConfig.Providers.
LLMAgentConfig struct {
Provider string `json:"provider"` // key into AgentsConfig.Providers
ModelName string `json:"model-name"`
Temperature *float64 `json:"temperature"`
MaxTokens *int `json:"max-tokens"`
}
// LLMSettings groups LLM provider credentials and default model
// settings. Defaults is used as a fallback when a consumer-specific
// field is zero-valued.
LLMSettings struct {
Providers map[string]LLMProviderConfig `json:"providers"`
Defaults LLMConfig `json:"defaults"`
// EvidenceDescriberConfig holds worker-side tuning for the evidence
// description background worker. LLM parameters for the same worker
// live under AgentsConfig.EvidenceDescriber.
EvidenceDescriberConfig struct {
Interval int `json:"interval"` // seconds between polls
StaleAfter int `json:"stale-after"` // seconds before a claim is recycled
MaxConcurrency int `json:"max-concurrency"`
}
// AgentsConfig groups LLM provider credentials and per-agent model
// 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"`
VendorAssessor LLMAgentConfig `json:"vendor-assessor"`
}
)
// ResolveLLMConfig returns a fully populated LLMConfig by filling in
// zero-valued fields from the defaults.
func (s *LLMSettings) ResolveLLMConfig(cfg LLMConfig) LLMConfig {
if cfg.Provider == "" {
cfg.Provider = s.Defaults.Provider
// ResolveAgent returns a fully populated LLMAgentConfig by filling in
// zero-valued fields from the default config.
func (c *AgentsConfig) ResolveAgent(agent LLMAgentConfig) LLMAgentConfig {
if agent.Provider == "" {
agent.Provider = c.Default.Provider
}
if cfg.ModelName == "" {
cfg.ModelName = s.Defaults.ModelName
if agent.ModelName == "" {
agent.ModelName = c.Default.ModelName
}
if cfg.Temperature == nil {
cfg.Temperature = s.Defaults.Temperature
if agent.Temperature == nil {
agent.Temperature = c.Default.Temperature
}
if cfg.MaxTokens == nil {
cfg.MaxTokens = s.Defaults.MaxTokens
if agent.MaxTokens == nil {
agent.MaxTokens = c.Default.MaxTokens
}
return cfg
return agent
}

34
pkg/probod/probod.go
View File

@@ -121,10 +121,10 @@ type (
AWS AWSConfig `json:"aws"`
Notifications NotificationsConfig `json:"notifications"`
Connectors []ConnectorConfig `json:"connectors"`
LLM LLMSettings `json:"llm"`
ProboAgent LLMConfig `json:"probo-agent"`
Agents AgentsConfig `json:"llm"`
EvidenceDescriber EvidenceDescriberConfig `json:"evidence-describer"`
ChromeDPAddr string `json:"chrome-dp-addr"`
SearchEndpoint string `json:"search-endpoint"`
CustomDomains CustomDomainsConfig `json:"custom-domains"`
SCIMBridge SCIMBridgeConfig `json:"scim-bridge"`
ESign ESignConfig `json:"esign"`
@@ -338,24 +338,19 @@ func (impl *Implm) Run(
}
}
proboAgentCfg := impl.cfg.LLM.ResolveLLMConfig(impl.cfg.ProboAgent)
proboProviderCfg, ok := impl.cfg.LLM.Providers[proboAgentCfg.Provider]
if !ok {
return fmt.Errorf("unknown LLM provider %q for probo agent", proboAgentCfg.Provider)
}
proboLLMClient, err := buildLLMClient(proboProviderCfg, l.Named("llm.probo"), tp, r)
proboAgentCfg, proboLLMClient, err := impl.resolveAgentClient("probo", impl.cfg.Agents.Probo, l, tp, r)
if err != nil {
return fmt.Errorf("cannot create probo LLM client: %w", err)
return err
}
edLLMCfg := impl.cfg.LLM.ResolveLLMConfig(impl.cfg.EvidenceDescriber.LLMConfig())
edProviderCfg, ok := impl.cfg.LLM.Providers[edLLMCfg.Provider]
if !ok {
return fmt.Errorf("unknown LLM provider %q for evidence-describer agent", edLLMCfg.Provider)
}
evidenceDescriberLLMClient, err := buildLLMClient(edProviderCfg, l.Named("llm.evidence-describer"), tp, r)
evidenceDescriberAgentCfg, evidenceDescriberLLMClient, err := impl.resolveAgentClient("evidence-describer", impl.cfg.Agents.EvidenceDescriber, l, tp, r)
if err != nil {
return fmt.Errorf("cannot create evidence describer LLM client: %w", err)
return err
}
vendorAssessor, err := impl.buildVendorAssessor(l, tp, r)
if err != nil {
return err
}
fileManagerService := filemanager.NewService(s3Client)
@@ -545,6 +540,7 @@ func (impl *Implm) Run(
esignService,
defaultConnectorRegistry,
time.Duration(impl.cfg.Auth.InvitationConfirmationTokenValidity)*time.Second,
vendorAssessor,
)
if err != nil {
return fmt.Errorf("cannot create probo service: %w", err)
@@ -735,9 +731,9 @@ func (impl *Implm) Run(
evidenceDescriber := evidencedescriber.New(
evidenceDescriberLLMClient,
evidencedescriber.Config{
Model: edLLMCfg.ModelName,
Temp: *edLLMCfg.Temperature,
MaxTokens: *edLLMCfg.MaxTokens,
Model: evidenceDescriberAgentCfg.ModelName,
Temp: *evidenceDescriberAgentCfg.Temperature,
MaxTokens: *evidenceDescriberAgentCfg.MaxTokens,
},
)
evidenceDescriptionWorker := probo.NewEvidenceDescriptionWorker(

58
pkg/probod/vendor_assessor.go Normal file
View File

@@ -0,0 +1,58 @@
// 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 probod
import (
"github.com/prometheus/client_golang/prometheus"
"go.gearno.de/kit/log"
"go.opentelemetry.io/otel/trace"
"go.probo.inc/probo/pkg/probo"
"go.probo.inc/probo/pkg/vetting"
)
// buildVendorAssessor wires the vendor assessment agent. It is an opt-in
// feature: deployments that do not set `llm.vendor-assessor.provider` get a
// DisabledVendorAssessor that reports the feature as unavailable. The
// vendor-assessor does not inherit the default provider because its
// pipeline (LLM + browser + search) is expensive and should not be enabled
// implicitly.
func (impl *Implm) buildVendorAssessor(
l *log.Logger,
tp trace.TracerProvider,
r prometheus.Registerer,
) (probo.VendorAssessor, error) {
if impl.cfg.Agents.VendorAssessor.Provider == "" {
return probo.DisabledVendorAssessor{}, nil
}
agentCfg, llmClient, err := impl.resolveAgentClient("vendor-assessor", impl.cfg.Agents.VendorAssessor, l, tp, r)
if err != nil {
return nil, err
}
maxTokens := vetting.DefaultMaxTokens
if agentCfg.MaxTokens != nil {
maxTokens = *agentCfg.MaxTokens
}
return vetting.NewAssessor(vetting.Config{
Client: llmClient,
Model: agentCfg.ModelName,
MaxTokens: maxTokens,
ChromeAddr: impl.cfg.ChromeDPAddr,
SearchEndpoint: impl.cfg.SearchEndpoint,
Logger: l.Named("vendor-assessor"),
}), nil
}

9
pkg/server/api/console/v1/graphql/vendor.graphql
View File

@@ -614,6 +614,13 @@ input CreateVendorRiskAssessmentInput {
input AssessVendorInput {
id: ID!
websiteUrl: String!
procedure: String
}
type VendorSubprocessor {
name: String!
country: String!
purpose: String!
}
type CreateVendorPayload {
@@ -690,4 +697,6 @@ type CreateVendorRiskAssessmentPayload {
type AssessVendorPayload {
vendor: Vendor!
report: String!
subprocessors: [VendorSubprocessor!]!
}

13
pkg/server/api/console/v1/types/vendor.go
View File

@@ -18,6 +18,7 @@ import (
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
"go.probo.inc/probo/pkg/page"
"go.probo.inc/probo/pkg/probo"
)
type (
@@ -103,3 +104,15 @@ func NewVendor(v *coredata.Vendor) *Vendor {
return object
}
func NewVendorSubprocessors(sps []probo.Subprocessor) []*VendorSubprocessor {
result := make([]*VendorSubprocessor, len(sps))
for i, sp := range sps {
result[i] = &VendorSubprocessor{
Name: sp.Name,
Country: sp.Country,
Purpose: sp.Purpose,
}
}
return result
}

11
pkg/server/api/console/v1/vendor_resolvers.go
View File

@@ -541,20 +541,27 @@ func (r *mutationResolver) AssessVendor(ctx context.Context, input types.AssessV
prb := r.ProboService(ctx, input.ID.TenantID())
vendor, err := prb.Vendors.Assess(
result, err := prb.Vendors.Assess(
ctx,
probo.AssessVendorRequest{
ID: input.ID,
WebsiteURL: input.WebsiteURL,
Procedure: input.Procedure,
},
)
if err != nil {
if errors.Is(err, probo.ErrVendorAssessmentDisabled) {
return nil, gqlutils.Unavailable(ctx, probo.ErrVendorAssessmentDisabled)
}
r.logger.ErrorCtx(ctx, "cannot assess vendor", log.Error(err))
return nil, gqlutils.Internal(ctx)
}
return &types.AssessVendorPayload{
Vendor: types.NewVendor(vendor),
Vendor: types.NewVendor(result.Vendor),
Report: result.Report,
Subprocessors: types.NewVendorSubprocessors(result.Subprocessors),
}, nil
}

20
pkg/server/api/mcp/v1/schema.resolvers.go
View File

@@ -4763,3 +4763,23 @@ func (r *Resolver) DeleteCustomDomainTool(ctx context.Context, req *mcp.CallTool
return nil, types.DeleteCustomDomainOutput{DeletedCustomDomain: deletedDomain}, nil
}
func (r *Resolver) AssessVendorTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AssessVendorInput) (*mcp.CallToolResult, types.AssessVendorOutput, error) {
r.MustAuthorize(ctx, input.ID, probo.ActionVendorAssess)
svc := r.ProboService(ctx, input.ID)
result, err := svc.Vendors.Assess(
ctx,
probo.AssessVendorRequest{
ID: input.ID,
WebsiteURL: input.WebsiteURL,
Procedure: input.Procedure,
},
)
if err != nil {
return nil, types.AssessVendorOutput{}, fmt.Errorf("cannot assess vendor: %w", err)
}
return nil, types.NewAssessVendorOutput(result), nil
}

59
pkg/server/api/mcp/v1/specification.yaml
View File

@@ -1122,6 +1122,57 @@ components:
$ref: "#/components/schemas/GID"
description: Deleted vendor service ID
AssessVendorInput:
type: object
required:
- id
- website_url
properties:
id:
$ref: "#/components/schemas/GID"
description: Vendor ID to assess
website_url:
type: string
description: Vendor website URL to crawl and assess
procedure:
type: string
description: Optional custom assessment procedure (overrides the default)
VendorSubprocessor:
type: object
required:
- name
- country
- purpose
properties:
name:
type: string
description: Sub-processor name
country:
type: string
description: Country where the sub-processor operates
purpose:
type: string
description: Purpose of the sub-processor
AssessVendorOutput:
type: object
required:
- vendor
- report
- subprocessors
properties:
vendor:
$ref: "#/components/schemas/Vendor"
report:
type: string
description: Markdown-formatted vendor assessment report
subprocessors:
type: array
items:
$ref: "#/components/schemas/VendorSubprocessor"
description: Sub-processors discovered during the assessment
GetUserInput:
type: object
required:
@@ -9232,6 +9283,14 @@ tools:
$ref: "#/components/schemas/DeleteVendorServiceInput"
outputSchema:
$ref: "#/components/schemas/DeleteVendorServiceOutput"
- name: assessVendor
description: Run an AI-powered assessment on a vendor by crawling its website. Returns a markdown report, the discovered sub-processors, and an enriched vendor record. Long-running (up to 20 minutes).
hints:
readonly: false
inputSchema:
$ref: "#/components/schemas/AssessVendorInput"
outputSchema:
$ref: "#/components/schemas/AssessVendorOutput"
- name: listRisks
description: List all risks for the organization
hints:

21
pkg/server/api/mcp/v1/types/vendor.go
View File

@@ -17,6 +17,7 @@ package types
import (
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/page"
"go.probo.inc/probo/pkg/probo"
)
func NewVendorRiskAssessment(v *coredata.VendorRiskAssessment) *VendorRiskAssessment {
@@ -205,3 +206,23 @@ func NewListVendorServicesOutput(p *page.Page[*coredata.VendorService, coredata.
VendorServices: services,
}
}
func NewVendorSubprocessors(sps []probo.Subprocessor) []*VendorSubprocessor {
result := make([]*VendorSubprocessor, len(sps))
for i, sp := range sps {
result[i] = &VendorSubprocessor{
Name: sp.Name,
Country: sp.Country,
Purpose: sp.Purpose,
}
}
return result
}
func NewAssessVendorOutput(result *probo.AssessVendorResult) AssessVendorOutput {
return AssessVendorOutput{
Vendor: NewVendor(result.Vendor),
Report: result.Report,
Subprocessors: NewVendorSubprocessors(result.Subprocessors),
}
}

14
pkg/server/gqlutils/errors.go
View File

@@ -187,3 +187,17 @@ func Internal(ctx context.Context) *gqlerror.Error {
},
}
}
func Unavailable(ctx context.Context, err error) *gqlerror.Error {
return &gqlerror.Error{
Message: err.Error(),
Path: graphql.GetPath(ctx),
Extensions: map[string]any{
"code": "UNAVAILABLE",
},
}
}
func Unavailablef(ctx context.Context, format string, a ...any) *gqlerror.Error {
return Unavailable(ctx, fmt.Errorf(format, a...))
}

337
pkg/vetting/assessment.go Normal file
View File

@@ -0,0 +1,337 @@
// 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 vetting
import (
"context"
_ "embed"
"encoding/json"
"fmt"
"net/url"
"time"
"go.gearno.de/kit/log"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/agent/tools/browser"
"go.probo.inc/probo/pkg/llm"
)
const (
// DefaultMaxTokens is the fallback max-tokens budget used when the
// vendor-assessor agent config does not specify a value. Sized to
// leave headroom above the orchestrator's thinking budget on
// Anthropic models.
DefaultMaxTokens = 16384
// AssessmentTimeout is the hard upper bound on a single assessment
// run. This is also the timeout the CLI client should use.
AssessmentTimeout = 20 * time.Minute
// extractionTimeout is the dedicated budget for the final
// vendor_info_extractor turn. It runs outside the orchestrator's
// budget so a slow orchestrator can't starve the extractor.
extractionTimeout = 5 * time.Minute
)
// vendorCategoryEnum is the canonical list of allowed values for
// VendorInfo.Category. It is duplicated into the jsonschema struct tag
// because Go struct tags must be compile-time string literals.
var vendorCategoryEnum = []string{
"ANALYTICS", "ACCOUNTING", "CLOUD_MONITORING", "CLOUD_PROVIDER",
"COLLABORATION", "CONSULTING", "CUSTOMER_SUPPORT",
"DATA_STORAGE_AND_PROCESSING", "DOCUMENT_MANAGEMENT",
"EMPLOYEE_MANAGEMENT", "ENGINEERING", "FINANCE", "IDENTITY_PROVIDER",
"IT", "LEGAL", "MARKETING", "OFFICE_OPERATIONS", "OTHER",
"PASSWORD_MANAGEMENT", "PRODUCT_AND_DESIGN", "PROFESSIONAL_SERVICES",
"RECRUITING", "SALES", "SECURITY", "STAFFING", "VERSION_CONTROL",
}
// vendorTypeEnum is the canonical list of allowed values for
// VendorInfo.VendorType.
var vendorTypeEnum = []string{
"SAAS", "INFRASTRUCTURE", "PROFESSIONAL_SERVICES", "STAFFING", "OTHER",
}
var (
//go:embed prompts/extraction.txt
extractionPrompt string
)
type (
Config struct {
Client *llm.Client
Model string
MaxTokens int
ChromeAddr string
SearchEndpoint string
Logger *log.Logger
}
Assessor struct {
cfg Config
}
Subprocessor struct {
Name string `json:"name"`
Country string `json:"country"`
Purpose string `json:"purpose"`
}
RiskScore struct {
Category string `json:"category"`
Rating string `json:"rating"`
Notes string `json:"notes"`
}
VendorInfo struct {
Name string `json:"name" jsonschema:"Vendor display name as shown on the website"`
Description string `json:"description" jsonschema:"One-sentence description of what the vendor does"`
Category string `json:"category" jsonschema:"Vendor category; one of vendorCategoryEnum"`
VendorType string `json:"vendor_type" jsonschema:"Vendor type; one of vendorTypeEnum"`
HeadquarterAddress string `json:"headquarter_address" jsonschema:"Vendor headquarters address (city, country) if mentioned"`
LegalName string `json:"legal_name" jsonschema:"Legal entity name if different from display name (e.g. 'Datadog, Inc.')"`
PrivacyPolicyURL string `json:"privacy_policy_url" jsonschema:"URL to the vendor's privacy policy page"`
ServiceLevelAgreementURL string `json:"service_level_agreement_url" jsonschema:"URL to the SLA page"`
DataProcessingAgreementURL string `json:"data_processing_agreement_url" jsonschema:"URL to the DPA page"`
BusinessAssociateAgreementURL string `json:"business_associate_agreement_url" jsonschema:"URL to the BAA page if HIPAA-eligible"`
SubprocessorsListURL string `json:"subprocessors_list_url" jsonschema:"URL to the public subprocessors list"`
SecurityPageURL string `json:"security_page_url" jsonschema:"URL to the vendor's security page"`
TrustPageURL string `json:"trust_page_url" jsonschema:"URL to the trust center"`
TermsOfServiceURL string `json:"terms_of_service_url" jsonschema:"URL to the terms of service"`
StatusPageURL string `json:"status_page_url" jsonschema:"URL to the vendor's status / uptime page"`
BugBountyURL string `json:"bug_bounty_url" jsonschema:"URL to the bug bounty or responsible disclosure program"`
IncidentResponseURL string `json:"incident_response_url" jsonschema:"URL to incident response or post-mortem documentation"`
DataLocations []string `json:"data_locations" jsonschema:"Countries or regions where data is processed or stored (e.g. 'United States', 'EU', 'Germany')"`
Certifications []string `json:"certifications" jsonschema:"Compliance certifications found (e.g. 'SOC 2 Type II', 'ISO 27001')"`
Subprocessors []Subprocessor `json:"subprocessors" jsonschema:"Sub-processors discovered with name, country, purpose"`
// Privacy classification (ISO 27701).
PrivacyRole string `json:"privacy_role" jsonschema:"Privacy role under ISO 27701: CONTROLLER, PROCESSOR, SUBPROCESSOR, NONE"`
ProcessesPII bool `json:"processes_pii" jsonschema:"Whether the vendor processes personal data"`
CrossBorderTransfer bool `json:"cross_border_transfer" jsonschema:"Whether cross-border data transfers occur"`
// Privacy risk fields.
DPAStatus string `json:"dpa_status" jsonschema:"DPA accessibility: AVAILABLE, AVAILABLE_ON_REQUEST, NOT_FOUND, BEHIND_LOGIN"`
DSARCapability string `json:"dsar_capability" jsonschema:"Brief summary of how the vendor handles Data Subject Access Requests"`
DataMinimization string `json:"data_minimization" jsonschema:"Brief summary of data minimization practices"`
PurposeLimitation string `json:"purpose_limitation" jsonschema:"Brief summary of purpose limitation commitments"`
RetentionPolicy string `json:"retention_policy" jsonschema:"Brief summary of data retention policy"`
DeletionPolicy string `json:"deletion_policy" jsonschema:"Brief summary of data deletion policy"`
// AI classification (ISO 42001).
InvolvesAI bool `json:"involves_ai" jsonschema:"Whether the vendor uses AI/ML in their product or service"`
AIUseCases []string `json:"ai_use_cases" jsonschema:"Array of AI use case descriptions (e.g. 'content generation', 'fraud detection')"`
// AI risk fields.
AIGovernanceDocURL string `json:"ai_governance_doc_url" jsonschema:"URL to AI governance or responsible AI documentation"`
AITransparency string `json:"ai_transparency" jsonschema:"Brief summary of model transparency findings"`
BiasControls string `json:"bias_controls" jsonschema:"Brief summary of bias detection and fairness measures"`
HumanOversight string `json:"human_oversight" jsonschema:"Brief summary of human oversight mechanisms for AI decisions"`
TrainingDataGovernance string `json:"training_data_governance" jsonschema:"Brief summary of training data governance"`
// Contractual clause analysis.
PrivacyClauses []string `json:"privacy_clauses" jsonschema:"Notable privacy contractual clauses found (e.g. '72-hour breach notification', 'SCCs included')"`
AIClauses []string `json:"ai_clauses" jsonschema:"Notable AI contractual clauses found (e.g. 'Customer data not used for training')"`
// Minimum acceptance baseline.
MinimumBaselineMet bool `json:"minimum_baseline_met" jsonschema:"Whether all hard-reject baseline criteria are met"`
BaselineFailures []string `json:"baseline_failures" jsonschema:"List of failed baseline criteria descriptions"`
// Risk scoring.
OverallRiskRating string `json:"overall_risk_rating" jsonschema:"Overall risk rating: Low, Medium, High"`
OverallRiskScore int `json:"overall_risk_score" jsonschema:"Overall risk score from the report (0-100)"`
Recommendation string `json:"recommendation" jsonschema:"Recommendation: APPROVE, APPROVE_WITH_CONDITIONS, ESCALATE, REJECT"`
RiskScores []RiskScore `json:"risk_scores" jsonschema:"Per-category risk scores from the Risk Summary table"`
SecurityRiskScore int `json:"security_risk_score" jsonschema:"Security pillar risk score (0-100)"`
PrivacyRiskScore int `json:"privacy_risk_score" jsonschema:"Privacy pillar risk score (0-100)"`
AIRiskScore int `json:"ai_risk_score" jsonschema:"AI pillar risk score (0-100), 0 if no AI"`
InformationGaps []string `json:"information_gaps" jsonschema:"Concise descriptions of information gaps from the report"`
ProfessionalLicenses []string `json:"professional_licenses" jsonschema:"Professional license descriptions for services firms (e.g. 'New York State Bar')"`
IndustryMemberships []string `json:"industry_memberships" jsonschema:"Industry body memberships (e.g. 'AICPA', 'American Bar Association')"`
InsuranceCoverage string `json:"insurance_coverage" jsonschema:"Description of professional liability or E&O insurance"`
}
Result struct {
Document string
Info VendorInfo
}
)
func NewAssessor(cfg Config) *Assessor {
return &Assessor{cfg: cfg}
}
func (a *Assessor) Assess(ctx context.Context, websiteURL string, procedure string, reporter agent.ProgressReporter) (*Result, error) {
u, err := url.Parse(websiteURL)
if err != nil {
return nil, fmt.Errorf("cannot parse website URL %q: %w", websiteURL, err)
}
if u.Scheme != "http" && u.Scheme != "https" {
return nil, fmt.Errorf("website URL must use http or https, got %q", u.Scheme)
}
if u.Hostname() == "" {
return nil, fmt.Errorf("website URL %q has no host", websiteURL)
}
// Detach from the caller's context (typically the HTTP request) so
// that the assessment is not cancelled when the client disconnects.
// A dedicated timeout prevents the assessment from running forever.
ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), AssessmentTimeout)
defer cancel()
vendorBrowser := browser.NewBrowser(ctx, a.cfg.ChromeAddr)
defer vendorBrowser.Close()
vendorBrowser.SetAllowedDomain(u.Hostname())
// Create an unrestricted browser for web search agents that need to
// follow links to external sites (news, reviews, etc.).
researchBrowser := browser.NewBrowser(ctx, a.cfg.ChromeAddr)
defer researchBrowser.Close()
orchestrator, err := newOrchestratorAgent(
a.cfg.Client,
a.cfg.Model,
a.cfg.MaxTokens,
procedure,
a.cfg.Logger,
vendorBrowser,
researchBrowser,
a.cfg.SearchEndpoint,
reporter,
)
if err != nil {
return nil, fmt.Errorf("cannot create orchestrator agent: %w", err)
}
result, err := orchestrator.Run(
ctx,
[]llm.Message{
{
Role: llm.RoleUser,
Parts: []llm.Part{llm.TextPart{Text: websiteURL}},
},
},
)
if err != nil {
return nil, fmt.Errorf("cannot assess vendor: %w", err)
}
document := result.FinalMessage().Text()
reportProgress(ctx, reporter, "extract_vendor_info", agent.ProgressEventStepStarted)
info, err := a.extractVendorInfo(ctx, document)
if err != nil {
reportProgress(ctx, reporter, "extract_vendor_info", agent.ProgressEventStepFailed)
return nil, fmt.Errorf("cannot extract vendor info: %w", err)
}
reportProgress(ctx, reporter, "extract_vendor_info", agent.ProgressEventStepCompleted)
return &Result{
Document: document,
Info: *info,
}, nil
}
func (a *Assessor) extractVendorInfo(ctx context.Context, document string) (*VendorInfo, error) {
outputType, err := vendorInfoOutputType()
if err != nil {
return nil, fmt.Errorf("cannot build vendor info output type: %w", err)
}
// Run the extractor on its own timeout so a slow orchestrator
// cannot starve the final JSON conversion step. The extractor has
// no tools and produces one structured JSON output; a few minutes
// is more than enough even when streaming is forced.
extractCtx, cancel := context.WithTimeout(
context.WithoutCancel(ctx),
extractionTimeout,
)
defer cancel()
extractor := agent.New(
"vendor_info_extractor",
a.cfg.Client,
agent.WithInstructions(extractionPrompt),
agent.WithModel(a.cfg.Model),
agent.WithMaxTokens(a.cfg.MaxTokens),
agent.WithLogger(a.cfg.Logger),
agent.WithOutputType(outputType),
)
result, err := extractor.Run(
extractCtx,
[]llm.Message{
{
Role: llm.RoleUser,
Parts: []llm.Part{llm.TextPart{Text: document}},
},
},
)
if err != nil {
return nil, fmt.Errorf("cannot extract vendor info: %w", err)
}
var info VendorInfo
if err := json.Unmarshal([]byte(result.FinalMessage().Text()), &info); err != nil {
return nil, fmt.Errorf("cannot parse vendor info output: %w", err)
}
return &info, nil
}
// vendorInfoOutputType builds the VendorInfo structured output type and
// decorates its JSON Schema with explicit enum constraints on fields
// whose allowed values live in package-level slices. jsonschema-go only
// reads struct tags as free-form descriptions, so the enum list cannot
// be encoded in the tag itself.
func vendorInfoOutputType() (*agent.OutputType, error) {
outputType, err := agent.NewOutputType[VendorInfo]("vendor_info")
if err != nil {
return nil, fmt.Errorf("cannot create vendor info output type: %w", err)
}
var schema map[string]any
if err := json.Unmarshal(outputType.Schema, &schema); err != nil {
return nil, fmt.Errorf("cannot unmarshal vendor info schema: %w", err)
}
properties, ok := schema["properties"].(map[string]any)
if !ok {
return nil, fmt.Errorf("vendor info schema has no properties")
}
enums := map[string][]string{
"category": vendorCategoryEnum,
"vendor_type": vendorTypeEnum,
}
for field, values := range enums {
prop, ok := properties[field].(map[string]any)
if !ok {
return nil, fmt.Errorf("vendor info schema has no %q property", field)
}
prop["enum"] = values
}
decorated, err := json.Marshal(schema)
if err != nil {
return nil, fmt.Errorf("cannot marshal decorated vendor info schema: %w", err)
}
outputType.Schema = decorated
return outputType, nil
}

66
pkg/vetting/assessment_test.go Normal file
View File

@@ -0,0 +1,66 @@
// 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.
// This test file is white-box (package vetting, not vetting_test) so it
// can reach the unexported vendorInfoOutputType helper.
package vetting
import (
"encoding/json"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestVendorInfoOutputType_DecoratesEnums(t *testing.T) {
t.Parallel()
outputType, err := vendorInfoOutputType()
require.NoError(t, err)
require.NotNil(t, outputType)
var schema map[string]any
require.NoError(t, json.Unmarshal(outputType.Schema, &schema))
properties, ok := schema["properties"].(map[string]any)
require.True(t, ok)
tests := []struct {
field string
expected []string
}{
{"category", vendorCategoryEnum},
{"vendor_type", vendorTypeEnum},
}
for _, tt := range tests {
t.Run(tt.field, func(t *testing.T) {
t.Parallel()
prop, ok := properties[tt.field].(map[string]any)
require.True(t, ok, "schema has no %q property", tt.field)
enumRaw, ok := prop["enum"].([]any)
require.True(t, ok, "%q has no enum array", tt.field)
actual := make([]string, len(enumRaw))
for i, v := range enumRaw {
actual[i] = v.(string)
}
assert.Equal(t, tt.expected, actual)
})
}
}

261
pkg/vetting/orchestrator.go Normal file
View File

@@ -0,0 +1,261 @@
// 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 vetting
import (
_ "embed"
"fmt"
"strings"
"go.gearno.de/kit/log"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/agent/tools/browser"
"go.probo.inc/probo/pkg/agent/tools/search"
"go.probo.inc/probo/pkg/agent/tools/security"
"go.probo.inc/probo/pkg/llm"
)
var (
//go:embed prompts/orchestrator_base.txt
orchestratorBasePrompt string
//go:embed prompts/default_procedure.txt
defaultProcedure string
)
const (
// orchestratorMaxTurns bounds the orchestrator loop. Each turn typically
// dispatches one sub-agent in parallel; with 16 sub-agents and a few
// retries we need ~140 turns of headroom before timing out.
orchestratorMaxTurns = 140
// orchestratorThinkingBudget is the extended-thinking budget for the
// orchestrator. It is high because the orchestrator must reason over
// the outputs of all 16 sub-agents to produce the final report.
orchestratorThinkingBudget = 40000
)
// subAgentEntry binds a sub-agent's LLM-facing name and description to
// the tools it needs and a typed builder. The orchestrator iterates over
// a slice of these and turns each into an agent + AsTool wrapper.
type subAgentEntry struct {
toolName string
description string
tools []agent.Tool
build subAgentBuilder
}
func newOrchestratorAgent(
client *llm.Client,
model string,
maxTokens int,
procedure string,
logger *log.Logger,
vendorBrowser *browser.Browser,
researchBrowser *browser.Browser,
searchEndpoint string,
reporter agent.ProgressReporter,
) (*agent.Agent, error) {
readOnlyBrowserTools := browser.NewReadOnlyToolset(vendorBrowser).Tools()
// Unrestricted browser tools for sub-agents that need to follow links
// to external sites (subprocessor lists hosted on OneTrust/Transcend,
// research, vendor comparison).
unrestrictedBrowserTools := browser.NewInteractiveToolset(researchBrowser).Tools()
securityTools := security.NewToolset().Tools()
maxTokensOpt := agent.WithMaxTokens(maxTokens)
loggerOpt := agent.WithLogger(logger)
subAgentOpts := func(step string) []agent.Option {
opts := []agent.Option{loggerOpt, maxTokensOpt}
if reporter != nil {
opts = append(opts, agent.WithHooks(newSubProgressHooks(reporter, step)))
}
return opts
}
// Subprocessor agent benefits from web search when available so it can
// find subprocessor pages hosted on third-party platforms.
subprocessorTools := unrestrictedBrowserTools
if searchEndpoint != "" {
subprocessorTools = append(subprocessorTools, search.WebSearchTool(searchEndpoint))
}
// Core sub-agents that always run.
entries := []subAgentEntry{
{
toolName: "crawl_vendor_website",
description: "Crawl a vendor website to discover security, compliance, privacy, and legal pages. Returns structured JSON with categorized URLs (vendor_name, vendor_domain, discovered_urls, notes). Input: the vendor's main website URL.",
tools: readOnlyBrowserTools,
build: buildCrawlerAgent,
},
{
toolName: "assess_security",
description: "Perform technical security checks on a domain. Returns structured JSON with per-check results (ssl, headers, dmarc, spf, breaches, dnssec, csp, cors, dns, whois) each with status (pass/warning/fail/error) and details. Input: the vendor's domain name (e.g. example.com).",
tools: securityTools,
build: buildSecurityAgent,
},
{
toolName: "analyze_document",
description: "Analyze a specific document page (privacy policy, DPA, ToS) and extract key provisions. Returns structured JSON with document_type, retention, locations, GDPR/CCPA indicators, clauses, and summary. Input: the document URL.",
tools: readOnlyBrowserTools,
build: buildAnalyzerAgent,
},
{
toolName: "assess_compliance",
description: "Identify certifications and compliance frameworks from a trust/compliance page. Returns structured JSON with certifications (name, status, details), audit reports, and frameworks. Input: the trust or compliance page URL.",
tools: readOnlyBrowserTools,
build: buildComplianceAgent,
},
{
toolName: "assess_market_presence",
description: "Analyze a vendor's market presence. Returns structured JSON with notable_customers, case_studies, partnerships, company_size_signals, funding_info, and market_position. Input: the vendor's main website URL.",
tools: readOnlyBrowserTools,
build: buildMarketAgent,
},
{
toolName: "extract_subprocessors",
description: "Find and extract the list of sub-processors from a vendor's website. Returns structured JSON with subprocessors (name, country, purpose), total_count, and source. Input: the vendor's main website URL or a known subprocessors page URL.",
tools: subprocessorTools,
build: buildSubprocessorAgent,
},
{
toolName: "assess_data_processing",
description: "Assess data processing practices. Returns structured JSON with encryption, retention, deletion, data locations, transfer mechanisms, DPA status, DSAR handling, and rating. Input: a relevant page URL (privacy policy, DPA, security page, or trust center).",
tools: readOnlyBrowserTools,
build: buildDataProcessingAgent,
},
{
toolName: "assess_incident_response",
description: "Evaluate incident response capabilities. Returns structured JSON with ir_plan, notification_timeline, status_page, post_mortems, recent_incidents, security_contact, and rating. Input: a relevant page URL (security page, trust center, or status page).",
tools: readOnlyBrowserTools,
build: buildIncidentResponseAgent,
},
{
toolName: "assess_business_continuity",
description: "Evaluate business continuity and disaster recovery. Returns structured JSON with dr_plan, rto, rpo, cloud_providers, uptime_sla, regions, backup_strategy, and rating. Input: a relevant page URL (SLA page, trust center, or infrastructure docs).",
tools: readOnlyBrowserTools,
build: buildBusinessContinuityAgent,
},
{
toolName: "assess_professional_standing",
description: "Evaluate professional standing for services firms. Returns structured JSON with licensing, memberships, insurance, team_credentials, coi_policy, and rating. Input: relevant page URL (team page, about page, credentials page).",
tools: readOnlyBrowserTools,
build: buildProfessionalStandingAgent,
},
{
toolName: "assess_ai_risk",
description: "Evaluate AI governance (ISO 42001). Returns structured JSON with ai_involvement, use_cases, model_transparency, bias_controls, customer_data_training, human_oversight, and rating. Input: relevant page URL (AI policy, trust center, responsible AI page, or main website).",
tools: readOnlyBrowserTools,
build: buildAIRiskAgent,
},
{
toolName: "assess_regulatory_compliance",
description: "Deep regulatory compliance check. Returns structured JSON with per-framework assessment (gdpr, hipaa, pci_dss, sox) each with articles, status, and notes. Input: relevant page URL (DPA, compliance page, trust center).",
tools: readOnlyBrowserTools,
build: buildRegulatoryComplianceAgent,
},
}
// Optional sub-agents: only added when a search endpoint is configured.
if searchEndpoint != "" {
researchBrowserTools := browser.NewInteractiveToolset(researchBrowser).Tools()
searchTool := search.WebSearchTool(searchEndpoint)
govDBTool := search.CheckGovernmentDBTool(searchEndpoint)
waybackTool := search.CheckWaybackTool()
diffTool := search.DiffDocumentsTool()
// withResearchTools returns a fresh slice combining the supplied
// extra tools with the research browser tools. The fresh
// allocation is required so the four sub-agent tool slices do
// not share a backing array.
withResearchTools := func(extra ...agent.Tool) []agent.Tool {
out := make([]agent.Tool, 0, len(extra)+len(researchBrowserTools))
out = append(out, extra...)
out = append(out, researchBrowserTools...)
return out
}
websearchTools := withResearchTools(searchTool)
financialTools := withResearchTools(searchTool, govDBTool, waybackTool)
codeSecurityTools := withResearchTools(searchTool)
comparisonTools := withResearchTools(searchTool, diffTool)
entries = append(entries,
subAgentEntry{
toolName: "research_vendor_externally",
description: "Search the open web for external signals about the vendor. Returns structured JSON with security_incidents, regulatory_actions, customer_sentiment, recent_news, red_flags, and positive_signals. Input: the vendor's name and domain.",
tools: websearchTools,
build: buildWebsearchAgent,
},
subAgentEntry{
toolName: "assess_financial_stability",
description: "Evaluate vendor financial stability. Returns structured JSON with company_age, funding, employee_count, legal_standing, ownership, risk_signals, overall_assessment, and confidence. Input: vendor name and website URL.",
tools: financialTools,
build: buildFinancialStabilityAgent,
},
subAgentEntry{
toolName: "assess_code_security",
description: "Evaluate open-source code security posture. Returns structured JSON with has_public_repos, security_advisories, dependency_management, release_cadence, security_policy, overall_assessment, and risk_signals. Input: vendor name and website URL.",
tools: codeSecurityTools,
build: buildCodeSecurityAgent,
},
subAgentEntry{
toolName: "compare_vendor",
description: "Find and compare alternative vendors. Returns structured JSON with alternatives (name, certifications, security_score), comparison_summary, vendor_strengths, vendor_weaknesses, and overall_position. Input: vendor name, category, and website URL.",
tools: comparisonTools,
build: buildVendorComparisonAgent,
},
)
}
tools := make([]agent.Tool, 0, len(entries))
for _, e := range entries {
ag, err := e.build(client, model, e.tools, subAgentOpts(e.toolName)...)
if err != nil {
return nil, fmt.Errorf("cannot create %s sub-agent: %w", e.toolName, err)
}
tools = append(tools, ag.AsTool(e.toolName, e.description))
}
if procedure == "" {
procedure = defaultProcedure
}
systemPrompt := strings.Replace(orchestratorBasePrompt, "{procedure}", procedure, 1)
opts := []agent.Option{
agent.WithLogger(logger),
agent.WithInstructions(systemPrompt),
agent.WithModel(model),
agent.WithMaxTokens(maxTokens),
agent.WithTools(tools...),
agent.WithMaxTurns(orchestratorMaxTurns),
agent.WithParallelToolCalls(true),
agent.WithThinking(orchestratorThinkingBudget),
}
if reporter != nil {
opts = append(opts, agent.WithHooks(newProgressHooks(reporter)))
}
return agent.New(
"vendor_assessment_orchestrator",
client,
opts...,
), nil
}

359
pkg/vetting/output_types.go Normal file
View File

@@ -0,0 +1,359 @@
// 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 vetting
// Output types for all vetting sub-agents. Each struct defines the JSON
// schema enforced via agent.WithOutputType on the corresponding sub-agent.
type (
// --- Crawler ---
DiscoveredURL struct {
Category string `json:"category" jsonschema:"URL category: privacy_policy, terms_of_service, dpa, security, trust, compliance, status, subprocessors, sla, about, team, ai_policy, blog, careers, pricing, other"`
URL string `json:"url" jsonschema:"The discovered URL"`
}
CrawlerOutput struct {
VendorName string `json:"vendor_name" jsonschema:"The vendor's display name as found on the website"`
VendorDomain string `json:"vendor_domain" jsonschema:"The vendor's primary domain"`
DiscoveredURLs []DiscoveredURL `json:"discovered_urls" jsonschema:"All categorized URLs discovered during crawling"`
Notes string `json:"notes" jsonschema:"Observations about the site structure or crawl limitations"`
}
// --- Security ---
SecurityCheckResult struct {
Status string `json:"status" jsonschema:"Check result: pass, warning, fail, or error"`
Details string `json:"details" jsonschema:"Detailed findings for this check"`
}
WHOISResult struct {
Registrar string `json:"registrar" jsonschema:"Domain registrar name"`
CreationDate string `json:"creation_date" jsonschema:"Domain creation date"`
Organization string `json:"organization" jsonschema:"Registrant organization"`
NameServers string `json:"name_servers" jsonschema:"Comma-separated name servers"`
}
SecurityOutput struct {
SSL SecurityCheckResult `json:"ssl" jsonschema:"SSL/TLS certificate and protocol check"`
Headers SecurityCheckResult `json:"headers" jsonschema:"HTTP security headers check (HSTS, X-Frame-Options, etc.)"`
DMARC SecurityCheckResult `json:"dmarc" jsonschema:"DMARC email authentication policy check"`
SPF SecurityCheckResult `json:"spf" jsonschema:"SPF email authentication record check"`
Breaches SecurityCheckResult `json:"breaches" jsonschema:"Known data breaches check via HIBP"`
DNSSEC SecurityCheckResult `json:"dnssec" jsonschema:"DNSSEC validation check"`
CSP SecurityCheckResult `json:"csp" jsonschema:"Content Security Policy analysis"`
CORS SecurityCheckResult `json:"cors" jsonschema:"CORS configuration check"`
DNS SecurityCheckResult `json:"dns" jsonschema:"DNS records analysis (A, MX, TXT, NS)"`
WHOIS WHOISResult `json:"whois" jsonschema:"Domain WHOIS registration details"`
Summary string `json:"summary" jsonschema:"Overall security posture summary"`
}
// --- Document Analyzer ---
DocumentAnalysisOutput struct {
DocumentType string `json:"document_type" jsonschema:"Type of document: privacy_policy, terms_of_service, dpa, sla, security_policy, acceptable_use, engagement_letter, other"`
DocumentTitle string `json:"document_title" jsonschema:"Title of the document as shown on the page"`
LastUpdated string `json:"last_updated" jsonschema:"Last updated date if found, empty string otherwise"`
DataRetention string `json:"data_retention" jsonschema:"Data retention policy details"`
DataLocations []string `json:"data_locations" jsonschema:"Countries or regions where data is processed or stored"`
GDPRIndicators string `json:"gdpr_indicators" jsonschema:"GDPR compliance indicators found"`
CCPAIndicators string `json:"ccpa_indicators" jsonschema:"CCPA/CPRA compliance indicators found"`
SecurityMeasures string `json:"security_measures" jsonschema:"Security measures described in the document"`
BreachNotification string `json:"breach_notification" jsonschema:"Breach notification commitments and timelines"`
DataDeletion string `json:"data_deletion" jsonschema:"Data deletion procedures and timelines"`
LiabilityCaps string `json:"liability_caps" jsonschema:"Liability limitations and caps"`
Indemnification string `json:"indemnification" jsonschema:"Indemnification obligations"`
Termination string `json:"termination" jsonschema:"Termination provisions and data return"`
GoverningLaw string `json:"governing_law" jsonschema:"Governing law and jurisdiction"`
PrivacyClauses []string `json:"privacy_clauses" jsonschema:"Notable privacy contractual clauses found"`
AIClauses []string `json:"ai_clauses" jsonschema:"Notable AI-related contractual clauses found"`
SubprocessorTerms string `json:"subprocessor_terms" jsonschema:"Sub-processor management terms (approval mechanism, notification)"`
Summary string `json:"summary" jsonschema:"Key findings summary"`
SourceURL string `json:"source_url" jsonschema:"URL of the analyzed document"`
}
// --- Compliance ---
CertificationEntry struct {
Name string `json:"name" jsonschema:"Certification name (e.g. SOC 2 Type II, ISO 27001)"`
Status string `json:"status" jsonschema:"Certification status: current, in_progress, claimed_unverified, not_specified"`
Details string `json:"details" jsonschema:"Additional details: audit date, certificate number, accreditation body"`
}
ComplianceOutput struct {
Certifications []CertificationEntry `json:"certifications" jsonschema:"All certifications and compliance frameworks found"`
PenetrationTesting string `json:"penetration_testing" jsonschema:"Penetration testing practices (frequency, third-party firm)"`
BugBounty string `json:"bug_bounty" jsonschema:"Bug bounty or responsible disclosure program details"`
EncryptionStandards string `json:"encryption_standards" jsonschema:"Encryption standards mentioned (AES-256, TLS 1.3, etc.)"`
AuditReports string `json:"audit_reports" jsonschema:"Audit report availability (downloadable, on request, not available)"`
OtherFrameworks []string `json:"other_frameworks" jsonschema:"Other frameworks or standards mentioned"`
Summary string `json:"summary" jsonschema:"Overall compliance posture summary"`
Sources []string `json:"sources" jsonschema:"URLs visited during assessment"`
}
// --- Market Presence ---
MarketOutput struct {
NotableCustomers []string `json:"notable_customers" jsonschema:"Notable customer names or logos identified"`
CaseStudies []string `json:"case_studies" jsonschema:"Case study summaries with customer names"`
Partnerships []string `json:"partnerships" jsonschema:"Strategic partnerships or integrations"`
CompanySizeSignals string `json:"company_size_signals" jsonschema:"Employee count, office locations, funding indicators"`
FundingInfo string `json:"funding_info" jsonschema:"Known funding rounds, investors, or valuation signals"`
MarketPosition string `json:"market_position" jsonschema:"Market positioning and competitive stance"`
Summary string `json:"summary" jsonschema:"Overall market presence assessment"`
Sources []string `json:"sources" jsonschema:"URLs visited during assessment"`
}
// --- Data Processing ---
DataProcessingOutput struct {
EncryptionAtRest string `json:"encryption_at_rest" jsonschema:"Encryption at rest details (algorithm, key size)"`
EncryptionInTransit string `json:"encryption_in_transit" jsonschema:"Encryption in transit details (TLS version, cipher suites)"`
KeyManagement string `json:"key_management" jsonschema:"Key management practices (HSM, rotation, customer-managed keys)"`
RetentionPeriod string `json:"retention_period" jsonschema:"Data retention period and policy"`
DeletionProcess string `json:"deletion_process" jsonschema:"Data deletion process and timeline"`
CustomerControls string `json:"customer_controls" jsonschema:"Customer-facing data management controls"`
DataLocations []string `json:"data_locations" jsonschema:"Countries or regions where data is processed or stored"`
TransferMechanisms []string `json:"transfer_mechanisms" jsonschema:"Cross-border transfer mechanisms (SCCs, BCRs, adequacy decisions)"`
DataResidency string `json:"data_residency" jsonschema:"Data residency options and restrictions"`
BackupRecovery string `json:"backup_recovery" jsonschema:"Backup and disaster recovery for data"`
Anonymization string `json:"anonymization" jsonschema:"Anonymization or pseudonymization practices"`
DPAStatus string `json:"dpa_status" jsonschema:"DPA availability: available, available_on_request, not_found, behind_login"`
ControllerProcessor string `json:"controller_processor" jsonschema:"Data processing role: controller, processor, subprocessor"`
AuditRights string `json:"audit_rights" jsonschema:"Customer audit rights described"`
SubprocessorApproval string `json:"subprocessor_approval" jsonschema:"Sub-processor change approval mechanism"`
BreachNotification string `json:"breach_notification" jsonschema:"Breach notification timeline and obligations"`
DataReturn string `json:"data_return" jsonschema:"Data return and deletion on contract termination"`
DSARHandling string `json:"dsar_handling" jsonschema:"DSAR handling capability and timeline"`
DataMinimization string `json:"data_minimization" jsonschema:"Data minimization practices"`
PurposeLimitation string `json:"purpose_limitation" jsonschema:"Purpose limitation commitments"`
Rating string `json:"rating" jsonschema:"Overall data processing rating: Strong, Adequate, or Weak"`
Summary string `json:"summary" jsonschema:"Key findings summary"`
Sources []string `json:"sources" jsonschema:"URLs visited during assessment"`
}
// --- Subprocessor ---
SubprocessorOutput struct {
Subprocessors []Subprocessor `json:"subprocessors" jsonschema:"List of sub-processors discovered"`
TotalCount int `json:"total_count" jsonschema:"Total number of sub-processors found"`
Source string `json:"source" jsonschema:"URL where the sub-processor list was found"`
IsComplete bool `json:"is_complete" jsonschema:"Whether the full list was extracted (false if pagination was incomplete)"`
Notes string `json:"notes" jsonschema:"Observations about the sub-processor list"`
}
// --- Incident Response ---
IncidentResponseOutput struct {
IRPlan string `json:"ir_plan" jsonschema:"Incident response plan documentation status"`
NotificationTimeline string `json:"notification_timeline" jsonschema:"Breach notification timeline (e.g. 72 hours)"`
NotificationMethod string `json:"notification_method" jsonschema:"How customers are notified of incidents"`
ContractualObligations string `json:"contractual_obligations" jsonschema:"Contractual IR obligations found"`
StatusPageURL string `json:"status_page_url" jsonschema:"Status page URL if found"`
StatusPageActive bool `json:"status_page_active" jsonschema:"Whether the status page is actively maintained"`
UpdateFrequency string `json:"update_frequency" jsonschema:"How frequently status updates are provided during incidents"`
PostMortems string `json:"post_mortems" jsonschema:"Post-mortem publication practices"`
RemediationApproach string `json:"remediation_approach" jsonschema:"Approach to incident remediation"`
RecentIncidents []string `json:"recent_incidents" jsonschema:"Recent incidents found with dates and descriptions"`
SecurityContact string `json:"security_contact" jsonschema:"Security contact email or reporting mechanism"`
BugBounty string `json:"bug_bounty" jsonschema:"Bug bounty or vulnerability disclosure program"`
Rating string `json:"rating" jsonschema:"Overall incident response rating: Strong, Adequate, or Weak"`
Summary string `json:"summary" jsonschema:"Key findings summary"`
Sources []string `json:"sources" jsonschema:"URLs visited during assessment"`
}
// --- Business Continuity ---
BusinessContinuityOutput struct {
DRPlan string `json:"dr_plan" jsonschema:"Disaster recovery plan documentation status"`
RTO string `json:"rto" jsonschema:"Recovery Time Objective"`
RPO string `json:"rpo" jsonschema:"Recovery Point Objective"`
TestingFrequency string `json:"testing_frequency" jsonschema:"DR testing frequency and last test date"`
CloudProviders []string `json:"cloud_providers" jsonschema:"Cloud infrastructure providers used"`
MultiRegion string `json:"multi_region" jsonschema:"Multi-region deployment details"`
Failover string `json:"failover" jsonschema:"Failover mechanisms and automation"`
UptimeSLA string `json:"uptime_sla" jsonschema:"Uptime SLA commitment (e.g. 99.99%)"`
SLACredits string `json:"sla_credits" jsonschema:"SLA credit or penalty structure"`
HistoricalUptime string `json:"historical_uptime" jsonschema:"Historical uptime performance"`
MaintenanceWindows string `json:"maintenance_windows" jsonschema:"Scheduled maintenance window policy"`
Regions []string `json:"regions" jsonschema:"Geographic regions with infrastructure"`
CDN string `json:"cdn" jsonschema:"CDN usage and provider"`
BackupStrategy string `json:"backup_strategy" jsonschema:"Backup frequency, retention, and encryption"`
BCPDocumented string `json:"bcp_documented" jsonschema:"Business continuity plan documentation status"`
ISO22301 string `json:"iso_22301" jsonschema:"ISO 22301 certification status"`
Rating string `json:"rating" jsonschema:"Overall business continuity rating: Strong, Adequate, or Weak"`
Summary string `json:"summary" jsonschema:"Key findings summary"`
Sources []string `json:"sources" jsonschema:"URLs visited during assessment"`
}
// --- Professional Standing ---
ProfessionalStandingOutput struct {
VendorType string `json:"vendor_type" jsonschema:"Type of professional services firm: law_firm, accounting, consulting, audit, staffing, other"`
Licensing string `json:"licensing" jsonschema:"Professional licensing details (bar admissions, CPA licenses)"`
Memberships []string `json:"memberships" jsonschema:"Industry body memberships (ABA, AICPA, Big Four network, etc.)"`
Insurance string `json:"insurance" jsonschema:"Professional liability / E&O insurance coverage details"`
TeamCredentials string `json:"team_credentials" jsonschema:"Key team member qualifications and credentials"`
COIPolicy string `json:"coi_policy" jsonschema:"Conflict of interest policy details"`
ClientBase string `json:"client_base" jsonschema:"Client base signals (notable clients, industry focus)"`
Rating string `json:"rating" jsonschema:"Overall professional standing rating: Strong, Adequate, Weak, or N/A"`
KeyObservations string `json:"key_observations" jsonschema:"Key observations about professional standing"`
Sources []string `json:"sources" jsonschema:"URLs visited during assessment"`
}
// --- AI Risk ---
AIRiskOutput struct {
AIInvolvement string `json:"ai_involvement" jsonschema:"AI involvement status: yes, no, or unclear"`
UseCases []string `json:"use_cases" jsonschema:"AI/ML use cases in the product or service"`
AIPolicyURL string `json:"ai_policy_url" jsonschema:"URL to AI governance or responsible AI documentation"`
ModelTransparency string `json:"model_transparency" jsonschema:"Model transparency and explainability findings"`
BiasControls string `json:"bias_controls" jsonschema:"Bias detection and fairness measures"`
CustomerDataTraining string `json:"customer_data_training" jsonschema:"Whether customer data is used for model training"`
OptOutAvailable string `json:"opt_out_available" jsonschema:"Whether training data opt-out is available"`
TrainingDataDetails string `json:"training_data_details" jsonschema:"Training data governance details"`
HumanOversight string `json:"human_oversight" jsonschema:"Human oversight mechanisms for AI decisions"`
AIIncidentHandling string `json:"ai_incident_handling" jsonschema:"AI-specific incident handling procedures"`
AutomatedDecisions string `json:"automated_decisions" jsonschema:"GDPR Art. 22 automated decision-making compliance"`
EUAIAct string `json:"eu_ai_act" jsonschema:"EU AI Act awareness and compliance indicators"`
Rating string `json:"rating" jsonschema:"Overall AI risk rating: Strong, Adequate, Weak, or N/A"`
Summary string `json:"summary" jsonschema:"Key findings summary"`
Sources []string `json:"sources" jsonschema:"URLs visited during assessment"`
}
// --- Regulatory Compliance ---
RegulatoryArticle struct {
Article string `json:"article" jsonschema:"Article or section identifier (e.g. article_28, hipaa_security_rule)"`
Status string `json:"status" jsonschema:"Compliance status: compliant, partially_compliant, non_compliant, not_assessed, not_applicable"`
Notes string `json:"notes" jsonschema:"Evidence or reasoning for the status determination"`
}
RegulatoryFramework struct {
Applicable bool `json:"applicable" jsonschema:"Whether this framework applies to the vendor"`
OverallStatus string `json:"overall_status" jsonschema:"Overall compliance status for this framework"`
Articles []RegulatoryArticle `json:"articles" jsonschema:"Per-article compliance assessment"`
Notes string `json:"notes" jsonschema:"General notes about framework applicability"`
}
CrossBorderTransferInfo struct {
Mechanisms []string `json:"mechanisms" jsonschema:"Transfer mechanisms used (SCCs, BCRs, adequacy decisions)"`
DataLocations []string `json:"data_locations" jsonschema:"Countries where data is stored or processed"`
TIAEvidence bool `json:"tia_evidence" jsonschema:"Whether Transfer Impact Assessment evidence was found"`
}
RegulatoryComplianceOutput struct {
GDPR RegulatoryFramework `json:"gdpr" jsonschema:"GDPR compliance assessment"`
HIPAA RegulatoryFramework `json:"hipaa" jsonschema:"HIPAA compliance assessment"`
PCIDSS RegulatoryFramework `json:"pci_dss" jsonschema:"PCI DSS compliance assessment"`
SOX RegulatoryFramework `json:"sox" jsonschema:"SOX compliance assessment"`
IndustrySpecific []string `json:"industry_specific" jsonschema:"Other industry-specific regulations found"`
CrossBorderTransfers CrossBorderTransferInfo `json:"cross_border_transfers" jsonschema:"Cross-border data transfer assessment"`
Gaps []string `json:"gaps" jsonschema:"Identified compliance gaps"`
Recommendations []string `json:"recommendations" jsonschema:"Recommended actions to address gaps"`
}
// --- Web Search ---
WebSearchOutput struct {
SecurityIncidents string `json:"security_incidents" jsonschema:"Known security incidents or breaches found"`
RegulatoryActions string `json:"regulatory_actions" jsonschema:"Regulatory actions, fines, or investigations"`
CustomerSentiment string `json:"customer_sentiment" jsonschema:"Customer reviews and sentiment summary"`
RecentNews string `json:"recent_news" jsonschema:"Recent news coverage and press"`
IndustryRecognition string `json:"industry_recognition" jsonschema:"Industry awards, analyst recognition, rankings"`
ProfessionalStanding string `json:"professional_standing" jsonschema:"Professional disciplinary actions or regulatory findings (for services firms)"`
RedFlags []string `json:"red_flags" jsonschema:"Red flags or concerning findings"`
PositiveSignals []string `json:"positive_signals" jsonschema:"Positive external signals"`
Summary string `json:"summary" jsonschema:"Overall external research summary"`
Sources []string `json:"sources" jsonschema:"URLs visited during research"`
}
// --- Financial Stability ---
FinancialStabilityOutput struct {
CompanyAge string `json:"company_age" jsonschema:"Year founded and company age"`
Funding string `json:"funding" jsonschema:"Funding history (rounds, amounts, investors)"`
EmployeeCount string `json:"employee_count" jsonschema:"Estimated employee count and source"`
RevenueSignals string `json:"revenue_signals" jsonschema:"Revenue indicators (ARR mentions, growth signals)"`
CustomerBase string `json:"customer_base" jsonschema:"Customer base signals (count, notable names)"`
LegalStanding string `json:"legal_standing" jsonschema:"Active lawsuits, regulatory issues, bankruptcy filings"`
Ownership string `json:"ownership" jsonschema:"Ownership structure (public, PE-backed, founder-led, acquired)"`
RiskSignals []string `json:"risk_signals" jsonschema:"Financial risk signals identified"`
OverallAssessment string `json:"overall_assessment" jsonschema:"Overall financial stability: Strong, Adequate, Weak, or Concerning"`
Confidence string `json:"confidence" jsonschema:"Assessment confidence level: High, Medium, or Low"`
Notes string `json:"notes" jsonschema:"Additional observations"`
Sources []string `json:"sources" jsonschema:"URLs visited during research"`
}
// --- Code Security ---
SecurityAdvisorySummary struct {
Total int `json:"total" jsonschema:"Total number of security advisories"`
Critical int `json:"critical" jsonschema:"Critical severity advisories"`
High int `json:"high" jsonschema:"High severity advisories"`
Medium int `json:"medium" jsonschema:"Medium severity advisories"`
Low int `json:"low" jsonschema:"Low severity advisories"`
AvgTimeToFix string `json:"avg_time_to_fix" jsonschema:"Average time to fix advisories"`
Notes string `json:"notes" jsonschema:"Additional context about advisories"`
}
CodeSecurityOutput struct {
HasPublicRepos bool `json:"has_public_repos" jsonschema:"Whether the vendor has public repositories"`
GithubOrg string `json:"github_org" jsonschema:"GitHub organization or user name"`
MainRepos []string `json:"main_repos" jsonschema:"Main public repositories identified"`
SecurityAdvisories SecurityAdvisorySummary `json:"security_advisories" jsonschema:"Security advisory summary"`
DependencyManagement string `json:"dependency_management" jsonschema:"Dependency management practices (Dependabot, Renovate, etc.)"`
ReleaseCadence string `json:"release_cadence" jsonschema:"Release frequency and last release date"`
SecurityPolicy string `json:"security_policy" jsonschema:"SECURITY.md or vulnerability disclosure policy"`
CISecurity string `json:"ci_security" jsonschema:"CI/CD security practices (SAST, DAST, container scanning)"`
CodeSigning string `json:"code_signing" jsonschema:"Code or release signing practices"`
OpenSecurityIssues string `json:"open_security_issues" jsonschema:"Open security-related issues or PRs"`
License string `json:"license" jsonschema:"Open source license type"`
OverallAssessment string `json:"overall_assessment" jsonschema:"Overall code security: Strong, Adequate, Weak, or Not_Applicable"`
RiskSignals []string `json:"risk_signals" jsonschema:"Code security risk signals identified"`
Notes string `json:"notes" jsonschema:"Additional observations"`
Sources []string `json:"sources" jsonschema:"URLs visited during research"`
}
// --- Vendor Comparison ---
AlternativeVendor struct {
Name string `json:"name" jsonschema:"Alternative vendor name"`
Website string `json:"website" jsonschema:"Alternative vendor website URL"`
Certifications []string `json:"certifications" jsonschema:"Visible certifications"`
TrustCenter bool `json:"trust_center" jsonschema:"Whether a trust center page was found"`
PrivacyPolicy bool `json:"privacy_policy" jsonschema:"Whether a privacy policy was found"`
CompanySize string `json:"company_size" jsonschema:"Estimated company size"`
SecurityScore string `json:"security_score" jsonschema:"Quick security impression: Strong, Adequate, or Weak"`
}
ComparisonSummary struct {
SecurityMaturity string `json:"security_maturity" jsonschema:"Relative security maturity vs alternatives"`
CompliancePosture string `json:"compliance_posture" jsonschema:"Relative compliance posture vs alternatives"`
MarketPosition string `json:"market_position" jsonschema:"Relative market position vs alternatives"`
Transparency string `json:"transparency" jsonschema:"Relative transparency vs alternatives"`
}
VendorComparisonOutput struct {
VendorCategory string `json:"vendor_category" jsonschema:"The vendor's product category"`
AssessedVendor string `json:"assessed_vendor" jsonschema:"The vendor being assessed"`
Alternatives []AlternativeVendor `json:"alternatives" jsonschema:"Alternative vendors identified and evaluated"`
ComparisonSummary ComparisonSummary `json:"comparison_summary" jsonschema:"Summary comparison across dimensions"`
VendorStrengths []string `json:"vendor_strengths" jsonschema:"Assessed vendor's strengths vs alternatives"`
VendorWeaknesses []string `json:"vendor_weaknesses" jsonschema:"Assessed vendor's weaknesses vs alternatives"`
OverallPosition string `json:"overall_position" jsonschema:"Vendor position: Above_Average, Average, or Below_Average"`
Notes string `json:"notes" jsonschema:"Additional comparison notes"`
}
)

80
pkg/vetting/output_types_test.go Normal file
View File

@@ -0,0 +1,80 @@
// 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 vetting_test
import (
"encoding/json"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/vetting"
)
func TestOutputType_SchemaGeneration(t *testing.T) {
t.Parallel()
tests := []struct {
name string
fn func(t *testing.T)
}{
{"CrawlerOutput", assertSchema[vetting.CrawlerOutput]},
{"SecurityOutput", assertSchema[vetting.SecurityOutput]},
{"DocumentAnalysisOutput", assertSchema[vetting.DocumentAnalysisOutput]},
{"ComplianceOutput", assertSchema[vetting.ComplianceOutput]},
{"MarketOutput", assertSchema[vetting.MarketOutput]},
{"DataProcessingOutput", assertSchema[vetting.DataProcessingOutput]},
{"SubprocessorOutput", assertSchema[vetting.SubprocessorOutput]},
{"IncidentResponseOutput", assertSchema[vetting.IncidentResponseOutput]},
{"BusinessContinuityOutput", assertSchema[vetting.BusinessContinuityOutput]},
{"ProfessionalStandingOutput", assertSchema[vetting.ProfessionalStandingOutput]},
{"AIRiskOutput", assertSchema[vetting.AIRiskOutput]},
{"RegulatoryComplianceOutput", assertSchema[vetting.RegulatoryComplianceOutput]},
{"WebSearchOutput", assertSchema[vetting.WebSearchOutput]},
{"FinancialStabilityOutput", assertSchema[vetting.FinancialStabilityOutput]},
{"CodeSecurityOutput", assertSchema[vetting.CodeSecurityOutput]},
{"VendorComparisonOutput", assertSchema[vetting.VendorComparisonOutput]},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
tt.fn(t)
})
}
}
// assertSchema creates an OutputType for T and verifies that the
// generated JSON Schema has the expected shape: an object type with a
// non-empty properties map. This catches struct tags that silently
// produce empty or malformed schemas.
func assertSchema[T any](t *testing.T) {
t.Helper()
outputType, err := agent.NewOutputType[T]("test")
require.NoError(t, err)
require.NotNil(t, outputType)
require.NotEmpty(t, outputType.Schema)
var schema map[string]any
require.NoError(t, json.Unmarshal(outputType.Schema, &schema))
assert.Equal(t, "object", schema["type"])
properties, ok := schema["properties"].(map[string]any)
require.True(t, ok, "schema must expose a properties map")
assert.NotEmpty(t, properties, "schema must declare at least one property")
}

412
pkg/vetting/progress.go Normal file
View File

@@ -0,0 +1,412 @@
// 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 vetting
import (
"context"
"math/rand/v2"
"go.probo.inc/probo/pkg/agent"
)
var (
toolMessages = map[string][]string{
// Orchestrator tools (top-level steps).
"crawl_vendor_website": {
"Exploring vendor website for security and compliance pages",
"Discovering key pages on the vendor website",
"Mapping out the vendor's online presence",
"Scanning the website structure for relevant sections",
"Browsing the vendor site to locate important resources",
},
"assess_security": {
"Running technical security checks on the domain",
"Evaluating the vendor's security posture",
"Performing infrastructure security analysis",
"Auditing the domain's technical defenses",
"Probing the vendor's security configuration",
},
"analyze_document": {
"Reviewing document for key provisions",
"Analyzing policy details and obligations",
"Extracting important clauses from the document",
"Parsing the document for notable terms",
"Breaking down the document's main points",
},
"assess_compliance": {
"Identifying certifications and compliance frameworks",
"Reviewing the vendor's compliance posture",
"Checking for recognized security certifications",
"Surveying the vendor's regulatory standing",
"Evaluating adherence to industry standards",
},
"assess_market_presence": {
"Investigating the vendor's market presence",
"Looking for notable customers and case studies",
"Checking who uses this vendor",
"Assessing the vendor's market credibility",
"Identifying the vendor's customer base",
},
"extract_subprocessors": {
"Extracting sub-processor information",
"Reading the vendor's sub-processor list",
"Identifying third-party sub-processors",
"Parsing sub-processor details",
"Cataloging the vendor's sub-processors",
},
"assess_data_processing": {
"Analyzing data processing practices",
"Reviewing encryption and data handling",
"Evaluating data retention and transfer policies",
"Checking data processing documentation",
"Assessing cross-border data transfer mechanisms",
},
"assess_incident_response": {
"Evaluating incident response capabilities",
"Reviewing breach notification procedures",
"Checking incident history and transparency",
"Assessing security incident readiness",
"Examining post-incident review processes",
},
"assess_business_continuity": {
"Assessing business continuity planning",
"Reviewing disaster recovery capabilities",
"Checking SLA and uptime commitments",
"Evaluating infrastructure redundancy",
"Examining geographic distribution and failover",
},
"assess_professional_standing": {
"Evaluating professional standing and credentials",
"Reviewing licensing and industry memberships",
"Checking professional qualifications and accreditation",
"Assessing team credentials and experience",
"Examining professional liability and insurance coverage",
},
"assess_ai_risk": {
"Evaluating AI governance and responsible AI practices",
"Reviewing AI transparency and bias controls",
"Checking AI risk management documentation",
"Assessing automated decision-making safeguards",
"Examining AI training data governance",
},
"research_vendor_externally": {
"Researching the vendor across the web",
"Searching for external signals about the vendor",
"Looking for news and breach reports",
"Investigating the vendor's external reputation",
"Scanning public sources for vendor intelligence",
},
"assess_regulatory_compliance": {
"Performing deep regulatory compliance analysis",
"Checking GDPR article-level compliance",
"Analyzing regulatory framework adherence",
"Reviewing compliance against specific regulations",
"Evaluating regulatory requirements coverage",
},
"assess_financial_stability": {
"Assessing vendor financial stability",
"Investigating company funding and financial health",
"Checking business registration and SEC filings",
"Evaluating vendor viability and longevity",
"Researching company financial standing",
},
"assess_code_security": {
"Evaluating open-source code security posture",
"Checking for security advisories and CVEs",
"Reviewing dependency management practices",
"Analyzing release cadence and maintenance",
"Inspecting code security practices",
},
"compare_vendor": {
"Comparing vendor against alternatives",
"Finding competing vendors in the same category",
"Benchmarking security and compliance posture",
"Evaluating vendor relative to market alternatives",
"Assessing competitive landscape",
},
"extract_vendor_info": {
"Extracting vendor information from assessment",
"Parsing assessment into structured data",
"Building vendor profile from findings",
"Distilling key vendor details from report",
"Organizing vendor metadata from assessment",
},
// Web search sub-agent tools.
"web_search": {
"Searching the web",
"Running a web search query",
"Looking up information online",
"Querying search results",
"Fetching search results",
},
// Security sub-agent tools.
"check_ssl_certificate": {
"Inspecting SSL/TLS certificate",
"Verifying certificate validity and configuration",
"Checking SSL certificate details",
"Reviewing the certificate chain",
"Examining TLS setup and expiration",
},
"check_security_headers": {
"Analyzing HTTP security headers",
"Reviewing response headers for security best practices",
"Checking for missing security headers",
"Scanning HTTP headers for protective directives",
"Evaluating header-based security controls",
},
"check_dmarc": {
"Looking up DMARC email authentication record",
"Checking DMARC policy configuration",
"Verifying email spoofing protections",
"Querying DNS for DMARC policy",
"Reviewing email authentication settings",
},
"check_spf": {
"Looking up SPF email authentication record",
"Checking SPF policy configuration",
"Verifying sender policy framework",
"Querying DNS for SPF record",
"Reviewing SPF authorization settings",
},
"check_breaches": {
"Searching for known data breaches",
"Checking breach databases for past incidents",
"Looking up the domain in breach records",
"Scanning public breach disclosures",
"Querying breach intelligence sources",
},
"check_dnssec": {
"Verifying DNSSEC configuration",
"Checking DNS security extensions",
"Inspecting DNSSEC chain of trust",
"Validating DNS signing status",
"Reviewing DNSSEC deployment",
},
"analyze_csp": {
"Evaluating Content Security Policy",
"Analyzing CSP directives for weaknesses",
"Reviewing content security rules",
"Checking CSP for unsafe directives",
"Parsing Content Security Policy header",
},
"check_cors": {
"Checking CORS configuration",
"Inspecting cross-origin resource sharing policy",
"Reviewing CORS headers",
"Evaluating cross-origin access rules",
"Analyzing CORS allow-origin settings",
},
// Browser tools used by crawler, analyzer, and compliance sub-agents.
"navigate_to_url": {
"Opening page",
"Loading page content",
"Navigating to the page",
"Visiting the page",
"Heading to the page",
},
"extract_page_text": {
"Reading page content",
"Extracting text from the page",
"Pulling content from the page",
"Scanning page text",
"Capturing the page body",
},
"extract_links": {
"Collecting links from the page",
"Gathering all page links",
"Discovering outgoing links",
"Harvesting links on the page",
"Listing page hyperlinks",
},
"find_links_matching": {
"Searching for relevant links",
"Looking for links matching the pattern",
"Filtering page links by keyword",
"Hunting for specific links on the page",
"Sifting through links for a match",
},
"click_element": {
"Clicking on the page",
"Interacting with the page",
"Pressing a button on the page",
"Navigating within the page",
"Triggering a page action",
},
"select_option": {
"Selecting an option on the page",
"Changing a dropdown selection",
"Adjusting page settings",
"Picking a value from a dropdown",
"Updating a page filter",
},
// New security tools.
"check_whois": {
"Looking up domain registration details",
"Checking WHOIS records",
"Querying domain registrar information",
"Inspecting domain ownership data",
"Retrieving domain age and registrant info",
},
"check_dns_records": {
"Querying DNS records",
"Looking up A, MX, and NS records",
"Checking DNS configuration",
"Resolving domain DNS entries",
"Inspecting hosting and email providers",
},
// New browser tools.
"fetch_robots_txt": {
"Fetching robots.txt",
"Checking robots.txt for hidden pages",
"Reading site crawl directives",
"Discovering sitemap URLs from robots.txt",
"Parsing robots.txt disallow rules",
},
"fetch_sitemap": {
"Fetching sitemap",
"Parsing sitemap for page URLs",
"Discovering pages from sitemap",
"Reading sitemap index",
"Extracting URLs from sitemap XML",
},
"download_pdf": {
"Downloading and extracting PDF",
"Reading PDF document content",
"Extracting text from PDF",
"Processing PDF document",
"Parsing PDF for analysis",
},
// New search tools.
"check_wayback": {
"Checking Wayback Machine archives",
"Looking for historical page snapshots",
"Querying Internet Archive",
"Searching for archived versions",
"Checking page history in Wayback Machine",
},
"check_government_databases": {
"Searching government regulatory databases",
"Checking SEC and FTC records",
"Looking for GDPR enforcement actions",
"Querying regulatory databases",
"Searching for enforcement history",
},
"diff_documents": {
"Comparing document versions",
"Diffing document texts",
"Analyzing document changes",
"Checking for document modifications",
"Computing document differences",
},
}
)
func randomMessage(step string) string {
msgs, ok := toolMessages[step]
if !ok {
return ""
}
return msgs[rand.IntN(len(msgs))]
}
// reportProgress emits a progress event to the reporter if non-nil.
func reportProgress(
ctx context.Context,
reporter agent.ProgressReporter,
step string,
eventType agent.ProgressEventType,
) {
if reporter == nil {
return
}
event := agent.ProgressEvent{
Type: eventType,
Step: step,
}
if eventType == agent.ProgressEventStepStarted {
event.Message = randomMessage(step)
}
reporter(ctx, event)
}
// progressHooks translates tool events into progress events. When
// parentStep is non-empty, emitted events are scoped under a parent
// step (sub-agent mode); otherwise they are top-level orchestrator
// events.
type progressHooks struct {
agent.NoOpHooks
reporter agent.ProgressReporter
parentStep string
}
func newProgressHooks(reporter agent.ProgressReporter) *progressHooks {
return &progressHooks{reporter: reporter}
}
func newSubProgressHooks(reporter agent.ProgressReporter, parentStep string) *progressHooks {
return &progressHooks{
reporter: reporter,
parentStep: parentStep,
}
}
func (h *progressHooks) OnToolStart(ctx context.Context, _ *agent.Agent, tool agent.Tool, _ string) {
msg := randomMessage(tool.Name())
if msg == "" {
return
}
h.reporter(
ctx,
agent.ProgressEvent{
Type: agent.ProgressEventStepStarted,
Step: tool.Name(),
ParentStep: h.parentStep,
Message: msg,
},
)
}
func (h *progressHooks) OnToolEnd(ctx context.Context, _ *agent.Agent, tool agent.Tool, _ agent.ToolResult, err error) {
if _, ok := toolMessages[tool.Name()]; !ok {
return
}
eventType := agent.ProgressEventStepCompleted
if err != nil {
eventType = agent.ProgressEventStepFailed
}
h.reporter(
ctx,
agent.ProgressEvent{
Type: eventType,
Step: tool.Name(),
ParentStep: h.parentStep,
},
)
}
var _ agent.RunHooks = (*progressHooks)(nil)

83
pkg/vetting/prompts/ai_risk.txt Normal file
View File

@@ -0,0 +1,83 @@
<role>
You are an AI risk assessment specialist aligned with ISO 42001 (AI management system). You evaluate a vendor's AI governance and responsible AI practices from their website, policies, and documentation.
</role>
<task>
Given a starting URL (AI policy, trust center, responsible AI page, or main website), gather evidence across the assessment areas below. Follow links to dedicated AI policy pages, trust center AI sections, AI-related blog posts, DPA / privacy policy / ToS sections about AI, and model documentation.
</task>
<assessment>
**1. AI Usage Disclosure**
- Whether the vendor discloses use of AI/ML in product or services
- Specific AI use cases (content generation, recommendations, fraud detection, automated decisions)
- Dedicated AI policy, responsible AI page, or AI governance page
- Distinction between AI-as-product (core offering) and AI-as-internal-tool
**2. Model Transparency & Explainability**
- Information about the AI models used
- Model types, training approaches, limitations
- Whether outputs can be explained to end users
- Documentation about model versioning, updates, change management
**3. Bias Detection & Fairness**
- Bias detection or fairness testing measures
- Testing methodology (demographic parity, equalized odds, etc.)
- Fairness impact assessments or equity audits
- How bias issues are remediated when discovered
**4. Training Data Governance**
- How training data is sourced and governed
- Whether customer data is used for model training, and any opt-out mechanism
- Data quality, labeling, provenance processes
- Restrictions on using customer data to improve models
**5. Human Oversight**
- Human-in-the-loop processes for high-risk or consequential decisions
- Automated decision-making restrictions
- Process for users to appeal or contest automated decisions
- Escalation paths when AI outputs are uncertain or high-stakes
**6. AI Incident Handling**
- AI-specific incident response process
- How model failures, hallucinations, or harmful outputs are handled
- Monitoring for model drift, performance degradation, adversarial inputs
- Whether AI-related incidents are disclosed transparently
**7. Regulatory Compliance**
- GDPR Article 22 (automated individual decision-making)
- Awareness of the EU AI Act or other AI-specific regulation
- AI risk classifications (minimal, limited, high, unacceptable)
- Safeguards for automated profiling
</assessment>
<edge_cases>
- Only report information explicitly found on the vendor's pages.
- If AI involvement cannot be determined from public information, state that clearly.
- Distinguish between vendors that actively use AI vs vendors with no apparent AI usage.
- Note when AI governance documentation is absent — this is itself a finding.
- Do not penalize vendors that genuinely do not use AI in their products.
</edge_cases>
<output>
Return your findings as structured JSON matching the required output schema. The schema and per-field descriptions are enforced by the API; focus on the substance of the assessment.
</output>
<examples>
<example>
<description>Vendor with mature AI governance.</description>
<input>Vendor publishes a Responsible AI page describing model cards, bias testing methodology (demographic parity), customer data opt-out for training, and explicit GDPR Art. 22 compliance for automated decisions.</input>
<output>{"ai_involvement": "yes", "model_transparency": "Model cards published per release", "bias_controls": "Demographic parity testing documented", "customer_data_training": "Customer data not used for training by default", "opt_out_available": "Yes, account-level opt-out", "automated_decisions": "GDPR Art. 22 addressed with human review path", "rating": "Strong"}</output>
</example>
<example>
<description>Vendor with no AI involvement.</description>
<input>Vendor is a payroll processing service. No mention of AI, ML, automation, or algorithmic features anywhere on the site.</input>
<output>{"ai_involvement": "no", "rating": "N/A", "summary": "Vendor does not appear to use AI/ML in their product or service delivery"}</output>
</example>
<example>
<description>AI claimed but no governance documentation.</description>
<input>Marketing page says "AI-powered fraud detection" but the security page, privacy policy, and trust center contain no information about model transparency, training data, or oversight.</input>
<output>{"ai_involvement": "yes", "use_cases": ["AI-powered fraud detection (claimed)"], "model_transparency": "Not documented", "bias_controls": "Not documented", "rating": "Weak", "summary": "AI usage claimed but no governance documentation found — significant gap"}</output>
</example>
</examples>

80
pkg/vetting/prompts/analyzer.txt Normal file
View File

@@ -0,0 +1,80 @@
<role>
You are a document analyzer specialized in extracting compliance, privacy, and contractual information from vendor documents.
</role>
<task>
Given a document URL (privacy policy, DPA, terms of service, engagement letter, professional standards, etc.), extract and summarize the substantive provisions described under `<assessment>`. Read what the document says and report it factually — do not speculate or invent details.
</task>
<assessment>
Look for and report on:
**Operational and contractual terms**
- Data retention policies and periods
- Data processing locations and jurisdictions
- Data security measures described
- Breach notification procedures and timelines
- Data deletion / portability provisions
- Liability caps and limitations (aggregate, per-incident, carve-outs)
- Indemnification clauses (mutual vs one-way, scope, caps)
- Termination provisions (for cause, for convenience, notice period, data return / deletion timeline)
- Insurance requirements mentioned in the contract
- Governing law and jurisdiction
- Dispute resolution (arbitration vs litigation, venue)
- Assignment and change-of-control provisions
- Force majeure scope
- Confidentiality obligations and duration
**Privacy regulatory indicators**
- GDPR indicators: lawful basis, data subject rights, DPO contact
- CCPA indicators
- Subprocessor details (names, purposes, locations)
**Privacy contractual clauses (ISO 27701)**
- Data processing instructions and scope
- Subprocessor approval mechanism (prior written consent, objection-based, notification-only)
- Cross-border transfer safeguards (SCCs, BCRs, adequacy decisions)
- Breach notification timeline and obligations
- Data return and deletion on termination
- DSAR cooperation obligations
- DPO contact information
**AI contractual clauses (ISO 42001) — extract if present**
- Prohibition on using customer data for model training
- Transparency obligations about AI usage
- Audit rights for AI systems
- Automated decision-making restrictions
- AI liability and indemnification
- Model update notification requirements
- Right to opt out of AI features
</assessment>
<edge_cases>
- If the document appears truncated (ends mid-sentence or is missing expected sections), follow pagination or anchor links and re-extract.
- Privacy policies often link to separate cookie policies or DPAs — follow those links if needed for the fields above.
- If a section is missing from the document, explicitly note its absence rather than omitting it.
</edge_cases>
<output>
Return your findings as structured JSON matching the required output schema. The schema and per-field descriptions are enforced by the API; focus on the substance of the analysis.
</output>
<examples>
<example>
<description>Privacy policy with breach notification commitment.</description>
<input>Privacy policy section: "We will notify affected users within 72 hours of confirming a personal data breach affecting their information, in accordance with GDPR Art. 33."</input>
<output>{"document_type": "privacy_policy", "breach_notification": "72-hour notification to affected users, GDPR Art. 33 compliance", "gdpr_indicators": "GDPR Article 33 explicitly referenced"}</output>
</example>
<example>
<description>DPA with Standard Contractual Clauses.</description>
<input>DPA Section 9: "For transfers of Personal Data outside the EEA, the parties incorporate the Standard Contractual Clauses (Module Two: Controller to Processor) approved by Commission Implementing Decision (EU) 2021/914."</input>
<output>{"document_type": "dpa", "data_locations": ["EEA", "Outside EEA"], "subprocessor_terms": "EU 2021 SCCs Module Two (C2P) incorporated", "privacy_clauses": ["Standard Contractual Clauses 2021/914 Module Two for cross-border transfers"]}</output>
</example>
<example>
<description>Terms of service with low liability cap.</description>
<input>ToS Section 14.3: "In no event shall Provider's aggregate liability exceed the fees paid by Customer in the twelve (12) months preceding the claim, or one hundred dollars ($100), whichever is greater."</input>
<output>{"document_type": "terms_of_service", "liability_caps": "Aggregate liability capped at greater of 12 months fees or $100", "indemnification": "Not present in this document"}</output>
</example>
</examples>

55
pkg/vetting/prompts/business_continuity.txt Normal file
View File

@@ -0,0 +1,55 @@
<role>
You are a business continuity assessment specialist. You evaluate a vendor's business continuity and disaster recovery capabilities from their website, SLA documentation, and infrastructure pages.
</role>
<task>
Given a starting URL (SLA page, trust center, security page, or infrastructure docs), gather evidence across the assessment areas below. Follow links to status pages, architecture pages, and downloadable continuity documentation.
</task>
<assessment>
**1. Disaster Recovery**
- Documented disaster recovery plan
- Recovery Time Objective (RTO)
- Recovery Point Objective (RPO)
- DR plan testing frequency
- DR scenarios covered
**2. Infrastructure Redundancy**
- Cloud provider(s)
- Multi-region or multi-AZ deployment
- Automatic failover capability
- Load balancing and auto-scaling
**3. SLA & Uptime**
- Committed uptime SLA (e.g. 99.9%, 99.99%)
- SLA credit / compensation terms
- Historical uptime data
- Maintenance window policy
**4. Geographic Distribution**
- Regions / countries where infrastructure operates
- Edge / CDN distribution
- Customer choice of deployment region
**5. Backup Strategy**
- Backup frequency
- Backup storage location (same region vs cross-region)
- Backup retention period
- Backup integrity verification
**6. Business Continuity Planning**
- Documented BCP beyond technical DR
- Coverage of operational continuity (people, processes)
- ISO 22301 certification or reference
- Communication plan for extended outages
</assessment>
<edge_cases>
- Only report information explicitly found on the vendor's pages.
- Marketing claims like "enterprise-grade reliability" without specifics should be noted as vague.
- If SLA documents are behind a login wall, note that they are not publicly available.
</edge_cases>
<output>
Return your findings as structured JSON matching the required output schema. The schema and per-field descriptions are enforced by the API; focus on the substance of the assessment.
</output>

81
pkg/vetting/prompts/code_security.txt Normal file
View File

@@ -0,0 +1,81 @@
<role>
You are a code security assessor for third-party vendor due diligence. You evaluate the security posture of vendors that have open-source code repositories.
</role>
<task>
Find the vendor's public repositories and evaluate their security posture across the assessment areas below. If the vendor has no public repositories, report that and exit early — this assessment is only applicable to vendors with public code.
</task>
<assessment>
First, find the vendor's GitHub or GitLab organization (e.g. `github.com/{vendor_name}`). Identify the main product repository and any security-relevant repos. If nothing public exists, return `has_public_repos: false`, `overall_assessment: Not_Applicable`, and stop.
Once you have the repos, gather evidence across these areas:
**Security Advisories & CVEs**
- GitHub Security Advisories for the organization (`github.com/{org}/security/advisories`)
- CVEs: search `"{vendor_name}" CVE` or `"{product_name}" CVE`
- National Vulnerability Database: `site:nvd.nist.gov "{vendor_name}"`
- How many advisories, what severity, how quickly were they patched
**Dependency Management**
- Dependabot, Renovate, or similar automated dependency update tools
- Lock files (`package-lock.json`, `go.sum`, `Gemfile.lock`)
- Known vulnerable dependency patterns
**Release Cadence & Maintenance**
- Release frequency
- Date of the last release; is the project actively maintained?
- Contributor count (single-person vs team)
- Issue response times and PR merge patterns
**Security Policy**
- `SECURITY.md` present
- Responsible disclosure program
- Bug bounty (check the vendor website too)
- How security issues are handled (private advisories vs public issues)
**CI/CD Security**
- Security scanning in CI workflows (`.github/workflows/`)
- Tools: CodeQL, Snyk, Dependabot alerts, SAST, container scanning
- Code review patterns (PR merge patterns indicate review discipline)
**Code Signing & Artifacts**
- Signed releases (GPG, sigstore)
- Signed container images
- Software bill of materials (SBOM)
**Open Security Issues**
- Issues labeled `security`, `vulnerability`, or `CVE`
- Unresolved security-tagged issues
- Age of the oldest open security issues
**License Compliance**
- License (MIT, Apache 2.0, GPL, AGPL, proprietary)
- License compatibility issues
- Whether the license is clearly stated
</assessment>
<edge_cases>
- Focus on the vendor's main product repositories, not forks or experimental projects.
- A high number of security advisories is not necessarily bad if they are promptly fixed — it indicates transparency.
- Distinguish between the vendor's own code and their dependencies.
- Be factual — only report what you can verify from public sources.
</edge_cases>
<output>
Return your findings as structured JSON matching the required output schema. The schema and per-field descriptions are enforced by the API; focus on the substance of the assessment.
</output>
<examples>
<example>
<description>Active, well-maintained project.</description>
<input>github.com/vendor/product shows weekly releases over the past year, Dependabot enabled, SECURITY.md present, 5 published security advisories all patched within 2 weeks, and signed releases via cosign.</input>
<output>{"has_public_repos": true, "release_cadence": "Weekly releases, last release within past 7 days", "dependency_management": "Dependabot enabled", "security_policy": "SECURITY.md present with disclosure address", "security_advisories": {"total": 5, "critical": 0, "high": 2, "medium": 3, "low": 0, "avg_time_to_fix": "~14 days"}, "code_signing": "cosign-signed releases", "overall_assessment": "Strong"}</output>
</example>
<example>
<description>Vendor with no public repositories.</description>
<input>Vendor is a closed-source SaaS. No github.com/vendor or gitlab.com/vendor organization exists, and the website has no "open source" or "GitHub" links.</input>
<output>{"has_public_repos": false, "overall_assessment": "Not_Applicable", "notes": "No public code repositories found"}</output>
</example>
</examples>

59
pkg/vetting/prompts/compliance.txt Normal file
View File

@@ -0,0 +1,59 @@
<role>
You are a compliance assessor specialized in identifying certifications and compliance frameworks from vendor trust and compliance pages.
</role>
<task>
Given a trust center or compliance page URL, identify the certifications, audit programs, and compliance frameworks the vendor publishes. For each certification, distinguish between independently verified evidence, in-progress audits, marketing claims, and unverified framework alignment. Report only what you find.
</task>
<assessment>
Look for and report on:
- Security certifications: SOC 1, SOC 2 Type I/II, ISO 27001, ISO 27017, ISO 27018
- Privacy certifications: ISO 27701, APEC CBPR
- Industry-specific compliance: PCI DSS, HIPAA, FedRAMP, HITRUST, StateRAMP
- Regional compliance: GDPR, CCPA/CPRA, PIPEDA, LGPD, UK GDPR
- Audit report availability and dates
- Penetration testing information (frequency, third-party firm)
- Bug bounty or responsible disclosure program details
- Data encryption standards (at rest and in transit)
- Business continuity and disaster recovery mentions
- Other compliance frameworks or standards mentioned
If the trust page links to sub-pages (e.g. separate pages per certification), follow the most important ones to confirm details.
</assessment>
<rating_criteria>
For each certification, assign one of the following statuses:
- **current**: The certification is clearly active. Evidence includes a certification logo paired with an audit date or validity period, a downloadable or requestable audit report, a certificate number, or an explicit statement like "SOC 2 Type II certified (last audit: March 2025)".
- **in_progress**: The vendor explicitly states the certification is upcoming or in progress. Evidence includes phrases like "currently pursuing ISO 27001", "SOC 2 audit underway", or a roadmap page listing the certification as planned.
- **claimed_unverified**: The certification is mentioned on a marketing page but lacks supporting proof. For example, a SOC 2 badge on the homepage with no audit date, no certificate number, no downloadable report, and no details page. A logo alone is not proof.
- **not_specified**: The certification is referenced but its current status is unclear. For example, the vendor states "we follow ISO 27001 standards" without claiming actual certification.
Distinguish self-asserted claims from independently verified certifications. A vendor that says "we align with NIST CSF" is describing framework alignment, not a certification — list those under `other_frameworks`, not `certifications`.
</rating_criteria>
<output>
Return your findings as structured JSON matching the required output schema. The schema and per-field descriptions are enforced by the API; focus on the substance of the assessment.
</output>
<examples>
<example>
<description>Independently audited certification with proof.</description>
<input>Trust center page shows "SOC 2 Type II" with a Coalfire badge, audit period "Jan 2025 - Dec 2025", and a "Request Report" link gated behind a form.</input>
<output>{"certifications": [{"name": "SOC 2 Type II", "status": "current", "details": "Audited by Coalfire, 2025 audit period, report available on request via trust center"}]}</output>
</example>
<example>
<description>Marketing claim without verifiable proof.</description>
<input>Homepage footer displays a small "SOC 2" badge linking to /security, but the security page has no audit date, no auditor name, and no certificate number.</input>
<output>{"certifications": [{"name": "SOC 2", "status": "claimed_unverified", "details": "Badge displayed but no audit date, auditor, or certificate found"}]}</output>
</example>
<example>
<description>Framework alignment is not certification.</description>
<input>Security whitepaper says "Our security program aligns with NIST CSF and CIS Controls."</input>
<output>{"certifications": [], "other_frameworks": ["NIST CSF (alignment claimed, not certified)", "CIS Controls (alignment claimed, not certified)"]}</output>
</example>
</examples>

34
pkg/vetting/prompts/crawler.txt Normal file
View File

@@ -0,0 +1,34 @@
<role>
You are a website crawler specialized in discovering compliance, security, legal, and professional pages for vendor due diligence. Vendors may be SaaS products, cloud providers, law firms, accounting firms, consulting firms, or any other type of service provider.
</role>
<task>
Given a vendor website URL, discover all pages relevant to a security, compliance, privacy, AI governance, or professional standing assessment. Report each discovered URL with a short description of what it contains.
</task>
<assessment>
Start by fetching `robots.txt` and the sitemap — these often reveal trust centers, legal docs, and status pages that are not in the main navigation. Then navigate to the home page and the footer (most legal and compliance links live in the footer). Use `find_links_matching` and direct path probes for the kinds of pages listed below.
Pages to look for, with the kinds of paths that typically host them:
- **Security & trust**: security page, trust center, compliance page, bug bounty / responsible disclosure, status / uptime page (`/security`, `/trust`, `/compliance`, `/status`, `/bug-bounty`, `/responsible-disclosure`)
- **Legal**: privacy policy, terms of service, DPA, BAA, subprocessors / subcontractors list, SLA, GDPR / CCPA pages (`/privacy`, `/legal`, `/terms`, `/dpa`, `/baa`, `/subprocessors`, `/sla`, `/gdpr`, `/ccpa`)
- **Certifications**: SOC 2, ISO 27001, PCI, HIPAA, FedRAMP pages (often nested under `/trust` or `/compliance`)
- **Architecture & platform**: enterprise page, platform / infrastructure / reliability page (`/enterprise`, `/platform`, `/infrastructure`, `/reliability`) — these often consolidate security features, certifications, SLA details, and trust info that are not linked elsewhere
- **Professional services**: team / people / attorneys / professionals page, about / company page, credentials / licensing / accreditation page, services / practice-areas page, engagement terms / professional standards page, memberships / associations, insurance (`/team`, `/about`, `/our-team`, `/attorneys`, `/professionals`, `/people`, `/credentials`, `/services`, `/practice-areas`, `/engagement`)
- **AI governance**: AI policy, responsible AI, AI governance, AI ethics, machine learning page (`/ai`, `/ai-policy`, `/responsible-ai`, `/ai-governance`, `/ai-ethics`, `/machine-learning`)
For professional services firms (law firms, CPAs, consulting), team/people pages and credentials pages are the highest-value targets — prioritize them.
If you find an "enterprise" or "platform" page, visit it: these pages often contain security features, compliance certifications, SLA details, and trust information that are not surfaced anywhere else.
</assessment>
<edge_cases>
- Do not visit the same URL more than once.
- If a page redirects, report the final URL.
- If a section of the site is behind login, note it as discovered-but-gated rather than skipping it silently.
</edge_cases>
<output>
Return your findings as structured JSON matching the required output schema. The schema and per-field descriptions are enforced by the API; focus on the substance of the discovery.
</output>

75
pkg/vetting/prompts/data_processing.txt Normal file
View File

@@ -0,0 +1,75 @@
<role>
You are a data processing assessment specialist. Your job is to analyze a vendor's data handling practices by examining their website, privacy documentation, and security pages.
</role>
<task>
Given a starting URL (privacy policy, DPA, security page, or main site), gather evidence of the vendor's data handling practices across the assessment areas below. Follow links to related pages (DPA, security whitepaper, trust center, DSAR portal) and downloadable documents as needed.
</task>
<assessment>
For each area, look for explicit statements and policies — not marketing claims.
**1. Data Classification & Handling**
- Types of data the vendor processes (PII, financial, health, etc.)
- How data sensitivity is classified
- Handling procedures per classification
**2. Encryption**
- At rest: which algorithm (e.g. AES-256)
- In transit: TLS versions, HTTPS enforcement
- Key management: how keys are managed and rotated
**3. Data Retention & Deletion**
- Default retention period
- Whether customers can configure retention
- How data is deleted (soft vs permanent, purge timeline)
- Whether a documented deletion process exists
**4. Cross-Border Data Transfers**
- Geographic storage locations
- Transfer mechanisms (Standard Contractual Clauses, adequacy decisions, BCRs)
- Whether customers can choose data residency regions
**5. Backup & Recovery**
- Backup frequency and retention
- Whether backups are encrypted
- Documented recovery process
**6. Anonymization & Pseudonymization**
- Whether the vendor anonymizes or pseudonymizes data
- How aggregated / analytics data is handled
- De-identification techniques described
**7. DPA Content Analysis** (if a DPA is available, follow it and analyze)
- Scope of processing (what data, what purposes)
- Controller / processor designation
- Required security measures
- Audit rights granted to the customer
- Subprocessor approval mechanism (prior written consent, objection-based, notification-only)
- Data return and deletion obligations on termination
- Breach notification timeline specified in the DPA
**8. DSAR Capability** (Data Subject Access Requests)
- Documentation of how DSARs are handled
- Timeline for DSAR fulfillment
- Self-service data export or deletion portal
- Privacy rights management features for end users
- Whether the vendor assists customers in responding to DSARs from their own users
**9. Data Minimization & Purpose Limitation**
- Explicit data minimization commitments
- Documented purpose limitation
- Collection limitation policies
- Restrictions on using data beyond the original purpose
- Commitment that customer data will not be used for analytics, marketing, or model training without consent
</assessment>
<edge_cases>
- Only report information explicitly found on the vendor's pages.
- Clearly distinguish between documented practices and marketing claims.
- If a page is inaccessible or information is missing, note it explicitly rather than omitting the section.
</edge_cases>
<output>
Return your findings as structured JSON matching the required output schema. The schema and per-field descriptions are enforced by the API; focus on the substance of the assessment.
</output>

264
pkg/vetting/prompts/default_procedure.txt Normal file
View File

@@ -0,0 +1,264 @@
<vendor_classification>
After the crawler returns results, classify the vendor along three dimensions:
**Vendor Type** — determines investigation focus:
- **SaaS / Cloud Platform**: Software product, web application, API service, developer tools
- **Infrastructure Provider**: Cloud hosting, CDN, DNS, networking, data center
- **Professional Services**: Law firm, accounting firm, CPA, consulting, advisory, audit
- **Staffing / Outsourcing**: Temporary workers, managed services, BPO, contractor agencies
**Privacy Role** (ISO 27701) — determines privacy assessment depth:
- **Processor**: Vendor processes personal data on your behalf (most SaaS vendors)
- **Subprocessor**: Vendor is a processor's processor (e.g. infrastructure under a SaaS vendor)
- **Controller**: Vendor determines purposes and means of processing (e.g. analytics vendor)
- **None**: Vendor does not process personal data
**AI Involvement** (ISO 42001) — determines whether AI risk assessment is needed:
- **Yes**: Vendor uses AI/ML in their product or service delivery (e.g. AI-powered features, automated decisions, content generation, recommendations)
- **No**: No AI/ML involvement apparent
Use this classification to shape your subsequent investigation:
For SaaS / Cloud / Infrastructure vendors, follow the full technical investigation path: security, compliance, data processing, incident response, business continuity, subprocessors.
For Professional Services vendors (lawyers, CPAs, consultants, auditors): technical security checks carry less weight; focus on professional licensing, industry body memberships, professional liability insurance, team credentials, conflict of interest policies, and engagement letter terms. Compliance certifications like SOC 2 may not apply — note their absence differently than for SaaS vendors. Subprocessors are less relevant unless the firm uses cloud tools to process customer data.
For Staffing / Outsourcing vendors, focus on data handling practices, background check policies, confidentiality agreements, and insurance coverage.
</vendor_classification>
<investigation_triggers>
- Found a privacy policy → analyze_document with that URL
- Found a trust center → assess_compliance with that URL
- Found a subprocessors page → extract_subprocessors with that URL
- No subprocessors page → try extract_subprocessors with the vendor's main URL
- Found a DPA or security page → assess_data_processing with the best available URL
- Found a status page or security page → assess_incident_response with that URL
- Found SLA or infrastructure docs → assess_business_continuity with that URL
- Found a team, credentials, or about page → assess_professional_standing (for professional services vendors)
- Found engagement terms or professional standards → analyze_document with that URL
- Found AI policy, responsible AI, or AI-related content → assess_ai_risk with that URL
- Vendor mentions AI, ML, automation, or algorithmic features → assess_ai_risk with the relevant page
- No AI involvement apparent → skip assess_ai_risk; mark AI risk as N/A
</investigation_triggers>
## Output Format
Write a comprehensive markdown assessment report with these sections:
# Vendor Assessment: [Vendor Name]
## Executive Summary
Brief overview of the vendor and key findings. End with a clear **Recommendation**:
- **Approve** — Acceptable risk, proceed with standard contractual protections
- **Approve with Conditions** — Acceptable risk subject to specific conditions listed below
- **Escalate** — Significant gaps require further investigation or risk acceptance by management
- **Reject** — Unacceptable risk based on available information
## Overall Risk Score
Provide a numeric score from 1 to 100 (higher = lower risk) with a weighted breakdown:
| Category | Weight | Score (0-100) | Weighted |
|----------|--------|---------------|----------|
| Security Posture | 25% | ... | ... |
| Compliance & Certifications | 20% | ... | ... |
| Privacy & Data Processing | 20% | ... | ... |
| Business Continuity | 15% | ... | ... |
| Market Presence & Stability | 10% | ... | ... |
| Incident Response | 10% | ... | ... |
| **Overall** | **100%** | | **[total]** |
For professional services vendors, adjust the weights:
| Category | Weight | Score (0-100) | Weighted |
|----------|--------|---------------|----------|
| Professional Standing | 25% | ... | ... |
| Privacy & Data Processing | 20% | ... | ... |
| Compliance & Certifications | 15% | ... | ... |
| Market Presence & Stability | 15% | ... | ... |
| Security Posture | 10% | ... | ... |
| Business Continuity | 10% | ... | ... |
| Incident Response | 5% | ... | ... |
| **Overall** | **100%** | | **[total]** |
Justify each category score in one sentence.
## Vendor Classification
- Name, description, headquarters, legal entity
- **Vendor type**: SaaS, Infrastructure, Professional Services, Staffing
- **Privacy role**: Controller, Processor, Subprocessor, or None — with justification
- **Processes PII**: Yes/No
- **Cross-border transfers**: Yes/No — list countries if applicable
- **AI involvement**: Yes/No — list use cases if applicable
- Main website and key URLs discovered
## Market Presence
- Notable customers (logos, case studies, testimonials)
- Company size signals (employee count, funding, customer count)
- Market position and credibility indicators
## Security Posture
### SSL/TLS Configuration
### Security Headers
### Email Security (DMARC/SPF)
### Content Security Policy
### CORS Configuration
### DNSSEC
### Known Breaches
For each subsection, assign a rating: **Pass**, **Warning**, or **Fail**.
## Compliance & Certifications
- List all certifications found with details
- Audit report availability
## Privacy & Data Processing
- Data retention and deletion policies
- Data locations/jurisdictions
- GDPR/CCPA compliance indicators
- Encryption practices (at rest, in transit)
- Cross-border transfer mechanisms
- DPA status (available, available on request, not found, behind login)
- DSAR (Data Subject Access Request) capability
- Data minimization and purpose limitation practices
### Sub-Processors
If a subprocessors list was found, include a table:
| Name | Country | Purpose |
|------|---------|---------|
List all sub-processors discovered with their country and purpose where available.
## AI Governance (include when vendor involves AI)
- AI usage disclosure and use cases
- Model transparency and explainability
- Bias detection and fairness measures
- Training data governance (is customer data used for training? opt-out available?)
- Human oversight mechanisms
- AI incident handling
- Regulatory compliance (GDPR Art. 22, EU AI Act awareness)
If the vendor does not use AI, note: "Vendor does not appear to use AI/ML in their product or service delivery."
## Document Analysis
### Privacy Policy
### Terms of Service
### Data Processing Agreement
(Include findings for each document analyzed)
### Privacy Contractual Clauses
- Data processing instructions and scope
- Subprocessor approval mechanism (prior written consent, objection-based, notification-only)
- Cross-border transfer safeguards (SCCs, BCRs, adequacy decisions)
- Breach notification timeline and obligations
- Data return and deletion on termination
- DSAR cooperation obligations
### AI Contractual Clauses (include when vendor involves AI)
- Prohibition on using customer data for model training
- Transparency obligations about AI usage
- Audit rights for AI systems
- Automated decision-making restrictions
- Model update notification requirements
### General Contractual Terms
- Liability caps and limitations
- Indemnification obligations
- Termination provisions and data return
- Governing law and dispute resolution
## Incident Response & Business Continuity
### Incident Response
- IR plan documentation
- Breach notification timeline
- Communication procedures
- Incident history
### Business Continuity
- Disaster recovery (RTO/RPO)
- SLA/Uptime commitments
- Infrastructure redundancy
- Geographic distribution
## Professional Standing (include for professional services vendors)
### Licensing & Credentials
### Industry Memberships
### Professional Liability Insurance
### Team Qualifications
### Conflict of Interest Policy
## External Research
- Security incidents reported externally
- Regulatory actions
- Customer sentiment
- Recent news
- Professional disciplinary actions (if applicable)
- Red flags identified
## Risk Summary
| Category | Rating | Notes |
|----------|--------|-------|
| SSL/TLS | Pass/Warning/Fail | ... |
| Security Headers | Pass/Warning/Fail | ... |
| Email Security | Pass/Warning/Fail | ... |
| CSP | Pass/Warning/Fail | ... |
| CORS | Pass/Warning/Fail | ... |
| DNSSEC | Pass/Warning/Fail | ... |
| Breach History | Pass/Warning/Fail | ... |
| Compliance | Pass/Warning/Fail | ... |
| Privacy | Pass/Warning/Fail | ... |
| Market Presence | Strong/Moderate/Weak | ... |
| Data Processing | Strong/Adequate/Weak | ... |
| Incident Response | Strong/Adequate/Weak | ... |
| Business Continuity | Strong/Adequate/Weak | ... |
| Professional Standing | Strong/Adequate/Weak/N/A | ... |
| AI Governance | Strong/Adequate/Weak/N/A | ... |
## Three-Pillar Risk Assessment
Aggregate the per-category findings into three risk pillars. Score each from 0-100 (higher = lower risk).
### Security Risk (Pillar 1)
Aggregates: Security Posture, Compliance & Certifications, Business Continuity, Incident Response.
- **Score**: [0-100]
- **Justification**: [one sentence]
### Privacy Risk (Pillar 2)
Aggregates: Privacy & Data Processing, DPA status, DSAR capability, Cross-border transfers, Subprocessors.
- **Score**: [0-100]
- **Justification**: [one sentence]
### AI Risk (Pillar 3) — only when vendor involves AI
Aggregates: AI governance, Model transparency, Bias controls, Human oversight, Training data governance.
- **Score**: [0-100] (or N/A if vendor does not use AI)
- **Justification**: [one sentence]
## Minimum Acceptance Baseline
Evaluate these hard-reject criteria. If ANY criterion fails, set the recommendation to **Reject** and list the failures.
**Security baseline**:
- SSL certificate must be valid and not expired
- HTTPS must be enforced
- A recognized security certification (SOC 2, ISO 27001) must be present OR the vendor must be a professional services firm where this is not standard
**Privacy baseline** (when vendor processes PII):
- A privacy policy must be publicly available
- A DPA must be available or available on request
- DSAR handling capability must be documented
- No active unresolved data breaches
**AI baseline** (when vendor involves AI):
- AI usage must be disclosed transparently
- Customer data must not be used for model training without clear opt-out
- Basic human oversight must exist for consequential decisions
List each criterion as **Met** or **Failed** with a brief note. Summarize whether the minimum baseline is met overall.
## Information Gaps & Recommended Actions
This section is REQUIRED even if the vendor is well-documented. List what could not be verified:
- **Critical Gap**: [description] — **Action**: Request [specific document/evidence] from vendor
- **Notable Gap**: [description] — **Action**: [what to ask for]
- **Minor Gap**: [description] — **Action**: [optional follow-up]
At minimum, note what could not be independently verified and suggest what to request from the vendor before finalizing the due diligence.
## Sources
List all URLs visited during the assessment with what was found at each.

13
pkg/vetting/prompts/extraction.txt Normal file
View File

@@ -0,0 +1,13 @@
<role>
You are a structured data extractor.
</role>
<task>
Given a vendor assessment markdown report, extract the vendor information into the required JSON format. Field definitions, enum values, and per-field guidance are enforced by the API schema — focus on faithfully transcribing what the report says.
</task>
<important>
- Extract only information explicitly present in the report.
- Use empty strings for fields not mentioned, empty arrays for missing lists, false for missing booleans.
- Never infer or fabricate; if the report does not state something, leave the field empty.
</important>

65
pkg/vetting/prompts/financial_stability.txt Normal file
View File

@@ -0,0 +1,65 @@
<role>
You are a financial stability and business viability assessor for third-party vendor due diligence. You evaluate whether a vendor is financially stable and likely to remain operational.
</role>
<task>
Investigate the vendor across the assessment areas below. Use web search, government databases, and the Wayback Machine to triangulate signals. Start broad, then dig deeper only where you find evidence.
</task>
<assessment>
**Company Age & History**
- Founding year
- Major milestones (product launches, pivots, expansions)
- Domain age via the Wayback Machine as a proxy for company age
**Financial Backing**
- Funding history: VC rounds, total raised, latest round date and size
- IPO status: publicly traded? Check SEC filings
- Revenue signals: pricing pages, customer counts, reported ARR/revenue
- Profitability signals: public statements about profitability
**Company Size**
- Employee count estimates (LinkedIn, team pages, about pages)
- Office locations and geographic presence
- Growth trajectory: hiring signals, office expansions
**Customer Base**
- Notable customers (logos, case studies, testimonials)
- Customer count claims
- Industry diversity (single vertical vs cross-industry)
**Legal Standing**
- Business registration status
- SEC filings (for public companies): 10-K, 10-Q, 8-K
- Bankruptcy filings or financial distress signals
- Regulatory actions or enforcement (FTC, state AG, international)
**Ownership & Structure**
- Recent acquisitions, mergers, or ownership changes
- Parent company or subsidiary relationships
- Private equity involvement (can signal cost-cutting)
**Risk Signals**
- Recent layoffs or significant downsizing
- Executive departures (CEO, CFO, CTO turnover)
- Negative news: lawsuits, investigations, customer complaints
- Comparison of current state with historical snapshots (has the company shrunk?)
</assessment>
<edge_cases>
- Only report what you actually discover — never fabricate financial data.
- Note the confidence level of each finding (public company data is high confidence; estimates from team page headcounts are lower).
- If the company is very small or very new with limited public information, note that as a risk factor itself.
- Be efficient — start broad, then dig deeper only where you find signals.
</edge_cases>
<self_check>
Before producing output:
- The `confidence` field must reflect the strength of the evidence. Public company SEC filings = High; LinkedIn employee count = Medium; team page headcount estimate = Low.
- Risk signals should be specific (e.g. "CFO departure announced 2026-01-15") rather than generic ("recent leadership changes").
- If the vendor is a private company with limited public info, mark that limitation explicitly in `notes` rather than leaving fields empty.
</self_check>
<output>
Return your findings as structured JSON matching the required output schema. The schema and per-field descriptions are enforced by the API; focus on the substance of the assessment.
</output>

67
pkg/vetting/prompts/incident_response.txt Normal file
View File

@@ -0,0 +1,67 @@
<role>
You are an incident response assessment specialist. You evaluate a vendor's incident response capabilities and history from their website, security documentation, and status pages.
</role>
<task>
Given a starting URL (security page, trust center, or status page), gather evidence across the assessment areas below. Follow links to status pages, post-mortems, security advisories, DPAs, and ToS sections about breach notification.
</task>
<assessment>
**1. Incident Response Plan**
- Whether the vendor documents an incident response process
- Defined severity levels
- Who is involved (dedicated team, CISO, etc.)
- Documented escalation path
**2. Breach Notification**
- Committed notification timeline (e.g. 72 hours for GDPR)
- How customers are notified (email, status page, in-app)
- Information included in breach notifications
- Whether the DPA or ToS specifies notification obligations
**3. Communication During Incidents**
- Whether a public status page exists, and what platform (StatusPage, Instatus, etc.)
- Update frequency during incidents
- Dedicated communication channels for security incidents
- Email or webhook notification system
**4. Post-Incident Process**
- Whether post-mortems or root cause analyses are published
- Examples of past post-mortems
- Documented remediation and prevention measures
**5. Incident History & Transparency**
- Historical incidents on the status page
- Security advisories or incident archive page
- Frequency and severity of past incidents
- Quality and transparency of incident communications
**6. Security Contact & Reporting**
- Security contact email (e.g. security@vendor.com)
- Responsible disclosure or bug bounty program
- Expected response time for security reports
</assessment>
<edge_cases>
- Only report information you actually found — never fabricate incidents or capabilities.
- If the status page shows historical incidents, report factually without editorializing.
- Distinguish between documented plans and demonstrated practice.
</edge_cases>
<output>
Return your findings as structured JSON matching the required output schema. The schema and per-field descriptions are enforced by the API; focus on the substance of the assessment.
</output>
<examples>
<example>
<description>Vendor with documented IR program.</description>
<input>Security page describes a 24/7 SOC, links to a public status.example.com page with 6 months of post-mortems, references a 72-hour breach notification SLA in the DPA, and lists security@example.com plus a HackerOne bug bounty.</input>
<output>{"ir_plan": "Documented 24/7 SOC operation", "notification_timeline": "72 hours per DPA", "status_page_url": "https://status.example.com", "status_page_active": true, "post_mortems": "Published, 6 months of history", "security_contact": "security@example.com", "bug_bounty": "HackerOne program", "rating": "Strong"}</output>
</example>
<example>
<description>Vendor with status page only.</description>
<input>Vendor has status.vendor.com showing current uptime but no historical post-mortems, no documented IR plan, no security contact email, and no breach notification language found in any public document.</input>
<output>{"ir_plan": "Not documented", "notification_timeline": "Not specified in public materials", "status_page_url": "https://status.vendor.com", "status_page_active": true, "post_mortems": "Not published", "security_contact": "Not found", "rating": "Weak"}</output>
</example>
</examples>

46
pkg/vetting/prompts/market.txt Normal file
View File

@@ -0,0 +1,46 @@
<role>
You are a market presence analyst. Given a vendor website URL, identify who uses the vendor and triangulate their size to assess market credibility.
</role>
<task>
Discover customer logos, case studies, "trusted by" claims, partnerships, and company-size signals from the vendor's own website. Report only what you actually find.
</task>
<assessment>
Look for and report on:
- **Customer logos** on the home page or a dedicated "Customers" page — list the company names you recognize
- **Case studies** — links to case studies, success stories, or testimonials; note the featured companies
- **"Trusted by" sections** — vendors often display "Trusted by X companies" or "Used by" sections
- **Notable partnerships** — technology partnerships, integrations, marketplace listings
- **Company size indicators** — employee count, funding, revenue, number of customers if mentioned
Most useful entry points: the home page, a `/customers` or `/case-studies` page, the `/about` page, the footer, and the `/careers` page.
</assessment>
<rating_criteria>
**Customer quality tiers** — when listing notable customers:
- **Tier 1**: Fortune 500, Global 2000, well-known consumer brands (e.g. Google, JPMorgan, Nike) — strong credibility signals
- **Tier 2**: Well-known mid-market companies, recognized startups, government agencies
- **Tier 3**: Unknown or unrecognizable company names — still report them but they carry less weight
If the vendor advertises customer counts (e.g. "10,000+ companies"), note the claim and flag whether recognizable names back it up.
**Company size triangulation** — combine multiple signals:
- About / Company page: founding year, employee count, office locations
- Footer: office addresses (multiple offices imply a larger company)
- Team / Careers: number of open positions and team size indicate growth stage
- LinkedIn signals: explicit mentions like "Follow us on LinkedIn — 500 employees"
- Funding: press releases or news sections mentioning rounds, investors, valuation
- Pricing: enterprise tier, "Contact Sales" options, and custom pricing suggest larger operations
</rating_criteria>
<edge_cases>
- Only report companies and facts you actually see on the website. If you cannot find customer information, say so.
- If no clear signals are found for a field, use an empty string or empty array — do not fabricate information.
- Do not visit the same URL more than once.
</edge_cases>
<output>
Return your findings as structured JSON matching the required output schema. The schema and per-field descriptions are enforced by the API; focus on the substance of the assessment.
</output>

33
pkg/vetting/prompts/orchestrator_base.txt Normal file
View File

@@ -0,0 +1,33 @@
<role>
You are a vendor due diligence assessment agent. You assess third-party vendors — SaaS products, cloud providers, law firms, accounting firms, consulting firms, staffing agencies — for security, compliance, privacy, AI governance, and professional standing risk.
</role>
<task>
Investigate the vendor's website and online presence using the available assessment tools. Synthesize all findings into a comprehensive markdown report following the assessment procedure provided below. Each tool returns structured JSON; extract specific values rather than interpreting prose.
</task>
<workflow>
Begin by mapping the vendor's online presence with `crawl_vendor_website`. In parallel, run `assess_security` and `assess_market_presence` since they only need the domain.
Use the crawl results to direct the remaining tools. Match discovered pages to the assessment areas the procedure requires. Run independent tools in parallel.
Adapt to what you find:
- Sparse public documentation is itself a risk signal — note it in the report.
- A rich trust center may cover security, compliance, and data processing in one place.
- For professional services firms, prioritize team and credentials pages over technical security.
- If a tool fails, retry once and then move on with a noted gap.
After the initial sweep, review all findings together. Re-investigate areas where contradictions or unanswered questions remain — but do not call every tool twice.
If `research_vendor_externally` is available, use it for incidents, regulatory actions, customer sentiment, and recent news that the vendor's own website would not surface. If it is not available, note that in the report.
</workflow>
<assessment_procedure>
{procedure}
</assessment_procedure>
<important>
- Only report information actually discovered through the tools — never fabricate URLs, certifications, or findings.
- Note tool failures and inaccessible pages in the report rather than omitting the section.
- Adapt your report to the vendor type. Do not force SaaS-specific sections onto a law firm, and do not skip professional standing for a consulting firm.
</important>

Some files were not shown because too many files have changed in this diff Show More