Files
probo/pkg/agent/agent_tool.go
Bryan Frimin d80a0fa7e4 Reject null input in agent tool parameter validation
JSON null unmarshals into an empty string, so the presence-only
key check let {"input":null} through, running the nested agent
with a blank user message.

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
2026-03-13 19:57:18 +01:00

121 lines
2.9 KiB
Go

// 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"
"encoding/json"
"fmt"
"go.probo.inc/probo/pkg/llm"
)
const DefaultMaxToolDepth = 16
type (
agentTool struct {
agent *Agent
toolName string
description string
schema json.RawMessage
}
agentToolParams struct {
Input string `json:"input" jsonschema:"The input to send to the agent"`
}
agentToolDepthKey struct{}
)
var (
agentToolParamsSchema = mustJSONSchemaFor[agentToolParams]()
)
func agentToolDepth(ctx context.Context) int {
if v, ok := ctx.Value(agentToolDepthKey{}).(int); ok {
return v
}
return 0
}
func newAgentTool(agent *Agent, name, description string) *agentTool {
return &agentTool{
agent: agent,
toolName: name,
description: description,
schema: agentToolParamsSchema,
}
}
func (t *agentTool) Name() string { return t.toolName }
func (t *agentTool) Definition() llm.Tool {
return llm.Tool{
Name: t.toolName,
Description: t.description,
Parameters: t.schema,
}
}
func (t *agentTool) Execute(ctx context.Context, arguments string) (ToolResult, error) {
depth := agentToolDepth(ctx)
if depth >= t.agent.maxToolDepth {
return ToolResult{}, &MaxToolDepthExceededError{MaxDepth: t.agent.maxToolDepth}
}
var fields map[string]json.RawMessage
if err := json.Unmarshal([]byte(arguments), &fields); err != nil {
return ToolResult{
Content: fmt.Sprintf("Invalid parameters: %s", err.Error()),
IsError: true,
}, nil
}
raw, ok := fields["input"]
if !ok || string(raw) == "null" {
return ToolResult{
Content: "Missing required parameters: input",
IsError: true,
}, nil
}
var params agentToolParams
if err := json.Unmarshal([]byte(arguments), &params); err != nil {
return ToolResult{
Content: fmt.Sprintf("Invalid parameters: %s", err.Error()),
IsError: true,
}, nil
}
ctx = context.WithValue(ctx, agentToolDepthKey{}, depth+1)
result, err := t.agent.Run(
ctx,
[]llm.Message{
{
Role: llm.RoleUser,
Parts: []llm.Part{
llm.TextPart{Text: params.Input},
},
},
},
)
if err != nil {
return ToolResult{}, err
}
return ToolResult{Content: result.FinalMessage().Text()}, nil
}