Skip to content

Commit ebf5a1c

Browse files
dgageotdgageot
authored
Merge pull request #320 from dgageot/more-validation
More validation
2 parents 08bfab2 + cf605d9 commit ebf5a1c

6 files changed

Lines changed: 166 additions & 52 deletions

File tree

‎pkg/config/testdata/invalid_path_v2.yaml‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
version: "2"
2+
3+
agents:
4+
root:
5+
model: openai/gpt-4o
6+
toolsets:
7+
- type: mcp
8+
path: invalid-tool-type

‎pkg/config/testdata/invalid_post_edit_v2.yaml‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
version: "2"
2+
3+
agents:
4+
root:
5+
model: openai/gpt-4o
6+
toolsets:
7+
- type: mcp
8+
post_edit:
9+
- path: /path

‎pkg/config/testdata/missing_memory_path_v2.yaml‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
version: "2"
2+
3+
agents:
4+
root:
5+
model: openai/gpt-4o
6+
toolsets:
7+
- type: memory

‎pkg/config/v2/types.go‎

Lines changed: 18 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,5 @@
11
package v2
22

3-
import (
4-
"errors"
5-
"strings"
6-
)
7-
83
// Config represents the entire configuration file
94
type Config struct {
105
Version string `json:"version,omitempty" yaml:"version,omitempty"`
@@ -72,58 +67,30 @@ type PostEditConfig struct {
7267

7368
// Toolset represents a tool configuration
7469
type Toolset struct {
75-
Type string `json:"type,omitempty" yaml:"type,omitempty"`
76-
Ref string `json:"ref,omitempty" yaml:"ref,omitempty"`
77-
Config any `json:"config,omitempty" yaml:"config,omitempty"`
78-
Command string `json:"command,omitempty" yaml:"command,omitempty"`
79-
Remote Remote `json:"remote,omitempty" yaml:"remote,omitempty"`
80-
Args []string `json:"args,omitempty" yaml:"args,omitempty"`
81-
Tools []string `json:"tools,omitempty" yaml:"tools,omitempty"`
82-
Env map[string]string `json:"env,omitempty" yaml:"env,omitempty"`
83-
84-
// For the think tool
85-
Shared bool `json:"shared,omitempty" yaml:"shared,omitempty"`
86-
// For the memory tool
87-
Path string `json:"path,omitempty" yaml:"path,omitempty"`
70+
Type string `json:"type,omitempty" yaml:"type,omitempty"`
71+
Tools []string `json:"tools,omitempty" yaml:"tools,omitempty"`
8872

89-
// For the script tool
90-
Shell map[string]ScriptShellToolConfig `json:"shell,omitempty" yaml:"shell,omitempty"`
73+
// For the `mcp` tool
74+
Command string `json:"command,omitempty" yaml:"command,omitempty"`
75+
Args []string `json:"args,omitempty" yaml:"args,omitempty"`
76+
Ref string `json:"ref,omitempty" yaml:"ref,omitempty"`
77+
Remote Remote `json:"remote,omitempty" yaml:"remote,omitempty"`
78+
Config any `json:"config,omitempty" yaml:"config,omitempty"`
9179

92-
// For the filesystem tool - post-edit commands
93-
PostEdit []PostEditConfig `json:"post_edit,omitempty" yaml:"post_edit,omitempty"`
94-
}
80+
// For `shell`, `script` or `mcp` tools
81+
Env map[string]string `json:"env,omitempty" yaml:"env,omitempty"`
9582

96-
// Ensure that either Command, Remote or Ref is set, but not all empty
97-
func (t *Toolset) validate() error {
98-
if len(t.Shell) > 0 && t.Type != "script" {
99-
return errors.New("shell can only be used with type 'script'")
100-
}
101-
if t.Type != "mcp" {
102-
return nil
103-
}
83+
// For the `todo` tool
84+
Shared bool `json:"shared,omitempty" yaml:"shared,omitempty"`
10485

105-
count := 0
106-
if t.Command != "" {
107-
count++
108-
}
109-
if t.Remote.URL != "" {
110-
count++
111-
}
112-
if t.Ref != "" {
113-
count++
114-
}
115-
if count == 0 {
116-
return errors.New("either command, remote or ref must be set")
117-
}
118-
if count > 1 {
119-
return errors.New("either command, remote or ref must be set, but only one of those")
120-
}
86+
// For the `memory` tool
87+
Path string `json:"path,omitempty" yaml:"path,omitempty"`
12188

122-
if t.Ref != "" && !strings.Contains(t.Ref, "docker:") {
123-
return errors.New("only docker refs are supported for MCP tools, e.g., 'docker:context7'")
124-
}
89+
// For the `script` tool
90+
Shell map[string]ScriptShellToolConfig `json:"shell,omitempty" yaml:"shell,omitempty"`
12591

126-
return nil
92+
// For the `filesystem` tool - post-edit commands
93+
PostEdit []PostEditConfig `json:"post_edit,omitempty" yaml:"post_edit,omitempty"`
12794
}
12895

12996
func (t *Toolset) UnmarshalYAML(unmarshal func(any) error) error {

‎pkg/config/v2/validate.go‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
package v2
2+
3+
import (
4+
"errors"
5+
"strings"
6+
)
7+
8+
func (t *Config) UnmarshalYAML(unmarshal func(any) error) error {
9+
type alias Config
10+
var tmp alias
11+
if err := unmarshal(&tmp); err != nil {
12+
return err
13+
}
14+
*t = Config(tmp)
15+
return t.validate()
16+
}
17+
18+
func (t *Config) validate() error {
19+
for _, agent := range t.Agents {
20+
for i := range agent.Toolsets {
21+
if err := agent.Toolsets[i].validate(); err != nil {
22+
return err
23+
}
24+
}
25+
}
26+
27+
return nil
28+
}
29+
30+
func (t *Toolset) validate() error {
31+
// Attributes used on the wrong toolset type.
32+
if len(t.Shell) > 0 && t.Type != "script" {
33+
return errors.New("shell can only be used with type 'script'")
34+
}
35+
if t.Path != "" && t.Type != "memory" {
36+
return errors.New("path can only be used with type 'memory'")
37+
}
38+
if len(t.PostEdit) > 0 && t.Type != "filesystem" {
39+
return errors.New("post_edit can only be used with type 'filesystem'")
40+
}
41+
if len(t.Env) > 0 && (t.Type != "shell" && t.Type != "script" && t.Type != "mcp") {
42+
return errors.New("env can only be used with type 'shell', 'script' or 'mcp'")
43+
}
44+
if t.Shared && t.Type != "todo" {
45+
return errors.New("shared can only be used with type 'todo'")
46+
}
47+
if t.Command != "" && t.Type != "mcp" {
48+
return errors.New("command can only be used with type 'mcp'")
49+
}
50+
if len(t.Args) > 0 && t.Type != "mcp" {
51+
return errors.New("args can only be used with type 'mcp'")
52+
}
53+
if t.Ref != "" && t.Type != "mcp" {
54+
return errors.New("ref can only be used with type 'mcp'")
55+
}
56+
if (t.Remote.URL != "" || t.Remote.TransportType != "" || len(t.Remote.Headers) > 0) && t.Type != "mcp" {
57+
return errors.New("remote can only be used with type 'mcp'")
58+
}
59+
if t.Config != nil && t.Type != "mcp" {
60+
return errors.New("config can only be used with type 'mcp'")
61+
}
62+
63+
switch t.Type {
64+
case "memory":
65+
if t.Path == "" {
66+
return errors.New("memory toolset requires a path to be set")
67+
}
68+
case "mcp":
69+
count := 0
70+
if t.Command != "" {
71+
count++
72+
}
73+
if t.Remote.URL != "" {
74+
count++
75+
}
76+
if t.Ref != "" {
77+
count++
78+
}
79+
if count == 0 {
80+
return errors.New("either command, remote or ref must be set")
81+
}
82+
if count > 1 {
83+
return errors.New("either command, remote or ref must be set, but only one of those")
84+
}
85+
86+
if t.Ref != "" && !strings.Contains(t.Ref, "docker:") {
87+
return errors.New("only docker refs are supported for MCP tools, e.g., 'docker:context7'")
88+
}
89+
}
90+
91+
return nil
92+
}

‎pkg/config/validation_test.go‎

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,8 @@ func TestValidatePathInDirectory(t *testing.T) {
6161

6262
for _, tt := range tests {
6363
t.Run(tt.name, func(t *testing.T) {
64+
t.Parallel()
65+
6466
result, err := ValidatePathInDirectory(tt.path, tt.allowedDir)
6567

6668
if tt.expectError {
@@ -79,7 +81,7 @@ func TestValidatePathInDirectory(t *testing.T) {
7981
}
8082
}
8183

82-
func TestLoadConfigSecure(t *testing.T) {
84+
func TestValidateMemoryPath(t *testing.T) {
8385
tmpDir, err := os.MkdirTemp("", "test_config_secure")
8486
require.NoError(t, err)
8587
defer os.RemoveAll(tmpDir)
@@ -106,3 +108,32 @@ agents:
106108
require.Error(t, err)
107109
require.Contains(t, err.Error(), "path validation failed")
108110
}
111+
112+
func TestValidationErrors(t *testing.T) {
113+
tests := []struct {
114+
name string
115+
path string
116+
}{
117+
{
118+
name: "memory toolset missing path",
119+
path: "testdata/missing_memory_path_v2.yaml",
120+
},
121+
{
122+
name: "path in non memory toolset",
123+
path: "testdata/invalid_path_v2.yaml",
124+
},
125+
{
126+
name: "post_edit in non filesystem toolset",
127+
path: "testdata/invalid_post_edit_v2.yaml",
128+
},
129+
}
130+
131+
for _, tt := range tests {
132+
t.Run(tt.name, func(t *testing.T) {
133+
t.Parallel()
134+
135+
_, err := loadConfig(tt.path)
136+
require.Error(t, err)
137+
})
138+
}
139+
}

0 commit comments

Comments
 (0)