diff --git a/internal/api/ai_tools.go b/internal/api/ai_tools.go index ed63240..5a6fff2 100644 --- a/internal/api/ai_tools.go +++ b/internal/api/ai_tools.go @@ -303,6 +303,145 @@ func (s *Server) aiToolRegistry() map[string]aiTool { return truncateToolOutput(content), nil }, }, + "write_deployment_file": { + Spec: ai.Tool{ + Name: "write_deployment_file", + Description: "Create or overwrite a text file inside a deployment's directory, for example a config or compose file. Requires write access to the deployment; the path cannot escape the deployment directory.", + Parameters: objSchema(map[string]interface{}{ + "deployment": strProp("Deployment name. Omit to use the session's deployment."), + "path": strProp("Path to the file relative to the deployment directory."), + "content": strProp("The full new contents of the file."), + }, "path", "content"), + }, + Run: func(s *Server, c *gin.Context, bound string, args map[string]interface{}) (string, error) { + name, err := s.toolAllowedDeploymentWrite(c, bound, args) + if err != nil { + return "", err + } + path := argString(args, "path") + if path == "" { + return "", fmt.Errorf("path is required") + } + content, ok := args["content"].(string) + if !ok { + return "", fmt.Errorf("content is required") + } + if blocked, reason, perr := s.protectedDeploymentActionBlocked(name, protectedActionUploadFile); perr != nil { + return "", fmt.Errorf("could not check protected mode: %w", perr) + } else if blocked { + return "", fmt.Errorf("%s", reason) + } + if err := s.filesManager.WriteFile(name, path, strings.NewReader(content)); err != nil { + return "", err + } + return fmt.Sprintf("Wrote %d bytes to %s in %s.", len(content), path, name), nil + }, + }, + "run_quick_action": { + Spec: ai.Tool{ + Name: "run_quick_action", + Description: "Run one of a deployment's configured quick actions by its id. Requires write access to the deployment.", + Parameters: objSchema(map[string]interface{}{ + "deployment": strProp("Deployment name. Omit to use the session's deployment."), + "action_id": strProp("The id of the quick action to run."), + }, "action_id"), + }, + Run: func(s *Server, c *gin.Context, bound string, args map[string]interface{}) (string, error) { + name, err := s.toolAllowedDeploymentWrite(c, bound, args) + if err != nil { + return "", err + } + actionID := argString(args, "action_id") + if actionID == "" { + return "", fmt.Errorf("action_id is required") + } + if blocked, reason, perr := s.protectedDeploymentActionBlocked(name, protectedActionQuickAction); perr != nil { + return "", fmt.Errorf("could not check protected mode: %w", perr) + } else if blocked { + return "", fmt.Errorf("%s", reason) + } + output, err := s.manager.ExecuteQuickAction(name, actionID) + if err != nil { + return "", err + } + redactor := ai.NewRedactor(s.deploymentSecretValues(name)) + redacted, _ := redactor.Redact(output) + return truncateToolOutput(redacted), nil + }, + }, + "control_deployment": { + Spec: ai.Tool{ + Name: "control_deployment", + Description: "Start, stop, or restart a whole deployment. Requires write access to the deployment.", + Parameters: objSchema(map[string]interface{}{ + "deployment": strProp("Deployment name. Omit to use the session's deployment."), + "action": map[string]interface{}{ + "type": "string", + "enum": []string{"start", "stop", "restart"}, + "description": "What to do: start, stop, or restart.", + }, + }, "action"), + }, + Run: func(s *Server, c *gin.Context, bound string, args map[string]interface{}) (string, error) { + name, err := s.toolAllowedDeploymentWrite(c, bound, args) + if err != nil { + return "", err + } + action := argString(args, "action") + var output string + switch action { + case "start": + output, err = s.manager.StartDeployment(name) + case "stop": + output, err = s.manager.StopDeployment(name) + case "restart": + output, err = s.manager.RestartDeployment(name) + default: + return "", fmt.Errorf("action must be start, stop, or restart") + } + if err != nil { + return "", err + } + return truncateToolOutput(fmt.Sprintf("Ran %s on deployment %s.\n%s", action, name, output)), nil + }, + }, + "get_security_events": { + Spec: ai.Tool{ + Name: "get_security_events", + Description: "List recent security events for a deployment (blocked requests, scanners, auth failures) so they can be summarized. Read-only.", + Parameters: objSchema(map[string]interface{}{ + "deployment": strProp("Deployment name. Omit to use the session's deployment."), + "limit": map[string]interface{}{"type": "integer", "description": "Maximum events to return (default 50)."}, + }), + }, + Run: func(s *Server, c *gin.Context, bound string, args map[string]interface{}) (string, error) { + name, err := s.toolAllowedDeployment(c, bound, args) + if err != nil { + return "", err + } + if s.securityManager == nil { + return "The security module is not enabled.", nil + } + limit := 50 + if v, ok := args["limit"].(float64); ok && v > 0 { + limit = int(v) + } + events, _, err := s.securityManager.GetEventsByDeployment(name, limit) + if err != nil { + return "", err + } + if len(events) == 0 { + return fmt.Sprintf("No security events for %s.", name), nil + } + var b strings.Builder + fmt.Fprintf(&b, "%d recent security events for %s:\n", len(events), name) + for _, e := range events { + fmt.Fprintf(&b, "- %s [%s] %s %s %d from %s: %s\n", + e.CreatedAt.Format(time.RFC3339), e.Severity, e.RequestMethod, e.RequestPath, e.StatusCode, e.SourceIP, e.Message) + } + return truncateToolOutput(b.String()), nil + }, + }, "exec_in_service": { Spec: ai.Tool{ Name: "exec_in_service", @@ -330,7 +469,9 @@ func (s *Server) aiToolRegistry() map[string]aiTool { if err != nil { return "", err } - if blocked, reason, perr := s.protectedDeploymentActionBlocked(name, protectedActionExec); perr == nil && blocked { + if blocked, reason, perr := s.protectedDeploymentActionBlocked(name, protectedActionExec); perr != nil { + return "", fmt.Errorf("could not check protected mode: %w", perr) + } else if blocked { return "", fmt.Errorf("%s", reason) } ctx, cancel := context.WithTimeout(c.Request.Context(), 20*time.Second) diff --git a/internal/api/ai_tools_new_test.go b/internal/api/ai_tools_new_test.go new file mode 100644 index 0000000..fa66b47 --- /dev/null +++ b/internal/api/ai_tools_new_test.go @@ -0,0 +1,109 @@ +package api + +import ( + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/flatrun/agent/internal/auth" + "github.com/flatrun/agent/internal/contextkeys" + "github.com/flatrun/agent/pkg/models" + "github.com/gin-gonic/gin" +) + +func toolCtx(actor *auth.ActorContext) *gin.Context { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/", nil) + if actor != nil { + c.Set(contextkeys.Actor, actor) + } + return c +} + +func TestWriteDeploymentFileTool(t *testing.T) { + s, tmpDir, _ := setupPlanTestServer(t) + createTestDeployment(t, tmpDir, "myapp", &models.ServiceMetadata{Name: "myapp"}) + + tool := s.aiToolRegistry()["write_deployment_file"] + out, err := tool.Run(s, toolCtx(nil), "myapp", map[string]interface{}{ + "path": "config/app.conf", + "content": "key = value\n", + }) + if err != nil { + t.Fatalf("write failed: %v", err) + } + if !strings.Contains(out, "app.conf") { + t.Errorf("unexpected output %q", out) + } + data, err := os.ReadFile(filepath.Join(tmpDir, "myapp", "config", "app.conf")) + if err != nil || string(data) != "key = value\n" { + t.Errorf("file not written correctly: %q err=%v", string(data), err) + } +} + +func TestWriteDeploymentFileRefusesTraversal(t *testing.T) { + s, tmpDir, _ := setupPlanTestServer(t) + createTestDeployment(t, tmpDir, "myapp", &models.ServiceMetadata{Name: "myapp"}) + + tool := s.aiToolRegistry()["write_deployment_file"] + if _, err := tool.Run(s, toolCtx(nil), "myapp", map[string]interface{}{ + "path": "../escape.txt", + "content": "x", + }); err == nil { + t.Fatal("a path escaping the deployment directory must be refused") + } +} + +func TestWriteDeploymentFileFailsClosedOnProtectedCheckError(t *testing.T) { + s, _, _ := setupPlanTestServer(t) + // No deployment is created, so the protected-mode check cannot load it and + // errors. A state-changing tool must refuse rather than proceed. + tool := s.aiToolRegistry()["write_deployment_file"] + _, err := tool.Run(s, toolCtx(nil), "ghost", map[string]interface{}{"path": "x", "content": "y"}) + if err == nil || !strings.Contains(err.Error(), "protected mode") { + t.Errorf("a failed protected-mode check must block the write, got %v", err) + } +} + +func TestMutatingToolsRequireWriteAccess(t *testing.T) { + s, tmpDir, _ := setupPlanTestServer(t) + createTestDeployment(t, tmpDir, "myapp", &models.ServiceMetadata{Name: "myapp"}) + + // A viewer with no write access to the deployment. + viewer := &auth.ActorContext{Role: auth.RoleViewer, User: &auth.User{ID: 1, Username: "v"}} + args := map[string]interface{}{"path": "x", "content": "y", "action_id": "a", "action": "restart"} + for _, name := range []string{"write_deployment_file", "run_quick_action", "control_deployment"} { + tool := s.aiToolRegistry()[name] + if _, err := tool.Run(s, toolCtx(viewer), "myapp", args); err == nil || !strings.Contains(err.Error(), "write access") { + t.Errorf("%s must require write access, got %v", name, err) + } + } +} + +func TestGetSecurityEventsWhenDisabled(t *testing.T) { + s, tmpDir, _ := setupPlanTestServer(t) + createTestDeployment(t, tmpDir, "myapp", &models.ServiceMetadata{Name: "myapp"}) + + tool := s.aiToolRegistry()["get_security_events"] + out, err := tool.Run(s, toolCtx(nil), "myapp", map[string]interface{}{}) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(out, "not enabled") { + t.Errorf("expected a disabled message, got %q", out) + } +} + +func TestControlDeploymentRejectsBadAction(t *testing.T) { + s, tmpDir, _ := setupPlanTestServer(t) + createTestDeployment(t, tmpDir, "myapp", &models.ServiceMetadata{Name: "myapp"}) + + tool := s.aiToolRegistry()["control_deployment"] + if _, err := tool.Run(s, toolCtx(nil), "myapp", map[string]interface{}{"action": "delete"}); err == nil || + !strings.Contains(err.Error(), "start, stop, or restart") { + t.Errorf("expected an invalid-action error, got %v", err) + } +}