diff --git a/frontend/package.json b/frontend/package.json index 69e28190..2e188202 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -38,7 +38,7 @@ "class-variance-authority": "^0.7.1", "clsx": "^2.1.1", "date-fns": "^4.1.0", - "dompurify": "^3.4.9", + "dompurify": "^3.4.11", "graphql": "^16.13.2", "lucide-react": "^0.511.0", "react": "^19.2.4", diff --git a/frontend/pnpm-lock.yaml b/frontend/pnpm-lock.yaml index bf93014f..91ff7991 100644 --- a/frontend/pnpm-lock.yaml +++ b/frontend/pnpm-lock.yaml @@ -75,8 +75,8 @@ importers: specifier: ^4.1.0 version: 4.1.0 dompurify: - specifier: ^3.4.9 - version: 3.4.10 + specifier: ^3.4.11 + version: 3.4.11 graphql: specifier: ^16.13.2 version: 16.13.2 @@ -2093,8 +2093,8 @@ packages: resolution: {integrity: sha512-35mSku4ZXK0vfCuHEDAwt55dg2jNajHZ1odvF+8SSr82EsZY4QmXfuWso8oEd8zRhVObSN18aM0CjSdoBX7zIw==} engines: {node: '>=0.10.0'} - dompurify@3.4.10: - resolution: {integrity: sha512-0xzNv0e7oYC6yyuOGZIABPM4qtg3QxLFniDNPP4ZP90wR8Yq3zgwpRbrNiT4N3IKqDbbYFEJLV+JWEs19aZ//w==} + dompurify@3.4.11: + resolution: {integrity: sha512-zhlUV12GsaRzMsf9q5M254YhA4+VuF0fG+QFqu6aYpoGlKtz+w8//jBcGVYBgQkR5GHjUomejY84AV+/uPbWdw==} dot-case@3.0.4: resolution: {integrity: sha512-Kv5nKlh6yRrdrGvxeJ2e5y2eRUpkUosIW4A2AS38zwSz27zu7ufDwQPi5Jhs3XAlGNetl3bmnGhQsMtkKJnj3w==} @@ -6084,7 +6084,7 @@ snapshots: dependencies: esutils: 2.0.3 - dompurify@3.4.10: + dompurify@3.4.11: optionalDependencies: '@types/trusted-types': 2.0.7 diff --git a/go.mod b/go.mod index 274f68cc..6dae36ec 100644 --- a/go.mod +++ b/go.mod @@ -12,19 +12,19 @@ require ( github.com/fatih/color v1.19.0 github.com/getsentry/sentry-go v0.44.1 github.com/go-chi/chi/v5 v5.2.5 - github.com/gollem-dev/gollem v0.26.0 - github.com/gollem-dev/tools/abusech v0.1.0 - github.com/gollem-dev/tools/bigquery v0.1.0 - github.com/gollem-dev/tools/github v0.1.0 - github.com/gollem-dev/tools/intune v0.1.0 - github.com/gollem-dev/tools/ipdb v0.1.0 - github.com/gollem-dev/tools/otx v0.1.0 - github.com/gollem-dev/tools/shodan v0.1.0 - github.com/gollem-dev/tools/slack v0.1.0 - github.com/gollem-dev/tools/urlscan v0.1.0 - github.com/gollem-dev/tools/vt v0.1.0 - github.com/gollem-dev/tools/webfetch v0.1.0 - github.com/gollem-dev/tools/whois v0.1.0 + github.com/gollem-dev/gollem v0.26.2-0.20260630002044-aed22083a5b0 + github.com/gollem-dev/tools/abusech v0.2.0 + github.com/gollem-dev/tools/bigquery v0.2.0 + github.com/gollem-dev/tools/github v0.3.0 + github.com/gollem-dev/tools/intune v0.2.0 + github.com/gollem-dev/tools/ipdb v0.2.0 + github.com/gollem-dev/tools/otx v0.2.0 + github.com/gollem-dev/tools/shodan v0.2.0 + github.com/gollem-dev/tools/slack v0.3.0 + github.com/gollem-dev/tools/urlscan v0.2.0 + github.com/gollem-dev/tools/vt v0.2.0 + github.com/gollem-dev/tools/webfetch v0.3.0 + github.com/gollem-dev/tools/whois v0.2.0 github.com/google/go-github/v74 v74.0.0 github.com/google/uuid v1.6.0 github.com/gorilla/websocket v1.5.3 @@ -205,13 +205,13 @@ require ( go.opentelemetry.io/otel/trace v1.43.0 // indirect go.yaml.in/yaml/v2 v2.4.4 // indirect go.yaml.in/yaml/v3 v3.0.4 // indirect - golang.org/x/crypto v0.51.0 // indirect + golang.org/x/crypto v0.52.0 // indirect golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f // indirect golang.org/x/mod v0.35.0 // indirect - golang.org/x/net v0.54.0 // indirect + golang.org/x/net v0.55.0 // indirect golang.org/x/oauth2 v0.36.0 // indirect - golang.org/x/sync v0.20.0 // indirect - golang.org/x/sys v0.44.0 // indirect + golang.org/x/sync v0.21.0 // indirect + golang.org/x/sys v0.45.0 // indirect golang.org/x/telemetry v0.0.0-20260409153401-be6f6cb8b1fa // indirect golang.org/x/term v0.43.0 // indirect golang.org/x/text v0.37.0 // indirect diff --git a/go.sum b/go.sum index 42c25e51..ab3358ca 100644 --- a/go.sum +++ b/go.sum @@ -218,32 +218,32 @@ github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63Y github.com/golang-jwt/jwt/v5 v5.3.1/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE= github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= -github.com/gollem-dev/gollem v0.26.0 h1:FIGZq1un1AWCXIU00VXdtJ3SPo5C2pwWj8jOc5z/yyY= -github.com/gollem-dev/gollem v0.26.0/go.mod h1:qEZEgRi5g/tGoNQyLgq8uWfM7mhzCMkjdP2IZLSGSxM= -github.com/gollem-dev/tools/abusech v0.1.0 h1:0EvwiABq6585q7sarWYkpNRyDvMIgIAYxAAT6wIQZt0= -github.com/gollem-dev/tools/abusech v0.1.0/go.mod h1:4xn+UrOECXLWiBTPSgZoOOVrsqo4P2wZ0u/umtbsRFc= -github.com/gollem-dev/tools/bigquery v0.1.0 h1:TL3v6IVrozstX/LbFsID4NFiP/I/OUDXojYzv+H2by8= -github.com/gollem-dev/tools/bigquery v0.1.0/go.mod h1:H+sKbHbwEpCP1+FCaToZ5ftiOpMMO1Wbi2DaL3E7AqE= -github.com/gollem-dev/tools/github v0.1.0 h1:HK3CKjYrC62gqhw7v+Oi39ai3yxWbkEqjHlN7hD42X4= -github.com/gollem-dev/tools/github v0.1.0/go.mod h1:LHt/fbIzFTU+eU2U5uXSgTTu6Yoco2Lkzc9DRMaT/zI= -github.com/gollem-dev/tools/intune v0.1.0 h1:Qe41jsuRDvQrXsSmk91q15hBCT68bhNacupKIeU2vhg= -github.com/gollem-dev/tools/intune v0.1.0/go.mod h1:WdgquLJMs6tPFhzISArygZ+AGr2KvpOPulhtp1gp+oY= -github.com/gollem-dev/tools/ipdb v0.1.0 h1:Df1j1PwcIQvVFvqI6LP09EtRqHAj10pLzPnQk8EKfoU= -github.com/gollem-dev/tools/ipdb v0.1.0/go.mod h1:Zz0w9yXKpBcimB+ut4I3jlRpXWV+Qs6QQqb2G6MnpuE= -github.com/gollem-dev/tools/otx v0.1.0 h1:+mKhHd+gikhh3dvNSmyvbSaiACIIMaZe1UoXwzlk2+8= -github.com/gollem-dev/tools/otx v0.1.0/go.mod h1:8cuXN15Bq/ILXN4CPACpCFIabe52pz3O76Z1iePHQEk= -github.com/gollem-dev/tools/shodan v0.1.0 h1:eWmb426Kg1koYlHCI11sJ6BnZUxcUKbdd0lEoWdKr4k= -github.com/gollem-dev/tools/shodan v0.1.0/go.mod h1:AvuHaiYDPmJMJiQOK7fClkJ00pMH+RscS6VjU38/zos= -github.com/gollem-dev/tools/slack v0.1.0 h1:8dmzH5c2jM1WZvwXeLGdvWE/TBQNkp6sWfM31byWwhQ= -github.com/gollem-dev/tools/slack v0.1.0/go.mod h1:g415ENkXWQRDToVCviuSHQ2BMiz0omrIQsjbmQwOahM= -github.com/gollem-dev/tools/urlscan v0.1.0 h1:AxMMuy9bsXFUSCWXaIWBesgKAIVYEGl2/Cjt3BvnTgg= -github.com/gollem-dev/tools/urlscan v0.1.0/go.mod h1:cZ7bUqclkpFCL2Z5pFKLWd6KWTWzf+iy69NvnFoewIo= -github.com/gollem-dev/tools/vt v0.1.0 h1:X2T0E1HT5aoU4BK6B/ZITwRvOLMiq6ttLAgHdVfPeEc= -github.com/gollem-dev/tools/vt v0.1.0/go.mod h1:iV1dpMTWMMQSZp7+dbUVzgb4Rd6eIJHLSxPEa8W8mKo= -github.com/gollem-dev/tools/webfetch v0.1.0 h1:jDcY+oQGH6zoKL/tBSzzgin0R9t60UY6xn3jwFMj/9A= -github.com/gollem-dev/tools/webfetch v0.1.0/go.mod h1:zTCZBmDrF3gFJhzES/gaL+pnYHa6BIAglc8E0GR8rrY= -github.com/gollem-dev/tools/whois v0.1.0 h1:YePfmuLx44mHFaGi7IOGKATKme81AE+dHqisdALGrfU= -github.com/gollem-dev/tools/whois v0.1.0/go.mod h1:U/iXnDoKtJmWZA7ma9bcrOVV4W2oaLwyuRP+vd+6AKY= +github.com/gollem-dev/gollem v0.26.2-0.20260630002044-aed22083a5b0 h1:8W3NXX/O2S1TQF7dTYQUER+H39cnny4V32DanPQyB+Y= +github.com/gollem-dev/gollem v0.26.2-0.20260630002044-aed22083a5b0/go.mod h1:20atdewh0lJcT3JV7eD3VeemZqwLFJFO7cQsBFnOtdk= +github.com/gollem-dev/tools/abusech v0.2.0 h1:VxLz8IRIYrajTb6YyVR7D5g9fIAvmV7nkh8K10JJCYw= +github.com/gollem-dev/tools/abusech v0.2.0/go.mod h1:sPSiX1y2ci+ne4SA+VSJ//rFMJ2kKj/ELIKDPF6CuIk= +github.com/gollem-dev/tools/bigquery v0.2.0 h1:MsOaLHPNK5baVvjr7v/5JME7EXdhsq/Glm/JUvUA8ig= +github.com/gollem-dev/tools/bigquery v0.2.0/go.mod h1:0EBYrrjDDNFs4YjZs9qrhyV5qzFHHnDAPYFH+mf34po= +github.com/gollem-dev/tools/github v0.3.0 h1:6+kbOQO9Aeli6em2P7F95y7teA3BmuyNNJ2HeXP5ZMY= +github.com/gollem-dev/tools/github v0.3.0/go.mod h1:wbU67QwtsWUh0EpvqUD2W5PO6Oz7NWsgU7l/oWy8xNY= +github.com/gollem-dev/tools/intune v0.2.0 h1:J9WNxKFY1AbvihYz4/tw1ZpwX6wkgq2SDejH7S9K1Qo= +github.com/gollem-dev/tools/intune v0.2.0/go.mod h1:7qRs1qRAfTHqMPlQDMsNq62qnUmUB6x8JiKqrPERo8g= +github.com/gollem-dev/tools/ipdb v0.2.0 h1:kCYEPwcz+63wkexZE5y9+fr3VZf+hykjknNJLON+TPk= +github.com/gollem-dev/tools/ipdb v0.2.0/go.mod h1:d/ZqlWrjeUWK/Fbt7ru9mhpfr+NIPVanelLV1wvV8Xo= +github.com/gollem-dev/tools/otx v0.2.0 h1:7x5kbtbu7Du9kHJGFvJukT7qB/zWG8ue7EqDa6JUVKI= +github.com/gollem-dev/tools/otx v0.2.0/go.mod h1:uWACbM3+s+aGDuKGRYcX46ZXV3Ux4PB5ys2cuWjW9fM= +github.com/gollem-dev/tools/shodan v0.2.0 h1:4BXM7SEUCWBuSI3w24UKr5q1CbIS9PyuNPclRHOIHgU= +github.com/gollem-dev/tools/shodan v0.2.0/go.mod h1:sppv1tNxn5tPWwJEatZm1rrDf0D+JZKAWQCZ2+D/1KI= +github.com/gollem-dev/tools/slack v0.3.0 h1:S82rwhur7ic1gFZ/kEeU1xvAoPwXaZScKW7CGP05WV0= +github.com/gollem-dev/tools/slack v0.3.0/go.mod h1:H083CL9pRnCCJfwAj++KOmmIjRgdRcR9MGCDmsmPCDg= +github.com/gollem-dev/tools/urlscan v0.2.0 h1:yIYcAuzN8uC6dziCfgfsKGfuwC1svpeNPldmqXtSKws= +github.com/gollem-dev/tools/urlscan v0.2.0/go.mod h1:F+6866NqaQ087pKFnpgKtlwuBqtVj/ggqW4HSrL3aEE= +github.com/gollem-dev/tools/vt v0.2.0 h1:czuLKgqSeClOV7sb3N+kXknK09xXFvlpyXURZ7P9Kmk= +github.com/gollem-dev/tools/vt v0.2.0/go.mod h1:cb8sKe5nF0uC3wqByU+Fzal8unqqXFRH3EGoQ3RLnCc= +github.com/gollem-dev/tools/webfetch v0.3.0 h1:h036gkB/QxXyY+xnNDKCzJ2SxqA13NvbJqurYemY2IA= +github.com/gollem-dev/tools/webfetch v0.3.0/go.mod h1:77FqZ/6fj6sp06InWNRTpRb5fNZxysb0f1c7JaUCk6k= +github.com/gollem-dev/tools/whois v0.2.0 h1:SRW0Qt1/w+DlH2WY0IH6APGW4kifaXmc/rAIZWMnsIQ= +github.com/gollem-dev/tools/whois v0.2.0/go.mod h1:MwjO3WlIrIVme+3zoDs4Czq8uZVd8TtebS17F3/Jb9o= github.com/google/flatbuffers v25.12.19+incompatible h1:haMV2JRRJCe1998HeW/p0X9UaMTK6SDo0ffLn2+DbLs= github.com/google/flatbuffers v25.12.19+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= @@ -488,21 +488,21 @@ go.yaml.in/yaml/v2 v2.4.4 h1:tuyd0P+2Ont/d6e2rl3be67goVK4R6deVxCUX5vyPaQ= go.yaml.in/yaml/v2 v2.4.4/go.mod h1:gMZqIpDtDqOfM0uNfy0SkpRhvUryYH0Z6wdMYcacYXQ= go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= -golang.org/x/crypto v0.51.0 h1:IBPXwPfKxY7cWQZ38ZCIRPI50YLeevDLlLnyC5wRGTI= -golang.org/x/crypto v0.51.0/go.mod h1:8AdwkbraGNABw2kOX6YFPs3WM22XqI4EXEd8g+x7Oc8= +golang.org/x/crypto v0.52.0 h1:RMs7fP2rXdep0CftQlK8Uf+kibLm7qkCcradZWYz988= +golang.org/x/crypto v0.52.0/go.mod h1:1QgfPxDqh0T2M/elOJtp9RvuR95kVjir0e6/BvEmGbc= golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f h1:W3F4c+6OLc6H2lb//N1q4WpJkhzJCK5J6kUi1NTVXfM= golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f/go.mod h1:J1xhfL/vlindoeF/aINzNzt2Bket5bjo9sdOYzOsU80= golang.org/x/mod v0.35.0 h1:Ww1D637e6Pg+Zb2KrWfHQUnH2dQRLBQyAtpr/haaJeM= golang.org/x/mod v0.35.0/go.mod h1:+GwiRhIInF8wPm+4AoT6L0FA1QWAad3OMdTRx4tFYlU= -golang.org/x/net v0.54.0 h1:2zJIZAxAHV/OHCDTCOHAYehQzLfSXuf/5SoL/Dv6w/w= -golang.org/x/net v0.54.0/go.mod h1:Sj4oj8jK6XmHpBZU/zWHw3BV3abl4Kvi+Ut7cQcY+cQ= +golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8= +golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww= golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= -golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= -golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM= +golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.44.0 h1:ildZl3J4uzeKP07r2F++Op7E9B29JRUy+a27EibtBTQ= -golang.org/x/sys v0.44.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= +golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/telemetry v0.0.0-20260409153401-be6f6cb8b1fa h1:efT73AJZfAAUV7SOip6pWGkwJDzIGiKBZGVzHYa+ve4= golang.org/x/telemetry v0.0.0-20260409153401-be6f6cb8b1fa/go.mod h1:kHjTxDEnAu6/Nl9lDkzjWpR+bmKfxeiRuSDlsMb70gE= golang.org/x/term v0.43.0 h1:S4RLU2sB31O/NCl+zFN9Aru9A/Cq2aqKpTZJ6B+DwT4= diff --git a/pkg/agents/bigquery/export_test.go b/pkg/agents/bigquery/export_test.go index 464e3b67..b8331be4 100644 --- a/pkg/agents/bigquery/export_test.go +++ b/pkg/agents/bigquery/export_test.go @@ -8,21 +8,13 @@ import ( func NewToolSetForTest(config *Config, projectID string, impersonateServiceAccount string) *toolSet { return &toolSet{ config: config, - tool: &internalTool{ - config: config, - projectID: projectID, - impersonateServiceAccount: impersonateServiceAccount, - }, + tool: newInternalTool(config, projectID, impersonateServiceAccount), } } // ExportNewInternalTool creates a new internalTool instance for testing func ExportNewInternalTool(config *Config, projectID string, impersonateServiceAccount string) gollem.ToolSet { - return &internalTool{ - config: config, - projectID: projectID, - impersonateServiceAccount: impersonateServiceAccount, - } + return newInternalTool(config, projectID, impersonateServiceAccount) } // ToolSpec is exported for testing diff --git a/pkg/agents/bigquery/factory.go b/pkg/agents/bigquery/factory.go index 04b00f5f..bd55261f 100644 --- a/pkg/agents/bigquery/factory.go +++ b/pkg/agents/bigquery/factory.go @@ -98,11 +98,7 @@ func (f *Factory) Configure(ctx context.Context) (interfaces.ToolSet, error) { // Create and return toolSet ts := &toolSet{ config: cfg, - tool: &internalTool{ - config: cfg, - projectID: f.projectID, - impersonateServiceAccount: f.impersonateServiceAccount, - }, + tool: newInternalTool(cfg, f.projectID, f.impersonateServiceAccount), } return ts, nil diff --git a/pkg/agents/bigquery/tool.go b/pkg/agents/bigquery/tool.go index 99a70459..1dde8c2d 100644 --- a/pkg/agents/bigquery/tool.go +++ b/pkg/agents/bigquery/tool.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "strings" "cloud.google.com/go/bigquery" "github.com/dustin/go-humanize" @@ -13,6 +14,7 @@ import ( "github.com/secmon-lab/warren/pkg/domain/types" "github.com/secmon-lab/warren/pkg/utils/logging" "github.com/secmon-lab/warren/pkg/utils/msg" + "github.com/secmon-lab/warren/pkg/utils/toolset" "google.golang.org/api/impersonate" "google.golang.org/api/iterator" "google.golang.org/api/option" @@ -23,26 +25,77 @@ type internalTool struct { config *Config projectID string impersonateServiceAccount string + + tools gollem.ToolSet } -// ID implements gollem.ToolSet -func (t *internalTool) ID() string { - return "bigquery_agent" +// Static tool descriptions. The bigquery_query description is built per-config +// in newInternalTool because it interpolates the available tables and scan +// limit; the rest are constant. +const ( + descBigQuerySchema = `Get detailed schema information for a specific BigQuery table. +Use this to understand available fields, data types, and nested structures before constructing queries. +Returns complete schema including nested RECORD fields.` + descGetRunbook = "Get the full SQL content of a runbook by its ID. Use this when you want to see or adapt a pre-written query template." +) + +// Typed inputs for each tool; the schema is inferred from these struct tags. +type bigqueryQueryInput struct { + SQL string `json:"sql" required:"true" description:"The SQL query to execute. Must be a valid BigQuery SQL query."` + Description string `json:"description" description:"Brief description of what this query is trying to achieve (optional, for logging/tracking)"` } -// Specs implements gollem.ToolSet -func (t *internalTool) Specs(ctx context.Context) ([]gollem.ToolSpec, error) { - // Build table descriptions for the prompt - tableDescriptions := "" - for i, table := range t.config.Tables { - tableDescriptions += fmt.Sprintf("\n%d. %s.%s.%s: %s", +type bigquerySchemaInput struct { + ProjectID string `json:"project_id" required:"true" description:"The project ID of the table"` + DatasetID string `json:"dataset_id" required:"true" description:"The dataset ID of the table"` + TableID string `json:"table_id" required:"true" description:"The table ID to get schema for"` +} + +type getRunbookInput struct { + RunbookID string `json:"runbook_id" required:"true" description:"The ID of the runbook to retrieve"` +} + +// Startup assertions: validate each tool's In/Out types form a valid schema at +// package init, so a malformed type fails immediately rather than at first use. +var ( + _ = gollem.MustToolSchema[bigqueryQueryInput, map[string]any]() + _ = gollem.MustToolSchema[bigquerySchemaInput, map[string]any]() + _ = gollem.MustToolSchema[getRunbookInput, map[string]any]() +) + +// newInternalTool creates an internalTool and builds its type-safe tool set. +// The get_runbook tool is included only when runbooks are configured, preserving +// the previous mode-dependent tool surface. +func newInternalTool(config *Config, projectID, impersonateServiceAccount string) *internalTool { + t := &internalTool{ + config: config, + projectID: projectID, + impersonateServiceAccount: impersonateServiceAccount, + } + + tools := []gollem.Tool{ + gollem.MustNewTool("bigquery_query", buildQueryDescription(config), t.executeQuery), + gollem.MustNewTool("bigquery_schema", descBigQuerySchema, t.getTableSchema), + } + if len(config.Runbooks) > 0 { + tools = append(tools, gollem.MustNewTool("get_runbook", descGetRunbook, t.getRunbook)) + } + t.tools = toolset.New(tools...) + + return t +} + +// buildQueryDescription renders the bigquery_query tool description, embedding +// the available tables and the configured scan-size limit. +func buildQueryDescription(config *Config) string { + var sb strings.Builder + for i, table := range config.Tables { + fmt.Fprintf(&sb, "\n%d. %s.%s.%s: %s", i+1, table.ProjectID, table.DatasetID, table.TableID, table.Description) } + tableDescriptions := sb.String() - specs := []gollem.ToolSpec{ - { - Name: "bigquery_query", - Description: fmt.Sprintf(`Execute a BigQuery SQL query to retrieve data. + return fmt.Sprintf(`Execute a BigQuery SQL query to retrieve data. Available tables:%s Important guidelines: @@ -58,88 +111,32 @@ Best practices: - Start with schema inspection if unfamiliar with table structure - Apply time range filters to reduce scan size - Use LIMIT to restrict rows appropriately`, - tableDescriptions, - humanize.Bytes(t.config.ScanSizeLimit), - ), - Parameters: map[string]*gollem.Parameter{ - "sql": { - Type: gollem.TypeString, - Description: "The SQL query to execute. Must be a valid BigQuery SQL query.", - Required: true, - }, - "description": { - Type: gollem.TypeString, - Description: "Brief description of what this query is trying to achieve (optional, for logging/tracking)", - }, - }, - }, - { - Name: "bigquery_schema", - Description: `Get detailed schema information for a specific BigQuery table. -Use this to understand available fields, data types, and nested structures before constructing queries. -Returns complete schema including nested RECORD fields.`, - Parameters: map[string]*gollem.Parameter{ - "project_id": { - Type: gollem.TypeString, - Description: "The project ID of the table", - Required: true, - }, - "dataset_id": { - Type: gollem.TypeString, - Description: "The dataset ID of the table", - Required: true, - }, - "table_id": { - Type: gollem.TypeString, - Description: "The table ID to get schema for", - Required: true, - }, - }, - }, - } + tableDescriptions, + humanize.Bytes(config.ScanSizeLimit), + ) +} - // Add get_runbook tool if runbooks are configured - if len(t.config.Runbooks) > 0 { - specs = append(specs, gollem.ToolSpec{ - Name: "get_runbook", - Description: "Get the full SQL content of a runbook by its ID. Use this when you want to see or adapt a pre-written query template.", - Parameters: map[string]*gollem.Parameter{ - "runbook_id": { - Type: gollem.TypeString, - Description: "The ID of the runbook to retrieve", - Required: true, - }, - }, - }) - } +// ID implements gollem.ToolSet +func (t *internalTool) ID() string { + return "bigquery_agent" +} - return specs, nil +// Specs implements gollem.ToolSet +func (t *internalTool) Specs(ctx context.Context) ([]gollem.ToolSpec, error) { + return t.tools.Specs(ctx) } // Run implements gollem.ToolSet func (t *internalTool) Run(ctx context.Context, name string, args map[string]any) (map[string]any, error) { - log := logging.From(ctx) - log.Debug("Internal tool run started", "function", name, "args_keys", getMapKeys(args)) - - switch name { - case "bigquery_query": - return t.executeQuery(ctx, args) - case "bigquery_schema": - return t.getTableSchema(ctx, args) - case "get_runbook": - return t.getRunbook(ctx, args) - default: - log.Debug("Unknown internal tool function", "name", name) - return nil, goerr.New("unknown function", goerr.V("name", name)) - } + return t.tools.Run(ctx, name, args) } -func (t *internalTool) executeQuery(ctx context.Context, args map[string]any) (map[string]any, error) { +func (t *internalTool) executeQuery(ctx context.Context, in bigqueryQueryInput) (map[string]any, error) { log := logging.From(ctx) log.Debug("Executing BigQuery query") - sql, ok := args["sql"].(string) - if !ok { + sql := in.SQL + if sql == "" { log.Debug("SQL parameter is missing or invalid") return nil, goerr.New("sql parameter is required") } @@ -312,22 +309,22 @@ func (t *internalTool) executeQuery(ctx context.Context, args map[string]any) (m } // getTableSchema retrieves detailed schema information for a BigQuery table -func (t *internalTool) getTableSchema(ctx context.Context, args map[string]any) (map[string]any, error) { +func (t *internalTool) getTableSchema(ctx context.Context, in bigquerySchemaInput) (map[string]any, error) { log := logging.From(ctx) log.Debug("Getting BigQuery table schema") - projectID, ok := args["project_id"].(string) - if !ok { + projectID := in.ProjectID + if projectID == "" { log.Debug("project_id parameter is missing or invalid") return nil, goerr.New("project_id parameter is required") } - datasetID, ok := args["dataset_id"].(string) - if !ok { + datasetID := in.DatasetID + if datasetID == "" { log.Debug("dataset_id parameter is missing or invalid") return nil, goerr.New("dataset_id parameter is required") } - tableID, ok := args["table_id"].(string) - if !ok { + tableID := in.TableID + if tableID == "" { log.Debug("table_id parameter is missing or invalid") return nil, goerr.New("table_id parameter is required") } @@ -458,12 +455,12 @@ func convertBigQueryValue(val bigquery.Value) any { } // getRunbook retrieves a runbook entry by ID -func (t *internalTool) getRunbook(ctx context.Context, args map[string]any) (map[string]any, error) { +func (t *internalTool) getRunbook(ctx context.Context, in getRunbookInput) (map[string]any, error) { log := logging.From(ctx) log.Debug("Getting runbook") - runbookIDStr, ok := args["runbook_id"].(string) - if !ok { + runbookIDStr := in.RunbookID + if runbookIDStr == "" { log.Debug("runbook_id parameter is missing or invalid") return nil, goerr.New("runbook_id parameter is required") } @@ -489,12 +486,3 @@ func (t *internalTool) getRunbook(ctx context.Context, args map[string]any) (map "sql": entry.SQLContent, }, nil } - -// getMapKeys returns the keys of a map as a slice -func getMapKeys(m map[string]any) []string { - keys := make([]string, 0, len(m)) - for k := range m { - keys = append(keys, k) - } - return keys -} diff --git a/pkg/agents/falcon/tool.go b/pkg/agents/falcon/tool.go index 95f77240..e19f9cf8 100644 --- a/pkg/agents/falcon/tool.go +++ b/pkg/agents/falcon/tool.go @@ -17,6 +17,7 @@ import ( "github.com/secmon-lab/warren/pkg/utils/logging" "github.com/secmon-lab/warren/pkg/utils/msg" "github.com/secmon-lab/warren/pkg/utils/safe" + "github.com/secmon-lab/warren/pkg/utils/toolset" ) // internalTool implements gollem.ToolSet for CrowdStrike Falcon API operations. @@ -24,208 +25,131 @@ type internalTool struct { tokenProvider *tokenProvider baseURL string httpClient *http.Client + + tools gollem.ToolSet } // newInternalTool creates a new internalTool for Falcon API calls. func newInternalTool(tp *tokenProvider, baseURL string) *internalTool { - return &internalTool{ + t := &internalTool{ tokenProvider: tp, baseURL: baseURL, httpClient: &http.Client{Timeout: 60 * time.Second}, } + + // Build the type-safe tool set. Each tool's schema is inferred from its + // typed input struct, replacing the hand-written ToolSpec literals and the + // args[...] map extraction. + t.tools = toolset.New( + gollem.MustNewTool("falcon_search_incidents", descSearchIncidents, t.searchIncidents), + gollem.MustNewTool("falcon_get_incidents", descGetIncidents, t.getIncidents), + gollem.MustNewTool("falcon_search_alerts", descSearchAlerts, t.searchAlerts), + gollem.MustNewTool("falcon_get_alerts", descGetAlerts, t.getAlerts), + gollem.MustNewTool("falcon_search_behaviors", descSearchBehaviors, t.searchBehaviors), + gollem.MustNewTool("falcon_get_behaviors", descGetBehaviors, t.getBehaviors), + gollem.MustNewTool("falcon_get_crowdscores", descGetCrowdScores, t.getCrowdScores), + gollem.MustNewTool("falcon_search_devices", descSearchDevices, t.searchDevices), + gollem.MustNewTool("falcon_get_devices", descGetDevices, t.getDevices), + gollem.MustNewTool("falcon_search_events", descSearchEvents, t.searchEvents), + ) + + return t } -func (t *internalTool) Specs(_ context.Context) ([]gollem.ToolSpec, error) { - return []gollem.ToolSpec{ - { - Name: "falcon_search_incidents", - Description: "Search for incident IDs using FQL (Falcon Query Language) filters. Returns a list of incident IDs that can be used with falcon_get_incidents to retrieve full details.", - Parameters: map[string]*gollem.Parameter{ - "filter": { - Type: gollem.TypeString, - Description: "FQL filter expression (e.g., \"status:'30'\", \"tags:'critical'\", \"start:>'2025-01-01'\")", - }, - "sort": { - Type: gollem.TypeString, - Description: "Sort expression (e.g., \"start.desc\", \"end.asc\")", - }, - "limit": { - Type: gollem.TypeNumber, - Description: "Maximum number of IDs to return (default: 100, max: 500)", - }, - "offset": { - Type: gollem.TypeNumber, - Description: "Pagination offset", - }, - }, - }, - { - Name: "falcon_get_incidents", - Description: "Get detailed information for specific incidents by their IDs. Returns full incident details including status, tactics, techniques, hosts, and users involved.", - Parameters: map[string]*gollem.Parameter{ - "ids": { - Type: gollem.TypeString, - Description: "Comma-separated incident IDs (e.g., \"inc:abc123:def456,inc:abc123:ghi789\")", - Required: true, - }, - }, - }, - { - Name: "falcon_search_alerts", - Description: "Search and retrieve alert details in one call using FQL filters with cursor-based pagination. Returns full alert objects including severity, tactic, technique, and device info.", - Parameters: map[string]*gollem.Parameter{ - "filter": { - Type: gollem.TypeString, - Description: "FQL filter expression (e.g., \"status:'new'\", \"severity:>50\", \"tactics:'Lateral Movement'\")", - }, - "sort": { - Type: gollem.TypeString, - Description: "Sort property (e.g., \"timestamp|desc\", \"severity|asc\")", - }, - "limit": { - Type: gollem.TypeNumber, - Description: "Maximum number of alerts to return (default: 100, max: 1000)", - }, - "after": { - Type: gollem.TypeString, - Description: "Cursor pagination token from a previous response for fetching next page", - }, - }, - }, - { - Name: "falcon_get_alerts", - Description: "Get detailed alert information by composite IDs. Use this when you already have specific alert IDs.", - Parameters: map[string]*gollem.Parameter{ - "composite_ids": { - Type: gollem.TypeString, - Description: "Comma-separated composite alert IDs", - Required: true, - }, - }, - }, - { - Name: "falcon_search_behaviors", - Description: "Search for behavior IDs using FQL filters. Returns behavior IDs that can be used with falcon_get_behaviors for full details.", - Parameters: map[string]*gollem.Parameter{ - "filter": { - Type: gollem.TypeString, - Description: "FQL filter expression", - }, - "limit": { - Type: gollem.TypeNumber, - Description: "Maximum number of IDs to return (default: 100, max: 500)", - }, - "offset": { - Type: gollem.TypeNumber, - Description: "Pagination offset", - }, - }, - }, - { - Name: "falcon_get_behaviors", - Description: "Get detailed behavior information by IDs. Returns behavior details including tactic, technique, severity, pattern, and associated device info.", - Parameters: map[string]*gollem.Parameter{ - "ids": { - Type: gollem.TypeString, - Description: "Comma-separated behavior IDs", - Required: true, - }, - }, - }, - { - Name: "falcon_get_crowdscores", - Description: "Get CrowdScore values for the environment. CrowdScore is an overall threat level indicator.", - Parameters: map[string]*gollem.Parameter{ - "filter": { - Type: gollem.TypeString, - Description: "FQL filter expression (e.g., \"timestamp:>'2025-01-01'\")", - }, - }, - }, - { - Name: "falcon_search_devices", - Description: "Search for device (host) IDs using FQL filters. Returns a list of device IDs that can be used with falcon_get_devices to retrieve full host details including OS, IP addresses, sensor version, and containment status.", - Parameters: map[string]*gollem.Parameter{ - "filter": { - Type: gollem.TypeString, - Description: "FQL filter expression (e.g., \"hostname:'*web*'\", \"platform_name:'Windows'\", \"last_seen:>='2025-01-01'\", \"external_ip:'10.0.0.*'\")", - }, - "sort": { - Type: gollem.TypeString, - Description: "Sort expression (e.g., \"hostname.asc\", \"last_seen.desc\")", - }, - "limit": { - Type: gollem.TypeNumber, - Description: "Maximum number of IDs to return (default: 100, max: 5000)", - }, - "offset": { - Type: gollem.TypeString, - Description: "Pagination offset token from a previous response", - }, - }, - }, - { - Name: "falcon_get_devices", - Description: "Get detailed device (host) information by device IDs. Returns full host details including hostname, OS, IP addresses, sensor version, tags, and containment status.", - Parameters: map[string]*gollem.Parameter{ - "ids": { - Type: gollem.TypeString, - Description: "Comma-separated device IDs (up to 5000, e.g., \"abc123def456,ghi789jkl012\")", - Required: true, - }, - }, - }, - { - Name: "falcon_search_events", - Description: "Search EDR telemetry events using CrowdStrike Query Language (CQL). This uses the Next-Gen SIEM Search API to query raw event data (process executions, network connections, file writes, DNS requests, etc.). The search runs asynchronously and this tool automatically polls until results are ready.", - Parameters: map[string]*gollem.Parameter{ - "query_string": { - Type: gollem.TypeString, - Description: "CQL query string (e.g., \"aid=abc123\", \"#event_simpleName=ProcessRollup2 AND FileName=cmd.exe\", \"ComputerName=workstation1 | tail(100)\")", - Required: true, - }, - "repository": { - Type: gollem.TypeString, - Description: "Repository to search. Values: \"search-all\" (all data, default), \"investigate_view\" (Falcon EDR), \"third-party\" (third-party data), \"falcon_for_it_view\" (IT Automation), \"forensics_view\" (Forensics)", - }, - "start": { - Type: gollem.TypeString, - Description: "Start time for the search (e.g., \"1d\" for last 1 day, \"24h\" for last 24 hours, \"2025-01-01T00:00:00Z\" for absolute time). Default: \"1d\"", - }, - "end": { - Type: gollem.TypeString, - Description: "End time for the search (e.g., \"now\", \"2025-01-02T00:00:00Z\"). Default: \"now\"", - }, - }, - }, - }, nil +// Tool descriptions. Kept as constants so the typed-tool registration in +// newInternalTool stays readable and the wire-level descriptions are unchanged. +const ( + descSearchIncidents = "Search for incident IDs using FQL (Falcon Query Language) filters. Returns a list of incident IDs that can be used with falcon_get_incidents to retrieve full details." + descGetIncidents = "Get detailed information for specific incidents by their IDs. Returns full incident details including status, tactics, techniques, hosts, and users involved." + descSearchAlerts = "Search and retrieve alert details in one call using FQL filters with cursor-based pagination. Returns full alert objects including severity, tactic, technique, and device info." + descGetAlerts = "Get detailed alert information by composite IDs. Use this when you already have specific alert IDs." + descSearchBehaviors = "Search for behavior IDs using FQL filters. Returns behavior IDs that can be used with falcon_get_behaviors for full details." + descGetBehaviors = "Get detailed behavior information by IDs. Returns behavior details including tactic, technique, severity, pattern, and associated device info." + descGetCrowdScores = "Get CrowdScore values for the environment. CrowdScore is an overall threat level indicator." + descSearchDevices = "Search for device (host) IDs using FQL filters. Returns a list of device IDs that can be used with falcon_get_devices to retrieve full host details including OS, IP addresses, sensor version, and containment status." + descGetDevices = "Get detailed device (host) information by device IDs. Returns full host details including hostname, OS, IP addresses, sensor version, tags, and containment status." + descSearchEvents = "Search EDR telemetry events using CrowdStrike Query Language (CQL). This uses the Next-Gen SIEM Search API to query raw event data (process executions, network connections, file writes, DNS requests, etc.). The search runs asynchronously and this tool automatically polls until results are ready." +) + +// Typed inputs for each tool. The schema is inferred from these struct tags by +// gollem.NewTool. Pagination counts use float64 (not int64) so the inferred JSON +// schema stays "number" — matching the wire-level type the LLM was given before +// this migration. +type searchIncidentsInput struct { + Filter string `json:"filter" description:"FQL filter expression (e.g., \"status:'30'\", \"tags:'critical'\", \"start:>'2025-01-01'\")"` + Sort string `json:"sort" description:"Sort expression (e.g., \"start.desc\", \"end.asc\")"` + Limit float64 `json:"limit" description:"Maximum number of IDs to return (default: 100, max: 500)"` + Offset float64 `json:"offset" description:"Pagination offset"` +} + +type getIncidentsInput struct { + IDs string `json:"ids" required:"true" description:"Comma-separated incident IDs (e.g., \"inc:abc123:def456,inc:abc123:ghi789\")"` +} + +type searchAlertsInput struct { + Filter string `json:"filter" description:"FQL filter expression (e.g., \"status:'new'\", \"severity:>50\", \"tactics:'Lateral Movement'\")"` + Sort string `json:"sort" description:"Sort property (e.g., \"timestamp|desc\", \"severity|asc\")"` + Limit float64 `json:"limit" description:"Maximum number of alerts to return (default: 100, max: 1000)"` + After string `json:"after" description:"Cursor pagination token from a previous response for fetching next page"` +} + +type getAlertsInput struct { + CompositeIDs string `json:"composite_ids" required:"true" description:"Comma-separated composite alert IDs"` +} + +type searchBehaviorsInput struct { + Filter string `json:"filter" description:"FQL filter expression"` + Limit float64 `json:"limit" description:"Maximum number of IDs to return (default: 100, max: 500)"` + Offset float64 `json:"offset" description:"Pagination offset"` +} + +type getBehaviorsInput struct { + IDs string `json:"ids" required:"true" description:"Comma-separated behavior IDs"` +} + +type getCrowdScoresInput struct { + Filter string `json:"filter" description:"FQL filter expression (e.g., \"timestamp:>'2025-01-01'\")"` +} + +type searchDevicesInput struct { + Filter string `json:"filter" description:"FQL filter expression (e.g., \"hostname:'*web*'\", \"platform_name:'Windows'\", \"last_seen:>='2025-01-01'\", \"external_ip:'10.0.0.*'\")"` + Sort string `json:"sort" description:"Sort expression (e.g., \"hostname.asc\", \"last_seen.desc\")"` + Limit float64 `json:"limit" description:"Maximum number of IDs to return (default: 100, max: 5000)"` + Offset string `json:"offset" description:"Pagination offset token from a previous response"` +} + +type getDevicesInput struct { + IDs string `json:"ids" required:"true" description:"Comma-separated device IDs (up to 5000, e.g., \"abc123def456,ghi789jkl012\")"` +} + +type searchEventsInput struct { + QueryString string `json:"query_string" required:"true" description:"CQL query string (e.g., \"aid=abc123\", \"#event_simpleName=ProcessRollup2 AND FileName=cmd.exe\", \"ComputerName=workstation1 | tail(100)\")"` + Repository string `json:"repository" description:"Repository to search. Values: \"search-all\" (all data, default), \"investigate_view\" (Falcon EDR), \"third-party\" (third-party data), \"falcon_for_it_view\" (IT Automation), \"forensics_view\" (Forensics)"` + Start string `json:"start" description:"Start time for the search (e.g., \"1d\" for last 1 day, \"24h\" for last 24 hours, \"2025-01-01T00:00:00Z\" for absolute time). Default: \"1d\""` + End string `json:"end" description:"End time for the search (e.g., \"now\", \"2025-01-02T00:00:00Z\"). Default: \"now\""` +} + +// Startup assertions: validate each tool's In/Out types form a valid schema at +// package init, so a malformed type fails immediately rather than at first use. +var ( + _ = gollem.MustToolSchema[searchIncidentsInput, map[string]any]() + _ = gollem.MustToolSchema[getIncidentsInput, map[string]any]() + _ = gollem.MustToolSchema[searchAlertsInput, map[string]any]() + _ = gollem.MustToolSchema[getAlertsInput, map[string]any]() + _ = gollem.MustToolSchema[searchBehaviorsInput, map[string]any]() + _ = gollem.MustToolSchema[getBehaviorsInput, map[string]any]() + _ = gollem.MustToolSchema[getCrowdScoresInput, map[string]any]() + _ = gollem.MustToolSchema[searchDevicesInput, map[string]any]() + _ = gollem.MustToolSchema[getDevicesInput, map[string]any]() + _ = gollem.MustToolSchema[searchEventsInput, map[string]any]() +) + +func (t *internalTool) Specs(ctx context.Context) ([]gollem.ToolSpec, error) { + return t.tools.Specs(ctx) } func (t *internalTool) Run(ctx context.Context, name string, args map[string]any) (map[string]any, error) { - switch name { - case "falcon_search_incidents": - return t.searchIncidents(ctx, args) - case "falcon_get_incidents": - return t.getIncidents(ctx, args) - case "falcon_search_alerts": - return t.searchAlerts(ctx, args) - case "falcon_get_alerts": - return t.getAlerts(ctx, args) - case "falcon_search_behaviors": - return t.searchBehaviors(ctx, args) - case "falcon_get_behaviors": - return t.getBehaviors(ctx, args) - case "falcon_search_devices": - return t.searchDevices(ctx, args) - case "falcon_get_devices": - return t.getDevices(ctx, args) - case "falcon_get_crowdscores": - return t.getCrowdScores(ctx, args) - case "falcon_search_events": - return t.searchEvents(ctx, args) - default: - return nil, goerr.New("unknown tool name", goerr.V("name", name)) - } + return t.tools.Run(ctx, name, args) } // doRequest performs an authenticated HTTP request to the CrowdStrike API. @@ -321,12 +245,17 @@ func (t *internalTool) doRequestOnce(ctx context.Context, method, path string, b } // searchIncidents searches for incident IDs using FQL filters. -func (t *internalTool) searchIncidents(ctx context.Context, args map[string]any) (map[string]any, error) { - filter, _ := args["filter"].(string) +func (t *internalTool) searchIncidents(ctx context.Context, in searchIncidentsInput) (map[string]any, error) { + filter := in.Filter msg.Trace(ctx, "🔍 Searching incidents (filter: `%s`)", filter) path := "/incidents/queries/incidents/v1" - params := buildQueryParams(args, "filter", "sort", "limit", "offset") + params := queryParams( + [2]string{"filter", filter}, + [2]string{"sort", in.Sort}, + [2]string{"limit", numParam(in.Limit)}, + [2]string{"offset", numParam(in.Offset)}, + ) if params != "" { path += "?" + params } @@ -340,9 +269,9 @@ func (t *internalTool) searchIncidents(ctx context.Context, args map[string]any) } // getIncidents retrieves incident details by IDs. -func (t *internalTool) getIncidents(ctx context.Context, args map[string]any) (map[string]any, error) { - ids, ok := args["ids"].(string) - if !ok || ids == "" { +func (t *internalTool) getIncidents(ctx context.Context, in getIncidentsInput) (map[string]any, error) { + ids := in.IDs + if ids == "" { return nil, goerr.New("ids is required") } @@ -360,22 +289,22 @@ func (t *internalTool) getIncidents(ctx context.Context, args map[string]any) (m } // searchAlerts searches and retrieves alert details using FQL filters. -func (t *internalTool) searchAlerts(ctx context.Context, args map[string]any) (map[string]any, error) { - filter, _ := args["filter"].(string) +func (t *internalTool) searchAlerts(ctx context.Context, in searchAlertsInput) (map[string]any, error) { + filter := in.Filter msg.Trace(ctx, "🔍 Searching alerts (filter: `%s`)", filter) body := make(map[string]any) if filter != "" { body["filter"] = filter } - if sort, ok := args["sort"].(string); ok && sort != "" { - body["sort"] = sort + if in.Sort != "" { + body["sort"] = in.Sort } - if limit, ok := args["limit"].(float64); ok { - body["limit"] = int(limit) + if in.Limit > 0 { + body["limit"] = int(in.Limit) } - if after, ok := args["after"].(string); ok && after != "" { - body["after"] = after + if in.After != "" { + body["after"] = in.After } result, err := t.doRequest(ctx, http.MethodPost, "/alerts/combined/alerts/v1", body) if err != nil { @@ -387,9 +316,9 @@ func (t *internalTool) searchAlerts(ctx context.Context, args map[string]any) (m } // getAlerts retrieves alert details by composite IDs. -func (t *internalTool) getAlerts(ctx context.Context, args map[string]any) (map[string]any, error) { - compositeIDs, ok := args["composite_ids"].(string) - if !ok || compositeIDs == "" { +func (t *internalTool) getAlerts(ctx context.Context, in getAlertsInput) (map[string]any, error) { + compositeIDs := in.CompositeIDs + if compositeIDs == "" { return nil, goerr.New("composite_ids is required") } @@ -407,12 +336,16 @@ func (t *internalTool) getAlerts(ctx context.Context, args map[string]any) (map[ } // searchBehaviors searches for behavior IDs using FQL filters. -func (t *internalTool) searchBehaviors(ctx context.Context, args map[string]any) (map[string]any, error) { - filter, _ := args["filter"].(string) +func (t *internalTool) searchBehaviors(ctx context.Context, in searchBehaviorsInput) (map[string]any, error) { + filter := in.Filter msg.Trace(ctx, "🔍 Searching behaviors (filter: `%s`)", filter) path := "/incidents/queries/behaviors/v1" - params := buildQueryParams(args, "filter", "limit", "offset") + params := queryParams( + [2]string{"filter", filter}, + [2]string{"limit", numParam(in.Limit)}, + [2]string{"offset", numParam(in.Offset)}, + ) if params != "" { path += "?" + params } @@ -426,9 +359,9 @@ func (t *internalTool) searchBehaviors(ctx context.Context, args map[string]any) } // getBehaviors retrieves behavior details by IDs. -func (t *internalTool) getBehaviors(ctx context.Context, args map[string]any) (map[string]any, error) { - ids, ok := args["ids"].(string) - if !ok || ids == "" { +func (t *internalTool) getBehaviors(ctx context.Context, in getBehaviorsInput) (map[string]any, error) { + ids := in.IDs + if ids == "" { return nil, goerr.New("ids is required") } @@ -446,12 +379,18 @@ func (t *internalTool) getBehaviors(ctx context.Context, args map[string]any) (m } // searchDevices searches for device (host) IDs using FQL filters. -func (t *internalTool) searchDevices(ctx context.Context, args map[string]any) (map[string]any, error) { - filter, _ := args["filter"].(string) +func (t *internalTool) searchDevices(ctx context.Context, in searchDevicesInput) (map[string]any, error) { + filter := in.Filter msg.Trace(ctx, "🔍 Searching devices (filter: `%s`)", filter) path := "/devices/queries/devices-scroll/v1" - params := buildQueryParams(args, "filter", "sort", "limit", "offset") + // Note: device offset is a scroll token (string), not a numeric offset. + params := queryParams( + [2]string{"filter", filter}, + [2]string{"sort", in.Sort}, + [2]string{"limit", numParam(in.Limit)}, + [2]string{"offset", in.Offset}, + ) if params != "" { path += "?" + params } @@ -465,9 +404,9 @@ func (t *internalTool) searchDevices(ctx context.Context, args map[string]any) ( } // getDevices retrieves device (host) details by IDs. -func (t *internalTool) getDevices(ctx context.Context, args map[string]any) (map[string]any, error) { - ids, ok := args["ids"].(string) - if !ok || ids == "" { +func (t *internalTool) getDevices(ctx context.Context, in getDevicesInput) (map[string]any, error) { + ids := in.IDs + if ids == "" { return nil, goerr.New("ids is required") } @@ -485,12 +424,12 @@ func (t *internalTool) getDevices(ctx context.Context, args map[string]any) (map } // getCrowdScores retrieves CrowdScore values. -func (t *internalTool) getCrowdScores(ctx context.Context, args map[string]any) (map[string]any, error) { - filter, _ := args["filter"].(string) +func (t *internalTool) getCrowdScores(ctx context.Context, in getCrowdScoresInput) (map[string]any, error) { + filter := in.Filter msg.Trace(ctx, "📊 Retrieving CrowdScores (filter: `%s`)", filter) path := "/incidents/combined/crowdscores/v1" - params := buildQueryParams(args, "filter") + params := queryParams([2]string{"filter", filter}) if params != "" { path += "?" + params } @@ -505,29 +444,29 @@ func (t *internalTool) getCrowdScores(ctx context.Context, args map[string]any) // searchEvents runs a CQL query via the Next-Gen SIEM Search API. // It creates a query job and polls until the job completes, returning all events. -func (t *internalTool) searchEvents(ctx context.Context, args map[string]any) (map[string]any, error) { +func (t *internalTool) searchEvents(ctx context.Context, in searchEventsInput) (map[string]any, error) { log := logging.From(ctx) - queryString, ok := args["query_string"].(string) - if !ok || queryString == "" { + queryString := in.QueryString + if queryString == "" { return nil, goerr.New("query_string is required") } repository := "search-all" - if repo, ok := args["repository"].(string); ok && repo != "" { - repository = repo + if in.Repository != "" { + repository = in.Repository } body := map[string]any{ "queryString": queryString, } - if start, ok := args["start"].(string); ok && start != "" { - body["start"] = start + if in.Start != "" { + body["start"] = in.Start } else { body["start"] = "1d" } - if end, ok := args["end"].(string); ok && end != "" { - body["end"] = end + if in.End != "" { + body["end"] = in.End } else { body["end"] = "now" } @@ -632,24 +571,29 @@ func (t *internalTool) searchEvents(ctx context.Context, args map[string]any) (m }, nil } -// buildQueryParams constructs URL query parameters from tool arguments. -func buildQueryParams(args map[string]any, keys ...string) string { +// queryParams constructs a URL query string from string key/value pairs, +// skipping empty values. Callers pre-format non-string values (e.g. numeric +// limits) before passing them in. +func queryParams(pairs ...[2]string) string { params := url.Values{} - for _, key := range keys { - if val, ok := args[key]; ok { - switch v := val.(type) { - case string: - if v != "" { - params.Set(key, v) - } - case float64: - params.Set(key, fmt.Sprintf("%d", int(v))) - } + for _, kv := range pairs { + if kv[1] != "" { + params.Set(kv[0], kv[1]) } } return params.Encode() } +// numParam formats a numeric pagination value for a query string. Zero is +// treated as "absent" (returns ""), matching the previous behavior where a +// missing limit/offset argument was simply omitted from the request. +func numParam(v float64) string { + if v <= 0 { + return "" + } + return fmt.Sprintf("%d", int(v)) +} + // splitAndTrim splits a comma-separated string and trims whitespace from each element. func splitAndTrim(s string) []string { parts := strings.Split(s, ",") diff --git a/pkg/agents/slack/export_test.go b/pkg/agents/slack/export_test.go index 4408e1b9..bc81c198 100644 --- a/pkg/agents/slack/export_test.go +++ b/pkg/agents/slack/export_test.go @@ -11,16 +11,13 @@ type InternalTool = internalTool // NewInternalToolForTest creates an internalTool for testing func NewInternalToolForTest(slackClient interfaces.SlackClient, maxLimit int) *internalTool { - return &internalTool{ - slackClient: slackClient, - maxLimit: maxLimit, - } + return newInternalTool(slackClient, maxLimit) } // NewToolSetForTest creates a toolSet instance for testing func NewToolSetForTest(slackClient interfaces.SlackClient) *toolSet { return &toolSet{ - tool: &internalTool{slackClient: slackClient}, + tool: newInternalTool(slackClient, 0), } } diff --git a/pkg/agents/slack/factory.go b/pkg/agents/slack/factory.go index df756e23..3d263dd9 100644 --- a/pkg/agents/slack/factory.go +++ b/pkg/agents/slack/factory.go @@ -37,7 +37,7 @@ func (f *Factory) Configure(ctx context.Context) (interfaces.ToolSet, error) { slackClient := slackSDK.New(f.oauthToken) ts := &toolSet{ - tool: &internalTool{slackClient: slackClient}, + tool: newInternalTool(slackClient, 0), } logging.From(ctx).Info("Slack Search Agent configured") diff --git a/pkg/agents/slack/tool.go b/pkg/agents/slack/tool.go index 1da9e4f4..654e462d 100644 --- a/pkg/agents/slack/tool.go +++ b/pkg/agents/slack/tool.go @@ -10,12 +10,31 @@ import ( "github.com/gollem-dev/gollem" "github.com/m-mizutani/goerr/v2" "github.com/secmon-lab/warren/pkg/domain/interfaces" + "github.com/secmon-lab/warren/pkg/utils/toolset" slackSDK "github.com/slack-go/slack" ) type internalTool struct { slackClient interfaces.SlackClient maxLimit int // Maximum number of results allowed from search + + tools gollem.ToolSet +} + +// newInternalTool creates an internalTool and builds its type-safe tool set. +func newInternalTool(slackClient interfaces.SlackClient, maxLimit int) *internalTool { + t := &internalTool{ + slackClient: slackClient, + maxLimit: maxLimit, + } + + t.tools = toolset.New( + gollem.MustNewTool("slack_search_messages", descSearchMessages, t.searchMessages), + gollem.MustNewTool("slack_get_thread_messages", descGetThreadMessages, t.getThreadMessages), + gollem.MustNewTool("slack_get_context_messages", descGetContextMessages, t.getContextMessages), + ) + + return t } // parseSlackTimestamp parses a Slack timestamp string (e.g., "1234567890.123456") @@ -53,102 +72,57 @@ func parseSlackTimestamp(tsStr string) time.Time { return time.Unix(sec, nsec) } +// Tool descriptions. Kept as constants so the typed-tool registration stays +// readable and the wire-level descriptions are unchanged. +const ( + descSearchMessages = "Search for messages in Slack workspace using the search.messages API" + descGetThreadMessages = "Get all messages in a thread" + descGetContextMessages = "Get messages before and after a specific message timestamp" +) + +// Typed inputs for each tool. Counts use float64 so the inferred JSON schema +// stays "number", matching the wire-level type given before this migration. +type searchMessagesInput struct { + Query string `json:"query" required:"true" description:"The search query (e.g., 'from:@user', 'in:general', 'has:link')"` + Sort string `json:"sort" description:"Sort order: 'score' (relevance) or 'timestamp' (newest first)"` + SortDir string `json:"sort_dir" description:"Sort direction: 'asc' or 'desc'"` + Count float64 `json:"count" description:"Number of results to return (default: 20, max: 100)"` + Page float64 `json:"page" description:"Page number for pagination (default: 1)"` + Highlight bool `json:"highlight" description:"Enable highlighting of search terms in results"` +} + +type getThreadMessagesInput struct { + Channel string `json:"channel" required:"true" description:"Channel ID"` + ThreadTS string `json:"thread_ts" required:"true" description:"Thread timestamp (ts of the parent message)"` + Limit float64 `json:"limit" description:"Maximum number of messages to return (default: 50, max: 200)"` +} + +type getContextMessagesInput struct { + Channel string `json:"channel" required:"true" description:"Channel ID"` + AroundTS string `json:"around_ts" required:"true" description:"Timestamp of the message to get context around"` + Before float64 `json:"before" description:"Number of messages before the timestamp (default: 10)"` + After float64 `json:"after" description:"Number of messages after the timestamp (default: 10)"` +} + +// Startup assertions: validate each tool's In/Out types form a valid schema at +// package init, so a malformed type fails immediately rather than at first use. +var ( + _ = gollem.MustToolSchema[searchMessagesInput, map[string]any]() + _ = gollem.MustToolSchema[getThreadMessagesInput, map[string]any]() + _ = gollem.MustToolSchema[getContextMessagesInput, map[string]any]() +) + func (t *internalTool) Specs(ctx context.Context) ([]gollem.ToolSpec, error) { - return []gollem.ToolSpec{ - { - Name: "slack_search_messages", - Description: "Search for messages in Slack workspace using the search.messages API", - Parameters: map[string]*gollem.Parameter{ - "query": { - Type: gollem.TypeString, - Description: "The search query (e.g., 'from:@user', 'in:general', 'has:link')", - Required: true, - }, - "sort": { - Type: gollem.TypeString, - Description: "Sort order: 'score' (relevance) or 'timestamp' (newest first)", - }, - "sort_dir": { - Type: gollem.TypeString, - Description: "Sort direction: 'asc' or 'desc'", - }, - "count": { - Type: gollem.TypeNumber, - Description: "Number of results to return (default: 20, max: 100)", - }, - "page": { - Type: gollem.TypeNumber, - Description: "Page number for pagination (default: 1)", - }, - "highlight": { - Type: gollem.TypeBoolean, - Description: "Enable highlighting of search terms in results", - }, - }, - }, - { - Name: "slack_get_thread_messages", - Description: "Get all messages in a thread", - Parameters: map[string]*gollem.Parameter{ - "channel": { - Type: gollem.TypeString, - Description: "Channel ID", - Required: true, - }, - "thread_ts": { - Type: gollem.TypeString, - Description: "Thread timestamp (ts of the parent message)", - Required: true, - }, - "limit": { - Type: gollem.TypeNumber, - Description: "Maximum number of messages to return (default: 50, max: 200)", - }, - }, - }, - { - Name: "slack_get_context_messages", - Description: "Get messages before and after a specific message timestamp", - Parameters: map[string]*gollem.Parameter{ - "channel": { - Type: gollem.TypeString, - Description: "Channel ID", - Required: true, - }, - "around_ts": { - Type: gollem.TypeString, - Description: "Timestamp of the message to get context around", - Required: true, - }, - "before": { - Type: gollem.TypeNumber, - Description: "Number of messages before the timestamp (default: 10)", - }, - "after": { - Type: gollem.TypeNumber, - Description: "Number of messages after the timestamp (default: 10)", - }, - }, - }, - }, nil + return t.tools.Specs(ctx) } func (t *internalTool) Run(ctx context.Context, name string, args map[string]any) (map[string]any, error) { - switch name { - case "slack_search_messages": - return t.searchMessages(ctx, args) - case "slack_get_thread_messages": - return t.getThreadMessages(ctx, args) - case "slack_get_context_messages": - return t.getContextMessages(ctx, args) - default: - return nil, goerr.New("unknown tool name", goerr.V("name", name)) - } + return t.tools.Run(ctx, name, args) } -func (t *internalTool) searchMessages(ctx context.Context, args map[string]any) (map[string]any, error) { - query, ok := args["query"].(string) - if !ok || query == "" { +func (t *internalTool) searchMessages(ctx context.Context, in searchMessagesInput) (map[string]any, error) { + query := in.Query + if query == "" { return nil, goerr.New("query is required") } @@ -157,21 +131,15 @@ func (t *internalTool) searchMessages(ctx context.Context, args map[string]any) Page: 1, } - if sort, ok := args["sort"].(string); ok { - params.Sort = sort - } - if sortDir, ok := args["sort_dir"].(string); ok { - params.SortDirection = sortDir - } - if count, ok := args["count"].(float64); ok { - params.Count = int(count) - } - if page, ok := args["page"].(float64); ok { - params.Page = int(page) + params.Sort = in.Sort + params.SortDirection = in.SortDir + if in.Count > 0 { + params.Count = int(in.Count) } - if highlight, ok := args["highlight"].(bool); ok { - params.Highlight = highlight + if in.Page > 0 { + params.Page = int(in.Page) } + params.Highlight = in.Highlight resp, err := t.slackClient.SearchMessagesContext(ctx, query, params) if err != nil { @@ -216,20 +184,20 @@ func (t *internalTool) searchMessages(ctx context.Context, args map[string]any) }, nil } -func (t *internalTool) getThreadMessages(ctx context.Context, args map[string]any) (map[string]any, error) { - channel, ok := args["channel"].(string) - if !ok || channel == "" { +func (t *internalTool) getThreadMessages(ctx context.Context, in getThreadMessagesInput) (map[string]any, error) { + channel := in.Channel + if channel == "" { return nil, goerr.New("channel is required") } - threadTS, ok := args["thread_ts"].(string) - if !ok || threadTS == "" { + threadTS := in.ThreadTS + if threadTS == "" { return nil, goerr.New("thread_ts is required") } limit := 50 - if l, ok := args["limit"].(float64); ok { - limit = int(l) + if in.Limit > 0 { + limit = int(in.Limit) } if limit > 200 { limit = 200 @@ -263,25 +231,26 @@ func (t *internalTool) getThreadMessages(ctx context.Context, args map[string]an }, nil } -func (t *internalTool) getContextMessages(ctx context.Context, args map[string]any) (map[string]any, error) { - channel, ok := args["channel"].(string) - if !ok || channel == "" { +func (t *internalTool) getContextMessages(ctx context.Context, in getContextMessagesInput) (map[string]any, error) { + channel := in.Channel + if channel == "" { return nil, goerr.New("channel is required") } - aroundTS, ok := args["around_ts"].(string) - if !ok || aroundTS == "" { + aroundTS := in.AroundTS + if aroundTS == "" { return nil, goerr.New("around_ts is required") } + // A zero/absent count falls back to the documented default of 10. before := 10 - if b, ok := args["before"].(float64); ok { - before = int(b) + if in.Before > 0 { + before = int(in.Before) } after := 10 - if a, ok := args["after"].(float64); ok { - after = int(a) + if in.After > 0 { + after = int(in.After) } // Parse the timestamp diff --git a/pkg/tool/base/action.go b/pkg/tool/base/action.go index 4325d385..aece9e81 100644 --- a/pkg/tool/base/action.go +++ b/pkg/tool/base/action.go @@ -3,13 +3,12 @@ package base import ( "context" "log/slog" - "reflect" "github.com/gollem-dev/gollem" - "github.com/m-mizutani/goerr/v2" "github.com/secmon-lab/warren/pkg/domain/interfaces" "github.com/secmon-lab/warren/pkg/domain/model/ticket" "github.com/secmon-lab/warren/pkg/domain/types" + "github.com/secmon-lab/warren/pkg/utils/toolset" "github.com/urfave/cli/v3" ) @@ -21,6 +20,8 @@ type Warren struct { ticketID types.TicketID slackUpdate SlackUpdateFunc llmClient gollem.LLMClient + + tools gollem.ToolSet } var _ interfaces.Tool = &Warren{} @@ -29,33 +30,6 @@ func (x *Warren) Helper() *cli.Command { return nil } -func getArg[T any](args map[string]any, key string) (T, error) { - var null T - val, ok := args[key] - if !ok { - return null, nil - } - - // Handle special case for numeric types from JSON (which come as float64) - if reflect.TypeOf(null).Kind() == reflect.Int64 { - if floatVal, ok := val.(float64); ok { - result := int64(floatVal) - return any(result).(T), nil - } - } - - typedVal, ok := val.(T) - if !ok { - return null, goerr.New("invalid parameter type", - goerr.V("key", key), - goerr.V("expected_type", reflect.TypeOf(null).String()), - goerr.V("actual_type", reflect.TypeOf(val).String()), - goerr.V("value", val)) - } - - return typedVal, nil -} - func New(repo interfaces.Repository, ticketID types.TicketID, opts ...func(*Warren)) *Warren { w := &Warren{ repo: repo, @@ -66,6 +40,18 @@ func New(repo interfaces.Repository, ticketID types.TicketID, opts ...func(*Warr opt(w) } + // Build the type-safe tool set once the dependencies from opts are wired in. + // Each tool's schema is inferred from its typed input struct, replacing the + // hand-written ToolSpec literals and the getArg map[string]any extraction. + w.tools = toolset.New( + gollem.MustNewTool(cmdGetAlerts, descGetAlerts, w.getAlerts), + gollem.MustNewTool(cmdFindNearestTicket, descFindNearestTicket, w.findNearestTicket), + gollem.MustNewTool(cmdSearchTicketsByWords, descSearchTicketsByWords, w.searchTicketsByWords), + gollem.MustNewTool(cmdUpdateFinding, descUpdateFinding, w.updateFinding), + gollem.MustNewTool(cmdGetTicketSessionMessages, descGetTicketSessionMessages, w.getTicketSessionMessages), + gollem.MustNewTool(cmdSearchSessionMessages, descSearchSessionMessages, w.searchSessionMessages), + ) + return w } @@ -131,136 +117,69 @@ func IgnorableTool(name string) bool { } } +// Tool descriptions. Kept as constants so the typed-tool registration in New +// stays readable and the wire-level descriptions remain unchanged. +const ( + descGetAlerts = "Get a set of alerts that is bound to the ticket with pagination support" + descFindNearestTicket = "Search the previous tickets that are similar to the current ticket" + descSearchTicketsByWords = "Search tickets using natural language query or keywords. Uses semantic similarity to find relevant tickets." + descUpdateFinding = "Update the finding information of the current ticket with analysis results" + descGetTicketSessionMessages = "Get chat messages (user inputs, AI responses, traces, plans, warnings) from every Session attached to the current ticket. Supersedes warren_get_ticket_comments by returning both human-authored messages and AI-produced outputs across Slack/Web/CLI channels in a unified shape." + descSearchSessionMessages = "Full-text search across every Session.Message attached to the current ticket. Returns the top matching messages (case-insensitive substring match). Use this when you need to look up prior discussion or investigation traces by keyword rather than by source/type." +) + +// Typed inputs for each tool. The schema (field names, types, required, +// descriptions) is inferred from these struct tags by gollem.NewTool. +type getAlertsInput struct { + Limit int64 `json:"limit" description:"Maximum number of alerts to return"` + Offset int64 `json:"offset" description:"Number of alerts to skip"` +} + +type findNearestTicketInput struct { + Limit int64 `json:"limit" description:"Maximum number of tickets to return"` + Duration int64 `json:"duration" description:"Duration of the ticket in days"` +} + +type searchTicketsByWordsInput struct { + Query string `json:"query" required:"true" description:"Search query using natural language or keywords to find similar tickets"` + Limit int64 `json:"limit" description:"Maximum number of tickets to return (default: 10)"` + Duration int64 `json:"duration" description:"Duration to search back in days (default: 30)"` +} + +type updateFindingInput struct { + Summary string `json:"summary" required:"true" description:"Summary of the investigation results analyzed by the agent"` + Severity string `json:"severity" required:"true" description:"Severity level of the finding. Must be one of: 'low', 'medium', 'high', 'critical'"` + Reason string `json:"reason" required:"true" description:"Detailed reasoning and justification for the severity assessment"` + Recommendation string `json:"recommendation" required:"true" description:"Recommended actions based on the analysis results"` +} + +type getTicketSessionMessagesInput struct { + Source string `json:"source" description:"Optional filter by Session source: 'slack', 'web', or 'cli'. Omit to include all sources."` + Type string `json:"type" description:"Optional filter by Message type: 'user', 'trace', 'plan', 'response', or 'warning'. Omit to include all types."` + Limit int64 `json:"limit" description:"Maximum number of messages to return (default: 50)"` + Offset int64 `json:"offset" description:"Number of messages to skip for pagination (default: 0)"` +} + +type searchSessionMessagesInput struct { + Query string `json:"query" required:"true" description:"Case-insensitive substring to search for across message content."` + Limit int64 `json:"limit" description:"Maximum number of messages to return (default: 50)"` +} + +// Startup assertions: validate each tool's In/Out types form a valid schema at +// package init, so a malformed type fails immediately rather than at first use. +var ( + _ = gollem.MustToolSchema[getAlertsInput, map[string]any]() + _ = gollem.MustToolSchema[findNearestTicketInput, map[string]any]() + _ = gollem.MustToolSchema[searchTicketsByWordsInput, map[string]any]() + _ = gollem.MustToolSchema[updateFindingInput, map[string]any]() + _ = gollem.MustToolSchema[getTicketSessionMessagesInput, map[string]any]() + _ = gollem.MustToolSchema[searchSessionMessagesInput, map[string]any]() +) + func (x *Warren) Specs(ctx context.Context) ([]gollem.ToolSpec, error) { - return []gollem.ToolSpec{ - { - Name: cmdGetAlerts, - Description: "Get a set of alerts that is bound to the ticket with pagination support", - Parameters: map[string]*gollem.Parameter{ - "limit": { - Type: gollem.TypeInteger, - Description: "Maximum number of alerts to return", - }, - "offset": { - Type: gollem.TypeInteger, - Description: "Number of alerts to skip", - }, - }, - }, - { - Name: cmdFindNearestTicket, - Description: "Search the previous tickets that are similar to the current ticket", - Parameters: map[string]*gollem.Parameter{ - "limit": { - Type: gollem.TypeInteger, - Description: "Maximum number of tickets to return", - }, - "duration": { - Type: gollem.TypeInteger, - Description: "Duration of the ticket in days", - }, - }, - }, - { - Name: cmdSearchTicketsByWords, - Description: "Search tickets using natural language query or keywords. Uses semantic similarity to find relevant tickets.", - Parameters: map[string]*gollem.Parameter{ - "query": { - Type: gollem.TypeString, - Description: "Search query using natural language or keywords to find similar tickets", - Required: true, - }, - "limit": { - Type: gollem.TypeInteger, - Description: "Maximum number of tickets to return (default: 10)", - }, - "duration": { - Type: gollem.TypeInteger, - Description: "Duration to search back in days (default: 30)", - }, - }, - }, - { - Name: cmdUpdateFinding, - Description: "Update the finding information of the current ticket with analysis results", - Parameters: map[string]*gollem.Parameter{ - "summary": { - Type: gollem.TypeString, - Description: "Summary of the investigation results analyzed by the agent", - Required: true, - }, - "severity": { - Type: gollem.TypeString, - Description: "Severity level of the finding. Must be one of: 'low', 'medium', 'high', 'critical'", - Required: true, - }, - "reason": { - Type: gollem.TypeString, - Description: "Detailed reasoning and justification for the severity assessment", - Required: true, - }, - "recommendation": { - Type: gollem.TypeString, - Description: "Recommended actions based on the analysis results", - Required: true, - }, - }, - }, - { - Name: cmdGetTicketSessionMessages, - Description: "Get chat messages (user inputs, AI responses, traces, plans, warnings) from every Session attached to the current ticket. Supersedes warren_get_ticket_comments by returning both human-authored messages and AI-produced outputs across Slack/Web/CLI channels in a unified shape.", - Parameters: map[string]*gollem.Parameter{ - "source": { - Type: gollem.TypeString, - Description: "Optional filter by Session source: 'slack', 'web', or 'cli'. Omit to include all sources.", - }, - "type": { - Type: gollem.TypeString, - Description: "Optional filter by Message type: 'user', 'trace', 'plan', 'response', or 'warning'. Omit to include all types.", - }, - "limit": { - Type: gollem.TypeInteger, - Description: "Maximum number of messages to return (default: 50)", - }, - "offset": { - Type: gollem.TypeInteger, - Description: "Number of messages to skip for pagination (default: 0)", - }, - }, - }, - { - Name: cmdSearchSessionMessages, - Description: "Full-text search across every Session.Message attached to the current ticket. Returns the top matching messages (case-insensitive substring match). Use this when you need to look up prior discussion or investigation traces by keyword rather than by source/type.", - Parameters: map[string]*gollem.Parameter{ - "query": { - Type: gollem.TypeString, - Description: "Case-insensitive substring to search for across message content.", - Required: true, - }, - "limit": { - Type: gollem.TypeInteger, - Description: "Maximum number of messages to return (default: 50)", - }, - }, - }, - }, nil + return x.tools.Specs(ctx) } func (x *Warren) Run(ctx context.Context, name string, args map[string]any) (map[string]any, error) { - switch name { - case cmdGetAlerts: - return x.getAlerts(ctx, args) - case cmdFindNearestTicket: - return x.findNearestTicket(ctx, args) - case cmdSearchTicketsByWords: - return x.searchTicketsByWords(ctx, args) - case cmdUpdateFinding: - return x.updateFinding(ctx, args) - case cmdGetTicketSessionMessages: - return x.getTicketSessionMessages(ctx, args) - case cmdSearchSessionMessages: - return x.searchSessionMessages(ctx, args) - default: - return nil, goerr.New("invalid function name", goerr.V("name", name)) - } + return x.tools.Run(ctx, name, args) } diff --git a/pkg/tool/base/alert.go b/pkg/tool/base/alert.go index 155057e8..07f13086 100644 --- a/pkg/tool/base/alert.go +++ b/pkg/tool/base/alert.go @@ -9,20 +9,9 @@ import ( "github.com/secmon-lab/warren/pkg/utils/errutil" ) -func (x *Warren) getAlerts(ctx context.Context, args map[string]any) (map[string]any, error) { - limit, err := getArg[int64](args, "limit") - if err != nil { - return nil, goerr.Wrap(err, "failed to get limit", - goerr.TV(errutil.ParameterKey, "limit"), - goerr.T(errutil.TagValidation)) - } - - offset, err := getArg[int64](args, "offset") - if err != nil { - return nil, goerr.Wrap(err, "failed to get offset", - goerr.TV(errutil.ParameterKey, "offset"), - goerr.T(errutil.TagValidation)) - } +func (x *Warren) getAlerts(ctx context.Context, in getAlertsInput) (map[string]any, error) { + limit := in.Limit + offset := in.Offset // Get current ticket currentTicket, err := x.repo.GetTicket(ctx, x.ticketID) diff --git a/pkg/tool/base/search_session_messages.go b/pkg/tool/base/search_session_messages.go index f8c437c3..7d548a89 100644 --- a/pkg/tool/base/search_session_messages.go +++ b/pkg/tool/base/search_session_messages.go @@ -19,25 +19,14 @@ import ( // ticket-scoped through the Repository method, we expose it as another // function on the ticket-scoped warren tool rather than a separate // package. -func (x *Warren) searchSessionMessages(ctx context.Context, args map[string]any) (map[string]any, error) { - query, err := getArg[string](args, "query") - if err != nil { - return nil, goerr.Wrap(err, "failed to get query", - goerr.TV(errutil.ParameterKey, "query"), - goerr.T(errutil.TagValidation)) - } +func (x *Warren) searchSessionMessages(ctx context.Context, in searchSessionMessagesInput) (map[string]any, error) { + query := in.Query if query == "" { return nil, goerr.New("query is required", goerr.TV(errutil.ParameterKey, "query"), goerr.T(errutil.TagValidation)) } - limitVal, err := getArg[int64](args, "limit") - if err != nil { - return nil, goerr.Wrap(err, "failed to get limit", - goerr.TV(errutil.ParameterKey, "limit"), - goerr.T(errutil.TagValidation)) - } - limit := int(limitVal) + limit := int(in.Limit) if limit <= 0 { limit = DefaultCommentsLimit } diff --git a/pkg/tool/base/ticket.go b/pkg/tool/base/ticket.go index 5630c2b9..eb6d95fb 100644 --- a/pkg/tool/base/ticket.go +++ b/pkg/tool/base/ticket.go @@ -13,20 +13,9 @@ import ( "github.com/secmon-lab/warren/pkg/utils/errutil" ) -func (x *Warren) findNearestTicket(ctx context.Context, args map[string]any) (map[string]any, error) { - limit, err := getArg[int64](args, "limit") - if err != nil { - return nil, goerr.Wrap(err, "failed to get limit", - goerr.TV(errutil.ParameterKey, "limit"), - goerr.T(errutil.TagValidation)) - } - - duration, err := getArg[int64](args, "duration") - if err != nil { - return nil, goerr.Wrap(err, "failed to get duration", - goerr.TV(errutil.ParameterKey, "duration"), - goerr.T(errutil.TagValidation)) - } +func (x *Warren) findNearestTicket(ctx context.Context, in findNearestTicketInput) (map[string]any, error) { + limit := in.Limit + duration := in.Duration // Get current ticket currentTicket, err := x.repo.GetTicket(ctx, x.ticketID) @@ -64,13 +53,8 @@ func (x *Warren) findNearestTicket(ctx context.Context, args map[string]any) (ma }, nil } -func (x *Warren) searchTicketsByWords(ctx context.Context, args map[string]any) (map[string]any, error) { - query, err := getArg[string](args, "query") - if err != nil { - return nil, goerr.Wrap(err, "failed to get query", - goerr.TV(errutil.ParameterKey, "query"), - goerr.T(errutil.TagValidation)) - } +func (x *Warren) searchTicketsByWords(ctx context.Context, in searchTicketsByWordsInput) (map[string]any, error) { + query := in.Query if query == "" { return nil, goerr.New("query is required", goerr.TV(errutil.ParameterKey, "query"), @@ -83,22 +67,12 @@ func (x *Warren) searchTicketsByWords(ctx context.Context, args map[string]any) goerr.T(errutil.TagInternal)) } - limit, err := getArg[int64](args, "limit") - if err != nil { - return nil, goerr.Wrap(err, "failed to get limit", - goerr.TV(errutil.ParameterKey, "limit"), - goerr.T(errutil.TagValidation)) - } + limit := in.Limit if limit <= 0 { limit = DefaultSearchTicketsLimit } - duration, err := getArg[int64](args, "duration") - if err != nil { - return nil, goerr.Wrap(err, "failed to get duration", - goerr.TV(errutil.ParameterKey, "duration"), - goerr.T(errutil.TagValidation)) - } + duration := in.Duration if duration <= 0 { duration = DefaultSearchTicketsDuration } @@ -131,49 +105,29 @@ func (x *Warren) searchTicketsByWords(ctx context.Context, args map[string]any) }, nil } -func (x *Warren) updateFinding(ctx context.Context, args map[string]any) (map[string]any, error) { - summary, err := getArg[string](args, "summary") - if err != nil { - return nil, goerr.Wrap(err, "failed to get summary", - goerr.TV(errutil.ParameterKey, "summary"), - goerr.T(errutil.TagValidation)) - } +func (x *Warren) updateFinding(ctx context.Context, in updateFindingInput) (map[string]any, error) { + summary := in.Summary if summary == "" { return nil, goerr.New("summary is required", goerr.TV(errutil.ParameterKey, "summary"), goerr.T(errutil.TagValidation)) } - severityStr, err := getArg[string](args, "severity") - if err != nil { - return nil, goerr.Wrap(err, "failed to get severity", - goerr.TV(errutil.ParameterKey, "severity"), - goerr.T(errutil.TagValidation)) - } + severityStr := in.Severity if severityStr == "" { return nil, goerr.New("severity is required", goerr.TV(errutil.ParameterKey, "severity"), goerr.T(errutil.TagValidation)) } - reason, err := getArg[string](args, "reason") - if err != nil { - return nil, goerr.Wrap(err, "failed to get reason", - goerr.TV(errutil.ParameterKey, "reason"), - goerr.T(errutil.TagValidation)) - } + reason := in.Reason if reason == "" { return nil, goerr.New("reason is required", goerr.TV(errutil.ParameterKey, "reason"), goerr.T(errutil.TagValidation)) } - recommendation, err := getArg[string](args, "recommendation") - if err != nil { - return nil, goerr.Wrap(err, "failed to get recommendation", - goerr.TV(errutil.ParameterKey, "recommendation"), - goerr.T(errutil.TagValidation)) - } + recommendation := in.Recommendation if recommendation == "" { return nil, goerr.New("recommendation is required", goerr.TV(errutil.ParameterKey, "recommendation"), diff --git a/pkg/tool/base/ticket_session_messages.go b/pkg/tool/base/ticket_session_messages.go index 3d31530c..fca27eff 100644 --- a/pkg/tool/base/ticket_session_messages.go +++ b/pkg/tool/base/ticket_session_messages.go @@ -17,34 +17,12 @@ import ( // warren_get_ticket_comments. It returns a strictly richer payload // (author information, AI output types, turn grouping) while covering // the same ground on the Slack/user-comment side. -func (x *Warren) getTicketSessionMessages(ctx context.Context, args map[string]any) (map[string]any, error) { - limitVal, err := getArg[int64](args, "limit") - if err != nil { - return nil, goerr.Wrap(err, "failed to get limit", - goerr.TV(errutil.ParameterKey, "limit"), - goerr.T(errutil.TagValidation)) - } - offsetVal, err := getArg[int64](args, "offset") - if err != nil { - return nil, goerr.Wrap(err, "failed to get offset", - goerr.TV(errutil.ParameterKey, "offset"), - goerr.T(errutil.TagValidation)) - } - sourceStr, err := getArg[string](args, "source") - if err != nil { - return nil, goerr.Wrap(err, "failed to get source", - goerr.TV(errutil.ParameterKey, "source"), - goerr.T(errutil.TagValidation)) - } - typeStr, err := getArg[string](args, "type") - if err != nil { - return nil, goerr.Wrap(err, "failed to get type", - goerr.TV(errutil.ParameterKey, "type"), - goerr.T(errutil.TagValidation)) - } +func (x *Warren) getTicketSessionMessages(ctx context.Context, in getTicketSessionMessagesInput) (map[string]any, error) { + sourceStr := in.Source + typeStr := in.Type - limit := int(limitVal) - offset := int(offsetVal) + limit := int(in.Limit) + offset := int(in.Offset) if limit <= 0 { limit = DefaultCommentsLimit } diff --git a/pkg/tool/bigquery/helper.go b/pkg/tool/bigquery/helper.go index 22820196..b905f66b 100644 --- a/pkg/tool/bigquery/helper.go +++ b/pkg/tool/bigquery/helper.go @@ -18,6 +18,7 @@ import ( "github.com/secmon-lab/warren/pkg/domain/model/prompt" "github.com/secmon-lab/warren/pkg/service/llm" "github.com/secmon-lab/warren/pkg/utils/logging" + "github.com/secmon-lab/warren/pkg/utils/toolset" "github.com/urfave/cli/v3" "google.golang.org/api/iterator" "gopkg.in/yaml.v3" @@ -331,14 +332,14 @@ func generateConfigWithFactory(ctx context.Context, cfg generateConfigInput, fac gollem.WithLogger(logger), gollem.WithContentBlockMiddleware(llm.NewCompactionMiddleware(llmClient, logger)), gollem.WithContentStreamMiddleware(llm.NewCompactionStreamMiddleware(llmClient)), - gollem.WithToolSets(&configGeneratorTools{ + gollem.WithToolSets(newConfigGeneratorTools(&configGeneratorTools{ bqClient: bqClient, scanLimit: cfg.ScanLimit, tableDatasetID: cfg.TableDatasetID, tableTableID: cfg.TableTableID, outputPath: outputPath, metadata: tableMetadata, - }), + })), ) if _, err := agent.Execute(ctx, gollem.Text("Generate configuration")); err != nil { @@ -393,49 +394,52 @@ type configGeneratorTools struct { tableTableID string outputPath string metadata *bigquery.TableMetadata + + tools gollem.ToolSet +} + +// configQueryInput is the typed argument for the bigquery_query helper tool. +type configQueryInput struct { + Query string `json:"query" required:"true" description:"The SQL query to execute"` +} + +// generateConfigToolInput is the typed argument for the generate_config helper +// tool. Config stays a map[string]any so the inferred schema remains an opaque +// object property (the concrete shape is described to the LLM via the system +// prompt), matching the wire schema before this migration. +type generateConfigToolInput struct { + Config map[string]any `json:"config" required:"true" description:"The complete configuration object"` +} + +// Startup assertions: validate each tool's In/Out types form a valid schema at +// package init, so a malformed type fails immediately rather than at first use. +var ( + _ = gollem.MustToolSchema[configQueryInput, map[string]any]() + _ = gollem.MustToolSchema[generateConfigToolInput, map[string]any]() +) + +// newConfigGeneratorTools builds the configGeneratorTools and its type-safe tool +// set. The bigquery_query description embeds the configured scan-size limit. +func newConfigGeneratorTools(t *configGeneratorTools) *configGeneratorTools { + queryDesc := fmt.Sprintf("Execute SQL query against BigQuery. Performs dry run to check scan size (limit: %s). Only use field names from the provided schema.", humanize.Bytes(t.scanLimit)) + t.tools = toolset.New( + gollem.MustNewTool("bigquery_query", queryDesc, t.executeBigQueryQuery), + gollem.MustNewTool("generate_config", "Generate the final YAML configuration. This validates the config against the actual schema and saves it to a file.", t.generateConfigOutput), + ) + return t } func (t *configGeneratorTools) Specs(ctx context.Context) ([]gollem.ToolSpec, error) { - return []gollem.ToolSpec{ - { - Name: "bigquery_query", - Description: fmt.Sprintf("Execute SQL query against BigQuery. Performs dry run to check scan size (limit: %s). Only use field names from the provided schema.", humanize.Bytes(t.scanLimit)), - Parameters: map[string]*gollem.Parameter{ - "query": { - Type: gollem.TypeString, - Description: "The SQL query to execute", - Required: true, - }, - }, - }, - { - Name: "generate_config", - Description: "Generate the final YAML configuration. This validates the config against the actual schema and saves it to a file.", - Parameters: map[string]*gollem.Parameter{ - "config": { - Type: gollem.TypeObject, - Description: "The complete configuration object", - Required: true, - }, - }, - }, - }, nil + return t.tools.Specs(ctx) } func (t *configGeneratorTools) Run(ctx context.Context, name string, args map[string]any) (map[string]any, error) { - switch name { - case "bigquery_query": - return t.executeBigQueryQuery(ctx, args) - case "generate_config": - return t.generateConfigOutput(ctx, args) - default: - return nil, goerr.New("invalid function name", goerr.V("name", name)) - } + return t.tools.Run(ctx, name, args) } -func (t *configGeneratorTools) executeBigQueryQuery(ctx context.Context, args map[string]any) (map[string]any, error) { - query, ok := args["query"].(string) - if !ok { +func (t *configGeneratorTools) executeBigQueryQuery(ctx context.Context, in configQueryInput) (map[string]any, error) { + query := in.Query + if query == "" { return nil, goerr.New("query parameter is required") } @@ -508,9 +512,9 @@ func (t *configGeneratorTools) executeBigQueryQuery(ctx context.Context, args ma }, nil } -func (t *configGeneratorTools) generateConfigOutput(ctx context.Context, args map[string]any) (map[string]any, error) { - config, ok := args["config"].(map[string]any) - if !ok { +func (t *configGeneratorTools) generateConfigOutput(ctx context.Context, in generateConfigToolInput) (map[string]any, error) { + config := in.Config + if config == nil { return nil, goerr.New("config parameter is required") } diff --git a/pkg/tool/github/action_test.go b/pkg/tool/github/action_test.go index 13b1d7f9..f8e87f0d 100644 --- a/pkg/tool/github/action_test.go +++ b/pkg/tool/github/action_test.go @@ -138,7 +138,7 @@ func TestGitHubSpecsDelegation(t *testing.T) { specs, err := action.Specs(context.Background()) gt.NoError(t, err) - gt.A(t, specs).Length(5) + gt.A(t, specs).Length(7) names := map[string]bool{} for _, s := range specs { @@ -150,6 +150,8 @@ func TestGitHubSpecsDelegation(t *testing.T) { "github_get_content", "github_list_commits", "github_get_blame", + "github_get_issue", + "github_get_pull_request", } { gt.Value(t, names[want]).Equal(true) } diff --git a/pkg/tool/knowledge/tool.go b/pkg/tool/knowledge/tool.go index 7b3595c7..720b62ce 100644 --- a/pkg/tool/knowledge/tool.go +++ b/pkg/tool/knowledge/tool.go @@ -13,6 +13,7 @@ import ( "github.com/secmon-lab/warren/pkg/domain/types" svcknowledge "github.com/secmon-lab/warren/pkg/service/knowledge" "github.com/secmon-lab/warren/pkg/utils/msg" + "github.com/secmon-lab/warren/pkg/utils/toolset" "github.com/urfave/cli/v3" ) @@ -34,6 +35,8 @@ type Tool struct { svc *svcknowledge.Service category types.KnowledgeCategory mode Mode + + tools gollem.ToolSet } var _ interfaces.Tool = &Tool{} @@ -41,11 +44,39 @@ var _ interfaces.Tool = &Tool{} // New creates a new knowledge v2 tool. // category is fixed per use case (e.g., "fact" for investigation context, "technique" for analysis procedures). func New(svc *svcknowledge.Service, category types.KnowledgeCategory, mode Mode) *Tool { - return &Tool{ + x := &Tool{ svc: svc, category: category, mode: mode, } + + // The available tool set depends on the mode: search-only agents get just + // knowledge_search, read-only agents also get knowledge_tag_list, and + // read-write agents get the full CRUD + tag-management surface. Building the + // set here (rather than branching inside Specs/Run) keeps the mode logic in + // one place while each tool stays type-safe via gollem.NewTool. + tools := []gollem.Tool{ + gollem.MustNewTool(cmdSearch, descSearch, x.search), + } + if mode != ModeSearchOnly { + tools = append(tools, gollem.MustNewTool(cmdTagList, descTagList, x.tagList)) + } + if mode == ModeReadWrite { + tools = append(tools, + gollem.MustNewTool(cmdSave, descSave, x.save), + gollem.MustNewTool(cmdDelete, descDelete, x.delete), + gollem.MustNewTool(cmdList, descList, x.list), + gollem.MustNewTool(cmdGet, descGet, x.get), + gollem.MustNewTool(cmdHistory, descHistory, x.history), + gollem.MustNewTool(cmdTagCreate, descTagCreate, x.tagCreate), + gollem.MustNewTool(cmdTagUpdate, descTagUpdate, x.tagUpdate), + gollem.MustNewTool(cmdTagDelete, descTagDelete, x.tagDelete), + gollem.MustNewTool(cmdTagMerge, descTagMerge, x.tagMerge), + ) + } + x.tools = toolset.New(tools...) + + return x } func (x *Tool) ID() string { @@ -106,234 +137,109 @@ const ( cmdTagMerge = "knowledge_tag_merge" ) -func (x *Tool) Specs(_ context.Context) ([]gollem.ToolSpec, error) { - searchSpec := gollem.ToolSpec{ - Name: cmdSearch, - Description: "Search for knowledge. Returns lightweight summaries (ID, title, tags only). Use knowledge_get to retrieve full details.", - Parameters: map[string]*gollem.Parameter{ - "query": { - Type: gollem.TypeString, - Description: "Natural language query to search for", - Required: true, - }, - "tags": { - Type: gollem.TypeArray, - Items: &gollem.Parameter{Type: gollem.TypeString}, - Description: "Tag IDs to filter by (at least one required)", - Required: true, - }, - }, - } - - if x.mode == ModeSearchOnly { - return []gollem.ToolSpec{searchSpec}, nil - } - - specs := []gollem.ToolSpec{ - searchSpec, - { - Name: cmdTagList, - Description: "List all available knowledge tags with their IDs and descriptions.", - }, - } - - if x.mode == ModeReadWrite { - specs = append(specs, - gollem.ToolSpec{ - Name: cmdSave, - Description: "Create or update a knowledge entry. Specify ID to update existing, omit for new.", - Parameters: map[string]*gollem.Parameter{ - "id": { - Type: gollem.TypeString, - Description: "Knowledge ID (for update). Omit for new entry.", - }, - "title": { - Type: gollem.TypeString, - Description: "Title of the knowledge entry (topic name)", - Required: true, - }, - "claim": { - Type: gollem.TypeString, - Description: "Markdown content with facts or techniques", - Required: true, - }, - "tags": { - Type: gollem.TypeArray, - Items: &gollem.Parameter{Type: gollem.TypeString}, - Description: "Tag IDs to associate (at least one required)", - Required: true, - }, - "message": { - Type: gollem.TypeString, - Description: "Reason for this change", - Required: true, - }, - "ticket_id": { - Type: gollem.TypeString, - Description: "Related ticket ID (optional)", - }, - }, - }, - gollem.ToolSpec{ - Name: cmdDelete, - Description: "Delete a knowledge entry that contains incorrect information. Records the reason in the log.", - Parameters: map[string]*gollem.Parameter{ - "id": { - Type: gollem.TypeString, - Description: "Knowledge ID to delete", - Required: true, - }, - "reason": { - Type: gollem.TypeString, - Description: "Reason for deletion", - Required: true, - }, - }, - }, - gollem.ToolSpec{ - Name: cmdList, - Description: "List knowledge entries (ID, title, tags only - lightweight for browsing). Use knowledge_get to retrieve full details.", - Parameters: map[string]*gollem.Parameter{ - "limit": { - Type: gollem.TypeInteger, - Description: "Maximum number of entries to return (default: 25, max: 100)", - }, - "offset": { - Type: gollem.TypeInteger, - Description: "Number of entries to skip for pagination (default: 0)", - }, - }, - }, - gollem.ToolSpec{ - Name: cmdGet, - Description: "Get full details (including claim) of one or more knowledge entries by ID.", - Parameters: map[string]*gollem.Parameter{ - "ids": { - Type: gollem.TypeArray, - Items: &gollem.Parameter{Type: gollem.TypeString}, - Description: "Knowledge IDs to retrieve (1-10 IDs)", - Required: true, - }, - }, - }, - gollem.ToolSpec{ - Name: cmdHistory, - Description: "Get change history of a knowledge entry.", - Parameters: map[string]*gollem.Parameter{ - "id": { - Type: gollem.TypeString, - Description: "Knowledge ID", - Required: true, - }, - }, - }, - gollem.ToolSpec{ - Name: cmdTagCreate, - Description: "Create a new tag.", - Parameters: map[string]*gollem.Parameter{ - "name": { - Type: gollem.TypeString, - Description: "Tag name (lowercase, short)", - Required: true, - }, - "description": { - Type: gollem.TypeString, - Description: "Description of what this tag represents", - }, - }, - }, - gollem.ToolSpec{ - Name: cmdTagUpdate, - Description: "Update an existing tag.", - Parameters: map[string]*gollem.Parameter{ - "id": { - Type: gollem.TypeString, - Description: "Tag ID", - Required: true, - }, - "name": { - Type: gollem.TypeString, - Description: "New tag name", - Required: true, - }, - "description": { - Type: gollem.TypeString, - Description: "New description", - }, - }, - }, - gollem.ToolSpec{ - Name: cmdTagDelete, - Description: "Delete a tag. Removes it from all knowledge entries that reference it.", - Parameters: map[string]*gollem.Parameter{ - "id": { - Type: gollem.TypeString, - Description: "Tag ID to delete", - Required: true, - }, - }, - }, - gollem.ToolSpec{ - Name: cmdTagMerge, - Description: "Merge two tags: replaces old_id with new_id in all knowledge entries, then deletes old_id.", - Parameters: map[string]*gollem.Parameter{ - "old_id": { - Type: gollem.TypeString, - Description: "Tag ID to be merged (will be deleted)", - Required: true, - }, - "new_id": { - Type: gollem.TypeString, - Description: "Tag ID to merge into (will be kept)", - Required: true, - }, - }, - }, - ) - } +// Tool descriptions. Kept as constants so the typed-tool registration in New +// stays readable and the wire-level descriptions are unchanged. +const ( + descSearch = "Search for knowledge. Returns lightweight summaries (ID, title, tags only). Use knowledge_get to retrieve full details." + descTagList = "List all available knowledge tags with their IDs and descriptions." + descSave = "Create or update a knowledge entry. Specify ID to update existing, omit for new." + descDelete = "Delete a knowledge entry that contains incorrect information. Records the reason in the log." + descList = "List knowledge entries (ID, title, tags only - lightweight for browsing). Use knowledge_get to retrieve full details." + descGet = "Get full details (including claim) of one or more knowledge entries by ID." + descHistory = "Get change history of a knowledge entry." + descTagCreate = "Create a new tag." + descTagUpdate = "Update an existing tag." + descTagDelete = "Delete a tag. Removes it from all knowledge entries that reference it." + descTagMerge = "Merge two tags: replaces old_id with new_id in all knowledge entries, then deletes old_id." +) + +// Typed inputs for each tool. The schema is inferred from these struct tags. +type searchInput struct { + Query string `json:"query" required:"true" description:"Natural language query to search for"` + Tags []string `json:"tags" required:"true" description:"Tag IDs to filter by (at least one required)"` +} + +// emptyInput is used for tools that take no arguments (knowledge_tag_list). +type emptyInput struct{} + +type saveInput struct { + ID string `json:"id" description:"Knowledge ID (for update). Omit for new entry."` + Title string `json:"title" required:"true" description:"Title of the knowledge entry (topic name)"` + Claim string `json:"claim" required:"true" description:"Markdown content with facts or techniques"` + Tags []string `json:"tags" required:"true" description:"Tag IDs to associate (at least one required)"` + Message string `json:"message" required:"true" description:"Reason for this change"` + TicketID string `json:"ticket_id" description:"Related ticket ID (optional)"` +} + +type deleteInput struct { + ID string `json:"id" required:"true" description:"Knowledge ID to delete"` + Reason string `json:"reason" required:"true" description:"Reason for deletion"` +} + +type listInput struct { + Limit int64 `json:"limit" description:"Maximum number of entries to return (default: 25, max: 100)"` + Offset int64 `json:"offset" description:"Number of entries to skip for pagination (default: 0)"` +} + +type getInput struct { + IDs []string `json:"ids" required:"true" description:"Knowledge IDs to retrieve (1-10 IDs)"` +} + +type historyInput struct { + ID string `json:"id" required:"true" description:"Knowledge ID"` +} + +type tagCreateInput struct { + Name string `json:"name" required:"true" description:"Tag name (lowercase, short)"` + Description string `json:"description" description:"Description of what this tag represents"` +} - return specs, nil +type tagUpdateInput struct { + ID string `json:"id" required:"true" description:"Tag ID"` + Name string `json:"name" required:"true" description:"New tag name"` + Description string `json:"description" description:"New description"` +} + +type tagDeleteInput struct { + ID string `json:"id" required:"true" description:"Tag ID to delete"` +} + +type tagMergeInput struct { + OldID string `json:"old_id" required:"true" description:"Tag ID to be merged (will be deleted)"` + NewID string `json:"new_id" required:"true" description:"Tag ID to merge into (will be kept)"` +} + +// Startup assertions: validate each tool's In/Out types form a valid schema at +// package init, so a malformed type fails immediately rather than at first use. +var ( + _ = gollem.MustToolSchema[searchInput, map[string]any]() + _ = gollem.MustToolSchema[emptyInput, map[string]any]() + _ = gollem.MustToolSchema[saveInput, map[string]any]() + _ = gollem.MustToolSchema[deleteInput, map[string]any]() + _ = gollem.MustToolSchema[listInput, map[string]any]() + _ = gollem.MustToolSchema[getInput, map[string]any]() + _ = gollem.MustToolSchema[historyInput, map[string]any]() + _ = gollem.MustToolSchema[tagCreateInput, map[string]any]() + _ = gollem.MustToolSchema[tagUpdateInput, map[string]any]() + _ = gollem.MustToolSchema[tagDeleteInput, map[string]any]() + _ = gollem.MustToolSchema[tagMergeInput, map[string]any]() +) + +func (x *Tool) Specs(ctx context.Context) ([]gollem.ToolSpec, error) { + return x.tools.Specs(ctx) } func (x *Tool) Run(ctx context.Context, name string, args map[string]any) (map[string]any, error) { - switch name { - case cmdSearch: - return x.search(ctx, args) - case cmdTagList: - return x.tagList(ctx) - case cmdSave: - return x.save(ctx, args) - case cmdDelete: - return x.delete(ctx, args) - case cmdList: - return x.list(ctx, args) - case cmdGet: - return x.get(ctx, args) - case cmdHistory: - return x.history(ctx, args) - case cmdTagCreate: - return x.tagCreate(ctx, args) - case cmdTagUpdate: - return x.tagUpdate(ctx, args) - case cmdTagDelete: - return x.tagDelete(ctx, args) - case cmdTagMerge: - return x.tagMerge(ctx, args) - default: - return nil, goerr.New("unknown command", goerr.V("name", name)) - } -} - -func (x *Tool) search(ctx context.Context, args map[string]any) (map[string]any, error) { - query, _ := args["query"].(string) + return x.tools.Run(ctx, name, args) +} + +func (x *Tool) search(ctx context.Context, in searchInput) (map[string]any, error) { + query := in.Query if query == "" { return nil, goerr.New("query is required") } - tagIDs, err := extractTagIDs(args, "tags") - if err != nil { - return nil, err - } + tagIDs := toTagIDs(in.Tags) if len(tagIDs) == 0 { return nil, goerr.New("at least one tag is required") } @@ -351,17 +257,14 @@ func (x *Tool) search(ctx context.Context, args map[string]any) (map[string]any, }, nil } -func (x *Tool) save(ctx context.Context, args map[string]any) (map[string]any, error) { - title, _ := args["title"].(string) - claim, _ := args["claim"].(string) - message, _ := args["message"].(string) - ticketID, _ := args["ticket_id"].(string) - id, _ := args["id"].(string) +func (x *Tool) save(ctx context.Context, in saveInput) (map[string]any, error) { + title := in.Title + claim := in.Claim + message := in.Message + ticketID := in.TicketID + id := in.ID - tagIDs, err := extractTagIDs(args, "tags") - if err != nil { - return nil, err - } + tagIDs := toTagIDs(in.Tags) k := &knowledgeModel.Knowledge{ ID: types.KnowledgeID(id), @@ -384,9 +287,9 @@ func (x *Tool) save(ctx context.Context, args map[string]any) (map[string]any, e }, nil } -func (x *Tool) delete(ctx context.Context, args map[string]any) (map[string]any, error) { - id, _ := args["id"].(string) - reason, _ := args["reason"].(string) +func (x *Tool) delete(ctx context.Context, in deleteInput) (map[string]any, error) { + id := in.ID + reason := in.Reason if id == "" { return nil, goerr.New("id is required") } @@ -403,16 +306,16 @@ func (x *Tool) delete(ctx context.Context, args map[string]any) (map[string]any, return map[string]any{"success": true}, nil } -func (x *Tool) list(ctx context.Context, args map[string]any) (map[string]any, error) { +func (x *Tool) list(ctx context.Context, in listInput) (map[string]any, error) { // Extract limit and offset parameters limit := 25 // default offset := 0 // default - if l, ok := args["limit"].(float64); ok { - limit = int(l) + if in.Limit > 0 { + limit = int(in.Limit) } - if o, ok := args["offset"].(float64); ok { - offset = int(o) + if in.Offset > 0 { + offset = int(in.Offset) } // Apply max limit @@ -480,30 +383,15 @@ func (x *Tool) list(ctx context.Context, args map[string]any) (map[string]any, e }, nil } -func (x *Tool) get(ctx context.Context, args map[string]any) (map[string]any, error) { - // Extract IDs array - idsRaw, ok := args["ids"] - if !ok { - return nil, goerr.New("ids parameter is required") - } - - idsArray, ok := idsRaw.([]any) - if !ok { - return nil, goerr.New("ids must be an array") - } - - if len(idsArray) == 0 { +func (x *Tool) get(ctx context.Context, in getInput) (map[string]any, error) { + if len(in.IDs) == 0 { return nil, goerr.New("at least one ID is required") } // Convert to KnowledgeID slice and deduplicate idMap := make(map[types.KnowledgeID]struct{}) var ids []types.KnowledgeID - for _, v := range idsArray { - s, ok := v.(string) - if !ok { - return nil, goerr.New("ID must be a string") - } + for _, s := range in.IDs { id := types.KnowledgeID(s) if _, exists := idMap[id]; !exists { idMap[id] = struct{}{} @@ -526,8 +414,8 @@ func (x *Tool) get(ctx context.Context, args map[string]any) (map[string]any, er }, nil } -func (x *Tool) history(ctx context.Context, args map[string]any) (map[string]any, error) { - id, _ := args["id"].(string) +func (x *Tool) history(ctx context.Context, in historyInput) (map[string]any, error) { + id := in.ID if id == "" { return nil, goerr.New("id is required") } @@ -551,7 +439,7 @@ func (x *Tool) history(ctx context.Context, args map[string]any) (map[string]any return map[string]any{"logs": entries, "count": len(entries)}, nil } -func (x *Tool) tagList(ctx context.Context) (map[string]any, error) { +func (x *Tool) tagList(ctx context.Context, _ emptyInput) (map[string]any, error) { tags, err := x.svc.ListTags(ctx) if err != nil { return nil, goerr.Wrap(err, "failed to list tags") @@ -569,9 +457,9 @@ func (x *Tool) tagList(ctx context.Context) (map[string]any, error) { return map[string]any{"tags": entries, "count": len(entries)}, nil } -func (x *Tool) tagCreate(ctx context.Context, args map[string]any) (map[string]any, error) { - name, _ := args["name"].(string) - description, _ := args["description"].(string) +func (x *Tool) tagCreate(ctx context.Context, in tagCreateInput) (map[string]any, error) { + name := in.Name + description := in.Description tag, err := x.svc.CreateTag(ctx, name, description) if err != nil { @@ -583,10 +471,10 @@ func (x *Tool) tagCreate(ctx context.Context, args map[string]any) (map[string]a return map[string]any{"success": true, "id": tag.ID.String()}, nil } -func (x *Tool) tagUpdate(ctx context.Context, args map[string]any) (map[string]any, error) { - id, _ := args["id"].(string) - name, _ := args["name"].(string) - description, _ := args["description"].(string) +func (x *Tool) tagUpdate(ctx context.Context, in tagUpdateInput) (map[string]any, error) { + id := in.ID + name := in.Name + description := in.Description if id == "" { return nil, goerr.New("id is required") } @@ -599,8 +487,8 @@ func (x *Tool) tagUpdate(ctx context.Context, args map[string]any) (map[string]a return map[string]any{"success": true, "id": tag.ID.String()}, nil } -func (x *Tool) tagDelete(ctx context.Context, args map[string]any) (map[string]any, error) { - id, _ := args["id"].(string) +func (x *Tool) tagDelete(ctx context.Context, in tagDeleteInput) (map[string]any, error) { + id := in.ID if id == "" { return nil, goerr.New("id is required") } @@ -614,9 +502,9 @@ func (x *Tool) tagDelete(ctx context.Context, args map[string]any) (map[string]a return map[string]any{"success": true}, nil } -func (x *Tool) tagMerge(ctx context.Context, args map[string]any) (map[string]any, error) { - oldID, _ := args["old_id"].(string) - newID, _ := args["new_id"].(string) +func (x *Tool) tagMerge(ctx context.Context, in tagMergeInput) (map[string]any, error) { + oldID := in.OldID + newID := in.NewID if oldID == "" || newID == "" { return nil, goerr.New("old_id and new_id are required") } @@ -630,27 +518,13 @@ func (x *Tool) tagMerge(ctx context.Context, args map[string]any) (map[string]an return map[string]any{"success": true}, nil } -// extractTagIDs extracts tag IDs from tool arguments. -func extractTagIDs(args map[string]any, key string) ([]types.KnowledgeTagID, error) { - raw, ok := args[key] - if !ok { - return nil, nil - } - - arr, ok := raw.([]any) - if !ok { - return nil, goerr.New("tags must be an array") - } - - ids := make([]types.KnowledgeTagID, 0, len(arr)) - for _, v := range arr { - s, ok := v.(string) - if !ok { - return nil, goerr.New("tag ID must be a string") - } +// toTagIDs converts decoded string tag IDs into the typed KnowledgeTagID slice. +func toTagIDs(tags []string) []types.KnowledgeTagID { + ids := make([]types.KnowledgeTagID, 0, len(tags)) + for _, s := range tags { ids = append(ids, types.KnowledgeTagID(s)) } - return ids, nil + return ids } // formatKnowledgeSummary formats knowledges as lightweight summaries (no claim). diff --git a/pkg/tool/slack/action_test.go b/pkg/tool/slack/action_test.go index bbada230..2eb71ceb 100644 --- a/pkg/tool/slack/action_test.go +++ b/pkg/tool/slack/action_test.go @@ -7,6 +7,7 @@ import ( "os" "testing" + "github.com/gollem-dev/gollem" "github.com/m-mizutani/gt" "github.com/secmon-lab/warren/pkg/tool/slack" "github.com/secmon-lab/warren/pkg/utils/errutil" @@ -93,9 +94,17 @@ func TestSlackMessageSearch_Specs(t *testing.T) { action := newAction(t, server.URL) specs, err := action.Specs(context.Background()) gt.NoError(t, err) - gt.A(t, specs).Length(1) - gt.Value(t, specs[0].Name).Equal("slack_message_search") - gt.Map(t, specs[0].Parameters).HasKey("query") + gt.A(t, specs).Length(2) + + byName := map[string]gollem.ToolSpec{} + for _, s := range specs { + byName[s.Name] = s + } + searchSpec, ok := byName["slack_message_search"] + gt.True(t, ok) + gt.Map(t, searchSpec.Parameters).HasKey("query") + _, ok = byName["slack_get_messages"] + gt.True(t, ok) } func TestSlackConfigure(t *testing.T) { diff --git a/pkg/tool/webfetch/README.md b/pkg/tool/webfetch/README.md index 041cab8c..f5175281 100644 --- a/pkg/tool/webfetch/README.md +++ b/pkg/tool/webfetch/README.md @@ -104,7 +104,7 @@ No external API keys are required when running in LLM-disabled mode (every call ## Limitations - **No URL allow/deny filtering.** Access control is delegated to the HITL approval dialog (LLM-disabled mode) and / or the LLM IPI screen (LLM-enabled mode). -- **No SSRF protection** beyond scheme restriction. Internal IPs and metadata endpoints are not blocked at this layer. +- **SSRF guard blocks private/internal IPs by default.** The underlying `github.com/gollem-dev/tools/webfetch` module refuses to fetch URLs that resolve to private, loopback, or link-local IP ranges (including cloud metadata endpoints), so internal-network targets are rejected before the request is sent. warren runs with this guard enabled and does not expose an opt-out, so `web_fetch` cannot reach internal hosts. - **UTF-8 assumed.** The `Content-Type` charset is not honored; non-UTF-8 pages may produce garbled output. - **No summarization.** The LLM only reformats — it does not condense the body. Large pages (after extraction) are passed to the upstream agent as-is. - **No retries.** Transient LLM or HTTP failures bubble up to the agent loop, which owns the retry strategy. diff --git a/pkg/usecase/refine.go b/pkg/usecase/refine.go index f38fad20..1c50df01 100644 --- a/pkg/usecase/refine.go +++ b/pkg/usecase/refine.go @@ -20,6 +20,7 @@ import ( "github.com/secmon-lab/warren/pkg/service/llm" "github.com/secmon-lab/warren/pkg/utils/clock" "github.com/secmon-lab/warren/pkg/utils/logging" + "github.com/secmon-lab/warren/pkg/utils/toolset" slackSDK "github.com/slack-go/slack" ) @@ -224,62 +225,51 @@ func (uc *UseCases) reviewSingleTicket(ctx context.Context, t *ticket.Ticket) er return nil } -// slackPostMessageTool implements gollem.ToolSet for posting messages to a ticket's Slack thread. -type slackPostMessageTool struct { - uc *UseCases - ticket *ticket.Ticket +// slackPostMessageInput is the typed argument for the slack_post_message tool. +// The schema (name, description, required) is inferred from these struct tags by +// gollem.NewTool, replacing the hand-written gollem.ToolSpec literal. +type slackPostMessageInput struct { + Message string `json:"message" required:"true" description:"The message to post in the ticket's Slack thread. Supports Slack mrkdwn format. Use <@USER_ID> for mentions."` } -func newSlackPostMessageTool(uc *UseCases, t *ticket.Ticket) *slackPostMessageTool { - return &slackPostMessageTool{uc: uc, ticket: t} -} - -// Specs implements gollem.ToolSet. -func (t *slackPostMessageTool) Specs(ctx context.Context) ([]gollem.ToolSpec, error) { - return []gollem.ToolSpec{ - { - Name: "slack_post_message", - Description: "Post a message to the ticket's Slack thread. Use Slack mrkdwn format. For mentions, use <@USER_ID> format.", - Parameters: map[string]*gollem.Parameter{ - "message": { - Type: gollem.TypeString, - Description: "The message to post in the ticket's Slack thread. Supports Slack mrkdwn format. Use <@USER_ID> for mentions.", - Required: true, - }, - }, - }, - }, nil -} - -// Run implements gollem.ToolSet. -func (t *slackPostMessageTool) Run(ctx context.Context, name string, args map[string]any) (map[string]any, error) { - if name != "slack_post_message" { - return nil, goerr.New("unknown tool", goerr.V("name", name)) - } - - message, ok := args["message"].(string) - if !ok || message == "" { - return map[string]any{"error": "message parameter is required and must be a non-empty string"}, nil - } +// Startup assertion: validate the tool's In/Out types form a valid schema at +// package init, so a malformed type fails immediately rather than at first use. +var _ = gollem.MustToolSchema[slackPostMessageInput, map[string]any]() + +// newSlackPostMessageTool builds a gollem.ToolSet exposing slack_post_message, +// which posts a follow-up message to the ticket's Slack thread during the refine +// agent loop. +func newSlackPostMessageTool(uc *UseCases, t *ticket.Ticket) gollem.ToolSet { + return toolset.New( + gollem.MustNewTool( + "slack_post_message", + "Post a message to the ticket's Slack thread. Use Slack mrkdwn format. For mentions, use <@USER_ID> format.", + func(ctx context.Context, in slackPostMessageInput) (map[string]any, error) { + if in.Message == "" { + return map[string]any{"error": "message parameter is required and must be a non-empty string"}, nil + } - logger := logging.From(ctx) + logger := logging.From(ctx) - if t.uc.slackService != nil && t.ticket.SlackThread != nil { - threadSvc := t.uc.slackService.NewThread(*t.ticket.SlackThread) - if err := threadSvc.PostComment(ctx, message); err != nil { - logger.Error("failed to post refine message to ticket thread", - "ticket_id", t.ticket.ID, "error", err) - return map[string]any{"error": fmt.Sprintf("failed to post message: %v", err)}, nil - } - return map[string]any{"status": "posted"}, nil - } + if uc.slackService != nil && t.SlackThread != nil { + threadSvc := uc.slackService.NewThread(*t.SlackThread) + if err := threadSvc.PostComment(ctx, in.Message); err != nil { + logger.Error("failed to post refine message to ticket thread", + "ticket_id", t.ID, "error", err) + return map[string]any{"error": fmt.Sprintf("failed to post message: %v", err)}, nil + } + return map[string]any{"status": "posted"}, nil + } - // CLI mode fallback - logger.Info("refine follow-up", - "ticket_id", t.ticket.ID, - "message", message, + // CLI mode fallback + logger.Info("refine follow-up", + "ticket_id", t.ID, + "message", in.Message, + ) + return map[string]any{"status": "logged (CLI mode)"}, nil + }, + ), ) - return map[string]any{"status": "logged (CLI mode)"}, nil } // consolidateUnboundAlerts finds unbound alerts that can be grouped and proposes consolidation. diff --git a/pkg/utils/toolset/toolset.go b/pkg/utils/toolset/toolset.go new file mode 100644 index 00000000..ad297ea3 --- /dev/null +++ b/pkg/utils/toolset/toolset.go @@ -0,0 +1,67 @@ +// Package toolset adapts a fixed list of individual gollem.Tool values into the +// gollem.ToolSet (Specs/Run) contract. +// +// gollem's type-safe constructor gollem.NewTool builds one gollem.Tool per +// handler, but Warren groups tools as a ToolSet: the planner selects tools by +// ToolSet ID and hands the whole group to a task agent. This adapter bridges the +// two so Warren can author type-safe tools with gollem.NewTool while still +// exposing the ToolSet grouping that interfaces.ToolSet requires, removing the +// hand-written Specs() literals and args[...] type assertions each tool used to +// carry. +package toolset + +import ( + "context" + + "github.com/gollem-dev/gollem" + "github.com/m-mizutani/goerr/v2" + "github.com/secmon-lab/warren/pkg/utils/errutil" +) + +// Set is a gollem.ToolSet backed by a fixed list of gollem.Tool values. Run +// dispatches by tool name; Specs returns each tool's spec in registration order +// so the wire-level ordering stays stable across calls. +type Set struct { + order []string + tools map[string]gollem.Tool +} + +var _ gollem.ToolSet = (*Set)(nil) + +// New builds a Set from the given tools. The tool list is fixed at construction +// time (static registration), so a duplicate tool name is a programming error +// and panics rather than silently shadowing a sibling — mirroring the +// MustNewTool convention used to build the tools passed in here. +func New(tools ...gollem.Tool) *Set { + s := &Set{ + order: make([]string, 0, len(tools)), + tools: make(map[string]gollem.Tool, len(tools)), + } + for _, t := range tools { + name := t.Spec().Name + if _, ok := s.tools[name]; ok { + panic(goerr.New("duplicate tool name in toolset", goerr.T(errutil.TagInvalidState), goerr.V("name", name))) + } + s.order = append(s.order, name) + s.tools[name] = t + } + return s +} + +// Specs implements gollem.ToolSet. +func (s *Set) Specs(_ context.Context) ([]gollem.ToolSpec, error) { + specs := make([]gollem.ToolSpec, 0, len(s.order)) + for _, name := range s.order { + specs = append(specs, s.tools[name].Spec()) + } + return specs, nil +} + +// Run implements gollem.ToolSet. +func (s *Set) Run(ctx context.Context, name string, args map[string]any) (map[string]any, error) { + t, ok := s.tools[name] + if !ok { + return nil, goerr.New("unknown tool", goerr.T(errutil.TagInvalidRequest), goerr.V("name", name)) + } + return t.Run(ctx, args) +} diff --git a/pkg/utils/toolset/toolset_test.go b/pkg/utils/toolset/toolset_test.go new file mode 100644 index 00000000..4edbf49d --- /dev/null +++ b/pkg/utils/toolset/toolset_test.go @@ -0,0 +1,86 @@ +package toolset_test + +import ( + "context" + "testing" + + "github.com/gollem-dev/gollem" + "github.com/m-mizutani/gt" + "github.com/secmon-lab/warren/pkg/utils/toolset" +) + +type echoInput struct { + Text string `json:"text" required:"true" description:"text to echo back"` +} + +type addInput struct { + A int64 `json:"a" description:"first addend"` + B int64 `json:"b" description:"second addend"` +} + +func newEchoTool() gollem.Tool { + return gollem.MustNewTool("echo", "Echo the given text", + func(_ context.Context, in echoInput) (map[string]any, error) { + return map[string]any{"echo": in.Text}, nil + }) +} + +func newAddTool() gollem.Tool { + return gollem.MustNewTool("add", "Add two integers", + func(_ context.Context, in addInput) (map[string]any, error) { + return map[string]any{"sum": in.A + in.B}, nil + }) +} + +func TestSet_SpecsKeepsRegistrationOrder(t *testing.T) { + ctx := context.Background() + set := toolset.New(newEchoTool(), newAddTool()) + + specs, err := set.Specs(ctx) + gt.NoError(t, err) + gt.A(t, specs).Length(2) + gt.Equal(t, specs[0].Name, "echo") + gt.Equal(t, specs[1].Name, "add") + + // The schema is inferred from the In struct: a required field stays required. + gt.True(t, specs[0].Parameters["text"].Required) +} + +func TestSet_RunDispatchesByName(t *testing.T) { + ctx := context.Background() + set := toolset.New(newEchoTool(), newAddTool()) + + out, err := set.Run(ctx, "echo", map[string]any{"text": "hello"}) + gt.NoError(t, err) + gt.Value(t, out["echo"]).Equal("hello") +} + +// TestSet_RunDecodesIntegers guards the int64 path: JSON numbers arrive as +// float64 over the wire, and NewTool must decode them into int64 fields without +// the manual float64 special-casing the old hand-written tools carried. +func TestSet_RunDecodesIntegers(t *testing.T) { + ctx := context.Background() + set := toolset.New(newAddTool()) + + out, err := set.Run(ctx, "add", map[string]any{"a": float64(2), "b": float64(40)}) + gt.NoError(t, err) + gt.Value(t, out["sum"]).Equal(int64(42)) +} + +func TestSet_RunUnknownToolErrors(t *testing.T) { + ctx := context.Background() + set := toolset.New(newEchoTool()) + + _, err := set.Run(ctx, "missing", map[string]any{}) + gt.Error(t, err) +} + +func TestSet_DuplicateNamePanics(t *testing.T) { + defer func() { + r := recover() + gt.NotNil(t, r) + }() + + toolset.New(newEchoTool(), newEchoTool()) + t.Fatal("expected panic on duplicate tool name") +}