From 941a1381c124e9cdb7966062d59639326d58dc2b Mon Sep 17 00:00:00 2001 From: nfebe Date: Sun, 19 Jul 2026 20:24:09 +0100 Subject: [PATCH] feat(ai): Let the assistant edit files, run actions, and control deployments The assistant's tools were read-only. It can now write a file in a deployment's directory, run a configured quick action, start, stop, or restart a deployment, and list a deployment's recent security events to summarize them. Every state-changing tool requires write access to the deployment and honours protected mode, so the assistant can only do what the operator behind it is allowed to do. File writes stay inside the deployment directory. --- internal/api/ai_tools.go | 143 +++++++++++++++++++++++++++++- internal/api/ai_tools_new_test.go | 109 +++++++++++++++++++++++ 2 files changed, 251 insertions(+), 1 deletion(-) create mode 100644 internal/api/ai_tools_new_test.go 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) + } +}