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

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
}

38
pkg/agent/tools/search/search.go Normal file
View File

@@ -0,0 +1,38 @@
// 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 (
"go.probo.inc/probo/pkg/agent"
)
// Toolset provides web search tools.
type Toolset struct {
endpoint string
}
// 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)
},
)
}