From f34b6e89ee7caa98c7af01e2e31e9c01514c06e1 Mon Sep 17 00:00:00 2001 From: Sebastian Florek Date: Thu, 20 Aug 2026 11:16:04 +0200 Subject: [PATCH 01/15] feat(toolquery): implement sandboxed Python execution - Added `RunPythonInput` and `RunPythonOutput` message types in `toolquery.proto` for Python script execution - Introduced Python worker logic using a limited subset of Monty, with controlled execution environment and resources - Enhanced `Server` with graceful shutdown and service cleanup to safely terminate background processes - Updated `ToolQueryService` to support Python script execution with context handling - Added `Closer` interface for services that require resource cleanup - Implemented protocol versioning and request handling for Python execution in the worker process - Extended Elixir definitions in toolquery.pb.ex to include Python RPC services - Updated workbench documentation to incorporate Python tool information --- .gitignore | 1 + charts/console/values.yaml | 4 +- go/cloud-query/Dockerfile | 2 +- go/cloud-query/README.md | 14 +- go/cloud-query/api/proto/toolquery.proto | 17 + go/cloud-query/cmd/main.go | 17 +- go/cloud-query/docs/api-reference.md | 20 ++ go/cloud-query/go.mod | 6 +- go/cloud-query/go.sum | 8 + .../internal/proto/toolquery/toolquery.pb.go | 179 ++++++++-- .../proto/toolquery/toolquery_grpc.pb.go | 38 +++ go/cloud-query/internal/server/server.go | 41 ++- go/cloud-query/internal/service/service.go | 8 + go/cloud-query/internal/service/toolquery.go | 65 +++- .../internal/service/toolquery_python_test.go | 40 +++ .../internal/tools/python/README.md | 63 ++++ .../internal/tools/python/common.go | 25 ++ .../internal/tools/python/errors.go | 100 ++++++ .../internal/tools/python/protocol.go | 138 ++++++++ .../internal/tools/python/python.go | 322 ++++++++++++++++++ .../internal/tools/python/python_test.go | 301 ++++++++++++++++ .../internal/tools/python/worker.go | 148 ++++++++ go/go.work | 2 +- lib/cloud_query/toolquery.pb.ex | 26 ++ lib/console/ai/tools/workbench/python.ex | 38 +++ lib/console/ai/workbench/engine.ex | 2 + .../ai/workbench/subagents/infrastructure.ex | 2 + .../ai/workbench/subagents/observability.ex | 3 +- priv/prompts/workbench/infrastructure.md.eex | 2 +- priv/prompts/workbench/job.md.eex | 10 +- priv/prompts/workbench/observability.md.eex | 2 +- priv/tools/workbench/python.json | 22 ++ .../ai/tools/workbench/python_test.exs | 53 +++ 33 files changed, 1660 insertions(+), 59 deletions(-) create mode 100644 go/cloud-query/internal/service/toolquery_python_test.go create mode 100644 go/cloud-query/internal/tools/python/README.md create mode 100644 go/cloud-query/internal/tools/python/common.go create mode 100644 go/cloud-query/internal/tools/python/errors.go create mode 100644 go/cloud-query/internal/tools/python/protocol.go create mode 100644 go/cloud-query/internal/tools/python/python.go create mode 100644 go/cloud-query/internal/tools/python/python_test.go create mode 100644 go/cloud-query/internal/tools/python/worker.go create mode 100644 lib/console/ai/tools/workbench/python.ex create mode 100644 priv/tools/workbench/python.json create mode 100644 test/console/ai/tools/workbench/python_test.exs diff --git a/.gitignore b/.gitignore index 46c0e1e059..819fd07075 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ /.idea /.vscode/**/* !/.vscode/settings.json +/.codex # System .DS_Store diff --git a/charts/console/values.yaml b/charts/console/values.yaml index 8293991376..da2cca7bd8 100644 --- a/charts/console/values.yaml +++ b/charts/console/values.yaml @@ -605,8 +605,10 @@ cloudQuery: resources: requests: - memory: 250Mi + memory: 512Mi cpu: 100m + limits: + memory: 1Gi livenessProbe: ~ diff --git a/go/cloud-query/Dockerfile b/go/cloud-query/Dockerfile index 0d3f90226e..5b11841cc3 100644 --- a/go/cloud-query/Dockerfile +++ b/go/cloud-query/Dockerfile @@ -1,4 +1,4 @@ -FROM golang:1.26.5 AS builder +FROM golang:1.26.6 AS builder WORKDIR /workspace diff --git a/go/cloud-query/README.md b/go/cloud-query/README.md index c68e0da155..412b842534 100644 --- a/go/cloud-query/README.md +++ b/go/cloud-query/README.md @@ -7,16 +7,20 @@ Cloud Query is a service part of the Plural Console ecosystem that provides clou - Query cloud resources across multiple providers - Embedded PostgreSQL database for data storage and retrieval through PostgreSQL FDW steampipe extension - gRPC API for integration with other services +- Sandboxed execution for Lua and Monty's limited Python subset - Containerized deployment for easy scaling CloudQuery provider support includes AWS, Azure, GCP, and VMware vSphere. ## Prerequisites -- Go 1.24.2 or higher +- Go 1.26.6 or higher - Docker (for containerized deployment) - Make +Running the complete service locally requires a writable temporary directory. +The Python worker extracts gomonty's embedded native library there on first use. + ## Getting Started ### Local Development @@ -89,6 +93,14 @@ Cloud-Query also exposes ToolQuery gRPC endpoints for observability tools (metri ToolQuery also supports cloud function invocation via `InvokeLambda` for AWS Lambda, GCP Cloud Run services (Gen2), and Azure Functions using canonical identifiers and cloud connection credentials. +### Sandboxed Python execution + +`RunPython` executes Monty's limited Python subset in a crash-isolated `cloud-query python-worker` subprocess. The request's optional JSON object is available as `input`; `output` starts as an empty dictionary and must remain a JSON-serializable dictionary. The response returns that dictionary as `result_json` and standard-output text from `print()` separately as `stdout`. + +This is not CPython. The sandbox has no host filesystem, environment, network, subprocess, shell, `pip`, third-party packages, or callback access. Two workers start with only `TMPDIR=/tmp`; each request gets a fresh gomonty REPL. Failed or canceled workers are killed and replaced, and healthy workers are periodically recycled. Source is limited to 64 KiB, input and result JSON to 1 MiB, stdout to 64 KiB, interpreter execution to 10 seconds, interpreter-managed memory to 64 MiB, and recursion to 200 frames. Two runs execute concurrently and up to 16 more wait in a FIFO queue. A full queue is rejected. A parent watchdog ends a run after 15 seconds or the caller's earlier deadline. + +The image embeds gomonty `v0.0.14` and its platform-specific glibc library, built against official Monty commit `c9802b5f30d11fecf9f153feb1dfdab3abda070e`. It contains no separate `monty` executable. Both pins are recorded in OCI labels. Monty's memory limit covers interpreter-managed allocations rather than total pod RSS; operators should measure the workload before reducing the default cloud-query memory allocation. + ### Tool Provider Credentials and Permissions - `Dynatrace`: diff --git a/go/cloud-query/api/proto/toolquery.proto b/go/cloud-query/api/proto/toolquery.proto index 29e9cbb977..c4e22f390c 100644 --- a/go/cloud-query/api/proto/toolquery.proto +++ b/go/cloud-query/api/proto/toolquery.proto @@ -300,6 +300,22 @@ message RunLuaOutput { string result_json = 1; } +message RunPythonInput { + // Python source. Write structured results to the global `output` dictionary. + string script = 1; + + // Optional JSON object exposed as the global `input` dictionary. Empty means `{}`. + string input_json = 2; +} + +message RunPythonOutput { + // JSON-encoded contents of the global `output` dictionary. + string result_json = 1; + + // Text emitted by print() during the user script. + string stdout = 2; +} + service ToolQuery { rpc Metrics(MetricsQueryInput) returns (MetricsQueryOutput) {} rpc MetricsSearch(MetricsSearchInput) returns (MetricsSearchOutput) {} @@ -308,4 +324,5 @@ service ToolQuery { rpc Traces(TracesQueryInput) returns (TracesQueryOutput) {} rpc InvokeLambda(InvokeLambdaInput) returns (InvokeLambdaOutput) {} rpc RunLua(RunLuaInput) returns (RunLuaOutput) {} + rpc RunPython(RunPythonInput) returns (RunPythonOutput) {} } diff --git a/go/cloud-query/cmd/main.go b/go/cloud-query/cmd/main.go index a47c08f509..4c216f2c5e 100644 --- a/go/cloud-query/cmd/main.go +++ b/go/cloud-query/cmd/main.go @@ -13,6 +13,7 @@ import ( "github.com/pluralsh/console/go/cloud-query/internal/pool" "github.com/pluralsh/console/go/cloud-query/internal/server" "github.com/pluralsh/console/go/cloud-query/internal/service" + pythontools "github.com/pluralsh/console/go/cloud-query/internal/tools/python" ) func startHealthzHandler() { @@ -27,9 +28,20 @@ func startHealthzHandler() { } func main() { + if len(os.Args) == 2 && os.Args[1] == "python-worker" { + if err := pythontools.RunWorker(os.Stdin, os.Stdout); err != nil { + os.Exit(1) + } + return + } + startHealthzHandler() - services := []service.Service{service.NewToolQueryService()} + toolQueryService, err := service.NewToolQueryService(context.Background()) + if err != nil { + klog.Fatalf("failed to initialize tool query service: %v", err) + } + services := []service.Service{toolQueryService} if args.DatabaseEnabled() { p, err := pool.NewConnectionPool(args.DatabaseConnectionTTL()) @@ -67,8 +79,5 @@ func handleShutdown(cancel context.CancelFunc, s *server.Server) { <-signalChan klog.Info("received shutdown signal, shutting down gracefully...") - s.Stop() cancel() - klog.Info("stopped gracefully") - os.Exit(0) } diff --git a/go/cloud-query/docs/api-reference.md b/go/cloud-query/docs/api-reference.md index fd4da9019e..73cbda1b40 100644 --- a/go/cloud-query/docs/api-reference.md +++ b/go/cloud-query/docs/api-reference.md @@ -405,6 +405,8 @@ service ToolQuery { rpc Logs(LogsQueryInput) returns (LogsQueryOutput) {} rpc Traces(TracesQueryInput) returns (TracesQueryOutput) {} rpc InvokeLambda(InvokeLambdaInput) returns (InvokeLambdaOutput) {} + rpc RunLua(RunLuaInput) returns (RunLuaOutput) {} + rpc RunPython(RunPythonInput) returns (RunPythonOutput) {} } ``` @@ -1507,6 +1509,24 @@ message JaegerTracesOptions { } ``` +## Run Python + +`RunPython` synchronously executes Monty's limited Python subset in a fresh logical sandbox session. `input_json` is optional but, when present, must encode an object. It is exposed as the global `input`; scripts write their structured response to the global `output` dictionary. Printed text is returned separately. + +```protobuf +message RunPythonInput { + string script = 1; + string input_json = 2; +} + +message RunPythonOutput { + string result_json = 1; + string stdout = 2; +} +``` + +The runtime exposes no host filesystem, environment, network, subprocess, shell, package installation, third-party package, or host-tool callback. It limits source to 64 KiB, input and result JSON to 1 MiB, stdout to 64 KiB, execution to 10 seconds, memory to 64 MiB, recursion to 200 frames, wall time to 15 seconds, and concurrency to two runs per process. Up to 16 additional requests wait in a bounded FIFO queue. It uses gomonty `v0.0.14`, built against official Monty commit `c9802b5f30d11fecf9f153feb1dfdab3abda070e`; it is not CPython. + ## Invoke Lambda `InvokeLambda` invokes serverless functions using canonical provider identifiers only. diff --git a/go/cloud-query/go.mod b/go/cloud-query/go.mod index c628c63a5e..c5524ec57a 100644 --- a/go/cloud-query/go.mod +++ b/go/cloud-query/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/cloud-query -go 1.26.5 +go 1.26.6 require ( github.com/Azure/azure-sdk-for-go/sdk/azcore v1.21.1 @@ -18,6 +18,7 @@ require ( github.com/aws/aws-sdk-go-v2/service/lambda v1.89.1 github.com/aws/aws-sdk-go-v2/service/sts v1.42.2 github.com/elastic/go-elasticsearch/v9 v9.3.1 + github.com/ewhauser/gomonty v0.0.14 github.com/gofrs/uuid v4.4.0+incompatible github.com/lib/pq v1.12.3 github.com/orcaman/concurrent-map/v2 v2.0.1 @@ -58,6 +59,7 @@ require ( github.com/aws/smithy-go v1.27.1 // indirect github.com/beorn7/perks v1.0.1 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/ebitengine/purego v0.10.0 // indirect github.com/elastic/elastic-transport-go/v8 v8.8.0 // indirect github.com/felixge/httpsnoop v1.0.4 // indirect github.com/go-logr/logr v1.4.3 // indirect @@ -80,6 +82,8 @@ require ( github.com/prometheus/client_model v0.6.2 // indirect github.com/prometheus/procfs v0.20.1 // indirect github.com/rogpeppe/go-internal v1.15.0 // indirect + github.com/vmihailenco/msgpack/v5 v5.4.1 // indirect + github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect github.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb // indirect github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 // indirect github.com/xeipuuv/gojsonschema v1.2.0 // indirect diff --git a/go/cloud-query/go.sum b/go/cloud-query/go.sum index 730935a9b3..e698498192 100644 --- a/go/cloud-query/go.sum +++ b/go/cloud-query/go.sum @@ -78,10 +78,14 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/ebitengine/purego v0.10.0 h1:QIw4xfpWT6GWTzaW5XEKy3HXoqrJGx1ijYHzTF0/ISU= +github.com/ebitengine/purego v0.10.0/go.mod h1:iIjxzd6CiRiOG0UyXP+V1+jWqUXVjPKLAI0mRfJZTmQ= github.com/elastic/elastic-transport-go/v8 v8.8.0 h1:7k1Ua+qluFr6p1jfJjGDl97ssJS/P7cHNInzfxgBQAo= github.com/elastic/elastic-transport-go/v8 v8.8.0/go.mod h1:YLHer5cj0csTzNFXoNQ8qhtGY1GTvSqPnKWKaqQE3Hk= github.com/elastic/go-elasticsearch/v9 v9.3.1 h1:v5A9uFw0nLFA0luD3xAqliBXbscfuhch409HIinfhKY= github.com/elastic/go-elasticsearch/v9 v9.3.1/go.mod h1:B5u4H2jo2/v0+PrgbmIUdEyHdenFyavWtjciAFl7TA0= +github.com/ewhauser/gomonty v0.0.14 h1:DM+iSZ/WJzl+huEH5SGONLl8CTooQrfVFbC0xY1ghO4= +github.com/ewhauser/gomonty v0.0.14/go.mod h1:XCLUVfUFX733MIidhDw3iLWwaVNxT23i9xK+ioAAOVc= github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg= github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= @@ -163,6 +167,10 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+ github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/vmihailenco/msgpack/v5 v5.4.1 h1:cQriyiUvjTwOHg8QZaPihLWeRAAVoCpE00IUPn0Bjt8= +github.com/vmihailenco/msgpack/v5 v5.4.1/go.mod h1:GaZTsDaehaPpQVyxrf5mtQlH+pc21PIudVV/E3rRQok= +github.com/vmihailenco/tagparser/v2 v2.0.0 h1:y09buUbR+b5aycVFQs/g70pqKVZNBmxwAhO7/IwNM9g= +github.com/vmihailenco/tagparser/v2 v2.0.0/go.mod h1:Wri+At7QHww0WTrCBeu4J6bNtoV6mEfg5OIWRZA9qds= github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f/go.mod h1:N2zxlSyiKSe5eX1tZViRH5QA0qijqEDrYZiPEAiq3wU= github.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb h1:zGWFAtiMcyryUHoUjUJX0/lt1H2+i2Ka2n+D3DImSNo= github.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb/go.mod h1:N2zxlSyiKSe5eX1tZViRH5QA0qijqEDrYZiPEAiq3wU= diff --git a/go/cloud-query/internal/proto/toolquery/toolquery.pb.go b/go/cloud-query/internal/proto/toolquery/toolquery.pb.go index d004253d7e..f68d0754cf 100644 --- a/go/cloud-query/internal/proto/toolquery/toolquery.pb.go +++ b/go/cloud-query/internal/proto/toolquery/toolquery.pb.go @@ -1,6 +1,6 @@ // Code generated by protoc-gen-go. DO NOT EDIT. // versions: -// protoc-gen-go v1.36.11 +// protoc-gen-go v1.36.11-devel // protoc v6.31.1 // source: toolquery.proto @@ -2918,6 +2918,114 @@ func (x *RunLuaOutput) GetResultJson() string { return "" } +type RunPythonInput struct { + state protoimpl.MessageState `protogen:"open.v1"` + // Python source. Write structured results to the global `output` dictionary. + Script string `protobuf:"bytes,1,opt,name=script,proto3" json:"script,omitempty"` + // Optional JSON object exposed as the global `input` dictionary. Empty means `{}`. + InputJson string `protobuf:"bytes,2,opt,name=input_json,json=inputJson,proto3" json:"input_json,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RunPythonInput) Reset() { + *x = RunPythonInput{} + mi := &file_toolquery_proto_msgTypes[44] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RunPythonInput) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RunPythonInput) ProtoMessage() {} + +func (x *RunPythonInput) ProtoReflect() protoreflect.Message { + mi := &file_toolquery_proto_msgTypes[44] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RunPythonInput.ProtoReflect.Descriptor instead. +func (*RunPythonInput) Descriptor() ([]byte, []int) { + return file_toolquery_proto_rawDescGZIP(), []int{44} +} + +func (x *RunPythonInput) GetScript() string { + if x != nil { + return x.Script + } + return "" +} + +func (x *RunPythonInput) GetInputJson() string { + if x != nil { + return x.InputJson + } + return "" +} + +type RunPythonOutput struct { + state protoimpl.MessageState `protogen:"open.v1"` + // JSON-encoded contents of the global `output` dictionary. + ResultJson string `protobuf:"bytes,1,opt,name=result_json,json=resultJson,proto3" json:"result_json,omitempty"` + // Text emitted by print() during the user script. + Stdout string `protobuf:"bytes,2,opt,name=stdout,proto3" json:"stdout,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RunPythonOutput) Reset() { + *x = RunPythonOutput{} + mi := &file_toolquery_proto_msgTypes[45] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RunPythonOutput) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RunPythonOutput) ProtoMessage() {} + +func (x *RunPythonOutput) ProtoReflect() protoreflect.Message { + mi := &file_toolquery_proto_msgTypes[45] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RunPythonOutput.ProtoReflect.Descriptor instead. +func (*RunPythonOutput) Descriptor() ([]byte, []int) { + return file_toolquery_proto_rawDescGZIP(), []int{45} +} + +func (x *RunPythonOutput) GetResultJson() string { + if x != nil { + return x.ResultJson + } + return "" +} + +func (x *RunPythonOutput) GetStdout() string { + if x != nil { + return x.Stdout + } + return "" +} + var File_toolquery_proto protoreflect.FileDescriptor const file_toolquery_proto_rawDesc = "" + @@ -3232,7 +3340,15 @@ const file_toolquery_proto_rawDesc = "" + "\x06script\x18\x01 \x01(\tR\x06script\"/\n" + "\fRunLuaOutput\x12\x1f\n" + "\vresult_json\x18\x01 \x01(\tR\n" + - "resultJson2\x9c\x04\n" + + "resultJson\"G\n" + + "\x0eRunPythonInput\x12\x16\n" + + "\x06script\x18\x01 \x01(\tR\x06script\x12\x1d\n" + + "\n" + + "input_json\x18\x02 \x01(\tR\tinputJson\"J\n" + + "\x0fRunPythonOutput\x12\x1f\n" + + "\vresult_json\x18\x01 \x01(\tR\n" + + "resultJson\x12\x16\n" + + "\x06stdout\x18\x02 \x01(\tR\x06stdout2\xe2\x04\n" + "\tToolQuery\x12H\n" + "\aMetrics\x12\x1c.toolquery.MetricsQueryInput\x1a\x1d.toolquery.MetricsQueryOutput\"\x00\x12P\n" + "\rMetricsSearch\x12\x1d.toolquery.MetricsSearchInput\x1a\x1e.toolquery.MetricsSearchOutput\"\x00\x12_\n" + @@ -3240,7 +3356,8 @@ const file_toolquery_proto_rawDesc = "" + "\x04Logs\x12\x19.toolquery.LogsQueryInput\x1a\x1a.toolquery.LogsQueryOutput\"\x00\x12E\n" + "\x06Traces\x12\x1b.toolquery.TracesQueryInput\x1a\x1c.toolquery.TracesQueryOutput\"\x00\x12M\n" + "\fInvokeLambda\x12\x1c.toolquery.InvokeLambdaInput\x1a\x1d.toolquery.InvokeLambdaOutput\"\x00\x12;\n" + - "\x06RunLua\x12\x16.toolquery.RunLuaInput\x1a\x17.toolquery.RunLuaOutput\"\x00BEZCgithub.com/pluralsh/console/go/cloud-query/internal/proto/toolqueryb\x06proto3" + "\x06RunLua\x12\x16.toolquery.RunLuaInput\x1a\x17.toolquery.RunLuaOutput\"\x00\x12D\n" + + "\tRunPython\x12\x19.toolquery.RunPythonInput\x1a\x1a.toolquery.RunPythonOutput\"\x00BEZCgithub.com/pluralsh/console/go/cloud-query/internal/proto/toolqueryb\x06proto3" var ( file_toolquery_proto_rawDescOnce sync.Once @@ -3254,7 +3371,7 @@ func file_toolquery_proto_rawDescGZIP() []byte { return file_toolquery_proto_rawDescData } -var file_toolquery_proto_msgTypes = make([]protoimpl.MessageInfo, 47) +var file_toolquery_proto_msgTypes = make([]protoimpl.MessageInfo, 49) var file_toolquery_proto_goTypes = []any{ (*ElasticConnection)(nil), // 0: toolquery.ElasticConnection (*OpensearchConnection)(nil), // 1: toolquery.OpensearchConnection @@ -3300,11 +3417,13 @@ var file_toolquery_proto_goTypes = []any{ (*InvokeLambdaOutput)(nil), // 41: toolquery.InvokeLambdaOutput (*RunLuaInput)(nil), // 42: toolquery.RunLuaInput (*RunLuaOutput)(nil), // 43: toolquery.RunLuaOutput - nil, // 44: toolquery.MetricPoint.LabelsEntry - nil, // 45: toolquery.LogEntry.LabelsEntry - nil, // 46: toolquery.TraceSpan.TagsEntry - (*timestamppb.Timestamp)(nil), // 47: google.protobuf.Timestamp - (*cloudquery.Connection)(nil), // 48: cloudquery.Connection + (*RunPythonInput)(nil), // 44: toolquery.RunPythonInput + (*RunPythonOutput)(nil), // 45: toolquery.RunPythonOutput + nil, // 46: toolquery.MetricPoint.LabelsEntry + nil, // 47: toolquery.LogEntry.LabelsEntry + nil, // 48: toolquery.TraceSpan.TagsEntry + (*timestamppb.Timestamp)(nil), // 49: google.protobuf.Timestamp + (*cloudquery.Connection)(nil), // 50: cloudquery.Connection } var file_toolquery_proto_depIdxs = []int32{ 0, // 0: toolquery.ToolConnection.elastic:type_name -> toolquery.ElasticConnection @@ -3318,8 +3437,8 @@ var file_toolquery_proto_depIdxs = []int32{ 10, // 8: toolquery.ToolConnection.azure:type_name -> toolquery.AzureConnection 6, // 9: toolquery.ToolConnection.jaeger:type_name -> toolquery.JaegerConnection 1, // 10: toolquery.ToolConnection.opensearch:type_name -> toolquery.OpensearchConnection - 47, // 11: toolquery.TimeRange.start:type_name -> google.protobuf.Timestamp - 47, // 12: toolquery.TimeRange.end:type_name -> google.protobuf.Timestamp + 49, // 11: toolquery.TimeRange.start:type_name -> google.protobuf.Timestamp + 49, // 12: toolquery.TimeRange.end:type_name -> google.protobuf.Timestamp 11, // 13: toolquery.MetricsQueryInput.connection:type_name -> toolquery.ToolConnection 12, // 14: toolquery.MetricsQueryInput.range:type_name -> toolquery.TimeRange 14, // 15: toolquery.MetricsQueryInput.options:type_name -> toolquery.MetricsOptions @@ -3334,8 +3453,8 @@ var file_toolquery_proto_depIdxs = []int32{ 21, // 24: toolquery.TracesQueryInput.options:type_name -> toolquery.TracesOptions 23, // 25: toolquery.TracesOptions.jaeger:type_name -> toolquery.JaegerTracesOptions 22, // 26: toolquery.JaegerTracesOptions.attributes:type_name -> toolquery.JaegerTraceQueryAttribute - 47, // 27: toolquery.MetricPoint.timestamp:type_name -> google.protobuf.Timestamp - 44, // 28: toolquery.MetricPoint.labels:type_name -> toolquery.MetricPoint.LabelsEntry + 49, // 27: toolquery.MetricPoint.timestamp:type_name -> google.protobuf.Timestamp + 46, // 28: toolquery.MetricPoint.labels:type_name -> toolquery.MetricPoint.LabelsEntry 24, // 29: toolquery.MetricsQueryOutput.metrics:type_name -> toolquery.MetricPoint 11, // 30: toolquery.MetricsSearchInput.connection:type_name -> toolquery.ToolConnection 27, // 31: toolquery.MetricsSearchInput.options:type_name -> toolquery.MetricsSearchOptions @@ -3345,14 +3464,14 @@ var file_toolquery_proto_depIdxs = []int32{ 32, // 35: toolquery.MetricsLabelSearchInput.options:type_name -> toolquery.MetricsLabelSearchOptions 33, // 36: toolquery.MetricsLabelSearchOptions.azure:type_name -> toolquery.AzureMetricsLabelSearchOptions 34, // 37: toolquery.MetricsLabelSearchOutput.results:type_name -> toolquery.MetricsLabelSearchResult - 47, // 38: toolquery.LogEntry.timestamp:type_name -> google.protobuf.Timestamp - 45, // 39: toolquery.LogEntry.labels:type_name -> toolquery.LogEntry.LabelsEntry + 49, // 38: toolquery.LogEntry.timestamp:type_name -> google.protobuf.Timestamp + 47, // 39: toolquery.LogEntry.labels:type_name -> toolquery.LogEntry.LabelsEntry 36, // 40: toolquery.LogsQueryOutput.logs:type_name -> toolquery.LogEntry - 47, // 41: toolquery.TraceSpan.start:type_name -> google.protobuf.Timestamp - 47, // 42: toolquery.TraceSpan.end:type_name -> google.protobuf.Timestamp - 46, // 43: toolquery.TraceSpan.tags:type_name -> toolquery.TraceSpan.TagsEntry + 49, // 41: toolquery.TraceSpan.start:type_name -> google.protobuf.Timestamp + 49, // 42: toolquery.TraceSpan.end:type_name -> google.protobuf.Timestamp + 48, // 43: toolquery.TraceSpan.tags:type_name -> toolquery.TraceSpan.TagsEntry 38, // 44: toolquery.TracesQueryOutput.spans:type_name -> toolquery.TraceSpan - 48, // 45: toolquery.InvokeLambdaInput.connection:type_name -> cloudquery.Connection + 50, // 45: toolquery.InvokeLambdaInput.connection:type_name -> cloudquery.Connection 13, // 46: toolquery.ToolQuery.Metrics:input_type -> toolquery.MetricsQueryInput 26, // 47: toolquery.ToolQuery.MetricsSearch:input_type -> toolquery.MetricsSearchInput 31, // 48: toolquery.ToolQuery.MetricsLabelSearch:input_type -> toolquery.MetricsLabelSearchInput @@ -3360,15 +3479,17 @@ var file_toolquery_proto_depIdxs = []int32{ 20, // 50: toolquery.ToolQuery.Traces:input_type -> toolquery.TracesQueryInput 40, // 51: toolquery.ToolQuery.InvokeLambda:input_type -> toolquery.InvokeLambdaInput 42, // 52: toolquery.ToolQuery.RunLua:input_type -> toolquery.RunLuaInput - 25, // 53: toolquery.ToolQuery.Metrics:output_type -> toolquery.MetricsQueryOutput - 30, // 54: toolquery.ToolQuery.MetricsSearch:output_type -> toolquery.MetricsSearchOutput - 35, // 55: toolquery.ToolQuery.MetricsLabelSearch:output_type -> toolquery.MetricsLabelSearchOutput - 37, // 56: toolquery.ToolQuery.Logs:output_type -> toolquery.LogsQueryOutput - 39, // 57: toolquery.ToolQuery.Traces:output_type -> toolquery.TracesQueryOutput - 41, // 58: toolquery.ToolQuery.InvokeLambda:output_type -> toolquery.InvokeLambdaOutput - 43, // 59: toolquery.ToolQuery.RunLua:output_type -> toolquery.RunLuaOutput - 53, // [53:60] is the sub-list for method output_type - 46, // [46:53] is the sub-list for method input_type + 44, // 53: toolquery.ToolQuery.RunPython:input_type -> toolquery.RunPythonInput + 25, // 54: toolquery.ToolQuery.Metrics:output_type -> toolquery.MetricsQueryOutput + 30, // 55: toolquery.ToolQuery.MetricsSearch:output_type -> toolquery.MetricsSearchOutput + 35, // 56: toolquery.ToolQuery.MetricsLabelSearch:output_type -> toolquery.MetricsLabelSearchOutput + 37, // 57: toolquery.ToolQuery.Logs:output_type -> toolquery.LogsQueryOutput + 39, // 58: toolquery.ToolQuery.Traces:output_type -> toolquery.TracesQueryOutput + 41, // 59: toolquery.ToolQuery.InvokeLambda:output_type -> toolquery.InvokeLambdaOutput + 43, // 60: toolquery.ToolQuery.RunLua:output_type -> toolquery.RunLuaOutput + 45, // 61: toolquery.ToolQuery.RunPython:output_type -> toolquery.RunPythonOutput + 54, // [54:62] is the sub-list for method output_type + 46, // [46:54] is the sub-list for method input_type 46, // [46:46] is the sub-list for extension type_name 46, // [46:46] is the sub-list for extension extendee 0, // [0:46] is the sub-list for field type_name @@ -3420,7 +3541,7 @@ func file_toolquery_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_toolquery_proto_rawDesc), len(file_toolquery_proto_rawDesc)), NumEnums: 0, - NumMessages: 47, + NumMessages: 49, NumExtensions: 0, NumServices: 1, }, diff --git a/go/cloud-query/internal/proto/toolquery/toolquery_grpc.pb.go b/go/cloud-query/internal/proto/toolquery/toolquery_grpc.pb.go index f990f45cbe..400a79eb5d 100644 --- a/go/cloud-query/internal/proto/toolquery/toolquery_grpc.pb.go +++ b/go/cloud-query/internal/proto/toolquery/toolquery_grpc.pb.go @@ -26,6 +26,7 @@ const ( ToolQuery_Traces_FullMethodName = "/toolquery.ToolQuery/Traces" ToolQuery_InvokeLambda_FullMethodName = "/toolquery.ToolQuery/InvokeLambda" ToolQuery_RunLua_FullMethodName = "/toolquery.ToolQuery/RunLua" + ToolQuery_RunPython_FullMethodName = "/toolquery.ToolQuery/RunPython" ) // ToolQueryClient is the client API for ToolQuery service. @@ -39,6 +40,7 @@ type ToolQueryClient interface { Traces(ctx context.Context, in *TracesQueryInput, opts ...grpc.CallOption) (*TracesQueryOutput, error) InvokeLambda(ctx context.Context, in *InvokeLambdaInput, opts ...grpc.CallOption) (*InvokeLambdaOutput, error) RunLua(ctx context.Context, in *RunLuaInput, opts ...grpc.CallOption) (*RunLuaOutput, error) + RunPython(ctx context.Context, in *RunPythonInput, opts ...grpc.CallOption) (*RunPythonOutput, error) } type toolQueryClient struct { @@ -119,6 +121,16 @@ func (c *toolQueryClient) RunLua(ctx context.Context, in *RunLuaInput, opts ...g return out, nil } +func (c *toolQueryClient) RunPython(ctx context.Context, in *RunPythonInput, opts ...grpc.CallOption) (*RunPythonOutput, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(RunPythonOutput) + err := c.cc.Invoke(ctx, ToolQuery_RunPython_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + // ToolQueryServer is the server API for ToolQuery service. // All implementations must embed UnimplementedToolQueryServer // for forward compatibility. @@ -130,6 +142,7 @@ type ToolQueryServer interface { Traces(context.Context, *TracesQueryInput) (*TracesQueryOutput, error) InvokeLambda(context.Context, *InvokeLambdaInput) (*InvokeLambdaOutput, error) RunLua(context.Context, *RunLuaInput) (*RunLuaOutput, error) + RunPython(context.Context, *RunPythonInput) (*RunPythonOutput, error) mustEmbedUnimplementedToolQueryServer() } @@ -161,6 +174,9 @@ func (UnimplementedToolQueryServer) InvokeLambda(context.Context, *InvokeLambdaI func (UnimplementedToolQueryServer) RunLua(context.Context, *RunLuaInput) (*RunLuaOutput, error) { return nil, status.Error(codes.Unimplemented, "method RunLua not implemented") } +func (UnimplementedToolQueryServer) RunPython(context.Context, *RunPythonInput) (*RunPythonOutput, error) { + return nil, status.Error(codes.Unimplemented, "method RunPython not implemented") +} func (UnimplementedToolQueryServer) mustEmbedUnimplementedToolQueryServer() {} func (UnimplementedToolQueryServer) testEmbeddedByValue() {} @@ -308,6 +324,24 @@ func _ToolQuery_RunLua_Handler(srv interface{}, ctx context.Context, dec func(in return interceptor(ctx, in, info, handler) } +func _ToolQuery_RunPython_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(RunPythonInput) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(ToolQueryServer).RunPython(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: ToolQuery_RunPython_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(ToolQueryServer).RunPython(ctx, req.(*RunPythonInput)) + } + return interceptor(ctx, in, info, handler) +} + // ToolQuery_ServiceDesc is the grpc.ServiceDesc for ToolQuery service. // It's only intended for direct use with grpc.RegisterService, // and not to be introspected or modified (even as a copy) @@ -343,6 +377,10 @@ var ToolQuery_ServiceDesc = grpc.ServiceDesc{ MethodName: "RunLua", Handler: _ToolQuery_RunLua_Handler, }, + { + MethodName: "RunPython", + Handler: _ToolQuery_RunPython_Handler, + }, }, Streams: []grpc.StreamDesc{}, Metadata: "toolquery.proto", diff --git a/go/cloud-query/internal/server/server.go b/go/cloud-query/internal/server/server.go index a2560d3cab..7c018e1abb 100644 --- a/go/cloud-query/internal/server/server.go +++ b/go/cloud-query/internal/server/server.go @@ -5,6 +5,7 @@ import ( "crypto/tls" "fmt" "net" + "sync" "time" "google.golang.org/grpc" @@ -23,6 +24,7 @@ type Server struct { server *grpc.Server services []service.Service stopped chan struct{} // Channel to signal when server is stopped + stopOnce sync.Once } // Start initializes and starts the gRPC server @@ -49,21 +51,32 @@ func (in *Server) Start(ctx context.Context) error { // Stop gracefully shuts down the server func (in *Server) Stop() { - klog.Info("stopping gRPC server") - go func() { - in.server.GracefulStop() - close(in.stopped) - }() + in.stopOnce.Do(func() { + klog.Info("stopping gRPC server") + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + go func() { + in.server.GracefulStop() + close(in.stopped) + }() + + select { + case <-in.stopped: + klog.Info("gRPC server stopped gracefully") + case <-ctx.Done(): + klog.Info("timeout waiting for server to stop, forcing shutdown") + in.server.Stop() + } - // Wait for a graceful stop with a timeout - timeout := time.After(10 * time.Second) - select { - case <-in.stopped: - klog.Info("gRPC server stopped gracefully") - case <-timeout: - klog.Info("timeout waiting for server to stop, forcing shutdown") - in.server.Stop() - } + for _, registered := range in.services { + if closer, ok := registered.(service.Closer); ok { + if err := closer.Close(ctx); err != nil { + klog.ErrorS(err, "failed to close service") + } + } + } + }) } func (in *Server) register() { diff --git a/go/cloud-query/internal/service/service.go b/go/cloud-query/internal/service/service.go index 8c4b92a8e3..d55cd87da2 100644 --- a/go/cloud-query/internal/service/service.go +++ b/go/cloud-query/internal/service/service.go @@ -1,6 +1,8 @@ package service import ( + "context" + "google.golang.org/grpc" ) @@ -8,3 +10,9 @@ import ( type Service interface { Install(server *grpc.Server) } + +// Closer is implemented by services that own background processes or other +// resources which must be released during server shutdown. +type Closer interface { + Close(context.Context) error +} diff --git a/go/cloud-query/internal/service/toolquery.go b/go/cloud-query/internal/service/toolquery.go index 9fc32ed4b4..3643ef7f93 100644 --- a/go/cloud-query/internal/service/toolquery.go +++ b/go/cloud-query/internal/service/toolquery.go @@ -16,11 +16,13 @@ import ( "github.com/pluralsh/console/go/cloud-query/internal/tools" lambdatools "github.com/pluralsh/console/go/cloud-query/internal/tools/lambda" luatools "github.com/pluralsh/console/go/cloud-query/internal/tools/lua" + pythontools "github.com/pluralsh/console/go/cloud-query/internal/tools/python" ) // ToolQueryService implements the toolquery.ToolQueryServer interface. type ToolQueryService struct { toolquery.UnimplementedToolQueryServer + python *pythontools.Runner } // Install registers the ToolQuery service with the gRPC server. @@ -29,9 +31,22 @@ func (in *ToolQueryService) Install(server *grpc.Server) { toolquery.RegisterToolQueryServer(server, in) } -// NewToolQueryService creates a new instance of the ToolQuery server. -func NewToolQueryService() Service { - return &ToolQueryService{} +// NewToolQueryService creates a ToolQuery server and verifies its sandboxed +// Python worker pool before the gRPC server starts accepting requests. +func NewToolQueryService(ctx context.Context) (Service, error) { + runner, err := pythontools.New(ctx) + if err != nil { + return nil, err + } + return &ToolQueryService{python: runner}, nil +} + +// Close reaps all Python worker subprocesses. +func (in *ToolQueryService) Close(ctx context.Context) error { + if in.python != nil { + return in.python.CloseContext(ctx) + } + return nil } func (in *ToolQueryService) Metrics(ctx context.Context, input *toolquery.MetricsQueryInput) (*toolquery.MetricsQueryOutput, error) { @@ -204,6 +219,50 @@ func (in *ToolQueryService) RunLua(ctx context.Context, input *toolquery.RunLuaI }, nil } +func (in *ToolQueryService) RunPython(ctx context.Context, input *toolquery.RunPythonInput) (*toolquery.RunPythonOutput, error) { + if input == nil { + return nil, status.Error(codes.InvalidArgument, "input is required") + } + if in.python == nil { + return nil, status.Error(codes.Unavailable, "python runtime is unavailable") + } + + output, err := in.python.Run(ctx, pythontools.RunInput{ + Script: input.GetScript(), + InputJSON: input.GetInputJson(), + }) + if err != nil { + if pythontools.ErrorCode(err) == pythontools.Internal { + klog.ErrorS(err, "python worker failed") + } + return nil, status.Error(pythonGRPCCode(err), err.Error()) + } + + return &toolquery.RunPythonOutput{ + ResultJson: output.ResultJSON, + Stdout: output.Stdout, + }, nil +} + +func pythonGRPCCode(err error) codes.Code { + switch pythontools.ErrorCode(err) { + case pythontools.InvalidArgument: + return codes.InvalidArgument + case pythontools.FailedPrecondition: + return codes.FailedPrecondition + case pythontools.Canceled: + return codes.Canceled + case pythontools.DeadlineExceeded: + return codes.DeadlineExceeded + case pythontools.ResourceExhausted: + return codes.ResourceExhausted + case pythontools.Unavailable: + return codes.Unavailable + default: + return codes.Internal + } +} + func (in *ToolQueryService) validateInput(connection *toolquery.ToolConnection, query string, timeRange *toolquery.TimeRange) error { if err := in.validateSearchInput(connection); err != nil { return err diff --git a/go/cloud-query/internal/service/toolquery_python_test.go b/go/cloud-query/internal/service/toolquery_python_test.go new file mode 100644 index 0000000000..3b65353292 --- /dev/null +++ b/go/cloud-query/internal/service/toolquery_python_test.go @@ -0,0 +1,40 @@ +package service + +import ( + "context" + "testing" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/pluralsh/console/go/cloud-query/internal/proto/toolquery" + pythontools "github.com/pluralsh/console/go/cloud-query/internal/tools/python" +) + +func TestRunPythonValidatesInputAndAvailability(t *testing.T) { + server := &ToolQueryService{} + + if _, err := server.RunPython(context.Background(), nil); status.Code(err) != codes.InvalidArgument { + t.Fatalf("nil input code = %s, want %s", status.Code(err), codes.InvalidArgument) + } + if _, err := server.RunPython(context.Background(), &toolquery.RunPythonInput{Script: "output = {}"}); status.Code(err) != codes.Unavailable { + t.Fatalf("missing runner code = %s, want %s", status.Code(err), codes.Unavailable) + } +} + +func TestPythonGRPCCode(t *testing.T) { + for runCode, want := range map[pythontools.Code]codes.Code{ + pythontools.InvalidArgument: codes.InvalidArgument, + pythontools.FailedPrecondition: codes.FailedPrecondition, + pythontools.Canceled: codes.Canceled, + pythontools.DeadlineExceeded: codes.DeadlineExceeded, + pythontools.ResourceExhausted: codes.ResourceExhausted, + pythontools.Unavailable: codes.Unavailable, + pythontools.Internal: codes.Internal, + } { + err := &pythontools.Error{Code: runCode, Msg: "safe"} + if got := pythonGRPCCode(err); got != want { + t.Errorf("pythonGRPCCode(%s) = %s, want %s", runCode, got, want) + } + } +} diff --git a/go/cloud-query/internal/tools/python/README.md b/go/cloud-query/internal/tools/python/README.md new file mode 100644 index 0000000000..2374eba57d --- /dev/null +++ b/go/cloud-query/internal/tools/python/README.md @@ -0,0 +1,63 @@ +# Sandboxed Python runner + +This package implements `RunPython` with Monty's limited Python subset. It is +not CPython. The parent starts four copies of `cloud-query python-worker` and +communicates with them over private stdin/stdout pipes. Each worker loads +gomonty's embedded native library and creates a fresh REPL for every request. + +## Execution contract + +Each run receives Python source and an optional JSON object. The object becomes +the global `input` dictionary. `output` starts as an empty dictionary and must +remain JSON serializable and object shaped. The response returns `output` and +the user script's standard output separately. + +Workers receive only `TMPDIR=/tmp`. The host supplies no filesystem, +environment, network, subprocess, shell, package, OS callback, external +function, or name-lookup capability. + +Limits: + +- Source: 64 KiB +- Input and result JSON: 1 MiB each +- Captured stdout: 64 KiB +- Monty execution: 10 seconds +- Parent wall clock: 15 seconds or the caller's earlier deadline +- Monty-managed memory: 64 MiB +- Recursion: 200 frames +- Active workers: four +- Waiting requests: 16 by default +- Worker recycling: 10 successful requests + +Monty's memory limit does not bound total worker or pod RSS. The parent provides +the hard crash and wall-time limit by killing and replacing a worker process. + +## Private protocol + +Protocol version 1 uses a four-byte big-endian length followed by strict JSON. +Frames are bounded before allocation, unknown fields and trailing JSON values +are rejected, and only `health` and `run` request types are accepted. Every +response must repeat the request ID and type. Malformed protocol data terminates +the worker. + +The parent admits requests to a bounded FIFO queue. Caller cancellation applies +while waiting and during execution. A worker is replaced after any execution +error, cancellation that reaches its exchange, protocol or transport failure, +or after ten successful responses. It does not run a health request between +successful executions. Shutdown stops admission, drains queued and active work +until the server shutdown context expires, then kills and reaps the remaining +workers. + +## Native runtime + +The Go module pins gomonty `v0.0.14`, which embeds glibc shared libraries for +Linux amd64 and arm64. That release builds against official Monty commit +`c9802b5f30d11fecf9f153feb1dfdab3abda070e`. The production image contains no +separate `monty` executable. + +Run focused validation with: + +```sh +cd go/cloud-query +go test ./internal/tools/python ./internal/service +``` diff --git a/go/cloud-query/internal/tools/python/common.go b/go/cloud-query/internal/tools/python/common.go new file mode 100644 index 0000000000..a45be7fc44 --- /dev/null +++ b/go/cloud-query/internal/tools/python/common.go @@ -0,0 +1,25 @@ +package python + +import ( + "encoding/json" + "strconv" + "strings" +) + +func validateRunInput(raw string) (string, error) { + if strings.TrimSpace(raw) == "" { + return "{}", nil + } + var value map[string]json.RawMessage + if err := json.Unmarshal([]byte(raw), &value); err != nil || value == nil { + return "", invalid("input_json must be a JSON object") + } + return raw, nil +} + +func pythonString(value string) string { return strconv.Quote(value) } + +func isJSONObject(value string) bool { + var object map[string]json.RawMessage + return json.Unmarshal([]byte(value), &object) == nil && object != nil +} diff --git a/go/cloud-query/internal/tools/python/errors.go b/go/cloud-query/internal/tools/python/errors.go new file mode 100644 index 0000000000..94f6d60492 --- /dev/null +++ b/go/cloud-query/internal/tools/python/errors.go @@ -0,0 +1,100 @@ +package python + +import ( + "context" + "errors" + "strings" +) + +const ( + maxExceptionSummaryRunes = 512 + runtimeUnavailableMessage = "python runtime is unavailable" +) + +// Code is a safe, stable error category suitable for mapping to an RPC status. +type Code string + +const ( + // InvalidArgument indicates invalid source, input, or output data. + InvalidArgument Code = "invalid_argument" + // FailedPrecondition indicates Python could not complete the requested operation. + FailedPrecondition Code = "failed_precondition" + // Canceled indicates the caller canceled execution. + Canceled Code = "canceled" + // DeadlineExceeded indicates execution exceeded its allotted time. + DeadlineExceeded Code = "deadline_exceeded" + // ResourceExhausted indicates a configured size, memory, or concurrency limit was exceeded. + ResourceExhausted Code = "resource_exhausted" + // Unavailable indicates no healthy Python worker is available. + Unavailable Code = "unavailable" + // Internal indicates an unexpected runner or worker failure. + Internal Code = "internal" +) + +// Error contains only a stable category and a safe summary. Protocol frames, +// child stderr, and host paths are intentionally never retained. +type Error struct { + Code Code + Msg string +} + +// Error returns the safe error summary. +func (e *Error) Error() string { return e.Msg } + +// ErrorCode returns a safe classification for err. Errors outside this package +// are intentionally treated as Internal so callers never infer host details. +func ErrorCode(err error) Code { + var runErr *Error + if errors.As(err, &runErr) { + return runErr.Code + } + return Internal +} + +func invalid(msg string) error { return &Error{Code: InvalidArgument, Msg: msg} } + +func runtimeError(kind, message string) error { + switch kind { + case "SyntaxError", "TypeError": + return invalid(pythonExceptionSummary(kind, message)) + case "MemoryError", "RecursionError": + return &Error{Code: ResourceExhausted, Msg: "python resource limit exceeded"} + case "TimeoutError": + return executionTimeoutError() + default: + return &Error{Code: FailedPrecondition, Msg: pythonExceptionSummary(kind, message)} + } +} + +// pythonExceptionSummary exposes the useful exception type and message while +// excluding tracebacks, worker diagnostics, protocol data, and unbounded text. +func pythonExceptionSummary(kind, message string) string { + kind = strings.TrimSpace(strings.ToValidUTF8(kind, "")) + message = strings.Join(strings.Fields(strings.ToValidUTF8(message, "")), " ") + if kind == "" { + kind = "Exception" + } + summary := "python " + kind + if message != "" { + summary += ": " + message + } + runes := []rune(summary) + if len(runes) > maxExceptionSummaryRunes { + summary = string(runes[:maxExceptionSummaryRunes-1]) + "…" + } + return summary +} + +func runtimeFailure() error { return &Error{Code: Internal, Msg: "python runtime failed"} } + +func executionTimeoutError() error { + return &Error{Code: DeadlineExceeded, Msg: "python execution timed out"} +} + +func executionContextError(err error) error { + if errors.Is(err, context.Canceled) { + return &Error{Code: Canceled, Msg: "python execution was canceled"} + } + + return executionTimeoutError() +} diff --git a/go/cloud-query/internal/tools/python/protocol.go b/go/cloud-query/internal/tools/python/protocol.go new file mode 100644 index 0000000000..07f4e68470 --- /dev/null +++ b/go/cloud-query/internal/tools/python/protocol.go @@ -0,0 +1,138 @@ +package python + +import ( + "bytes" + "context" + "encoding/binary" + "encoding/json" + "errors" + "io" + "sync" +) + +const ( + protocolVersion = 1 + requestHealth = "health" + requestRun = "run" +) + +type protocolRequest struct { + Version int `json:"version"` + Type string `json:"type"` + ID string `json:"id"` + Script string `json:"script,omitempty"` + InputJSON string `json:"input_json,omitempty"` +} + +type protocolResponse struct { + Version int `json:"version"` + Type string `json:"type"` + ID string `json:"id"` + ResultJSON string `json:"result_json,omitempty"` + Stdout string `json:"stdout,omitempty"` + Code Code `json:"code,omitempty"` + Message string `json:"message,omitempty"` +} + +func readFrame(in io.Reader, into any) error { + var header [4]byte + if _, err := io.ReadFull(in, header[:]); err != nil { + return err + } + length := binary.BigEndian.Uint32(header[:]) + if length == 0 || length > maxProtocolFrameLen { + return errors.New("invalid protocol frame") + } + body := make([]byte, length) + if _, err := io.ReadFull(in, body); err != nil { + return err + } + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(into); err != nil { + return errors.New("invalid protocol payload") + } + if err := decoder.Decode(&struct{}{}); !errors.Is(err, io.EOF) { + return errors.New("invalid protocol payload") + } + return nil +} + +func writeFrame(out io.Writer, value any) error { + body, err := json.Marshal(value) + if err != nil || len(body) == 0 || len(body) > maxProtocolFrameLen { + return errors.New("invalid protocol response") + } + frame := make([]byte, 4+len(body)) + binary.BigEndian.PutUint32(frame, uint32(len(body))) + copy(frame[4:], body) + _, err = io.Copy(out, bytes.NewReader(frame)) + return err +} + +type childWorker struct { + stdin io.WriteCloser + stdout io.ReadCloser + kill func() + mu sync.Mutex + stopOnce sync.Once +} + +func (w *childWorker) stop() { w.stopOnce.Do(w.kill) } + +func (w *childWorker) exchange(ctx context.Context, request protocolRequest) (protocolResponse, error) { + w.mu.Lock() + defer w.mu.Unlock() + result := make(chan struct { + response protocolResponse + err error + }, 1) + go func() { + if err := writeFrame(w.stdin, request); err != nil { + result <- struct { + response protocolResponse + err error + }{err: err} + return + } + var response protocolResponse + result <- struct { + response protocolResponse + err error + }{response: response, err: readFrame(w.stdout, &response)} + }() + select { + case result := <-result: + if result.err != nil || result.response.Version != protocolVersion || result.response.ID != request.ID || result.response.Type != request.Type { + return protocolResponse{}, runtimeFailure() + } + if (result.response.Code == "") != (result.response.Message == "") { + return protocolResponse{}, runtimeFailure() + } + if request.Type == requestHealth && (result.response.ResultJSON != "" || result.response.Stdout != "") { + return protocolResponse{}, runtimeFailure() + } + if result.response.Code != "" && (result.response.ResultJSON != "" || result.response.Stdout != "") { + return protocolResponse{}, runtimeFailure() + } + if request.Type == requestRun && result.response.Code == "" && (len(result.response.ResultJSON) > MaxResultBytes || len(result.response.Stdout) > MaxStdoutBytes || !isJSONObject(result.response.ResultJSON)) { + return protocolResponse{}, runtimeFailure() + } + return result.response, nil + case <-ctx.Done(): + w.stop() + return protocolResponse{}, executionContextError(ctx.Err()) + } +} + +func protocolError(response protocolResponse) error { + if response.Code == "" || response.Message == "" { + return runtimeFailure() + } + switch response.Code { + case InvalidArgument, FailedPrecondition, Canceled, DeadlineExceeded, ResourceExhausted, Unavailable, Internal: + return &Error{Code: response.Code, Msg: response.Message} + default: + return runtimeFailure() + } +} diff --git a/go/cloud-query/internal/tools/python/python.go b/go/cloud-query/internal/tools/python/python.go new file mode 100644 index 0000000000..2587cce78b --- /dev/null +++ b/go/cloud-query/internal/tools/python/python.go @@ -0,0 +1,322 @@ +// Package python runs Monty in crash-isolated self-spawned workers. +package python + +import ( + "context" + "fmt" + "os" + "os/exec" + "strings" + "sync" + "time" +) + +const ( + MaxSourceBytes = 64 << 10 + MaxInputBytes = 1 << 20 + MaxResultBytes = 1 << 20 + MaxStdoutBytes = 64 << 10 + + defaultWorkers = 4 + defaultQueueSize = 16 + maxCheckouts = 10 + executionTimeout = 10 * time.Second + wallTimeout = 15 * time.Second + maxMemoryBytes = 64 << 20 + maxRecursionDepth = 200 + maxProtocolFrameLen = MaxInputBytes + MaxSourceBytes + MaxResultBytes + MaxStdoutBytes +) + +type RunInput struct{ Script, InputJSON string } +type RunOutput struct{ ResultJSON, Stdout string } + +// Config controls the parent pool. Zero numeric fields use the package defaults. +type Config struct { + Workers int + QueueSize int + MaxSuccessfulRuns int + Executable string + Arguments []string +} + +type Option func(*Config) + +// WithBinaryPath remains available for tests. Production defaults to the +// current executable, which must dispatch python-worker to RunWorker. +func WithBinaryPath(path string) Option { return func(c *Config) { c.Executable = path } } +func withArguments(arguments ...string) Option { + return func(c *Config) { c.Arguments = append([]string(nil), arguments...) } +} +func WithConfig(config Config) Option { return func(c *Config) { *c = config } } + +type job struct { + ctx context.Context + input RunInput + response chan runResult +} +type runResult struct { + output *RunOutput + err error +} + +type Runner struct { + config Config + jobs chan *job + mu sync.Mutex + closed bool + healthy int + workers map[*childWorker]struct{} + workerWG sync.WaitGroup + replaceWG sync.WaitGroup + closeOnce sync.Once + closedDone chan struct{} +} + +func New(ctx context.Context, options ...Option) (*Runner, error) { + executable, err := os.Executable() + if err != nil { + return nil, &Error{Code: Unavailable, Msg: runtimeUnavailableMessage} + } + cfg := Config{Workers: defaultWorkers, QueueSize: defaultQueueSize, MaxSuccessfulRuns: maxCheckouts, Executable: executable, Arguments: []string{"python-worker"}} + for _, option := range options { + option(&cfg) + } + if cfg.Workers == 0 { + cfg.Workers = defaultWorkers + } + if cfg.QueueSize == 0 { + cfg.QueueSize = defaultQueueSize + } + if cfg.MaxSuccessfulRuns == 0 { + cfg.MaxSuccessfulRuns = maxCheckouts + } + if cfg.Executable == "" { + cfg.Executable = executable + } + if cfg.Arguments == nil { + cfg.Arguments = []string{"python-worker"} + } + if cfg.Workers <= 0 || cfg.QueueSize < 0 || cfg.MaxSuccessfulRuns <= 0 || strings.TrimSpace(cfg.Executable) == "" { + return nil, &Error{Code: Unavailable, Msg: runtimeUnavailableMessage} + } + r := &Runner{config: cfg, jobs: make(chan *job, cfg.QueueSize), workers: make(map[*childWorker]struct{}), closedDone: make(chan struct{})} + for range cfg.Workers { + if err := r.startWorker(ctx); err != nil { + r.Close() + return nil, &Error{Code: Unavailable, Msg: runtimeUnavailableMessage} + } + } + return r, nil +} + +func (r *Runner) startWorker(ctx context.Context) error { + ctx, cancel := context.WithTimeout(ctx, wallTimeout) + defer cancel() + + w, err := r.newWorker() + if err != nil { + return err + } + r.mu.Lock() + if r.closed { + r.mu.Unlock() + w.stop() + return fmt.Errorf("runner is closed") + } + r.workers[w] = struct{}{} + r.mu.Unlock() + healthy := false + defer func() { + if healthy { + return + } + w.stop() + r.mu.Lock() + delete(r.workers, w) + r.mu.Unlock() + }() + + response, err := w.exchange(ctx, protocolRequest{Version: protocolVersion, Type: requestHealth, ID: requestHealth}) + if err != nil || response.Code != "" { + return fmt.Errorf("unhealthy worker") + } + r.mu.Lock() + if r.closed { + r.mu.Unlock() + return fmt.Errorf("runner is closed") + } + r.healthy++ + r.workerWG.Add(1) + healthy = true + r.mu.Unlock() + go r.serveWorker(w) + return nil +} + +func (r *Runner) newWorker() (*childWorker, error) { + cmd := exec.Command(r.config.Executable, r.config.Arguments...) + cmd.Env = []string{"TMPDIR=/tmp"} + stdin, err := cmd.StdinPipe() + if err != nil { + return nil, err + } + stdout, err := cmd.StdoutPipe() + if err != nil { + return nil, err + } + if err := cmd.Start(); err != nil { + return nil, err + } + w := &childWorker{stdin: stdin, stdout: stdout} + w.kill = func() { + _ = stdin.Close() + _ = stdout.Close() + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + _ = cmd.Wait() + } + return w, nil +} + +func (r *Runner) serveWorker(w *childWorker) { + defer r.workerWG.Done() + defer func() { + w.stop() + r.mu.Lock() + delete(r.workers, w) + r.healthy-- + closed := r.closed + if !closed { + r.replaceWG.Add(1) + } + r.mu.Unlock() + if !closed { + go func() { + defer r.replaceWG.Done() + r.replaceWorker() + }() + } + }() + successes := 0 + for request := range r.jobs { + if err := request.ctx.Err(); err != nil { + request.response <- runResult{err: executionContextError(err)} + continue + } + ctx, cancel := context.WithTimeout(request.ctx, wallTimeout) + response, err := w.exchange(ctx, protocolRequest{Version: protocolVersion, Type: requestRun, ID: fmt.Sprintf("%d", time.Now().UnixNano()), Script: request.input.Script, InputJSON: request.input.InputJSON}) + cancel() + if err != nil { + request.response <- runResult{err: err} + return + } + if response.Code != "" { + request.response <- runResult{err: protocolError(response)} + return + } + request.response <- runResult{output: &RunOutput{ResultJSON: response.ResultJSON, Stdout: response.Stdout}} + successes++ + if successes >= r.config.MaxSuccessfulRuns { + return + } + } +} + +func (r *Runner) replaceWorker() { + ticker := time.NewTicker(time.Second) + defer ticker.Stop() + for { + r.mu.Lock() + closed := r.closed + r.mu.Unlock() + if closed { + return + } + if r.startWorker(context.Background()) == nil { + return + } + select { + case <-r.closedDone: + return + case <-ticker.C: + } + } +} + +func (r *Runner) Run(ctx context.Context, input RunInput) (*RunOutput, error) { + if len(strings.TrimSpace(input.Script)) == 0 { + return nil, invalid("script is required") + } + if len(input.Script) > MaxSourceBytes { + return nil, invalid("script exceeds the source limit") + } + if len(input.InputJSON) > MaxInputBytes { + return nil, invalid("input exceeds the input limit") + } + inputJSON, err := validateRunInput(input.InputJSON) + if err != nil { + return nil, err + } + if err := ctx.Err(); err != nil { + return nil, executionContextError(err) + } + request := &job{ctx: ctx, input: RunInput{Script: input.Script, InputJSON: inputJSON}, response: make(chan runResult, 1)} + r.mu.Lock() + if r.closed || r.healthy == 0 { + r.mu.Unlock() + return nil, &Error{Code: Unavailable, Msg: runtimeUnavailableMessage} + } + select { + case r.jobs <- request: + r.mu.Unlock() + case <-ctx.Done(): + r.mu.Unlock() + return nil, executionContextError(ctx.Err()) + default: + r.mu.Unlock() + return nil, &Error{Code: ResourceExhausted, Msg: "python queue is full"} + } + select { + case result := <-request.response: + return result.output, result.err + case <-ctx.Done(): + return nil, executionContextError(ctx.Err()) + } +} + +func (r *Runner) Close() { _ = r.CloseContext(context.Background()) } + +// CloseContext stops admissions and lets queued and active requests finish +// until ctx expires. On expiry it kills workers and fails remaining queued work. +func (r *Runner) CloseContext(ctx context.Context) error { + r.closeOnce.Do(func() { + r.mu.Lock() + r.closed = true + close(r.jobs) + r.mu.Unlock() + go func() { + r.workerWG.Wait() + r.replaceWG.Wait() + close(r.closedDone) + }() + }) + select { + case <-r.closedDone: + return nil + case <-ctx.Done(): + r.mu.Lock() + workers := make([]*childWorker, 0, len(r.workers)) + for w := range r.workers { + workers = append(workers, w) + } + r.mu.Unlock() + for _, w := range workers { + w.stop() + } + for request := range r.jobs { + request.response <- runResult{err: &Error{Code: Unavailable, Msg: runtimeUnavailableMessage}} + } + return ctx.Err() + } +} diff --git a/go/cloud-query/internal/tools/python/python_test.go b/go/cloud-query/internal/tools/python/python_test.go new file mode 100644 index 0000000000..52f3b69157 --- /dev/null +++ b/go/cloud-query/internal/tools/python/python_test.go @@ -0,0 +1,301 @@ +package python + +import ( + "bytes" + "context" + "encoding/binary" + "errors" + "fmt" + "io" + "os" + "strings" + "testing" + "time" +) + +func TestMain(m *testing.M) { + if len(os.Args) == 2 && os.Args[1] == "python-worker" { + if err := RunWorker(os.Stdin, os.Stdout); err != nil { + os.Exit(1) + } + os.Exit(0) + } + if len(os.Args) == 2 && os.Args[1] == "fake-python-worker" { + if err := runFakeWorker(os.Stdin, os.Stdout); err != nil { + os.Exit(1) + } + os.Exit(0) + } + os.Exit(m.Run()) +} + +func TestValidateRunInput(t *testing.T) { + for _, raw := range []string{"[]", "1", "{"} { + if _, err := validateRunInput(raw); ErrorCode(err) != InvalidArgument { + t.Fatalf("validateRunInput(%q) = %v", raw, err) + } + } + if got, err := validateRunInput(""); err != nil || got != "{}" { + t.Fatalf("validateRunInput(empty) = %q, %v", got, err) + } +} + +func TestPoolDefaults(t *testing.T) { + if defaultWorkers != 4 { + t.Fatalf("defaultWorkers = %d, want 4", defaultWorkers) + } + if maxCheckouts != 10 { + t.Fatalf("maxCheckouts = %d, want 10", maxCheckouts) + } +} + +func TestDefaultQueueSize(t *testing.T) { + runner, err := newFakeRunner(t, Config{Workers: 1, MaxSuccessfulRuns: 1}) + if err != nil { + t.Fatalf("New() error: %v", err) + } + defer runner.Close() + if got := cap(runner.jobs); got != 16 { + t.Fatalf("queue capacity = %d, want 16", got) + } +} + +func TestWorkerReusesTenSuccessfulRunsWithoutHealthChecks(t *testing.T) { + runner, err := newFakeRunner(t, Config{Workers: 1}) + if err != nil { + t.Fatalf("New() error: %v", err) + } + defer runner.Close() + + var firstPID string + for run := 0; run < maxCheckouts; run++ { + output, err := runner.Run(context.Background(), RunInput{Script: "success"}) + if err != nil { + t.Fatalf("run %d error: %v", run+1, err) + } + if run == 0 { + firstPID = output.ResultJSON + } else if output.ResultJSON != firstPID { + t.Fatalf("run %d used %s, want original worker %s", run+1, output.ResultJSON, firstPID) + } + } + + output, err := runAfterReplacement(t, runner, RunInput{Script: "success"}) + if err != nil { + t.Fatalf("replacement run error: %v", err) + } + if output.ResultJSON == firstPID { + t.Fatalf("worker was not recycled after %d successful runs", maxCheckouts) + } +} + +func TestWorkerResponseErrorTriggersReplacement(t *testing.T) { + runner, err := newFakeRunner(t, Config{Workers: 1}) + if err != nil { + t.Fatalf("New() error: %v", err) + } + defer runner.Close() + + first, err := runner.Run(context.Background(), RunInput{Script: "success"}) + if err != nil { + t.Fatalf("initial run error: %v", err) + } + if _, err := runner.Run(context.Background(), RunInput{Script: "resource-error"}); ErrorCode(err) != ResourceExhausted { + t.Fatalf("resource response error = %v", err) + } + second, err := runAfterReplacement(t, runner, RunInput{Script: "success"}) + if err != nil { + t.Fatalf("replacement run error: %v", err) + } + if second.ResultJSON == first.ResultJSON { + t.Fatal("worker was reused after a resource-exhausted response") + } +} + +func TestProtocolRejectsUnknownAndTrailingValues(t *testing.T) { + for _, body := range []string{ + `{"version":1,"type":"health","id":"1","extra":true}`, + `{"version":1,"type":"health","id":"1"} {}`, + } { + var framed bytes.Buffer + if err := writeRawFrame(&framed, []byte(body)); err != nil { + t.Fatal(err) + } + var request protocolRequest + if err := readFrame(&framed, &request); err == nil { + t.Fatalf("readFrame(%s) accepted invalid payload", body) + } + } +} + +func TestRunWorkerHealth(t *testing.T) { + configureWorkerEnvironment(t) + var input, output bytes.Buffer + if err := writeFrame(&input, protocolRequest{Version: protocolVersion, Type: requestHealth, ID: "1"}); err != nil { + t.Fatal(err) + } + if err := RunWorker(&input, &output); err != nil { + t.Fatalf("RunWorker() error: %v", err) + } + var response protocolResponse + if err := readFrame(&output, &response); err != nil { + t.Fatal(err) + } + if response.Code != "" || response.ID != "1" { + t.Fatalf("health response = %#v", response) + } +} + +func TestRunWorkerUsesFreshStateAndSeparatesStdout(t *testing.T) { + configureWorkerEnvironment(t) + var input, output bytes.Buffer + requests := []protocolRequest{ + {Version: protocolVersion, Type: requestRun, ID: "first", Script: "secret = 42\noutput = {'sum': input['value'] + 1}\nprint('ok')", InputJSON: `{"value": 1}`}, + {Version: protocolVersion, Type: requestRun, ID: "second", Script: "output = {'secret': secret}", InputJSON: `{}`}, + } + for _, request := range requests { + if err := writeFrame(&input, request); err != nil { + t.Fatal(err) + } + } + if err := RunWorker(&input, &output); err != nil { + t.Fatalf("RunWorker() error: %v", err) + } + + var first protocolResponse + if err := readFrame(&output, &first); err != nil { + t.Fatal(err) + } + if first.Code != "" || first.ResultJSON != `{"sum": 2}` || first.Stdout != "ok\n" { + t.Fatalf("first response = %#v", first) + } + var second protocolResponse + if err := readFrame(&output, &second); err != nil { + t.Fatal(err) + } + if second.Code != FailedPrecondition || second.ResultJSON != "" || second.Stdout != "" { + t.Fatalf("second response = %#v", second) + } +} + +func TestRunRejectsBoundsBeforeAdmission(t *testing.T) { + runner := &Runner{jobs: make(chan *job), healthy: 1} + _, err := runner.Run(context.Background(), RunInput{Script: strings.Repeat("x", MaxSourceBytes+1)}) + if ErrorCode(err) != InvalidArgument { + t.Fatalf("source error = %v", err) + } + _, err = runner.Run(context.Background(), RunInput{Script: "pass", InputJSON: strings.Repeat("x", MaxInputBytes+1)}) + if ErrorCode(err) != InvalidArgument { + t.Fatalf("input error = %v", err) + } +} + +func TestRunnerSelfSpawnsAndQueuesUntilCapacityReturns(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + runner, err := New(ctx, WithConfig(Config{Workers: 1, QueueSize: 1, MaxSuccessfulRuns: 100})) + if err != nil { + t.Fatalf("New() error: %v", err) + } + defer runner.Close() + + firstCtx, cancelFirst := context.WithCancel(ctx) + first := make(chan error, 1) + go func() { + _, err := runner.Run(firstCtx, RunInput{Script: "while True:\n pass"}) + first <- err + }() + time.Sleep(50 * time.Millisecond) + + second := make(chan runResult, 1) + go func() { + output, err := runner.Run(ctx, RunInput{Script: "output = {'value': 42}"}) + second <- runResult{output: output, err: err} + }() + deadline := time.Now().Add(time.Second) + for len(runner.jobs) != 1 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if len(runner.jobs) != 1 { + t.Fatal("second request did not enter the queue") + } + if _, err := runner.Run(ctx, RunInput{Script: "output = {}"}); ErrorCode(err) != ResourceExhausted { + t.Fatalf("full queue error = %v", err) + } + + cancelFirst() + if err := <-first; ErrorCode(err) != Canceled { + t.Fatalf("canceled run error = %v", err) + } + result := <-second + if result.err != nil || result.output == nil || result.output.ResultJSON != `{"value": 42}` { + t.Fatalf("queued run = %#v, %v", result.output, result.err) + } +} + +func writeRawFrame(out *bytes.Buffer, body []byte) error { + frame := make([]byte, 4) + binary.BigEndian.PutUint32(frame, uint32(len(body))) + _, err := out.Write(append(frame, body...)) + return err +} + +func configureWorkerEnvironment(t *testing.T) { + t.Helper() + t.Setenv("HOME", "") + t.Setenv("XDG_CACHE_HOME", "") + t.Setenv("TMPDIR", "/tmp") +} + +func newFakeRunner(t *testing.T, config Config) (*Runner, error) { + t.Helper() + config.Executable = os.Args[0] + config.Arguments = []string{"fake-python-worker"} + return New(context.Background(), WithConfig(config)) +} + +func runAfterReplacement(t *testing.T, runner *Runner, input RunInput) (*RunOutput, error) { + t.Helper() + deadline := time.Now().Add(time.Second) + for { + output, err := runner.Run(context.Background(), input) + if ErrorCode(err) != Unavailable || time.Now().After(deadline) { + return output, err + } + time.Sleep(time.Millisecond) + } +} + +func runFakeWorker(in *os.File, out *os.File) error { + ran := false + for { + var request protocolRequest + if err := readFrame(in, &request); err != nil { + if errors.Is(err, io.EOF) { + return nil + } + return err + } + response := protocolResponse{Version: protocolVersion, Type: request.Type, ID: request.ID} + switch request.Type { + case requestHealth: + if ran { + response.Code = Internal + response.Message = "unexpected health request" + } + case requestRun: + ran = true + if request.Script == "resource-error" { + response.Code = ResourceExhausted + response.Message = "python resource limit exceeded" + } else { + response.ResultJSON = fmt.Sprintf(`{"pid":%d}`, os.Getpid()) + } + default: + return fmt.Errorf("unexpected request type %q", request.Type) + } + if err := writeFrame(out, response); err != nil { + return err + } + } +} diff --git a/go/cloud-query/internal/tools/python/worker.go b/go/cloud-query/internal/tools/python/worker.go new file mode 100644 index 0000000000..5ecf17bb6c --- /dev/null +++ b/go/cloud-query/internal/tools/python/worker.go @@ -0,0 +1,148 @@ +package python + +import ( + "context" + "errors" + "io" + "strings" + + monty "github.com/ewhauser/gomonty" +) + +// RunWorker serves the private parent protocol. cmd/main must call it when its +// first argument is python-worker, before initializing the normal service. +func RunWorker(in io.Reader, out io.Writer) error { + for { + var request protocolRequest + if err := readFrame(in, &request); err != nil { + if errors.Is(err, io.EOF) { + return nil + } + return err + } + response := protocolResponse{Version: protocolVersion, Type: request.Type, ID: request.ID} + if request.Version != protocolVersion || request.ID == "" { + return errors.New("invalid protocol request") + } + switch request.Type { + case requestHealth: + if request.Script != "" || request.InputJSON != "" { + return errors.New("invalid protocol request") + } + if err := healthCheck(); err != nil { + response.Code, response.Message = errorFields(err) + } + case requestRun: + if len(request.Script) == 0 || len(request.Script) > MaxSourceBytes || len(request.InputJSON) > MaxInputBytes { + response.Code, response.Message = errorFields(invalid("invalid python request")) + break + } + inputJSON, err := validateRunInput(request.InputJSON) + if err != nil { + response.Code, response.Message = errorFields(err) + break + } + output, err := runMonty(request.Script, inputJSON) + if err != nil { + response.Code, response.Message = errorFields(err) + } else { + response.ResultJSON, response.Stdout = output.ResultJSON, output.Stdout + } + default: + return errors.New("invalid protocol request") + } + if err := writeFrame(out, response); err != nil { + return err + } + } +} + +func healthCheck() error { + _, err := monty.NewRepl(monty.ReplOptions{ScriptName: "workbench.py", Limits: montyLimits()}) + if err != nil { + return runtimeFailure() + } + return nil +} + +func runMonty(script, inputJSON string) (*RunOutput, error) { + repl, err := monty.NewRepl(monty.ReplOptions{ScriptName: "workbench.py", Limits: montyLimits()}) + if err != nil { + return nil, runtimeFailure() + } + ctx, cancel := context.WithTimeout(context.Background(), executionTimeout) + defer cancel() + if _, err := repl.FeedRun(ctx, "import json as __workbench_json\ninput = __workbench_json.loads("+pythonString(inputJSON)+")\noutput = {}", monty.FeedOptions{}); err != nil { + return nil, montyError(err) + } + var stdout strings.Builder + stdoutLimit := false + print := func(stream, text string) { + if stream != "stdout" { + return + } + if stdout.Len()+len(text) > MaxStdoutBytes { + stdoutLimit = true + return + } + stdout.WriteString(text) + } + if _, err := repl.FeedRun(ctx, script, monty.FeedOptions{Print: print}); err != nil { + return nil, montyError(err) + } + if stdoutLimit { + return nil, &Error{Code: ResourceExhausted, Msg: "python stdout exceeds the stdout limit"} + } + value, err := repl.FeedRun(ctx, "__workbench_json.dumps(output)", monty.FeedOptions{}) + if err != nil { + return nil, montyError(err) + } + result, ok := value.Raw().(string) + if !ok { + return nil, invalid("output must be a JSON object") + } + if len(result) > MaxResultBytes { + return nil, &Error{Code: ResourceExhausted, Msg: "python result exceeds the result limit"} + } + if !isJSONObject(result) { + return nil, invalid("output must be a JSON object") + } + return &RunOutput{ResultJSON: result, Stdout: stdout.String()}, nil +} + +func montyLimits() *monty.ResourceLimits { + return &monty.ResourceLimits{MaxDuration: executionTimeout, MaxMemory: maxMemoryBytes, MaxRecursionDepth: maxRecursionDepth} +} + +func montyError(err error) error { + if errors.Is(err, context.DeadlineExceeded) { + return executionTimeoutError() + } + var syntax *monty.SyntaxError + if errors.As(err, &syntax) { + return invalid("python code is invalid") + } + var runtime *monty.RuntimeError + if errors.As(err, &runtime) { + kind, detail, ok := strings.Cut(strings.TrimSpace(runtime.Error()), ":") + if !ok { + kind, detail = "Exception", runtime.Error() + } + if kind == "ResourceError" { + if strings.Contains(strings.ToLower(detail), "time") || strings.Contains(strings.ToLower(detail), "duration") { + return executionTimeoutError() + } + return &Error{Code: ResourceExhausted, Msg: "python resource limit exceeded"} + } + return runtimeError(kind, detail) + } + return runtimeFailure() +} + +func errorFields(err error) (Code, string) { + var typed *Error + if errors.As(err, &typed) { + return typed.Code, typed.Msg + } + return Internal, "python runtime failed" +} diff --git a/go/go.work b/go/go.work index 1e16b66388..078d88fe76 100644 --- a/go/go.work +++ b/go/go.work @@ -1,4 +1,4 @@ -go 1.26.5 +go 1.26.6 use ( ./ai-proxy diff --git a/lib/cloud_query/toolquery.pb.ex b/lib/cloud_query/toolquery.pb.ex index 7cf14c6a84..1c96cf6319 100644 --- a/lib/cloud_query/toolquery.pb.ex +++ b/lib/cloud_query/toolquery.pb.ex @@ -639,6 +639,30 @@ defmodule Toolquery.RunLuaOutput do field :result_json, 1, type: :string, json_name: "resultJson" end +defmodule Toolquery.RunPythonInput do + @moduledoc false + + use Protobuf, + full_name: "toolquery.RunPythonInput", + protoc_gen_elixir_version: "0.16.0", + syntax: :proto3 + + field :script, 1, type: :string + field :input_json, 2, type: :string, json_name: "inputJson" +end + +defmodule Toolquery.RunPythonOutput do + @moduledoc false + + use Protobuf, + full_name: "toolquery.RunPythonOutput", + protoc_gen_elixir_version: "0.16.0", + syntax: :proto3 + + field :result_json, 1, type: :string, json_name: "resultJson" + field :stdout, 2, type: :string +end + defmodule Toolquery.ToolQuery.Service do @moduledoc false @@ -657,6 +681,8 @@ defmodule Toolquery.ToolQuery.Service do rpc :InvokeLambda, Toolquery.InvokeLambdaInput, Toolquery.InvokeLambdaOutput rpc :RunLua, Toolquery.RunLuaInput, Toolquery.RunLuaOutput + + rpc :RunPython, Toolquery.RunPythonInput, Toolquery.RunPythonOutput end defmodule Toolquery.ToolQuery.Stub do diff --git a/lib/console/ai/tools/workbench/python.ex b/lib/console/ai/tools/workbench/python.ex new file mode 100644 index 0000000000..6f1f3311b9 --- /dev/null +++ b/lib/console/ai/tools/workbench/python.ex @@ -0,0 +1,38 @@ +defmodule Console.AI.Tools.Workbench.Python do + use Console.AI.Tools.Workbench.Base + alias CloudQuery.Client + alias Toolquery.ToolQuery.Stub + alias Toolquery.{RunPythonInput, RunPythonOutput} + + embedded_schema do + field :explanation, :string + field :code, :string + field :input, :map + end + + @json_schema Console.priv_file!("tools/workbench/python.json") |> Jason.decode!() + + def name(), do: "workbench_python" + + def description(), + do: + "Execute a sandboxed Monty Python script for precise computation. It supports a limited Python subset only and cannot access the filesystem, environment, network, subprocesses, pip, packages, or host tools." + + def json_schema(), do: @json_schema + + def changeset(model, attrs) do + model + |> cast(attrs, [:explanation, :code, :input]) + |> validate_required([:explanation, :code]) + end + + def implement(%__MODULE__{code: code, input: input}) do + with {:ok, client} <- Client.connect(), + {:ok, input_json} <- Jason.encode(input || %{}), + request = %RunPythonInput{script: code, input_json: input_json}, + {:ok, %RunPythonOutput{result_json: result_json, stdout: stdout}} <- Stub.run_python(client, request, Client.cloud_query_rpc_opts()), + {:ok, result} <- Jason.decode(result_json) do + {:ok, %{result: result, stdout: stdout}} + end + end +end diff --git a/lib/console/ai/workbench/engine.ex b/lib/console/ai/workbench/engine.ex index 8410e9e733..ec5add692e 100644 --- a/lib/console/ai/workbench/engine.ex +++ b/lib/console/ai/workbench/engine.ex @@ -29,6 +29,7 @@ defmodule Console.AI.Workbench.Engine do } alias Console.AI.Tools.Workbench.{ Lua, + Python, Complete, Subagents, Subagent, @@ -330,6 +331,7 @@ defmodule Console.AI.Workbench.Engine do %Subagent{subagents: subagents}, %FetchNotes{job: job}, Lua, + Python, Notes, Complete, ] ++ type_tools(job) diff --git a/lib/console/ai/workbench/subagents/infrastructure.ex b/lib/console/ai/workbench/subagents/infrastructure.ex index 39c00b0056..7bd5692deb 100644 --- a/lib/console/ai/workbench/subagents/infrastructure.ex +++ b/lib/console/ai/workbench/subagents/infrastructure.ex @@ -9,6 +9,7 @@ defmodule Console.AI.Workbench.Subagents.Infrastructure do Scratchpad, History, Lua, + Python, Infrastructure.KubeGet, Infrastructure.KubeList, Infrastructure.Cluster, @@ -75,6 +76,7 @@ defmodule Console.AI.Workbench.Subagents.Infrastructure do %Skill{skills: skills}, Scratchpad, Lua, + Python, %History{job: job, activities: activities}, Result ]) diff --git a/lib/console/ai/workbench/subagents/observability.ex b/lib/console/ai/workbench/subagents/observability.ex index d2805ad297..1bd1663e7a 100644 --- a/lib/console/ai/workbench/subagents/observability.ex +++ b/lib/console/ai/workbench/subagents/observability.ex @@ -1,7 +1,7 @@ defmodule Console.AI.Workbench.Subagents.Observability do use Console.AI.Workbench.Subagents.Base alias Console.Schema.{Workbench, WorkbenchJob, WorkbenchJobActivity, WorkbenchTool, User} - alias Console.AI.Tools.Workbench.{ObservabilityResult, Skills, Skill, Lua, History, Infrastructure.PodLogs, Scratchpad} + alias Console.AI.Tools.Workbench.{ObservabilityResult, Skills, Skill, Lua, Python, History, Infrastructure.PodLogs, Scratchpad} alias Console.AI.Tools.Workbench.Observability.{Metrics, MetricsSearch, MetricsLabelSearch, Logs, Traces, Plrl} alias Console.AI.Tools.Workbench.Integration.Sentry.Tools, as: SentryTools alias Console.AI.Workbench.{Environment, MCP} @@ -52,6 +52,7 @@ defmodule Console.AI.Workbench.Subagents.Observability do Scratchpad, ObservabilityResult, Lua, + Python, %History{job: job, activities: activities} ]) end diff --git a/priv/prompts/workbench/infrastructure.md.eex b/priv/prompts/workbench/infrastructure.md.eex index 91fe77ed0d..c744fc2a60 100644 --- a/priv/prompts/workbench/infrastructure.md.eex +++ b/priv/prompts/workbench/infrastructure.md.eex @@ -76,7 +76,7 @@ You are producing output for a human user, and should expect them to want to rea 2. In general summarize_component has the ability to dive deep into a k8s object, but you'll need the Plural-specific context from a service_search or stack_search to gather them. 3. We can offer the ability to directly query kubernetes, but you need to be precise with your inputs and *always* include a Plural cluster handle to make it work. It's also possible user RBAC policies block your query. 4. If you're searching for a kubernetes resource and don't know exactly the namespace/name it might have, you can mix `service_search` with standard k8s tool calls since it has the ability to probe kubernetes objects just with a finely crafted prompt. Both are useful tools in the toolkit. -5. You have a `workbench_lua` tool available to you to do deterministic algorithmic calculations. You should be leveraging that instead of trying to infer based on prior tool calls and summaries, using exact raw data in the lua code used. +5. You have `workbench_lua` and `workbench_python` tools available to do deterministic algorithmic calculations. You should be leveraging either instead of trying to infer based on prior tool calls and summaries, using exact raw data in the supported script used. 6. You are also given a workbench_history tool which can be used to search past work outside of this subagent. Use this to grab additional context that might not be present in your prompt, but don't rely on it if the prompt is sufficient. Also since many tools require a time anchor, the current time is <%= Timex.now() |> Timex.format!("{ISO:Extended}") %>. diff --git a/priv/prompts/workbench/job.md.eex b/priv/prompts/workbench/job.md.eex index c7ee4cfd33..38bd6052b8 100644 --- a/priv/prompts/workbench/job.md.eex +++ b/priv/prompts/workbench/job.md.eex @@ -155,7 +155,7 @@ do not run them simultaneously, since they won't be able to communicate with eac In general, you can call any of the subagents and tools as many times as you need to accomplish your task. -For any tool you delegate to, you should have them give as close to exact raw data as possible, rather than computing on it. When you need an exact, deterministic computation, always use the `workbench_lua` tool with the appropriate Lua code to get the result for you. +For any tool you delegate to, you should have them give as close to exact raw data as possible, rather than computing on it. When you need an exact, deterministic computation, always use either the `workbench_lua` or `workbench_python` tool with an appropriate supported script to get the result for you. For instance, @@ -169,6 +169,14 @@ output.total = total (you must always set output to see the returned values). +For Python, `input` is the optional JSON object supplied to the tool and `output` must remain a JSON-serializable dictionary: + +```python +output = {"total": sum(input["values"])} +``` + +Monty is a limited Python subset, not CPython; use Lua when its supported libraries are a better fit, and do not expect files, environment variables, network access, subprocesses, pip, third-party packages, or host tools from Python. + That said, you *should* avoid repeatedly calling the notes tool unless you're meaningfully updating your working theory, plan, or conclusion. This will spam users and is unnneeded. In general, you should be focused on always making progress, don't repeatedly call tools that get no where important. diff --git a/priv/prompts/workbench/observability.md.eex b/priv/prompts/workbench/observability.md.eex index 84b8b115b4..17bfad6387 100644 --- a/priv/prompts/workbench/observability.md.eex +++ b/priv/prompts/workbench/observability.md.eex @@ -21,7 +21,7 @@ data or use queries relative to the current time. 1. When querying observability tools, you should give at least some time range to the query. A 5 minute window minimum is good practice. 2. Feel free to query the same data source multiple times to probe. If you're investigating an alert, assume it's a real thing and search. 3. When searching logs, you can always search for contextual logs by filtering on specific facets (eg pod name that isolates to a specific running container), and using time ranges around the log in question. -4. You have a `workbench_tool` tool available to you to do deterministic algorithmic calculations. You should be leveraging that instead of trying to infer based on prior tool calls and summaries, using exact raw data in the lua code used. +4. You have `workbench_lua` and `workbench_python` tools available to do deterministic algorithmic calculations. You should be leveraging either instead of trying to infer based on prior tool calls and summaries, using exact raw data in the supported script used. 5. You are also given a workbench_history tool which can be used to search past work outside of this subagent. Use this to grab additional context that might not be present in your prompt, but don't rely on it if the prompt is sufficient. Also since many tools require a time anchor, the current time is <%= Timex.now() |> Timex.format!("{ISO:Extended}") %>. diff --git a/priv/tools/workbench/python.json b/priv/tools/workbench/python.json new file mode 100644 index 0000000000..187136a61a --- /dev/null +++ b/priv/tools/workbench/python.json @@ -0,0 +1,22 @@ +{ + "type": "object", + "description": "Execute a sandboxed Monty Python script for precise computation. Monty supports a limited Python subset for expressions, variables, collections, control flow, functions, and selected built-ins. The JSON object `input` is available to the script as a global dictionary; write structured results to the global `output` dictionary and use print() for diagnostic output. Filesystem, environment, network, subprocesses, pip, packages, and host tools are unavailable.", + "properties": { + "code": { + "type": "string", + "description": "Monty-compatible Python source code. Read the optional global input dictionary, assign JSON-serializable fields on the global output dictionary, and use print() for diagnostics." + }, + "explanation": { + "type": "string", + "description": "Explanation of why you're running this script, such as 'I need to calculate the exact p95 latency from these samples.'." + }, + "input": { + "type": "object", + "description": "Optional JSON object exposed to the script as the global input dictionary." + } + }, + "required": [ + "code", + "explanation" + ] +} diff --git a/test/console/ai/tools/workbench/python_test.exs b/test/console/ai/tools/workbench/python_test.exs new file mode 100644 index 0000000000..be800a26f4 --- /dev/null +++ b/test/console/ai/tools/workbench/python_test.exs @@ -0,0 +1,53 @@ +defmodule Console.AI.Tools.Workbench.PythonTest do + use Console.DataCase, async: false + use Mimic + + alias CloudQuery.Client + alias Console.AI.Tools.Workbench.Python + alias Toolquery.{RunPythonOutput} + alias Toolquery.ToolQuery.Stub + + setup :set_mimic_global + + test "exposes the Python tool contract" do + assert Python.name() == "workbench_python" + assert %{"input" => %{"type" => "object"}} = Python.json_schema()["properties"] + assert Python.json_schema()["required"] == ["code", "explanation"] + end + + test "requires an explanation and Python code" do + changeset = Python.changeset(%Python{}, %{}) + + refute changeset.valid? + assert %{code: ["can't be blank"], explanation: ["can't be blank"]} = errors_on(changeset) + end + + test "runs Monty Python with JSON input and returns decoded output and stdout" do + expect(Client, :connect, fn -> {:ok, :channel} end) + + expect(Stub, :run_python, fn :channel, request, opts -> + assert request.script == "output['total'] = input['first'] + input['second']" + assert request.input_json == ~s({"first":20,"second":22}) + assert opts == Client.cloud_query_rpc_opts() + + {:ok, %RunPythonOutput{result_json: ~s({"total":42}), stdout: "calculated total\\n"}} + end) + + assert {:ok, %{result: %{"total" => 42}, stdout: "calculated total\n"}} = + Python.implement(%Python{ + code: "output['total'] = input['first'] + input['second']", + input: %{"first" => 20, "second" => 22} + }) + end + + test "defaults omitted input to an empty JSON object" do + expect(Client, :connect, fn -> {:ok, :channel} end) + + expect(Stub, :run_python, fn :channel, request, _opts -> + assert request.input_json == "{}" + {:ok, %RunPythonOutput{result_json: "{}", stdout: ""}} + end) + + assert {:ok, %{result: %{}, stdout: ""}} = Python.implement(%Python{code: "output = {}"}) + end +end From acc0d7dbca66a5131a93ba77e5b3c53c6af95cf4 Mon Sep 17 00:00:00 2001 From: Sebastian Florek Date: Thu, 20 Aug 2026 16:14:09 +0200 Subject: [PATCH 02/15] refactor(python): restructure Python tool implementation - Removed obsolete Python protocol and worker logic - Introduced new worker and contract packages for improved modularity - Added `internal/worker/os.go` for handling OS-level operations in a sandboxed environment - Enhanced `ToolQueryService` to utilize the refactored Python runner interface - Simplified error handling and error message generation - Improved code organization by separating logic into distinct packages --- go/cloud-query/cmd/main.go | 2 +- go/cloud-query/internal/service/toolquery.go | 39 +- .../internal/service/toolquery_python_test.go | 73 +++- .../internal/tools/python/README.md | 50 ++- .../internal/tools/python/common.go | 25 -- .../internal/tools/python/config.go | 78 ++++ .../internal/tools/python/config_test.go | 95 ++++ .../internal/tools/python/errors.go | 102 +---- .../internal/tools/python/grpc/status.go | 32 ++ .../internal/tools/python/grpc/status_test.go | 43 ++ .../tools/python/internal/contract/errors.go | 141 ++++++ .../python/internal/contract/errors_test.go | 18 + .../tools/python/internal/contract/types.go | 48 +++ .../python/internal/contract/validation.go | 44 ++ .../tools/python/internal/pool/config.go | 44 ++ .../tools/python/internal/pool/pool.go | 408 ++++++++++++++++++ .../tools/python/internal/pool/pool_test.go | 284 ++++++++++++ .../tools/python/internal/pool/process.go | 102 +++++ .../tools/python/internal/protocol/codec.go | 244 +++++++++++ .../python/internal/protocol/codec_test.go | 150 +++++++ .../tools/python/internal/worker/os.go | 86 ++++ .../tools/python/internal/worker/runtime.go | 165 +++++++ .../tools/python/internal/worker/server.go | 76 ++++ .../python/internal/worker/server_test.go | 134 ++++++ .../internal/tools/python/protocol.go | 138 ------ .../internal/tools/python/python.go | 322 -------------- .../internal/tools/python/python_test.go | 293 +------------ .../internal/tools/python/runner.go | 34 ++ go/cloud-query/internal/tools/python/types.go | 29 ++ .../internal/tools/python/worker.go | 146 +------ 30 files changed, 2393 insertions(+), 1052 deletions(-) delete mode 100644 go/cloud-query/internal/tools/python/common.go create mode 100644 go/cloud-query/internal/tools/python/config.go create mode 100644 go/cloud-query/internal/tools/python/config_test.go create mode 100644 go/cloud-query/internal/tools/python/grpc/status.go create mode 100644 go/cloud-query/internal/tools/python/grpc/status_test.go create mode 100644 go/cloud-query/internal/tools/python/internal/contract/errors.go create mode 100644 go/cloud-query/internal/tools/python/internal/contract/errors_test.go create mode 100644 go/cloud-query/internal/tools/python/internal/contract/types.go create mode 100644 go/cloud-query/internal/tools/python/internal/contract/validation.go create mode 100644 go/cloud-query/internal/tools/python/internal/pool/config.go create mode 100644 go/cloud-query/internal/tools/python/internal/pool/pool.go create mode 100644 go/cloud-query/internal/tools/python/internal/pool/pool_test.go create mode 100644 go/cloud-query/internal/tools/python/internal/pool/process.go create mode 100644 go/cloud-query/internal/tools/python/internal/protocol/codec.go create mode 100644 go/cloud-query/internal/tools/python/internal/protocol/codec_test.go create mode 100644 go/cloud-query/internal/tools/python/internal/worker/os.go create mode 100644 go/cloud-query/internal/tools/python/internal/worker/runtime.go create mode 100644 go/cloud-query/internal/tools/python/internal/worker/server.go create mode 100644 go/cloud-query/internal/tools/python/internal/worker/server_test.go delete mode 100644 go/cloud-query/internal/tools/python/protocol.go delete mode 100644 go/cloud-query/internal/tools/python/python.go create mode 100644 go/cloud-query/internal/tools/python/runner.go create mode 100644 go/cloud-query/internal/tools/python/types.go diff --git a/go/cloud-query/cmd/main.go b/go/cloud-query/cmd/main.go index 4c216f2c5e..c0c27efb24 100644 --- a/go/cloud-query/cmd/main.go +++ b/go/cloud-query/cmd/main.go @@ -29,7 +29,7 @@ func startHealthzHandler() { func main() { if len(os.Args) == 2 && os.Args[1] == "python-worker" { - if err := pythontools.RunWorker(os.Stdin, os.Stdout); err != nil { + if err := pythontools.NewWorker().Run(os.Stdin, os.Stdout); err != nil { os.Exit(1) } return diff --git a/go/cloud-query/internal/service/toolquery.go b/go/cloud-query/internal/service/toolquery.go index 3643ef7f93..5a121e7698 100644 --- a/go/cloud-query/internal/service/toolquery.go +++ b/go/cloud-query/internal/service/toolquery.go @@ -17,12 +17,13 @@ import ( lambdatools "github.com/pluralsh/console/go/cloud-query/internal/tools/lambda" luatools "github.com/pluralsh/console/go/cloud-query/internal/tools/lua" pythontools "github.com/pluralsh/console/go/cloud-query/internal/tools/python" + pythongrpc "github.com/pluralsh/console/go/cloud-query/internal/tools/python/grpc" ) // ToolQueryService implements the toolquery.ToolQueryServer interface. type ToolQueryService struct { toolquery.UnimplementedToolQueryServer - python *pythontools.Runner + python pythontools.Runner } // Install registers the ToolQuery service with the gRPC server. @@ -34,19 +35,17 @@ func (in *ToolQueryService) Install(server *grpc.Server) { // NewToolQueryService creates a ToolQuery server and verifies its sandboxed // Python worker pool before the gRPC server starts accepting requests. func NewToolQueryService(ctx context.Context) (Service, error) { - runner, err := pythontools.New(ctx) + runner, err := pythontools.NewRunner(ctx, pythontools.RunnerConfig{}) if err != nil { return nil, err } + return &ToolQueryService{python: runner}, nil } // Close reaps all Python worker subprocesses. func (in *ToolQueryService) Close(ctx context.Context) error { - if in.python != nil { - return in.python.CloseContext(ctx) - } - return nil + return in.python.CloseContext(ctx) } func (in *ToolQueryService) Metrics(ctx context.Context, input *toolquery.MetricsQueryInput) (*toolquery.MetricsQueryOutput, error) { @@ -223,19 +222,16 @@ func (in *ToolQueryService) RunPython(ctx context.Context, input *toolquery.RunP if input == nil { return nil, status.Error(codes.InvalidArgument, "input is required") } - if in.python == nil { - return nil, status.Error(codes.Unavailable, "python runtime is unavailable") - } - output, err := in.python.Run(ctx, pythontools.RunInput{ Script: input.GetScript(), InputJSON: input.GetInputJson(), }) if err != nil { - if pythontools.ErrorCode(err) == pythontools.Internal { + if pythontools.CodeOf(err) == pythontools.Internal { klog.ErrorS(err, "python worker failed") } - return nil, status.Error(pythonGRPCCode(err), err.Error()) + + return nil, pythongrpc.StatusError(err) } return &toolquery.RunPythonOutput{ @@ -244,25 +240,6 @@ func (in *ToolQueryService) RunPython(ctx context.Context, input *toolquery.RunP }, nil } -func pythonGRPCCode(err error) codes.Code { - switch pythontools.ErrorCode(err) { - case pythontools.InvalidArgument: - return codes.InvalidArgument - case pythontools.FailedPrecondition: - return codes.FailedPrecondition - case pythontools.Canceled: - return codes.Canceled - case pythontools.DeadlineExceeded: - return codes.DeadlineExceeded - case pythontools.ResourceExhausted: - return codes.ResourceExhausted - case pythontools.Unavailable: - return codes.Unavailable - default: - return codes.Internal - } -} - func (in *ToolQueryService) validateInput(connection *toolquery.ToolConnection, query string, timeRange *toolquery.TimeRange) error { if err := in.validateSearchInput(connection); err != nil { return err diff --git a/go/cloud-query/internal/service/toolquery_python_test.go b/go/cloud-query/internal/service/toolquery_python_test.go index 3b65353292..83f5774c65 100644 --- a/go/cloud-query/internal/service/toolquery_python_test.go +++ b/go/cloud-query/internal/service/toolquery_python_test.go @@ -2,6 +2,9 @@ package service import ( "context" + "errors" + "os" + "strings" "testing" "google.golang.org/grpc/codes" @@ -11,30 +14,64 @@ import ( pythontools "github.com/pluralsh/console/go/cloud-query/internal/tools/python" ) -func TestRunPythonValidatesInputAndAvailability(t *testing.T) { - server := &ToolQueryService{} +func TestMain(m *testing.M) { + if len(os.Args) == 2 && os.Args[1] == "python-worker" { + if err := pythontools.NewWorker().Run(os.Stdin, os.Stdout); err != nil { + os.Exit(1) + } + + os.Exit(0) + } + + os.Exit(m.Run()) +} + +type failingPythonRunner struct{ err error } + +func (r failingPythonRunner) Run(context.Context, pythontools.RunInput) (*pythontools.RunOutput, error) { + return nil, r.err +} + +func (failingPythonRunner) CloseContext(context.Context) error { return nil } + +func TestRunPythonValidatesInput(t *testing.T) { + server := &ToolQueryService{python: failingPythonRunner{err: errors.New("runner must not be called")}} if _, err := server.RunPython(context.Background(), nil); status.Code(err) != codes.InvalidArgument { t.Fatalf("nil input code = %s, want %s", status.Code(err), codes.InvalidArgument) } - if _, err := server.RunPython(context.Background(), &toolquery.RunPythonInput{Script: "output = {}"}); status.Code(err) != codes.Unavailable { - t.Fatalf("missing runner code = %s, want %s", status.Code(err), codes.Unavailable) +} + +func TestRunPythonHidesInternalCause(t *testing.T) { + server := &ToolQueryService{python: failingPythonRunner{err: errors.New("private process detail")}} + + _, err := server.RunPython(context.Background(), &toolquery.RunPythonInput{Script: "output = {}"}) + if got := status.Convert(err).Message(); got != "python runtime failed" { + t.Fatalf("RunPython() message = %q, want safe summary", got) } } -func TestPythonGRPCCode(t *testing.T) { - for runCode, want := range map[pythontools.Code]codes.Code{ - pythontools.InvalidArgument: codes.InvalidArgument, - pythontools.FailedPrecondition: codes.FailedPrecondition, - pythontools.Canceled: codes.Canceled, - pythontools.DeadlineExceeded: codes.DeadlineExceeded, - pythontools.ResourceExhausted: codes.ResourceExhausted, - pythontools.Unavailable: codes.Unavailable, - pythontools.Internal: codes.Internal, - } { - err := &pythontools.Error{Code: runCode, Msg: "safe"} - if got := pythonGRPCCode(err); got != want { - t.Errorf("pythonGRPCCode(%s) = %s, want %s", runCode, got, want) - } +func TestRunPythonPreservesSandboxErrorMessage(t *testing.T) { + runner, err := pythontools.NewRunner(context.Background(), pythontools.RunnerConfig{ + Workers: 1, + QueueSize: 1, + MaxSuccessfulRunsBeforeRecycle: 1, + }) + if err != nil { + t.Fatal(err) + } + defer runner.CloseContext(context.Background()) + + server := &ToolQueryService{python: runner} + _, err = server.RunPython(context.Background(), &toolquery.RunPythonInput{ + Script: "import socket\nsocket.create_connection(('example.com', 443))\noutput = {'reachable': True}", + }) + + got := status.Convert(err) + if got.Code() != codes.FailedPrecondition { + t.Fatalf("status code = %s", got.Code()) + } + if !strings.Contains(strings.ToLower(got.Message()), "socket") { + t.Fatalf("status message = %q", got.Message()) } } diff --git a/go/cloud-query/internal/tools/python/README.md b/go/cloud-query/internal/tools/python/README.md index 2374eba57d..8be87ac69b 100644 --- a/go/cloud-query/internal/tools/python/README.md +++ b/go/cloud-query/internal/tools/python/README.md @@ -1,9 +1,29 @@ # Sandboxed Python runner -This package implements `RunPython` with Monty's limited Python subset. It is -not CPython. The parent starts four copies of `cloud-query python-worker` and -communicates with them over private stdin/stdout pipes. Each worker loads -gomonty's embedded native library and creates a fresh REPL for every request. +This package provides the runner used by the `RunPython` service RPC with +Monty's limited Python subset. It is not CPython. The parent starts four copies +of `cloud-query python-worker` and communicates with them over private +stdin/stdout pipes. Each worker loads gomonty's embedded native library and +creates a fresh REPL for every request. + +## Architecture + +The root package is a facade for runner and worker construction, stable error +codes, and safe public error messages. Its implementation is split into four +private modules: + +- `contract` owns execution types, fixed limits, validation, and errors. +- `protocol` owns strict version 1 framing and message validation. +- `worker` joins the request server to the Monty runtime. +- `pool` owns isolated subprocesses, queueing, recycling, and shutdown. + +Dependencies flow inward: `pool` and `worker` use `protocol`; those modules use +`contract`; the root package joins them. Each worker receives a replacement +environment containing only `TMPDIR=/tmp` by default. + +Errors have a stable code, a bounded public summary, and a private diagnostic +capped at 64 KiB. Use `PublicMessage` before returning an error to an untrusted +caller. Protocol data and process diagnostics must remain internal. ## Execution contract @@ -12,20 +32,23 @@ the global `input` dictionary. `output` starts as an empty dictionary and must remain JSON serializable and object shaped. The response returns `output` and the user script's standard output separately. -Workers receive only `TMPDIR=/tmp`. The host supplies no filesystem, -environment, network, subprocess, shell, package, OS callback, external -function, or name-lookup capability. +Workers receive only `TMPDIR=/tmp`. The host exposes UTC clock callbacks for +`datetime.now()` and `date.today()` only; `datetime.now()` returns a naive UTC +value, and non-UTC timezone arguments are rejected. It supplies no filesystem, +environment, network, subprocess, shell, package, external function, or +name-lookup capability. Limits: - Source: 64 KiB - Input and result JSON: 1 MiB each - Captured stdout: 64 KiB +- Private diagnostics: 64 KiB - Monty execution: 10 seconds - Parent wall clock: 15 seconds or the caller's earlier deadline - Monty-managed memory: 64 MiB - Recursion: 200 frames -- Active workers: four +- Active workers: 4 - Waiting requests: 16 by default - Worker recycling: 10 successful requests @@ -36,9 +59,12 @@ the hard crash and wall-time limit by killing and replacing a worker process. Protocol version 1 uses a four-byte big-endian length followed by strict JSON. Frames are bounded before allocation, unknown fields and trailing JSON values -are rejected, and only `health` and `run` request types are accepted. Every -response must repeat the request ID and type. Malformed protocol data terminates -the worker. +are rejected, and only `health` and `run` request kinds are accepted. Every +response must repeat the request ID and kind. Error responses carry a stable +code, a bounded sanitized public summary, and a separate bounded private +diagnostic. The parent validates and sanitizes the public summary again before +returning it, while the private diagnostic remains internal. Malformed protocol +data terminates the worker. The parent admits requests to a bounded FIFO queue. Caller cancellation applies while waiting and during execution. A worker is replaced after any execution @@ -59,5 +85,5 @@ Run focused validation with: ```sh cd go/cloud-query -go test ./internal/tools/python ./internal/service +go test -race ./internal/tools/python/... ./internal/service ``` diff --git a/go/cloud-query/internal/tools/python/common.go b/go/cloud-query/internal/tools/python/common.go deleted file mode 100644 index a45be7fc44..0000000000 --- a/go/cloud-query/internal/tools/python/common.go +++ /dev/null @@ -1,25 +0,0 @@ -package python - -import ( - "encoding/json" - "strconv" - "strings" -) - -func validateRunInput(raw string) (string, error) { - if strings.TrimSpace(raw) == "" { - return "{}", nil - } - var value map[string]json.RawMessage - if err := json.Unmarshal([]byte(raw), &value); err != nil || value == nil { - return "", invalid("input_json must be a JSON object") - } - return raw, nil -} - -func pythonString(value string) string { return strconv.Quote(value) } - -func isJSONObject(value string) bool { - var object map[string]json.RawMessage - return json.Unmarshal([]byte(value), &object) == nil && object != nil -} diff --git a/go/cloud-query/internal/tools/python/config.go b/go/cloud-query/internal/tools/python/config.go new file mode 100644 index 0000000000..0a24a7930b --- /dev/null +++ b/go/cloud-query/internal/tools/python/config.go @@ -0,0 +1,78 @@ +package python + +import ( + "os" + + "github.com/samber/lo" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" +) + +const ( + defaultWorkers = 4 + defaultQueueSize = 16 + defaultRecycle = 10 +) + +var currentExecutable = os.Executable + +// WorkerProcessConfig configures the command used for each isolated worker. +// Environment replaces the inherited environment instead of extending it. +type WorkerProcessConfig struct { + Executable string + Arguments []string + Environment []string +} + +// RunnerConfig configures worker capacity, queueing, recycling, and process +// startup. Zero-valued capacity settings use the package defaults. +type RunnerConfig struct { + Workers int + QueueSize int + MaxSuccessfulRunsBeforeRecycle int + WorkerProcess WorkerProcessConfig +} + +func resolveRunnerConfig(overlay RunnerConfig) (RunnerConfig, error) { + config := RunnerConfig{ + Workers: defaultWorkers, + QueueSize: defaultQueueSize, + MaxSuccessfulRunsBeforeRecycle: defaultRecycle, + WorkerProcess: WorkerProcessConfig{ + Arguments: []string{"python-worker"}, + Environment: []string{"TMPDIR=/tmp"}, + }, + } + if overlay.Workers != 0 { + config.Workers = overlay.Workers + } + + if overlay.QueueSize != 0 { + config.QueueSize = overlay.QueueSize + } + + if overlay.MaxSuccessfulRunsBeforeRecycle != 0 { + config.MaxSuccessfulRunsBeforeRecycle = overlay.MaxSuccessfulRunsBeforeRecycle + } + + if overlay.WorkerProcess.Executable != "" { + config.WorkerProcess.Executable = overlay.WorkerProcess.Executable + } else { + executable, err := currentExecutable() + if err != nil { + return RunnerConfig{}, contract.UnavailableError("locating the python worker executable", err) + } + + config.WorkerProcess.Executable = executable + } + + if overlay.WorkerProcess.Arguments != nil { + config.WorkerProcess.Arguments = lo.Clone(overlay.WorkerProcess.Arguments) + } + + if overlay.WorkerProcess.Environment != nil { + config.WorkerProcess.Environment = lo.Clone(overlay.WorkerProcess.Environment) + } + + return config, nil +} diff --git a/go/cloud-query/internal/tools/python/config_test.go b/go/cloud-query/internal/tools/python/config_test.go new file mode 100644 index 0000000000..8d9a4cc5a7 --- /dev/null +++ b/go/cloud-query/internal/tools/python/config_test.go @@ -0,0 +1,95 @@ +package python + +import ( + "errors" + "testing" +) + +func TestResolveRunnerConfigDefaultsAndOverlays(t *testing.T) { + previous := currentExecutable + currentExecutable = func() (string, error) { return "/test/python", nil } + t.Cleanup(func() { currentExecutable = previous }) + + defaults, err := resolveRunnerConfig(RunnerConfig{}) + if err != nil { + t.Fatal(err) + } + + if defaults.Workers != 4 || defaults.QueueSize != 16 || defaults.MaxSuccessfulRunsBeforeRecycle != 10 { + t.Fatalf("defaults = %#v", defaults) + } + if defaults.WorkerProcess.Executable != "/test/python" { + t.Fatalf("executable = %q", defaults.WorkerProcess.Executable) + } + if got := defaults.WorkerProcess.Arguments; len(got) != 1 || got[0] != "python-worker" { + t.Fatalf("arguments = %#v", got) + } + if got := defaults.WorkerProcess.Environment; len(got) != 1 || got[0] != "TMPDIR=/tmp" { + t.Fatalf("worker defaults = %#v", defaults.WorkerProcess) + } + + arguments := []string{"child", "worker"} + environment := []string{"TMPDIR=/sandbox"} + resolved, err := resolveRunnerConfig(RunnerConfig{ + Workers: 2, + QueueSize: 3, + MaxSuccessfulRunsBeforeRecycle: 4, + WorkerProcess: WorkerProcessConfig{ + Executable: "/custom/python", + Arguments: arguments, + Environment: environment, + }, + }) + if err != nil { + t.Fatal(err) + } + + arguments[0] = "changed" + environment[0] = "changed=value" + if resolved.Workers != 2 || resolved.QueueSize != 3 || resolved.MaxSuccessfulRunsBeforeRecycle != 4 { + t.Fatalf("numeric overlay = %#v", resolved) + } + if resolved.WorkerProcess.Executable != "/custom/python" { + t.Fatalf("executable overlay = %q", resolved.WorkerProcess.Executable) + } + if resolved.WorkerProcess.Arguments[0] != "child" || + resolved.WorkerProcess.Environment[0] != "TMPDIR=/sandbox" { + t.Fatalf("overlay = %#v", resolved) + } + + empty, err := resolveRunnerConfig(RunnerConfig{WorkerProcess: WorkerProcessConfig{Arguments: []string{}, Environment: []string{}}}) + if err != nil || empty.WorkerProcess.Arguments == nil || empty.WorkerProcess.Environment == nil { + t.Fatalf("explicit empty slices = %#v, %v", empty.WorkerProcess, err) + } + + if len(empty.WorkerProcess.Arguments) != 0 || len(empty.WorkerProcess.Environment) != 0 { + t.Fatalf("explicit empty slices = %#v, %v", empty.WorkerProcess, err) + } +} + +func TestNewRunnerRejectsInvalidValues(t *testing.T) { + previous := currentExecutable + currentExecutable = func() (string, error) { return "/test/python", nil } + t.Cleanup(func() { currentExecutable = previous }) + for _, config := range []RunnerConfig{ + {Workers: -1}, + {QueueSize: -1}, + {MaxSuccessfulRunsBeforeRecycle: -1}, + {WorkerProcess: WorkerProcessConfig{Executable: " "}}, + {WorkerProcess: WorkerProcessConfig{Environment: []string{"missing-equals"}}}, + } { + if _, err := NewRunner(t.Context(), config); CodeOf(err) != InvalidArgument { + t.Fatalf("config %#v: %v", config, err) + } + } + + currentExecutable = func() (string, error) { return "", errors.New("missing executable") } + if _, err := resolveRunnerConfig(RunnerConfig{}); CodeOf(err) != Unavailable { + t.Fatalf("executable error = %v", err) + } + + custom, err := resolveRunnerConfig(RunnerConfig{WorkerProcess: WorkerProcessConfig{Executable: "/custom/python"}}) + if err != nil || custom.WorkerProcess.Executable != "/custom/python" { + t.Fatalf("custom executable with unavailable default = %#v, %v", custom.WorkerProcess, err) + } +} diff --git a/go/cloud-query/internal/tools/python/errors.go b/go/cloud-query/internal/tools/python/errors.go index 94f6d60492..fe15f2d1cd 100644 --- a/go/cloud-query/internal/tools/python/errors.go +++ b/go/cloud-query/internal/tools/python/errors.go @@ -1,100 +1,10 @@ package python -import ( - "context" - "errors" - "strings" -) +import "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" -const ( - maxExceptionSummaryRunes = 512 - runtimeUnavailableMessage = "python runtime is unavailable" -) +// CodeOf returns err's stable runner error code, or Internal for unclassified errors. +func CodeOf(err error) Code { return contract.CodeOf(err) } -// Code is a safe, stable error category suitable for mapping to an RPC status. -type Code string - -const ( - // InvalidArgument indicates invalid source, input, or output data. - InvalidArgument Code = "invalid_argument" - // FailedPrecondition indicates Python could not complete the requested operation. - FailedPrecondition Code = "failed_precondition" - // Canceled indicates the caller canceled execution. - Canceled Code = "canceled" - // DeadlineExceeded indicates execution exceeded its allotted time. - DeadlineExceeded Code = "deadline_exceeded" - // ResourceExhausted indicates a configured size, memory, or concurrency limit was exceeded. - ResourceExhausted Code = "resource_exhausted" - // Unavailable indicates no healthy Python worker is available. - Unavailable Code = "unavailable" - // Internal indicates an unexpected runner or worker failure. - Internal Code = "internal" -) - -// Error contains only a stable category and a safe summary. Protocol frames, -// child stderr, and host paths are intentionally never retained. -type Error struct { - Code Code - Msg string -} - -// Error returns the safe error summary. -func (e *Error) Error() string { return e.Msg } - -// ErrorCode returns a safe classification for err. Errors outside this package -// are intentionally treated as Internal so callers never infer host details. -func ErrorCode(err error) Code { - var runErr *Error - if errors.As(err, &runErr) { - return runErr.Code - } - return Internal -} - -func invalid(msg string) error { return &Error{Code: InvalidArgument, Msg: msg} } - -func runtimeError(kind, message string) error { - switch kind { - case "SyntaxError", "TypeError": - return invalid(pythonExceptionSummary(kind, message)) - case "MemoryError", "RecursionError": - return &Error{Code: ResourceExhausted, Msg: "python resource limit exceeded"} - case "TimeoutError": - return executionTimeoutError() - default: - return &Error{Code: FailedPrecondition, Msg: pythonExceptionSummary(kind, message)} - } -} - -// pythonExceptionSummary exposes the useful exception type and message while -// excluding tracebacks, worker diagnostics, protocol data, and unbounded text. -func pythonExceptionSummary(kind, message string) string { - kind = strings.TrimSpace(strings.ToValidUTF8(kind, "")) - message = strings.Join(strings.Fields(strings.ToValidUTF8(message, "")), " ") - if kind == "" { - kind = "Exception" - } - summary := "python " + kind - if message != "" { - summary += ": " + message - } - runes := []rune(summary) - if len(runes) > maxExceptionSummaryRunes { - summary = string(runes[:maxExceptionSummaryRunes-1]) + "…" - } - return summary -} - -func runtimeFailure() error { return &Error{Code: Internal, Msg: "python runtime failed"} } - -func executionTimeoutError() error { - return &Error{Code: DeadlineExceeded, Msg: "python execution timed out"} -} - -func executionContextError(err error) error { - if errors.Is(err, context.Canceled) { - return &Error{Code: Canceled, Msg: "python execution was canceled"} - } - - return executionTimeoutError() -} +// PublicMessage returns the safe summary suitable for untrusted callers. It +// never includes private worker or process diagnostics. +func PublicMessage(err error) string { return contract.PublicMessage(err) } diff --git a/go/cloud-query/internal/tools/python/grpc/status.go b/go/cloud-query/internal/tools/python/grpc/status.go new file mode 100644 index 0000000000..da8d74c76d --- /dev/null +++ b/go/cloud-query/internal/tools/python/grpc/status.go @@ -0,0 +1,32 @@ +package grpc + +import ( + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python" +) + +// StatusError translates a runner error into an RPC-safe gRPC status error. +func StatusError(err error) error { + return status.Error(statusCode(python.CodeOf(err)), python.PublicMessage(err)) +} + +func statusCode(code python.Code) codes.Code { + switch code { + case python.InvalidArgument: + return codes.InvalidArgument + case python.FailedPrecondition: + return codes.FailedPrecondition + case python.Canceled: + return codes.Canceled + case python.DeadlineExceeded: + return codes.DeadlineExceeded + case python.ResourceExhausted: + return codes.ResourceExhausted + case python.Unavailable: + return codes.Unavailable + default: + return codes.Internal + } +} diff --git a/go/cloud-query/internal/tools/python/grpc/status_test.go b/go/cloud-query/internal/tools/python/grpc/status_test.go new file mode 100644 index 0000000000..e917925235 --- /dev/null +++ b/go/cloud-query/internal/tools/python/grpc/status_test.go @@ -0,0 +1,43 @@ +package grpc + +import ( + "errors" + "testing" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python" + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" +) + +func TestStatusCode(t *testing.T) { + for code, want := range map[python.Code]codes.Code{ + python.InvalidArgument: codes.InvalidArgument, + python.FailedPrecondition: codes.FailedPrecondition, + python.Canceled: codes.Canceled, + python.DeadlineExceeded: codes.DeadlineExceeded, + python.ResourceExhausted: codes.ResourceExhausted, + python.Unavailable: codes.Unavailable, + python.Internal: codes.Internal, + } { + if got := statusCode(code); got != want { + t.Errorf("statusCode(%s) = %s, want %s", code, got, want) + } + } +} + +func TestStatusErrorPreservesPublicMessage(t *testing.T) { + err := contract.Failed( + "python ImportError: socket is unavailable", + errors.New("private sandbox diagnostic"), + ) + + got := status.Convert(StatusError(err)) + if got.Code() != codes.FailedPrecondition { + t.Fatalf("status code = %s", got.Code()) + } + if got.Message() != "python ImportError: socket is unavailable" { + t.Fatalf("status message = %q", got.Message()) + } +} diff --git a/go/cloud-query/internal/tools/python/internal/contract/errors.go b/go/cloud-query/internal/tools/python/internal/contract/errors.go new file mode 100644 index 0000000000..0fef58a678 --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/contract/errors.go @@ -0,0 +1,141 @@ +package contract + +import ( + "context" + "errors" + "strings" +) + +const maxPublicRunes = 512 +const detailTruncationMarker = "\n[python diagnostic truncated]" + +type runError struct { + code Code + summary string + cause error +} + +func (e *runError) Error() string { + if e.cause == nil { + return e.summary + } + return e.summary + ": " + e.cause.Error() +} + +func (e *runError) Unwrap() error { return e.cause } + +// CodeOf returns err's stable code, or Internal for an unclassified error. +func CodeOf(err error) Code { + var run *runError + if errors.As(err, &run) { + return run.code + } + return Internal +} + +// PublicMessage returns err's safe public summary without its diagnostic cause. +func PublicMessage(err error) string { + var run *runError + if errors.As(err, &run) { + return run.summary + } + return "python runtime failed" +} + +// Detail returns err's bounded diagnostic for private logs and protocol traffic. +func Detail(err error) string { return truncateDetail(err.Error()) } + +// New creates an error with a stable code, sanitized public summary, and +// bounded private diagnostic cause. +func New(code Code, summary string, cause error) error { + return &runError{ + code: code, + summary: public(summary), + cause: boundedCause(cause), + } +} + +// Invalid creates an InvalidArgument error. +func Invalid(summary string, cause error) error { + return New(InvalidArgument, summary, cause) +} + +// Failed creates a FailedPrecondition error. +func Failed(summary string, cause error) error { + return New(FailedPrecondition, summary, cause) +} + +// CanceledError creates a Canceled error with the standard public summary. +func CanceledError(cause error) error { + return New(Canceled, "python execution was canceled", cause) +} + +// Deadline creates a DeadlineExceeded error with the standard public summary. +func Deadline(cause error) error { + return New(DeadlineExceeded, "python execution timed out", cause) +} + +// Exhausted creates a ResourceExhausted error. +func Exhausted(summary string, cause error) error { + return New(ResourceExhausted, summary, cause) +} + +// UnavailableError creates an Unavailable error. +func UnavailableError(summary string, cause error) error { + return New(Unavailable, summary, cause) +} + +// InternalError creates an Internal error with the standard public summary. +func InternalError(cause error) error { + return New(Internal, "python runtime failed", cause) +} + +// ContextError maps a context cancellation or deadline error to a runner code. +func ContextError(err error) error { + if errors.Is(err, context.Canceled) { + return CanceledError(err) + } + return Deadline(err) +} + +func public(summary string) string { + summary = strings.Join(strings.Fields(strings.ToValidUTF8(summary, "")), " ") + if summary == "" { + return "python runtime failed" + } + runes := []rune(summary) + if len(runes) > maxPublicRunes { + return string(runes[:maxPublicRunes-1]) + "…" + } + return summary +} + +func boundedCause(cause error) error { + if cause == nil { + return nil + } + if len(cause.Error()) <= MaxDetailBytes { + return cause + } + return truncatedCause{detail: truncateDetail(cause.Error()), cause: cause} +} + +func truncateDetail(detail string) string { + if len(detail) <= MaxDetailBytes { + return detail + } + limit := MaxDetailBytes - len(detailTruncationMarker) + if limit < 0 { + limit = 0 + } + return detail[:limit] + detailTruncationMarker +} + +type truncatedCause struct { + detail string + cause error +} + +func (e truncatedCause) Error() string { return e.detail } + +func (e truncatedCause) Unwrap() error { return e.cause } diff --git a/go/cloud-query/internal/tools/python/internal/contract/errors_test.go b/go/cloud-query/internal/tools/python/internal/contract/errors_test.go new file mode 100644 index 0000000000..4ae97e5e5f --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/contract/errors_test.go @@ -0,0 +1,18 @@ +package contract + +import ( + "errors" + "strings" + "testing" +) + +func TestErrorKeepsCauseAndBoundsDisclosure(t *testing.T) { + cause := errors.New(strings.Repeat("x", MaxDetailBytes+1)) + err := Invalid(strings.Repeat("summary ", 100), cause) + if CodeOf(err) != InvalidArgument || len([]rune(PublicMessage(err))) > 512 { + t.Fatalf("unexpected classified error: %q", PublicMessage(err)) + } + if !errors.Is(err, cause) || !strings.HasSuffix(Detail(err), detailTruncationMarker) { + t.Fatal("cause or detail truncation was lost") + } +} diff --git a/go/cloud-query/internal/tools/python/internal/contract/types.go b/go/cloud-query/internal/tools/python/internal/contract/types.go new file mode 100644 index 0000000000..f492e734e0 --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/contract/types.go @@ -0,0 +1,48 @@ +package contract + +const ( + // MaxSourceBytes bounds Python source accepted for one execution. + MaxSourceBytes = 64 << 10 + // MaxInputBytes bounds the JSON input accepted for one execution. + MaxInputBytes = 1 << 20 + // MaxResultBytes bounds the JSON result returned for one execution. + MaxResultBytes = 1 << 20 + // MaxStdoutBytes bounds stdout captured for one execution. + MaxStdoutBytes = 64 << 10 + // MaxDetailBytes bounds private diagnostics retained or sent on the protocol. + MaxDetailBytes = 64 << 10 + // MaxPublicMessageBytes bounds a sanitized error summary sent by a worker. + MaxPublicMessageBytes = 4 << 10 +) + +// RunInput contains validated source and the JSON object supplied as input. +type RunInput struct { + Script string + InputJSON string +} + +// RunOutput contains the JSON object assigned to output and captured stdout. +type RunOutput struct { + ResultJSON string + Stdout string +} + +// Code identifies a stable class of execution failure. +type Code string + +const ( + // InvalidArgument reports invalid source, input, or output values. + InvalidArgument Code = "invalid_argument" + // FailedPrecondition reports a Python execution failure after validation. + FailedPrecondition Code = "failed_precondition" + // Canceled reports cancellation by the caller. + Canceled Code = "canceled" + // DeadlineExceeded reports an execution deadline that elapsed. + DeadlineExceeded Code = "deadline_exceeded" + // ResourceExhausted reports a configured capacity or resource limit. + ResourceExhausted Code = "resource_exhausted" + // Unavailable reports that no runtime worker can accept the request. + Unavailable Code = "unavailable" + // Internal reports an unclassified runner or worker failure. + Internal Code = "internal" +) diff --git a/go/cloud-query/internal/tools/python/internal/contract/validation.go b/go/cloud-query/internal/tools/python/internal/contract/validation.go new file mode 100644 index 0000000000..48b2de98cf --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/contract/validation.go @@ -0,0 +1,44 @@ +package contract + +import ( + "encoding/json" + "strings" +) + +// NormalizeInput returns raw after verifying it is a bounded JSON object. An +// empty input becomes an empty JSON object. +func NormalizeInput(raw string) (string, error) { + if strings.TrimSpace(raw) == "" { + return "{}", nil + } + if len(raw) > MaxInputBytes { + return "", Invalid("input exceeds the input limit", nil) + } + var object map[string]json.RawMessage + if err := json.Unmarshal([]byte(raw), &object); err != nil || object == nil { + return "", Invalid("input_json must be a JSON object", err) + } + return raw, nil +} + +// ValidateRun validates source and normalizes the JSON object supplied as input. +func ValidateRun(input RunInput) (RunInput, error) { + if strings.TrimSpace(input.Script) == "" { + return RunInput{}, Invalid("script is required", nil) + } + if len(input.Script) > MaxSourceBytes { + return RunInput{}, Invalid("script exceeds the source limit", nil) + } + normalized, err := NormalizeInput(input.InputJSON) + if err != nil { + return RunInput{}, err + } + input.InputJSON = normalized + return input, nil +} + +// IsJSONObject reports whether raw contains one JSON object and no other value. +func IsJSONObject(raw string) bool { + var object map[string]json.RawMessage + return json.Unmarshal([]byte(raw), &object) == nil && object != nil +} diff --git a/go/cloud-query/internal/tools/python/internal/pool/config.go b/go/cloud-query/internal/tools/python/internal/pool/config.go new file mode 100644 index 0000000000..249f17ee32 --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/pool/config.go @@ -0,0 +1,44 @@ +package pool + +import ( + "strings" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" +) + +// ProcessConfig configures the command used for an isolated worker process. +// Environment replaces the parent process environment. +type ProcessConfig struct { + Executable string + Arguments []string + Environment []string +} + +// Config configures a Runner's worker capacity, queue, recycling, and process. +type Config struct { + Workers int + QueueSize int + MaxSuccessfulRunsBeforeRecycle int + Process ProcessConfig +} + +func validate(config Config) error { + if config.Workers <= 0 { + return contract.Invalid("python runner workers must be greater than zero", nil) + } + if config.QueueSize < 0 { + return contract.Invalid("python runner queue size must not be negative", nil) + } + if config.MaxSuccessfulRunsBeforeRecycle <= 0 { + return contract.Invalid("python runner recycle limit must be greater than zero", nil) + } + if strings.TrimSpace(config.Process.Executable) == "" { + return contract.Invalid("python worker executable is required", nil) + } + for _, entry := range config.Process.Environment { + if !strings.Contains(entry, "=") { + return contract.Invalid("python worker environment entries must use NAME=value", nil) + } + } + return nil +} diff --git a/go/cloud-query/internal/tools/python/internal/pool/pool.go b/go/cloud-query/internal/tools/python/internal/pool/pool.go new file mode 100644 index 0000000000..53b22be54d --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/pool/pool.go @@ -0,0 +1,408 @@ +package pool + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "time" + + "k8s.io/klog/v2" + + "github.com/pluralsh/console/go/cloud-query/internal/log" + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/protocol" +) + +const wallTimeout = 15 * time.Second + +// Runner owns a bounded set of isolated worker processes. It replaces workers +// after failures and recycles them after the configured number of successes. +type Runner struct { + config Config + factory processFactory + jobs chan job + mu sync.Mutex + closed bool + processes map[process]struct{} + wg sync.WaitGroup + closeOnce sync.Once + done chan struct{} + requestIDs atomic.Uint64 + workerIDs atomic.Uint64 + retry func() <-chan time.Time +} + +type job struct { + requestID string + queuedAt time.Time + ctx context.Context + input contract.RunInput + response chan result +} + +type result struct { + output *contract.RunOutput + err error +} + +// New starts config.Workers health-checked worker processes. If startup fails, +// it shuts down workers already started before returning the error. +func New(ctx context.Context, config Config) (*Runner, error) { + return newRunner(ctx, config, newExecProcess) +} + +func newRunner(ctx context.Context, config Config, factory processFactory) (*Runner, error) { + if err := validate(config); err != nil { + return nil, err + } + + klog.V(log.LogLevelInfo).InfoS( + "initializing python worker pool", + "workers", config.Workers, + "queue_size", config.QueueSize, + "recycle_limit", config.MaxSuccessfulRunsBeforeRecycle, + ) + + runner := &Runner{ + config: config, + factory: factory, + jobs: make(chan job, config.QueueSize), + processes: map[process]struct{}{}, + done: make(chan struct{}), + retry: func() <-chan time.Time { + return time.After(time.Second) + }, + } + + for range config.Workers { + if err := runner.start(ctx); err != nil { + _ = runner.CloseContext(context.Background()) + return nil, contract.UnavailableError("python runtime is unavailable", err) + } + } + + klog.V(log.LogLevelInfo).InfoS("python worker pool is ready", "workers", config.Workers) + + return runner, nil +} + +func (r *Runner) start(ctx context.Context) error { + ctx, cancel := context.WithTimeout(ctx, wallTimeout) + defer cancel() + + // Worker IDs identify parent-side process lifecycles without exposing command details. + workerID := r.workerIDs.Add(1) + klog.V(log.LogLevelVerbose).InfoS("spawning python worker", "worker_id", workerID) + process, err := r.factory(r.config.Process) + if err != nil { + logWorkerFailure("failed to start python worker process", workerID, err) + return err + } + klog.V(log.LogLevelVerbose).InfoS("python worker process started", "worker_id", workerID) + + request := protocol.Request{ + Version: protocol.Version, + Kind: protocol.Health, + ID: fmt.Sprintf("%d", r.requestIDs.Add(1)), + } + + response, err := process.exchange(ctx, request) + if err != nil || response.Error != nil { + process.stop() + if err != nil { + logWorkerFailure("python worker health check failed", workerID, err) + return err + } + err = protocol.Error(response) + logWorkerFailure("python worker health check failed", workerID, err) + return err + } + klog.V(log.LogLevelVerbose).InfoS("python worker health check passed", "worker_id", workerID) + + r.mu.Lock() + if r.closed { + r.mu.Unlock() + process.stop() + return contract.UnavailableError("python runtime is unavailable", nil) + } + + r.processes[process] = struct{}{} + r.wg.Add(1) + r.mu.Unlock() + + go r.serve(process, workerID) + return nil +} + +// Run validates and enqueues input, then waits for its worker response or ctx. +// It rejects new work after shutdown begins or when the queue is full. +func (r *Runner) Run(ctx context.Context, input contract.RunInput) (*contract.RunOutput, error) { + input, err := contract.ValidateRun(input) + if err != nil { + return nil, err + } + + if err := ctx.Err(); err != nil { + return nil, contract.ContextError(err) + } + + job := job{ + requestID: fmt.Sprintf("%d", r.requestIDs.Add(1)), + queuedAt: time.Now(), + ctx: ctx, + input: input, + response: make(chan result, 1), + } + + r.mu.Lock() + if r.closed || len(r.processes) == 0 { + r.mu.Unlock() + return nil, contract.UnavailableError("python runtime is unavailable", nil) + } + + select { + case r.jobs <- job: + klog.V(log.LogLevelExtended).InfoS( + "python request queued", + "request_id", job.requestID, + "queued_jobs", len(r.jobs), + "queue_capacity", cap(r.jobs), + ) + r.mu.Unlock() + case <-ctx.Done(): + r.mu.Unlock() + return nil, contract.ContextError(ctx.Err()) + default: + r.mu.Unlock() + return nil, contract.Exhausted("python queue is full", nil) + } + + select { + case answer := <-job.response: + return answer.output, answer.err + case <-ctx.Done(): + return nil, contract.ContextError(ctx.Err()) + } +} + +func (r *Runner) serve(process process, workerID uint64) { + retirement := "shutdown" + var retirementErr error + defer func() { + process.stop() + + r.mu.Lock() + delete(r.processes, process) + closed := r.closed + r.mu.Unlock() + + if closed { + retirement = "shutdown" + } + + fields := []any{"worker_id", workerID, "reason", retirement} + if retirementErr != nil { + fields = append(fields, + "error_code", contract.CodeOf(retirementErr), + "error", contract.PublicMessage(retirementErr), + ) + } + + klog.V(log.LogLevelExtended).InfoS("retiring python worker", fields...) + if !closed { + r.replace() + } + + r.wg.Done() + }() + + successes := 0 + for job := range r.jobs { + klog.V(log.LogLevelExtended).InfoS( + "python request taken from queue", + "request_id", job.requestID, + "worker_id", workerID, + "queued_for", time.Since(job.queuedAt), + "queued_jobs", len(r.jobs), + ) + + if err := job.ctx.Err(); err != nil { + klog.V(log.LogLevelExtended).InfoS( + "python request canceled before execution", + "request_id", job.requestID, + "worker_id", workerID, + "error_code", contract.CodeOf(contract.ContextError(err)), + ) + job.response <- result{err: contract.ContextError(err)} + continue + } + + ctx, cancel := context.WithTimeout(job.ctx, wallTimeout) + request := protocol.Request{ + Version: protocol.Version, + Kind: protocol.Run, + ID: job.requestID, + Script: job.input.Script, + InputJSON: job.input.InputJSON, + } + + startedAt := time.Now() + klog.V(log.LogLevelExtended).InfoS( + "python request execution started", + "request_id", job.requestID, + "worker_id", workerID, + ) + + response, err := process.exchange(ctx, request) + cancel() + if err != nil { + logExecutionFinished(job.requestID, workerID, startedAt, err) + job.response <- result{err: err} + retirement = retirementReason(err) + retirementErr = err + return + } + if response.Error != nil { + err := protocol.Error(response) + logExecutionFinished(job.requestID, workerID, startedAt, err) + job.response <- result{err: err} + retirement = "remote_python_error" + retirementErr = err + return + } + + logExecutionFinished(job.requestID, workerID, startedAt, nil) + job.response <- result{output: &contract.RunOutput{ + ResultJSON: response.ResultJSON, + Stdout: response.Stdout, + }} + + successes++ + if successes >= r.config.MaxSuccessfulRunsBeforeRecycle { + retirement = "successful_recycle_limit" + return + } + } +} + +func logExecutionFinished(requestID string, workerID uint64, startedAt time.Time, err error) { + status := "success" + if err != nil { + status = "error" + } + + fields := []any{ + "request_id", requestID, + "worker_id", workerID, + "duration", time.Since(startedAt), + "status", status, + } + if err != nil { + fields = append(fields, "error_code", contract.CodeOf(err)) + } + + klog.V(log.LogLevelExtended).InfoS("python request execution finished", fields...) +} + +func (r *Runner) replace() { + for { + r.mu.Lock() + closed := r.closed + r.mu.Unlock() + + if closed { + return + } + klog.V(log.LogLevelExtended).InfoS("attempting python worker replacement") + if err := r.start(context.Background()); err == nil { + klog.V(log.LogLevelExtended).InfoS("python worker replacement succeeded") + return + } + + select { + case <-r.done: + return + case <-r.retry(): + } + } +} + +// CloseContext stops admitting jobs, waits for worker shutdown, and kills +// remaining processes if ctx expires. It is safe to call more than once. +func (r *Runner) CloseContext(ctx context.Context) error { + r.closeOnce.Do(func() { + r.mu.Lock() + r.closed = true + processCount := len(r.processes) + queuedJobs := len(r.jobs) + close(r.jobs) + r.mu.Unlock() + klog.V(log.LogLevelInfo).InfoS( + "starting python worker pool shutdown", + "processes", processCount, + "queued_jobs", queuedJobs, + ) + go func() { + r.wg.Wait() + klog.V(log.LogLevelInfo).InfoS("python worker pool shutdown completed") + close(r.done) + }() + }) + + select { + case <-r.done: + return nil + case <-ctx.Done(): + r.mu.Lock() + processes := make([]process, 0, len(r.processes)) + for process := range r.processes { + processes = append(processes, process) + } + r.mu.Unlock() + + for _, process := range processes { + process.stop() + } + + drainedJobs := 0 + for draining := true; draining; { + select { + case job, ok := <-r.jobs: + if !ok { + draining = false + continue + } + job.response <- result{err: contract.UnavailableError("python runtime is unavailable", nil)} + drainedJobs++ + default: + draining = false + } + } + klog.ErrorS( + nil, + "python worker pool shutdown deadline exceeded; forcing teardown", + "processes", len(processes), + "queued_jobs", drainedJobs, + ) + return ctx.Err() + } +} + +func retirementReason(err error) string { + if contract.CodeOf(err) == contract.Canceled { + return "cancellation" + } + return "execution_or_transport_failure" +} + +func logWorkerFailure(message string, workerID uint64, err error) { + fields := []any{ + "error_code", contract.CodeOf(err), + "error", contract.PublicMessage(err), + } + if workerID != 0 { + fields = append([]any{"worker_id", workerID}, fields...) + } + klog.ErrorS(nil, message, fields...) +} diff --git a/go/cloud-query/internal/tools/python/internal/pool/pool_test.go b/go/cloud-query/internal/tools/python/internal/pool/pool_test.go new file mode 100644 index 0000000000..1abe62ddf6 --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/pool/pool_test.go @@ -0,0 +1,284 @@ +package pool + +import ( + "context" + "errors" + "fmt" + "runtime" + "sync" + "testing" + "time" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/protocol" +) + +type fakeProcess struct { + id int + stopped bool + mu sync.Mutex + seen []string + block chan struct{} +} + +func (p *fakeProcess) exchange(ctx context.Context, request protocol.Request) (protocol.Response, error) { + p.mu.Lock() + p.seen = append(p.seen, request.ID) + p.mu.Unlock() + if request.Kind == protocol.Health { + return protocol.Response{Version: protocol.Version, Kind: request.Kind, ID: request.ID}, nil + } + if request.Script == "error" { + return protocol.Response{ + Version: protocol.Version, + Kind: request.Kind, + ID: request.ID, + Error: &protocol.WireError{ + Code: contract.FailedPrecondition, + PublicMessage: "python ValueError: bad value", + Detail: "private", + }, + }, nil + } + if request.Script == "block" { + select { + case <-p.block: + case <-ctx.Done(): + return protocol.Response{}, contract.ContextError(ctx.Err()) + } + } + return protocol.Response{ + Version: protocol.Version, + Kind: request.Kind, + ID: request.ID, + ResultJSON: fmt.Sprintf(`{"worker":%d}`, p.id), + }, nil +} + +func (p *fakeProcess) stop() { + p.mu.Lock() + p.stopped = true + p.mu.Unlock() +} + +type fakeFactory struct { + mu sync.Mutex + processes []*fakeProcess + failAt int +} + +type controlledProcess struct { + started chan struct{} + stopped chan struct{} + once sync.Once +} + +func (p *controlledProcess) exchange(ctx context.Context, request protocol.Request) (protocol.Response, error) { + if request.Kind == protocol.Health { + return protocol.Response{Version: protocol.Version, Kind: request.Kind, ID: request.ID}, nil + } + select { + case <-p.started: + default: + close(p.started) + } + select { + case <-p.stopped: + return protocol.Response{}, contract.InternalError(errors.New("process stopped")) + case <-ctx.Done(): + return protocol.Response{}, contract.ContextError(ctx.Err()) + } +} + +func (p *controlledProcess) stop() { + p.once.Do(func() { close(p.stopped) }) +} + +func (f *fakeFactory) count() int { + f.mu.Lock() + defer f.mu.Unlock() + + return len(f.processes) +} + +func (f *fakeFactory) new(ProcessConfig) (process, error) { + f.mu.Lock() + defer f.mu.Unlock() + if f.failAt > 0 && len(f.processes)+1 == f.failAt { + return nil, errors.New("start failed") + } + p := &fakeProcess{id: len(f.processes) + 1, block: make(chan struct{})} + f.processes = append(f.processes, p) + return p, nil +} + +func testConfig() Config { + return Config{ + Workers: 1, + QueueSize: 1, + MaxSuccessfulRunsBeforeRecycle: 2, + Process: ProcessConfig{ + Executable: "fake", + Arguments: []string{"python-worker"}, + Environment: []string{"TMPDIR=/tmp"}, + }, + } +} + +func TestPoolStartupRecycleAndRemoteFailure(t *testing.T) { + factory := &fakeFactory{} + runner, err := newRunner(context.Background(), testConfig(), factory.new) + if err != nil { + t.Fatal(err) + } + defer runner.CloseContext(context.Background()) + + first, err := runner.Run(context.Background(), contract.RunInput{Script: "ok"}) + if err != nil { + t.Fatal(err) + } + if _, err := runner.Run(context.Background(), contract.RunInput{Script: "ok"}); err != nil { + t.Fatal(err) + } + deadline := time.Now().Add(time.Second) + for factory.count() < 2 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + third, err := runner.Run(context.Background(), contract.RunInput{Script: "ok"}) + if err != nil || third.ResultJSON == first.ResultJSON { + t.Fatalf("recycle: %#v %v", third, err) + } + if _, err := runner.Run(context.Background(), contract.RunInput{Script: "error"}); contract.CodeOf(err) != contract.FailedPrecondition { + t.Fatalf("remote error: %v", err) + } +} + +func TestPoolFailsEagerStartupAndRejectsAfterClose(t *testing.T) { + factory := &fakeFactory{failAt: 2} + config := Config{ + Workers: 2, + QueueSize: 1, + MaxSuccessfulRunsBeforeRecycle: 1, + Process: ProcessConfig{ + Executable: "fake", + Arguments: []string{"python-worker"}, + Environment: []string{"TMPDIR=/tmp"}, + }, + } + if _, err := newRunner(context.Background(), config, factory.new); contract.CodeOf(err) != contract.Unavailable { + t.Fatalf("startup: %v", err) + } + runner, err := newRunner(context.Background(), testConfig(), (&fakeFactory{}).new) + if err != nil { + t.Fatal(err) + } + if err := runner.CloseContext(context.Background()); err != nil { + t.Fatal(err) + } + if _, err := runner.Run(context.Background(), contract.RunInput{Script: "ok"}); contract.CodeOf(err) != contract.Unavailable { + t.Fatalf("admission: %v", err) + } +} + +func TestRunCancellationRetiresActiveProcess(t *testing.T) { + proc := &controlledProcess{started: make(chan struct{}), stopped: make(chan struct{})} + runner, err := newRunner(context.Background(), testConfig(), func(ProcessConfig) (process, error) { return proc, nil }) + if err != nil { + t.Fatal(err) + } + defer runner.CloseContext(context.Background()) + ctx, cancel := context.WithCancel(context.Background()) + result := make(chan error, 1) + go func() { _, err := runner.Run(ctx, contract.RunInput{Script: "block"}); result <- err }() + <-proc.started + cancel() + if err := <-result; contract.CodeOf(err) != contract.Canceled { + t.Fatalf("cancellation = %v", err) + } + select { + case <-proc.stopped: + case <-time.After(time.Second): + t.Fatal("active process was not stopped") + } +} + +func TestCloseDeadlineStopsActiveProcessAndDrainsQueue(t *testing.T) { + proc := &controlledProcess{started: make(chan struct{}), stopped: make(chan struct{})} + runner, err := newRunner(context.Background(), testConfig(), func(ProcessConfig) (process, error) { return proc, nil }) + if err != nil { + t.Fatal(err) + } + active := make(chan error, 1) + go func() { _, err := runner.Run(context.Background(), contract.RunInput{Script: "block"}); active <- err }() + <-proc.started + queued := make(chan result, 1) + runner.jobs <- job{ctx: context.Background(), input: contract.RunInput{Script: "queued"}, response: queued} + deadline, cancel := context.WithTimeout(context.Background(), time.Millisecond) + defer cancel() + if err := runner.CloseContext(deadline); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("close = %v", err) + } + if answer := <-queued; contract.CodeOf(answer.err) != contract.Unavailable { + t.Fatalf("queued = %v", answer.err) + } + if err := <-active; contract.CodeOf(err) != contract.Internal { + t.Fatalf("active = %v", err) + } + if _, err := runner.Run(context.Background(), contract.RunInput{Script: "later"}); contract.CodeOf(err) != contract.Unavailable { + t.Fatalf("admission = %v", err) + } +} + +func TestReplacementRetriesUntilCapacityRecovers(t *testing.T) { + var mu sync.Mutex + attempts := 0 + processes := make(chan *fakeProcess, 4) + factory := func(ProcessConfig) (process, error) { + mu.Lock() + defer mu.Unlock() + attempts++ + if attempts > 1 && attempts < 4 { + return nil, errors.New("replacement failed") + } + + process := &fakeProcess{id: attempts, block: make(chan struct{})} + processes <- process + return process, nil + } + + runner, err := newRunner(context.Background(), testConfig(), factory) + if err != nil { + t.Fatal(err) + } + defer runner.CloseContext(context.Background()) + <-processes + runner.retry = func() <-chan time.Time { ready := make(chan time.Time); close(ready); return ready } + if _, err := runner.Run(context.Background(), contract.RunInput{Script: "error"}); contract.CodeOf(err) != contract.FailedPrecondition { + t.Fatalf("run = %v", err) + } + <-processes + waitForCapacity(t, runner) + if _, err := runner.Run(context.Background(), contract.RunInput{Script: "ok"}); err != nil { + t.Fatalf("recovered = %v", err) + } + mu.Lock() + got := attempts + mu.Unlock() + if got != 4 { + t.Fatalf("replacement attempts = %d, want 4", got) + } +} + +func waitForCapacity(t *testing.T, runner *Runner) { + t.Helper() + for range 100_000 { + runner.mu.Lock() + ready := len(runner.processes) > 0 + runner.mu.Unlock() + if ready { + return + } + runtime.Gosched() + } + t.Fatal("replacement did not restore capacity") +} diff --git a/go/cloud-query/internal/tools/python/internal/pool/process.go b/go/cloud-query/internal/tools/python/internal/pool/process.go new file mode 100644 index 0000000000..dc4bfb4bae --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/pool/process.go @@ -0,0 +1,102 @@ +package pool + +import ( + "context" + "fmt" + "io" + "os/exec" + "sync" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/protocol" +) + +type process interface { + exchange(context.Context, protocol.Request) (protocol.Response, error) + stop() +} + +type processFactory func(ProcessConfig) (process, error) + +type execProcess struct { + stdin io.WriteCloser + stdout io.ReadCloser + kill func() + codec protocol.Codec + mu sync.Mutex + once sync.Once +} + +func newExecProcess(config ProcessConfig) (process, error) { + cmd := exec.Command(config.Executable, config.Arguments...) + cmd.Env = append([]string(nil), config.Environment...) + + stdin, err := cmd.StdinPipe() + if err != nil { + return nil, fmt.Errorf("opening python worker stdin: %w", err) + } + + stdout, err := cmd.StdoutPipe() + if err != nil { + return nil, fmt.Errorf("opening python worker stdout: %w", err) + } + + if err = cmd.Start(); err != nil { + return nil, fmt.Errorf("starting python worker process: %w", err) + } + + p := &execProcess{ + stdin: stdin, + stdout: stdout, + codec: protocol.NewCodec(protocol.MaxFrameSize), + } + + p.kill = func() { + _ = stdin.Close() + _ = stdout.Close() + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + _ = cmd.Wait() + } + + return p, nil +} + +func (p *execProcess) stop() { + p.once.Do(p.kill) +} + +func (p *execProcess) exchange(ctx context.Context, request protocol.Request) (protocol.Response, error) { + p.mu.Lock() + defer p.mu.Unlock() + + result := make(chan struct { + response protocol.Response + err error + }, 1) + go func() { + if err := p.codec.WriteRequest(p.stdin, request); err != nil { + result <- struct { + response protocol.Response + err error + }{err: err} + return + } + response, err := p.codec.ReadResponse(p.stdout, request) + result <- struct { + response protocol.Response + err error + }{response, err} + }() + select { + case result := <-result: + if result.err != nil { + return protocol.Response{}, contract.InternalError(result.err) + } + return result.response, nil + case <-ctx.Done(): + p.stop() + return protocol.Response{}, contract.ContextError(ctx.Err()) + } +} diff --git a/go/cloud-query/internal/tools/python/internal/protocol/codec.go b/go/cloud-query/internal/tools/python/internal/protocol/codec.go new file mode 100644 index 0000000000..3def9e254a --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/protocol/codec.go @@ -0,0 +1,244 @@ +package protocol + +import ( + "bytes" + "encoding/binary" + "encoding/json" + "errors" + "fmt" + "io" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" +) + +// Version identifies the private parent-to-worker protocol version. +const Version = 1 + +// MaxFrameSize bounds a framed protocol payload before the codec allocates it. +const MaxFrameSize = contract.MaxSourceBytes + + contract.MaxInputBytes + + contract.MaxResultBytes + + contract.MaxStdoutBytes + + contract.MaxDetailBytes + + 16<<10 +const ( + // Health requests a worker startup or liveness check. + Health = "health" + // Run requests one restricted Python execution. + Run = "run" +) + +// Request is a validated parent-to-worker protocol message. +type Request struct { + Version int `json:"version"` + Kind string `json:"kind"` + ID string `json:"id"` + Script string `json:"script,omitempty"` + InputJSON string `json:"input_json,omitempty"` +} + +// Response is a validated worker-to-parent protocol message. +type Response struct { + Version int `json:"version"` + Kind string `json:"kind"` + ID string `json:"id"` + ResultJSON string `json:"result_json,omitempty"` + Stdout string `json:"stdout,omitempty"` + Error *WireError `json:"error,omitempty"` +} + +// WireError carries a stable code, sanitized public summary, and bounded private +// diagnostic over the worker protocol. Detail must not reach untrusted callers. +type WireError struct { + Code contract.Code `json:"code"` + PublicMessage string `json:"public_message"` + Detail string `json:"detail"` +} + +// Codec reads and writes strict, bounded, length-prefixed protocol frames. +type Codec struct{ maxFrameSize int } + +// NewCodec creates a codec that rejects frames larger than maxFrameSize. +func NewCodec(maxFrameSize int) Codec { return Codec{maxFrameSize: maxFrameSize} } + +// ReadRequest reads and validates one request frame. +func (c Codec) ReadRequest(in io.Reader) (Request, error) { + var request Request + if err := c.read(in, &request); err != nil { + return Request{}, err + } + if err := c.validRequest(request); err != nil { + return Request{}, err + } + return request, nil +} + +// WriteRequest validates and writes one request frame. +func (c Codec) WriteRequest(out io.Writer, request Request) error { + if err := c.validRequest(request); err != nil { + return err + } + return c.write(out, request) +} + +// ReadResponse reads and validates a response matched to request. +func (c Codec) ReadResponse(in io.Reader, request Request) (Response, error) { + var response Response + if err := c.read(in, &response); err != nil { + return Response{}, err + } + if err := c.validResponse(request, response); err != nil { + return Response{}, err + } + return response, nil +} + +// WriteResponse validates and writes a response matched to request. +func (c Codec) WriteResponse(out io.Writer, request Request, response Response) error { + if err := c.validResponse(request, response); err != nil { + return err + } + return c.write(out, response) +} + +func (c Codec) read(in io.Reader, into any) error { + var header [4]byte + if _, err := io.ReadFull(in, header[:]); err != nil { + return err + } + + length := int(binary.BigEndian.Uint32(header[:])) + if length == 0 || length > c.maxFrameSize { + return errors.New("invalid python protocol frame") + } + + body := make([]byte, length) + if _, err := io.ReadFull(in, body); err != nil { + return err + } + + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(into); err != nil { + return errors.New("invalid python protocol payload") + } + if err := decoder.Decode(&struct{}{}); !errors.Is(err, io.EOF) { + return errors.New("invalid python protocol payload") + } + return nil +} + +func (c Codec) write(out io.Writer, value any) error { + body, err := json.Marshal(value) + if err != nil || len(body) == 0 || len(body) > c.maxFrameSize { + return errors.New("invalid python protocol frame") + } + + frame := make([]byte, 4+len(body)) + binary.BigEndian.PutUint32(frame, uint32(len(body))) + copy(frame[4:], body) + + for len(frame) > 0 { + n, err := out.Write(frame) + if err != nil { + return err + } + if n <= 0 { + return io.ErrShortWrite + } + frame = frame[n:] + } + return nil +} + +func (c Codec) validRequest(request Request) error { + if request.Version != Version || request.ID == "" { + return errors.New("invalid python protocol request") + } + + switch request.Kind { + case Health: + if request.Script != "" || request.InputJSON != "" { + return errors.New("invalid python protocol request") + } + case Run: + if request.Script == "" || + len(request.Script) > contract.MaxSourceBytes || + len(request.InputJSON) > contract.MaxInputBytes { + return errors.New("invalid python protocol request") + } + if _, err := contract.NormalizeInput(request.InputJSON); err != nil { + return errors.New("invalid python protocol request") + } + default: + return errors.New("invalid python protocol request") + } + return nil +} + +func (c Codec) validResponse(request Request, response Response) error { + if response.Version != Version || + response.Kind != request.Kind || + response.ID != request.ID { + return errors.New("invalid python protocol response") + } + + if response.Error != nil { + if !c.known(response.Error.Code) || + response.Error.PublicMessage == "" || + len(response.Error.PublicMessage) > contract.MaxPublicMessageBytes || + response.Error.Detail == "" || + len(response.Error.Detail) > contract.MaxDetailBytes || + response.ResultJSON != "" || + response.Stdout != "" { + return errors.New("invalid python protocol response") + } + return nil + } + + if request.Kind == Health { + if response.ResultJSON != "" || response.Stdout != "" { + return errors.New("invalid python protocol response") + } + return nil + } + + if response.ResultJSON == "" || + len(response.ResultJSON) > contract.MaxResultBytes || + len(response.Stdout) > contract.MaxStdoutBytes || + !contract.IsJSONObject(response.ResultJSON) { + return errors.New("invalid python protocol response") + } + return nil +} + +func (Codec) known(code contract.Code) bool { + switch code { + case contract.InvalidArgument, + contract.FailedPrecondition, + contract.Canceled, + contract.DeadlineExceeded, + contract.ResourceExhausted, + contract.Unavailable, + contract.Internal: + return true + } + return false +} + +// Error converts a worker error response into an error with a safe public +// summary and private bounded diagnostic. +func Error(response Response) error { + if response.Error == nil { + return nil + } + if !(Codec{}).known(response.Error.Code) { + return contract.InternalError(fmt.Errorf("unknown python worker error code")) + } + + return contract.New( + response.Error.Code, + response.Error.PublicMessage, + errors.New(response.Error.Detail), + ) +} diff --git a/go/cloud-query/internal/tools/python/internal/protocol/codec_test.go b/go/cloud-query/internal/tools/python/internal/protocol/codec_test.go new file mode 100644 index 0000000000..c5f69f4c43 --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/protocol/codec_test.go @@ -0,0 +1,150 @@ +package protocol + +import ( + "bytes" + "encoding/binary" + "encoding/json" + "io" + "strings" + "testing" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" +) + +func TestCodecRejectsMalformedAndInvalidResponses(t *testing.T) { + codec := NewCodec(128) + invalidBodies := [][]byte{ + nil, + []byte(`{"version":1,"kind":"health","id":"x","extra":true}`), + []byte(`{"version":2,"kind":"health","id":"x"}`), + } + for _, body := range invalidBodies { + var input bytes.Buffer + var header [4]byte + binary.BigEndian.PutUint32(header[:], uint32(len(body))) + input.Write(header[:]) + input.Write(body) + if _, err := codec.ReadRequest(&input); err == nil { + t.Fatal("accepted invalid request") + } + } + request := Request{Version: Version, Kind: Run, ID: "x", Script: "output={}", InputJSON: "{}"} + response := Response{ + Version: Version, + Kind: Run, + ID: "x", + Error: &WireError{ + Code: contract.Internal, + PublicMessage: "python runtime failed", + Detail: "private", + }, + } + if err := codec.WriteResponse(&bytes.Buffer{}, request, response); err != nil { + t.Fatal(err) + } + if Public := Error(response); contract.PublicMessage(Public) != "python runtime failed" { + t.Fatal("private detail leaked") + } +} + +func TestCodecFrameAndSemanticInvariants(t *testing.T) { + codec := NewCodec(MaxFrameSize) + request := Request{Version: Version, Kind: Run, ID: "request", Script: "output={}", InputJSON: "{}"} + for _, body := range [][]byte{ + []byte(`{"version":1,"kind":"unknown","id":"request"}`), + []byte(`{"version":1,"kind":"run","id":"request","script":"","input_json":"{}"}`), + []byte(`{"version":1,"kind":"health","id":"request","script":"x"}`), + []byte(`{"version":1,"kind":"run","id":"request","script":"output={}","input_json":"{}"} {}`), + } { + if _, err := codec.ReadRequest(frame(body)); err == nil { + t.Fatalf("accepted request %s", body) + } + } + for _, response := range []Response{ + {Version: 2, Kind: Run, ID: request.ID, ResultJSON: "{}"}, + {Version: Version, Kind: Health, ID: request.ID, ResultJSON: "{}"}, + {Version: Version, Kind: Run, ID: "other", ResultJSON: "{}"}, + { + Version: Version, + Kind: Run, + ID: request.ID, + ResultJSON: "{}", + Error: &WireError{ + Code: contract.Internal, + PublicMessage: "python runtime failed", + Detail: "x", + }, + }, + {Version: Version, Kind: Run, ID: request.ID, Error: &WireError{Code: contract.Internal}}, + { + Version: Version, + Kind: Run, + ID: request.ID, + Error: &WireError{ + Code: contract.Internal, + PublicMessage: "python runtime failed", + Detail: strings.Repeat("x", contract.MaxDetailBytes+1), + }, + }, + {Version: Version, Kind: Run, ID: request.ID, ResultJSON: "[]"}, + {Version: Version, Kind: Run, ID: request.ID, ResultJSON: strings.Repeat("x", contract.MaxResultBytes+1)}, + } { + var encoded bytes.Buffer + body, _ := jsonMarshal(response) + encoded.Write(frameBytes(body)) + if _, err := codec.ReadResponse(&encoded, request); err == nil { + t.Fatalf("accepted response %#v", response) + } + } + for _, input := range []io.Reader{bytes.NewReader([]byte{0, 0, 0, 0}), bytes.NewReader([]byte{0, 0, 0, 1, '{'}), frame([]byte(`{`))} { + if _, err := codec.ReadRequest(input); err == nil { + t.Fatal("accepted malformed frame") + } + } + oversized := make([]byte, MaxFrameSize+1) + binary.BigEndian.PutUint32(oversized[:4], uint32(MaxFrameSize+1)) + if _, err := codec.ReadRequest(bytes.NewReader(oversized[:4])); err == nil { + t.Fatal("accepted oversized frame") + } +} + +type shortWriter struct{ writes int } + +func (w *shortWriter) Write(p []byte) (int, error) { + w.writes++ + if w.writes == 1 { + return 1, nil + } + return 0, nil +} + +func TestCodecHandlesShortWritesAndPrivateRemoteDetail(t *testing.T) { + codec := NewCodec(MaxFrameSize) + request := Request{Version: Version, Kind: Health, ID: "x"} + if err := codec.WriteRequest(&shortWriter{}, request); err != io.ErrShortWrite { + t.Fatalf("short write: %v", err) + } + err := Error(Response{Error: &WireError{ + Code: contract.Internal, + PublicMessage: "python runtime failed", + Detail: "secret path /tmp/private", + }}) + if strings.Contains(contract.PublicMessage(err), "secret") { + t.Fatal("remote detail leaked") + } +} + +func frame(body []byte) io.Reader { return bytes.NewReader(frameBytes(body)) } + +func frameBytes(body []byte) []byte { + var header [4]byte + binary.BigEndian.PutUint32(header[:], uint32(len(body))) + return append(header[:], body...) +} + +func jsonMarshal(value any) ([]byte, error) { + var buffer bytes.Buffer + encoder := json.NewEncoder(&buffer) + err := encoder.Encode(value) + return bytes.TrimSpace(buffer.Bytes()), err +} diff --git a/go/cloud-query/internal/tools/python/internal/worker/os.go b/go/cloud-query/internal/tools/python/internal/worker/os.go new file mode 100644 index 0000000000..0d0b882b43 --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/worker/os.go @@ -0,0 +1,86 @@ +package worker + +import ( + "context" + "time" + + monty "github.com/ewhauser/gomonty" +) + +const ( + osDateToday monty.OSFunction = "date.today" + osDateTimeNow monty.OSFunction = "datetime.now" +) + +// handleOS exposes the only host callbacks available to sandboxed Python. +// Clock values are always derived from UTC; filesystem, environment, and all +// other OS callbacks remain unavailable. +func (m *montyRuntime) handleOS(_ context.Context, call monty.OSCall) (monty.Result, error) { + switch call.Function { + case osDateToday: + if len(call.Args) != 0 || len(call.Kwargs) != 0 { + return m.invalidClockArguments(call.Function) + } + return monty.Return(monty.DateValue(m.dateOf(m.utcNow()))), nil + case osDateTimeNow: + return m.dateTimeNow(call) + default: + message := "OS function " + string(call.Function) + " is not available" + return monty.Raise(monty.Exception{Type: "NotImplementedError", Arg: &message}), nil + } +} + +func (m *montyRuntime) dateTimeNow(call monty.OSCall) (monty.Result, error) { + if len(call.Kwargs) != 0 || len(call.Args) > 1 { + return m.invalidClockArguments(call.Function) + } + + now := m.dateTimeOf(m.utcNow()) + if len(call.Args) == 0 || call.Args[0].Raw() == nil { + return monty.Return(monty.DateTimeValue(now)), nil + } + + timezone, ok := call.Args[0].TimeZone() + if !ok || timezone.OffsetSeconds != 0 { + message := "datetime.now only supports UTC" + return monty.Raise(monty.Exception{Type: "ValueError", Arg: &message}), nil + } + + zeroOffset := int32(0) + utcName := "UTC" + now.OffsetSeconds = &zeroOffset + now.TimezoneName = &utcName + return monty.Return(monty.DateTimeValue(now)), nil +} + +func (m *montyRuntime) utcNow() time.Time { + if m.now == nil { + return time.Now().UTC() + } + return m.now().UTC() +} + +func (m *montyRuntime) dateOf(value time.Time) monty.Date { + return monty.Date{ + Year: int32(value.Year()), + Month: uint8(value.Month()), + Day: uint8(value.Day()), + } +} + +func (m *montyRuntime) dateTimeOf(value time.Time) monty.DateTime { + return monty.DateTime{ + Year: int32(value.Year()), + Month: uint8(value.Month()), + Day: uint8(value.Day()), + Hour: uint8(value.Hour()), + Minute: uint8(value.Minute()), + Second: uint8(value.Second()), + Microsecond: uint32(time.Duration(value.Nanosecond()) / time.Microsecond), + } +} + +func (m *montyRuntime) invalidClockArguments(function monty.OSFunction) (monty.Result, error) { + message := string(function) + " received unsupported arguments" + return monty.Raise(monty.Exception{Type: "TypeError", Arg: &message}), nil +} diff --git a/go/cloud-query/internal/tools/python/internal/worker/runtime.go b/go/cloud-query/internal/tools/python/internal/worker/runtime.go new file mode 100644 index 0000000000..a57f6b2855 --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/worker/runtime.go @@ -0,0 +1,165 @@ +package worker + +import ( + "context" + "errors" + "strconv" + "strings" + "time" + + monty "github.com/ewhauser/gomonty" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" +) + +const executionTimeout = 10 * time.Second +const maxMemoryBytes = 64 << 20 +const maxRecursionDepth = 200 + +type montyRuntime struct { + now func() time.Time +} + +func newMontyRuntime() *montyRuntime { return &montyRuntime{now: time.Now} } + +func (m *montyRuntime) Health() error { + _, err := monty.NewRepl(monty.ReplOptions{ + ScriptName: "workbench.py", + Limits: m.limits(), + }) + if err != nil { + return contract.InternalError(err) + } + return nil +} + +func (m *montyRuntime) Run(ctx context.Context, script, inputJSON string) (*contract.RunOutput, error) { + input, err := contract.ValidateRun(contract.RunInput{Script: script, InputJSON: inputJSON}) + if err != nil { + return nil, err + } + + repl, err := monty.NewRepl(monty.ReplOptions{ + ScriptName: "workbench.py", + Limits: m.limits(), + }) + if err != nil { + return nil, contract.InternalError(err) + } + + ctx, cancel := context.WithTimeout(ctx, executionTimeout) + defer cancel() + + initialization := "import json as __workbench_json\n" + + "input = __workbench_json.loads(" + strconv.Quote(input.InputJSON) + ")\n" + + "output = {}" + + if _, err = repl.FeedRun(ctx, initialization, m.feedOptions()); err != nil { + return nil, m.mapError(err) + } + + var stdout strings.Builder + tooLarge := false + printFn := func(stream, text string) { + if stream != "stdout" { + return + } + if stdout.Len()+len(text) > contract.MaxStdoutBytes { + tooLarge = true + return + } + stdout.WriteString(text) + } + + options := m.feedOptions() + options.Print = printFn + if _, err = repl.FeedRun(ctx, input.Script, options); err != nil { + return nil, m.mapError(err) + } + + if tooLarge { + return nil, contract.Exhausted("python stdout exceeds the stdout limit", nil) + } + + value, err := repl.FeedRun(ctx, "__workbench_json.dumps(output)", m.feedOptions()) + if err != nil { + return nil, m.mapError(err) + } + + result, ok := value.Raw().(string) + if !ok || !contract.IsJSONObject(result) { + return nil, contract.Invalid("output must be a JSON object", nil) + } + + if len(result) > contract.MaxResultBytes { + return nil, contract.Exhausted("python result exceeds the result limit", nil) + } + + return &contract.RunOutput{ + ResultJSON: result, + Stdout: stdout.String(), + }, nil +} + +func (m *montyRuntime) feedOptions() monty.FeedOptions { + return monty.FeedOptions{OS: m.handleOS} +} + +func (m *montyRuntime) limits() *monty.ResourceLimits { + return &monty.ResourceLimits{ + MaxDuration: executionTimeout, + MaxMemory: maxMemoryBytes, + MaxRecursionDepth: maxRecursionDepth, + } +} + +func (m *montyRuntime) mapError(err error) error { + if errors.Is(err, context.DeadlineExceeded) { + return contract.Deadline(err) + } + + var syntax *monty.SyntaxError + if errors.As(err, &syntax) { + return contract.Invalid("python code is invalid", err) + } + + var runtime *monty.RuntimeError + if !errors.As(err, &runtime) { + return contract.InternalError(err) + } + + kind, detail, ok := strings.Cut(strings.TrimSpace(runtime.Error()), ":") + if !ok { + kind, detail = "Exception", runtime.Error() + } + + switch kind { + case "SyntaxError", "TypeError": + return contract.Invalid(m.toErrorSummary(kind, detail), err) + case "MemoryError", "RecursionError": + return contract.Exhausted("python resource limit exceeded", err) + case "TimeoutError": + return contract.Deadline(err) + case "ResourceError": + if strings.Contains(strings.ToLower(detail), "time") || strings.Contains(strings.ToLower(detail), "duration") { + return contract.Deadline(err) + } + return contract.Exhausted("python resource limit exceeded", err) + default: + return contract.Failed(m.toErrorSummary(kind, detail), err) + } +} + +func (m *montyRuntime) toErrorSummary(kind, detail string) string { + kind = strings.TrimSpace(kind) + detail = strings.Join(strings.Fields(detail), " ") + + if kind == "" { + kind = "Exception" + } + + if detail == "" { + return "python " + kind + } + return "python " + kind + ": " + detail +} diff --git a/go/cloud-query/internal/tools/python/internal/worker/server.go b/go/cloud-query/internal/tools/python/internal/worker/server.go new file mode 100644 index 0000000000..050bc169fc --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/worker/server.go @@ -0,0 +1,76 @@ +package worker + +import ( + "context" + "errors" + "io" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/protocol" +) + +type runtime interface { + Health() error + Run(context.Context, string, string) (*contract.RunOutput, error) +} + +// Server processes private protocol requests using one restricted runtime. +type Server struct { + runtime runtime + codec protocol.Codec +} + +// NewServer creates a server backed by the restricted Monty runtime. +func NewServer() *Server { + return newServer(newMontyRuntime()) +} + +func newServer(runtime runtime) *Server { + return &Server{runtime: runtime, codec: protocol.NewCodec(protocol.MaxFrameSize)} +} + +// Run processes protocol requests until input reaches EOF or a protocol or I/O +// failure occurs. It returns only private diagnostics to its caller. +func (s *Server) Run(in io.Reader, out io.Writer) error { + for { + request, err := s.codec.ReadRequest(in) + if errors.Is(err, io.EOF) { + return nil + } + if err != nil { + return contract.InternalError(err) + } + + response := protocol.Response{ + Version: protocol.Version, + Kind: request.Kind, + ID: request.ID, + } + + switch request.Kind { + case protocol.Health: + err = s.runtime.Health() + case protocol.Run: + var result *contract.RunOutput + result, err = s.runtime.Run(context.Background(), request.Script, request.InputJSON) + if err == nil { + response.ResultJSON = result.ResultJSON + response.Stdout = result.Stdout + } + default: + err = contract.InternalError(errors.New("unreachable protocol kind")) + } + + if err != nil { + response.Error = &protocol.WireError{ + Code: contract.CodeOf(err), + PublicMessage: contract.PublicMessage(err), + Detail: contract.Detail(err), + } + } + + if err := s.codec.WriteResponse(out, request, response); err != nil { + return contract.InternalError(err) + } + } +} diff --git a/go/cloud-query/internal/tools/python/internal/worker/server_test.go b/go/cloud-query/internal/tools/python/internal/worker/server_test.go new file mode 100644 index 0000000000..c739cfec98 --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/worker/server_test.go @@ -0,0 +1,134 @@ +package worker + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "strings" + "testing" + "time" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/protocol" +) + +type fakeRuntime struct { + health error + output *contract.RunOutput + run error + calls int +} + +func (f *fakeRuntime) Health() error { return f.health } + +func (f *fakeRuntime) Run(_ context.Context, _, _ string) (*contract.RunOutput, error) { + f.calls++ + return f.output, f.run +} + +func TestServerDispatchesAndKeepsErrorDetailPrivate(t *testing.T) { + codec := protocol.NewCodec(protocol.MaxFrameSize) + runtime := &fakeRuntime{output: &contract.RunOutput{ResultJSON: "{}"}} + var in, out bytes.Buffer + request := protocol.Request{Version: protocol.Version, Kind: protocol.Run, ID: "1", Script: "output={}", InputJSON: "{}"} + if err := codec.WriteRequest(&in, request); err != nil { + t.Fatal(err) + } + if err := newServer(runtime).Run(&in, &out); err != nil { + t.Fatal(err) + } + response, err := codec.ReadResponse(&out, request) + if err != nil || response.ResultJSON != "{}" || runtime.calls != 1 { + t.Fatalf("response=%#v err=%v", response, err) + } + runtime.run = contract.Failed("python execution failed", errors.New("private host detail")) + in.Reset() + out.Reset() + _ = codec.WriteRequest(&in, request) + if err := newServer(runtime).Run(&in, &out); err != nil { + t.Fatal(err) + } + response, err = codec.ReadResponse(&out, request) + if err != nil || response.Error == nil || contract.PublicMessage(protocol.Error(response)) != "python execution failed" { + t.Fatalf("error response=%#v err=%v", response, err) + } +} + +func TestMontyRuntimeFreshStateAndStdout(t *testing.T) { + runtime := newMontyRuntime() + first, err := runtime.Run(context.Background(), "secret = 42\noutput = {'sum': input['value'] + 1}\nprint('ok')", `{"value":1}`) + if err != nil || first.ResultJSON != "{\"sum\": 2}" || first.Stdout != "ok\n" { + t.Fatalf("first=%#v err=%v", first, err) + } + if _, err := runtime.Run(context.Background(), "output = {'secret': secret}", "{}"); contract.CodeOf(err) != contract.FailedPrecondition { + t.Fatalf("state leaked: %v", err) + } +} + +func TestMontyRuntimeUsesUTCClockOnly(t *testing.T) { + runtime := &montyRuntime{ + now: func() time.Time { + return time.Date(2026, time.April, 5, 6, 7, 8, 123456000, time.FixedZone("CEST", 2*60*60)) + }, + } + + output, err := runtime.Run(context.Background(), ` +from datetime import date, datetime +now = datetime.now() +today = date.today() +output = {"now": now.isoformat(), "today": today.isoformat()} +`, "{}") + if err != nil { + t.Fatal(err) + } + + var result map[string]string + if err := json.Unmarshal([]byte(output.ResultJSON), &result); err != nil { + t.Fatal(err) + } + if want := map[string]string{ + "now": "2026-04-05T04:07:08.123456", + "today": "2026-04-05", + }; !equalStringMaps(result, want) { + t.Fatalf("result = %#v, want %#v", result, want) + } + + for name, script := range map[string]string{ + "environment": "import os\noutput = {'value': os.getenv('SECRET')}", + "filesystem": "from pathlib import Path\noutput = {'value': Path('/tmp/test').exists()}", + "non-UTC clock": ` +from datetime import datetime, timedelta, timezone +output = {"value": datetime.now(timezone(timedelta(hours=1))).isoformat()} +`, + } { + t.Run(name, func(t *testing.T) { + _, err := runtime.Run(context.Background(), script, "{}") + if contract.CodeOf(err) != contract.FailedPrecondition { + t.Fatalf("err = %v", err) + } + + message := contract.PublicMessage(err) + if name == "non-UTC clock" && !strings.Contains(message, "only supports UTC") { + t.Fatalf("err = %v", err) + } + if name != "non-UTC clock" && !strings.Contains(message, "NotImplementedError") { + t.Fatalf("err = %v", err) + } + }) + } +} + +func equalStringMaps(got, want map[string]string) bool { + if len(got) != len(want) { + return false + } + + for key, wantValue := range want { + if got[key] != wantValue { + return false + } + } + + return true +} diff --git a/go/cloud-query/internal/tools/python/protocol.go b/go/cloud-query/internal/tools/python/protocol.go deleted file mode 100644 index 07f4e68470..0000000000 --- a/go/cloud-query/internal/tools/python/protocol.go +++ /dev/null @@ -1,138 +0,0 @@ -package python - -import ( - "bytes" - "context" - "encoding/binary" - "encoding/json" - "errors" - "io" - "sync" -) - -const ( - protocolVersion = 1 - requestHealth = "health" - requestRun = "run" -) - -type protocolRequest struct { - Version int `json:"version"` - Type string `json:"type"` - ID string `json:"id"` - Script string `json:"script,omitempty"` - InputJSON string `json:"input_json,omitempty"` -} - -type protocolResponse struct { - Version int `json:"version"` - Type string `json:"type"` - ID string `json:"id"` - ResultJSON string `json:"result_json,omitempty"` - Stdout string `json:"stdout,omitempty"` - Code Code `json:"code,omitempty"` - Message string `json:"message,omitempty"` -} - -func readFrame(in io.Reader, into any) error { - var header [4]byte - if _, err := io.ReadFull(in, header[:]); err != nil { - return err - } - length := binary.BigEndian.Uint32(header[:]) - if length == 0 || length > maxProtocolFrameLen { - return errors.New("invalid protocol frame") - } - body := make([]byte, length) - if _, err := io.ReadFull(in, body); err != nil { - return err - } - decoder := json.NewDecoder(bytes.NewReader(body)) - decoder.DisallowUnknownFields() - if err := decoder.Decode(into); err != nil { - return errors.New("invalid protocol payload") - } - if err := decoder.Decode(&struct{}{}); !errors.Is(err, io.EOF) { - return errors.New("invalid protocol payload") - } - return nil -} - -func writeFrame(out io.Writer, value any) error { - body, err := json.Marshal(value) - if err != nil || len(body) == 0 || len(body) > maxProtocolFrameLen { - return errors.New("invalid protocol response") - } - frame := make([]byte, 4+len(body)) - binary.BigEndian.PutUint32(frame, uint32(len(body))) - copy(frame[4:], body) - _, err = io.Copy(out, bytes.NewReader(frame)) - return err -} - -type childWorker struct { - stdin io.WriteCloser - stdout io.ReadCloser - kill func() - mu sync.Mutex - stopOnce sync.Once -} - -func (w *childWorker) stop() { w.stopOnce.Do(w.kill) } - -func (w *childWorker) exchange(ctx context.Context, request protocolRequest) (protocolResponse, error) { - w.mu.Lock() - defer w.mu.Unlock() - result := make(chan struct { - response protocolResponse - err error - }, 1) - go func() { - if err := writeFrame(w.stdin, request); err != nil { - result <- struct { - response protocolResponse - err error - }{err: err} - return - } - var response protocolResponse - result <- struct { - response protocolResponse - err error - }{response: response, err: readFrame(w.stdout, &response)} - }() - select { - case result := <-result: - if result.err != nil || result.response.Version != protocolVersion || result.response.ID != request.ID || result.response.Type != request.Type { - return protocolResponse{}, runtimeFailure() - } - if (result.response.Code == "") != (result.response.Message == "") { - return protocolResponse{}, runtimeFailure() - } - if request.Type == requestHealth && (result.response.ResultJSON != "" || result.response.Stdout != "") { - return protocolResponse{}, runtimeFailure() - } - if result.response.Code != "" && (result.response.ResultJSON != "" || result.response.Stdout != "") { - return protocolResponse{}, runtimeFailure() - } - if request.Type == requestRun && result.response.Code == "" && (len(result.response.ResultJSON) > MaxResultBytes || len(result.response.Stdout) > MaxStdoutBytes || !isJSONObject(result.response.ResultJSON)) { - return protocolResponse{}, runtimeFailure() - } - return result.response, nil - case <-ctx.Done(): - w.stop() - return protocolResponse{}, executionContextError(ctx.Err()) - } -} - -func protocolError(response protocolResponse) error { - if response.Code == "" || response.Message == "" { - return runtimeFailure() - } - switch response.Code { - case InvalidArgument, FailedPrecondition, Canceled, DeadlineExceeded, ResourceExhausted, Unavailable, Internal: - return &Error{Code: response.Code, Msg: response.Message} - default: - return runtimeFailure() - } -} diff --git a/go/cloud-query/internal/tools/python/python.go b/go/cloud-query/internal/tools/python/python.go deleted file mode 100644 index 2587cce78b..0000000000 --- a/go/cloud-query/internal/tools/python/python.go +++ /dev/null @@ -1,322 +0,0 @@ -// Package python runs Monty in crash-isolated self-spawned workers. -package python - -import ( - "context" - "fmt" - "os" - "os/exec" - "strings" - "sync" - "time" -) - -const ( - MaxSourceBytes = 64 << 10 - MaxInputBytes = 1 << 20 - MaxResultBytes = 1 << 20 - MaxStdoutBytes = 64 << 10 - - defaultWorkers = 4 - defaultQueueSize = 16 - maxCheckouts = 10 - executionTimeout = 10 * time.Second - wallTimeout = 15 * time.Second - maxMemoryBytes = 64 << 20 - maxRecursionDepth = 200 - maxProtocolFrameLen = MaxInputBytes + MaxSourceBytes + MaxResultBytes + MaxStdoutBytes -) - -type RunInput struct{ Script, InputJSON string } -type RunOutput struct{ ResultJSON, Stdout string } - -// Config controls the parent pool. Zero numeric fields use the package defaults. -type Config struct { - Workers int - QueueSize int - MaxSuccessfulRuns int - Executable string - Arguments []string -} - -type Option func(*Config) - -// WithBinaryPath remains available for tests. Production defaults to the -// current executable, which must dispatch python-worker to RunWorker. -func WithBinaryPath(path string) Option { return func(c *Config) { c.Executable = path } } -func withArguments(arguments ...string) Option { - return func(c *Config) { c.Arguments = append([]string(nil), arguments...) } -} -func WithConfig(config Config) Option { return func(c *Config) { *c = config } } - -type job struct { - ctx context.Context - input RunInput - response chan runResult -} -type runResult struct { - output *RunOutput - err error -} - -type Runner struct { - config Config - jobs chan *job - mu sync.Mutex - closed bool - healthy int - workers map[*childWorker]struct{} - workerWG sync.WaitGroup - replaceWG sync.WaitGroup - closeOnce sync.Once - closedDone chan struct{} -} - -func New(ctx context.Context, options ...Option) (*Runner, error) { - executable, err := os.Executable() - if err != nil { - return nil, &Error{Code: Unavailable, Msg: runtimeUnavailableMessage} - } - cfg := Config{Workers: defaultWorkers, QueueSize: defaultQueueSize, MaxSuccessfulRuns: maxCheckouts, Executable: executable, Arguments: []string{"python-worker"}} - for _, option := range options { - option(&cfg) - } - if cfg.Workers == 0 { - cfg.Workers = defaultWorkers - } - if cfg.QueueSize == 0 { - cfg.QueueSize = defaultQueueSize - } - if cfg.MaxSuccessfulRuns == 0 { - cfg.MaxSuccessfulRuns = maxCheckouts - } - if cfg.Executable == "" { - cfg.Executable = executable - } - if cfg.Arguments == nil { - cfg.Arguments = []string{"python-worker"} - } - if cfg.Workers <= 0 || cfg.QueueSize < 0 || cfg.MaxSuccessfulRuns <= 0 || strings.TrimSpace(cfg.Executable) == "" { - return nil, &Error{Code: Unavailable, Msg: runtimeUnavailableMessage} - } - r := &Runner{config: cfg, jobs: make(chan *job, cfg.QueueSize), workers: make(map[*childWorker]struct{}), closedDone: make(chan struct{})} - for range cfg.Workers { - if err := r.startWorker(ctx); err != nil { - r.Close() - return nil, &Error{Code: Unavailable, Msg: runtimeUnavailableMessage} - } - } - return r, nil -} - -func (r *Runner) startWorker(ctx context.Context) error { - ctx, cancel := context.WithTimeout(ctx, wallTimeout) - defer cancel() - - w, err := r.newWorker() - if err != nil { - return err - } - r.mu.Lock() - if r.closed { - r.mu.Unlock() - w.stop() - return fmt.Errorf("runner is closed") - } - r.workers[w] = struct{}{} - r.mu.Unlock() - healthy := false - defer func() { - if healthy { - return - } - w.stop() - r.mu.Lock() - delete(r.workers, w) - r.mu.Unlock() - }() - - response, err := w.exchange(ctx, protocolRequest{Version: protocolVersion, Type: requestHealth, ID: requestHealth}) - if err != nil || response.Code != "" { - return fmt.Errorf("unhealthy worker") - } - r.mu.Lock() - if r.closed { - r.mu.Unlock() - return fmt.Errorf("runner is closed") - } - r.healthy++ - r.workerWG.Add(1) - healthy = true - r.mu.Unlock() - go r.serveWorker(w) - return nil -} - -func (r *Runner) newWorker() (*childWorker, error) { - cmd := exec.Command(r.config.Executable, r.config.Arguments...) - cmd.Env = []string{"TMPDIR=/tmp"} - stdin, err := cmd.StdinPipe() - if err != nil { - return nil, err - } - stdout, err := cmd.StdoutPipe() - if err != nil { - return nil, err - } - if err := cmd.Start(); err != nil { - return nil, err - } - w := &childWorker{stdin: stdin, stdout: stdout} - w.kill = func() { - _ = stdin.Close() - _ = stdout.Close() - if cmd.Process != nil { - _ = cmd.Process.Kill() - } - _ = cmd.Wait() - } - return w, nil -} - -func (r *Runner) serveWorker(w *childWorker) { - defer r.workerWG.Done() - defer func() { - w.stop() - r.mu.Lock() - delete(r.workers, w) - r.healthy-- - closed := r.closed - if !closed { - r.replaceWG.Add(1) - } - r.mu.Unlock() - if !closed { - go func() { - defer r.replaceWG.Done() - r.replaceWorker() - }() - } - }() - successes := 0 - for request := range r.jobs { - if err := request.ctx.Err(); err != nil { - request.response <- runResult{err: executionContextError(err)} - continue - } - ctx, cancel := context.WithTimeout(request.ctx, wallTimeout) - response, err := w.exchange(ctx, protocolRequest{Version: protocolVersion, Type: requestRun, ID: fmt.Sprintf("%d", time.Now().UnixNano()), Script: request.input.Script, InputJSON: request.input.InputJSON}) - cancel() - if err != nil { - request.response <- runResult{err: err} - return - } - if response.Code != "" { - request.response <- runResult{err: protocolError(response)} - return - } - request.response <- runResult{output: &RunOutput{ResultJSON: response.ResultJSON, Stdout: response.Stdout}} - successes++ - if successes >= r.config.MaxSuccessfulRuns { - return - } - } -} - -func (r *Runner) replaceWorker() { - ticker := time.NewTicker(time.Second) - defer ticker.Stop() - for { - r.mu.Lock() - closed := r.closed - r.mu.Unlock() - if closed { - return - } - if r.startWorker(context.Background()) == nil { - return - } - select { - case <-r.closedDone: - return - case <-ticker.C: - } - } -} - -func (r *Runner) Run(ctx context.Context, input RunInput) (*RunOutput, error) { - if len(strings.TrimSpace(input.Script)) == 0 { - return nil, invalid("script is required") - } - if len(input.Script) > MaxSourceBytes { - return nil, invalid("script exceeds the source limit") - } - if len(input.InputJSON) > MaxInputBytes { - return nil, invalid("input exceeds the input limit") - } - inputJSON, err := validateRunInput(input.InputJSON) - if err != nil { - return nil, err - } - if err := ctx.Err(); err != nil { - return nil, executionContextError(err) - } - request := &job{ctx: ctx, input: RunInput{Script: input.Script, InputJSON: inputJSON}, response: make(chan runResult, 1)} - r.mu.Lock() - if r.closed || r.healthy == 0 { - r.mu.Unlock() - return nil, &Error{Code: Unavailable, Msg: runtimeUnavailableMessage} - } - select { - case r.jobs <- request: - r.mu.Unlock() - case <-ctx.Done(): - r.mu.Unlock() - return nil, executionContextError(ctx.Err()) - default: - r.mu.Unlock() - return nil, &Error{Code: ResourceExhausted, Msg: "python queue is full"} - } - select { - case result := <-request.response: - return result.output, result.err - case <-ctx.Done(): - return nil, executionContextError(ctx.Err()) - } -} - -func (r *Runner) Close() { _ = r.CloseContext(context.Background()) } - -// CloseContext stops admissions and lets queued and active requests finish -// until ctx expires. On expiry it kills workers and fails remaining queued work. -func (r *Runner) CloseContext(ctx context.Context) error { - r.closeOnce.Do(func() { - r.mu.Lock() - r.closed = true - close(r.jobs) - r.mu.Unlock() - go func() { - r.workerWG.Wait() - r.replaceWG.Wait() - close(r.closedDone) - }() - }) - select { - case <-r.closedDone: - return nil - case <-ctx.Done(): - r.mu.Lock() - workers := make([]*childWorker, 0, len(r.workers)) - for w := range r.workers { - workers = append(workers, w) - } - r.mu.Unlock() - for _, w := range workers { - w.stop() - } - for request := range r.jobs { - request.response <- runResult{err: &Error{Code: Unavailable, Msg: runtimeUnavailableMessage}} - } - return ctx.Err() - } -} diff --git a/go/cloud-query/internal/tools/python/python_test.go b/go/cloud-query/internal/tools/python/python_test.go index 52f3b69157..88254808a9 100644 --- a/go/cloud-query/internal/tools/python/python_test.go +++ b/go/cloud-query/internal/tools/python/python_test.go @@ -1,27 +1,15 @@ package python import ( - "bytes" "context" - "encoding/binary" - "errors" - "fmt" - "io" "os" "strings" "testing" - "time" ) func TestMain(m *testing.M) { if len(os.Args) == 2 && os.Args[1] == "python-worker" { - if err := RunWorker(os.Stdin, os.Stdout); err != nil { - os.Exit(1) - } - os.Exit(0) - } - if len(os.Args) == 2 && os.Args[1] == "fake-python-worker" { - if err := runFakeWorker(os.Stdin, os.Stdout); err != nil { + if err := NewWorker().Run(os.Stdin, os.Stdout); err != nil { os.Exit(1) } os.Exit(0) @@ -29,273 +17,38 @@ func TestMain(m *testing.M) { os.Exit(m.Run()) } -func TestValidateRunInput(t *testing.T) { - for _, raw := range []string{"[]", "1", "{"} { - if _, err := validateRunInput(raw); ErrorCode(err) != InvalidArgument { - t.Fatalf("validateRunInput(%q) = %v", raw, err) - } - } - if got, err := validateRunInput(""); err != nil || got != "{}" { - t.Fatalf("validateRunInput(empty) = %q, %v", got, err) - } -} - -func TestPoolDefaults(t *testing.T) { - if defaultWorkers != 4 { - t.Fatalf("defaultWorkers = %d, want 4", defaultWorkers) - } - if maxCheckouts != 10 { - t.Fatalf("maxCheckouts = %d, want 10", maxCheckouts) - } -} - -func TestDefaultQueueSize(t *testing.T) { - runner, err := newFakeRunner(t, Config{Workers: 1, MaxSuccessfulRuns: 1}) +func TestRunnerSelfSpawnsWorker(t *testing.T) { + runner, err := NewRunner(context.Background(), RunnerConfig{Workers: 1, QueueSize: 1, MaxSuccessfulRunsBeforeRecycle: 1}) if err != nil { - t.Fatalf("New() error: %v", err) - } - defer runner.Close() - if got := cap(runner.jobs); got != 16 { - t.Fatalf("queue capacity = %d, want 16", got) - } -} - -func TestWorkerReusesTenSuccessfulRunsWithoutHealthChecks(t *testing.T) { - runner, err := newFakeRunner(t, Config{Workers: 1}) - if err != nil { - t.Fatalf("New() error: %v", err) - } - defer runner.Close() - - var firstPID string - for run := 0; run < maxCheckouts; run++ { - output, err := runner.Run(context.Background(), RunInput{Script: "success"}) - if err != nil { - t.Fatalf("run %d error: %v", run+1, err) - } - if run == 0 { - firstPID = output.ResultJSON - } else if output.ResultJSON != firstPID { - t.Fatalf("run %d used %s, want original worker %s", run+1, output.ResultJSON, firstPID) - } - } - - output, err := runAfterReplacement(t, runner, RunInput{Script: "success"}) - if err != nil { - t.Fatalf("replacement run error: %v", err) - } - if output.ResultJSON == firstPID { - t.Fatalf("worker was not recycled after %d successful runs", maxCheckouts) - } -} - -func TestWorkerResponseErrorTriggersReplacement(t *testing.T) { - runner, err := newFakeRunner(t, Config{Workers: 1}) - if err != nil { - t.Fatalf("New() error: %v", err) - } - defer runner.Close() - - first, err := runner.Run(context.Background(), RunInput{Script: "success"}) - if err != nil { - t.Fatalf("initial run error: %v", err) - } - if _, err := runner.Run(context.Background(), RunInput{Script: "resource-error"}); ErrorCode(err) != ResourceExhausted { - t.Fatalf("resource response error = %v", err) - } - second, err := runAfterReplacement(t, runner, RunInput{Script: "success"}) - if err != nil { - t.Fatalf("replacement run error: %v", err) - } - if second.ResultJSON == first.ResultJSON { - t.Fatal("worker was reused after a resource-exhausted response") - } -} - -func TestProtocolRejectsUnknownAndTrailingValues(t *testing.T) { - for _, body := range []string{ - `{"version":1,"type":"health","id":"1","extra":true}`, - `{"version":1,"type":"health","id":"1"} {}`, - } { - var framed bytes.Buffer - if err := writeRawFrame(&framed, []byte(body)); err != nil { - t.Fatal(err) - } - var request protocolRequest - if err := readFrame(&framed, &request); err == nil { - t.Fatalf("readFrame(%s) accepted invalid payload", body) - } - } -} - -func TestRunWorkerHealth(t *testing.T) { - configureWorkerEnvironment(t) - var input, output bytes.Buffer - if err := writeFrame(&input, protocolRequest{Version: protocolVersion, Type: requestHealth, ID: "1"}); err != nil { - t.Fatal(err) - } - if err := RunWorker(&input, &output); err != nil { - t.Fatalf("RunWorker() error: %v", err) - } - var response protocolResponse - if err := readFrame(&output, &response); err != nil { - t.Fatal(err) - } - if response.Code != "" || response.ID != "1" { - t.Fatalf("health response = %#v", response) - } -} - -func TestRunWorkerUsesFreshStateAndSeparatesStdout(t *testing.T) { - configureWorkerEnvironment(t) - var input, output bytes.Buffer - requests := []protocolRequest{ - {Version: protocolVersion, Type: requestRun, ID: "first", Script: "secret = 42\noutput = {'sum': input['value'] + 1}\nprint('ok')", InputJSON: `{"value": 1}`}, - {Version: protocolVersion, Type: requestRun, ID: "second", Script: "output = {'secret': secret}", InputJSON: `{}`}, - } - for _, request := range requests { - if err := writeFrame(&input, request); err != nil { - t.Fatal(err) - } - } - if err := RunWorker(&input, &output); err != nil { - t.Fatalf("RunWorker() error: %v", err) - } - - var first protocolResponse - if err := readFrame(&output, &first); err != nil { t.Fatal(err) } - if first.Code != "" || first.ResultJSON != `{"sum": 2}` || first.Stdout != "ok\n" { - t.Fatalf("first response = %#v", first) - } - var second protocolResponse - if err := readFrame(&output, &second); err != nil { - t.Fatal(err) - } - if second.Code != FailedPrecondition || second.ResultJSON != "" || second.Stdout != "" { - t.Fatalf("second response = %#v", second) - } -} - -func TestRunRejectsBoundsBeforeAdmission(t *testing.T) { - runner := &Runner{jobs: make(chan *job), healthy: 1} - _, err := runner.Run(context.Background(), RunInput{Script: strings.Repeat("x", MaxSourceBytes+1)}) - if ErrorCode(err) != InvalidArgument { - t.Fatalf("source error = %v", err) - } - _, err = runner.Run(context.Background(), RunInput{Script: "pass", InputJSON: strings.Repeat("x", MaxInputBytes+1)}) - if ErrorCode(err) != InvalidArgument { - t.Fatalf("input error = %v", err) + defer runner.CloseContext(context.Background()) + output, err := runner.Run(context.Background(), RunInput{Script: "output = {'value': 42}"}) + if err != nil || output.ResultJSON != "{\"value\": 42}" { + t.Fatalf("run = %#v, %v", output, err) } } -func TestRunnerSelfSpawnsAndQueuesUntilCapacityReturns(t *testing.T) { - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - runner, err := New(ctx, WithConfig(Config{Workers: 1, QueueSize: 1, MaxSuccessfulRuns: 100})) +func TestRunnerPreservesSandboxErrorMessage(t *testing.T) { + runner, err := NewRunner(context.Background(), RunnerConfig{ + Workers: 1, + QueueSize: 1, + MaxSuccessfulRunsBeforeRecycle: 1, + }) if err != nil { - t.Fatalf("New() error: %v", err) + t.Fatal(err) } - defer runner.Close() - - firstCtx, cancelFirst := context.WithCancel(ctx) - first := make(chan error, 1) - go func() { - _, err := runner.Run(firstCtx, RunInput{Script: "while True:\n pass"}) - first <- err - }() - time.Sleep(50 * time.Millisecond) + defer runner.CloseContext(context.Background()) - second := make(chan runResult, 1) - go func() { - output, err := runner.Run(ctx, RunInput{Script: "output = {'value': 42}"}) - second <- runResult{output: output, err: err} - }() - deadline := time.Now().Add(time.Second) - for len(runner.jobs) != 1 && time.Now().Before(deadline) { - time.Sleep(time.Millisecond) - } - if len(runner.jobs) != 1 { - t.Fatal("second request did not enter the queue") - } - if _, err := runner.Run(ctx, RunInput{Script: "output = {}"}); ErrorCode(err) != ResourceExhausted { - t.Fatalf("full queue error = %v", err) + _, err = runner.Run(context.Background(), RunInput{ + Script: "import socket\nsocket.create_connection(('example.com', 443))\noutput = {'reachable': True}", + }) + if err == nil { + t.Fatal("sandboxed socket import unexpectedly succeeded") } - cancelFirst() - if err := <-first; ErrorCode(err) != Canceled { - t.Fatalf("canceled run error = %v", err) - } - result := <-second - if result.err != nil || result.output == nil || result.output.ResultJSON != `{"value": 42}` { - t.Fatalf("queued run = %#v, %v", result.output, result.err) - } -} - -func writeRawFrame(out *bytes.Buffer, body []byte) error { - frame := make([]byte, 4) - binary.BigEndian.PutUint32(frame, uint32(len(body))) - _, err := out.Write(append(frame, body...)) - return err -} - -func configureWorkerEnvironment(t *testing.T) { - t.Helper() - t.Setenv("HOME", "") - t.Setenv("XDG_CACHE_HOME", "") - t.Setenv("TMPDIR", "/tmp") -} - -func newFakeRunner(t *testing.T, config Config) (*Runner, error) { - t.Helper() - config.Executable = os.Args[0] - config.Arguments = []string{"fake-python-worker"} - return New(context.Background(), WithConfig(config)) -} - -func runAfterReplacement(t *testing.T, runner *Runner, input RunInput) (*RunOutput, error) { - t.Helper() - deadline := time.Now().Add(time.Second) - for { - output, err := runner.Run(context.Background(), input) - if ErrorCode(err) != Unavailable || time.Now().After(deadline) { - return output, err - } - time.Sleep(time.Millisecond) - } -} - -func runFakeWorker(in *os.File, out *os.File) error { - ran := false - for { - var request protocolRequest - if err := readFrame(in, &request); err != nil { - if errors.Is(err, io.EOF) { - return nil - } - return err - } - response := protocolResponse{Version: protocolVersion, Type: request.Type, ID: request.ID} - switch request.Type { - case requestHealth: - if ran { - response.Code = Internal - response.Message = "unexpected health request" - } - case requestRun: - ran = true - if request.Script == "resource-error" { - response.Code = ResourceExhausted - response.Message = "python resource limit exceeded" - } else { - response.ResultJSON = fmt.Sprintf(`{"pid":%d}`, os.Getpid()) - } - default: - return fmt.Errorf("unexpected request type %q", request.Type) - } - if err := writeFrame(out, response); err != nil { - return err - } + message := PublicMessage(err) + if message == "python execution failed" || !strings.Contains(strings.ToLower(message), "socket") { + t.Fatalf("sandbox error message = %q", message) } } diff --git a/go/cloud-query/internal/tools/python/runner.go b/go/cloud-query/internal/tools/python/runner.go new file mode 100644 index 0000000000..6c2c8e6a0c --- /dev/null +++ b/go/cloud-query/internal/tools/python/runner.go @@ -0,0 +1,34 @@ +package python + +import ( + "context" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/pool" +) + +// Runner accepts Python executions and manages its worker processes. +type Runner interface { + // Run validates and executes input, respecting ctx while queued and running. + Run(context.Context, RunInput) (*RunOutput, error) + // CloseContext stops admission and waits for workers until ctx expires. + CloseContext(context.Context) error +} + +// NewRunner starts a configured set of isolated Python workers. +func NewRunner(ctx context.Context, config RunnerConfig) (Runner, error) { + config, err := resolveRunnerConfig(config) + if err != nil { + return nil, err + } + + return pool.New(ctx, pool.Config{ + Workers: config.Workers, + QueueSize: config.QueueSize, + MaxSuccessfulRunsBeforeRecycle: config.MaxSuccessfulRunsBeforeRecycle, + Process: pool.ProcessConfig{ + Executable: config.WorkerProcess.Executable, + Arguments: config.WorkerProcess.Arguments, + Environment: config.WorkerProcess.Environment, + }, + }) +} diff --git a/go/cloud-query/internal/tools/python/types.go b/go/cloud-query/internal/tools/python/types.go new file mode 100644 index 0000000000..deac542ecc --- /dev/null +++ b/go/cloud-query/internal/tools/python/types.go @@ -0,0 +1,29 @@ +package python + +import "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" + +// RunInput contains the source and JSON object supplied to a Python execution. +type RunInput = contract.RunInput + +// RunOutput contains the JSON object assigned to output and captured stdout. +type RunOutput = contract.RunOutput + +// Code identifies a stable category of Python runner failure. +type Code = contract.Code + +const ( + // InvalidArgument reports invalid source, input, or output values. + InvalidArgument = contract.InvalidArgument + // FailedPrecondition reports an execution failure after input validation. + FailedPrecondition = contract.FailedPrecondition + // Canceled reports cancellation by the caller. + Canceled = contract.Canceled + // DeadlineExceeded reports an execution deadline that elapsed. + DeadlineExceeded = contract.DeadlineExceeded + // ResourceExhausted reports a configured capacity or resource limit. + ResourceExhausted = contract.ResourceExhausted + // Unavailable reports that no Python runtime worker can accept the request. + Unavailable = contract.Unavailable + // Internal reports an unclassified runner or worker failure. + Internal = contract.Internal +) diff --git a/go/cloud-query/internal/tools/python/worker.go b/go/cloud-query/internal/tools/python/worker.go index 5ecf17bb6c..816ca14153 100644 --- a/go/cloud-query/internal/tools/python/worker.go +++ b/go/cloud-query/internal/tools/python/worker.go @@ -1,148 +1,16 @@ package python import ( - "context" - "errors" "io" - "strings" - monty "github.com/ewhauser/gomonty" + internalworker "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/worker" ) -// RunWorker serves the private parent protocol. cmd/main must call it when its -// first argument is python-worker, before initializing the normal service. -func RunWorker(in io.Reader, out io.Writer) error { - for { - var request protocolRequest - if err := readFrame(in, &request); err != nil { - if errors.Is(err, io.EOF) { - return nil - } - return err - } - response := protocolResponse{Version: protocolVersion, Type: request.Type, ID: request.ID} - if request.Version != protocolVersion || request.ID == "" { - return errors.New("invalid protocol request") - } - switch request.Type { - case requestHealth: - if request.Script != "" || request.InputJSON != "" { - return errors.New("invalid protocol request") - } - if err := healthCheck(); err != nil { - response.Code, response.Message = errorFields(err) - } - case requestRun: - if len(request.Script) == 0 || len(request.Script) > MaxSourceBytes || len(request.InputJSON) > MaxInputBytes { - response.Code, response.Message = errorFields(invalid("invalid python request")) - break - } - inputJSON, err := validateRunInput(request.InputJSON) - if err != nil { - response.Code, response.Message = errorFields(err) - break - } - output, err := runMonty(request.Script, inputJSON) - if err != nil { - response.Code, response.Message = errorFields(err) - } else { - response.ResultJSON, response.Stdout = output.ResultJSON, output.Stdout - } - default: - return errors.New("invalid protocol request") - } - if err := writeFrame(out, response); err != nil { - return err - } - } +// Worker serves the private parent-to-worker protocol on a pair of streams. +type Worker interface { + // Run processes requests until the input reaches EOF or an I/O error occurs. + Run(io.Reader, io.Writer) error } -func healthCheck() error { - _, err := monty.NewRepl(monty.ReplOptions{ScriptName: "workbench.py", Limits: montyLimits()}) - if err != nil { - return runtimeFailure() - } - return nil -} - -func runMonty(script, inputJSON string) (*RunOutput, error) { - repl, err := monty.NewRepl(monty.ReplOptions{ScriptName: "workbench.py", Limits: montyLimits()}) - if err != nil { - return nil, runtimeFailure() - } - ctx, cancel := context.WithTimeout(context.Background(), executionTimeout) - defer cancel() - if _, err := repl.FeedRun(ctx, "import json as __workbench_json\ninput = __workbench_json.loads("+pythonString(inputJSON)+")\noutput = {}", monty.FeedOptions{}); err != nil { - return nil, montyError(err) - } - var stdout strings.Builder - stdoutLimit := false - print := func(stream, text string) { - if stream != "stdout" { - return - } - if stdout.Len()+len(text) > MaxStdoutBytes { - stdoutLimit = true - return - } - stdout.WriteString(text) - } - if _, err := repl.FeedRun(ctx, script, monty.FeedOptions{Print: print}); err != nil { - return nil, montyError(err) - } - if stdoutLimit { - return nil, &Error{Code: ResourceExhausted, Msg: "python stdout exceeds the stdout limit"} - } - value, err := repl.FeedRun(ctx, "__workbench_json.dumps(output)", monty.FeedOptions{}) - if err != nil { - return nil, montyError(err) - } - result, ok := value.Raw().(string) - if !ok { - return nil, invalid("output must be a JSON object") - } - if len(result) > MaxResultBytes { - return nil, &Error{Code: ResourceExhausted, Msg: "python result exceeds the result limit"} - } - if !isJSONObject(result) { - return nil, invalid("output must be a JSON object") - } - return &RunOutput{ResultJSON: result, Stdout: stdout.String()}, nil -} - -func montyLimits() *monty.ResourceLimits { - return &monty.ResourceLimits{MaxDuration: executionTimeout, MaxMemory: maxMemoryBytes, MaxRecursionDepth: maxRecursionDepth} -} - -func montyError(err error) error { - if errors.Is(err, context.DeadlineExceeded) { - return executionTimeoutError() - } - var syntax *monty.SyntaxError - if errors.As(err, &syntax) { - return invalid("python code is invalid") - } - var runtime *monty.RuntimeError - if errors.As(err, &runtime) { - kind, detail, ok := strings.Cut(strings.TrimSpace(runtime.Error()), ":") - if !ok { - kind, detail = "Exception", runtime.Error() - } - if kind == "ResourceError" { - if strings.Contains(strings.ToLower(detail), "time") || strings.Contains(strings.ToLower(detail), "duration") { - return executionTimeoutError() - } - return &Error{Code: ResourceExhausted, Msg: "python resource limit exceeded"} - } - return runtimeError(kind, detail) - } - return runtimeFailure() -} - -func errorFields(err error) (Code, string) { - var typed *Error - if errors.As(err, &typed) { - return typed.Code, typed.Msg - } - return Internal, "python runtime failed" -} +// NewWorker constructs the worker implementation used by the python-worker command. +func NewWorker() Worker { return internalworker.NewServer() } From 2fb11cd66f5d6b63899692956215cabe46438891 Mon Sep 17 00:00:00 2001 From: Sebastian Florek Date: Thu, 20 Aug 2026 16:53:11 +0200 Subject: [PATCH 03/15] feat(python): update runtime limits and add tests - Increased Monty execution time to 60 seconds - Updated wall timeout to 65 seconds and Monty-managed memory to 100 MiB - Reduced recursion limit to 100 frames - Updated README and API reference documentation to reflect new limits - Added tests for wall timeout and Monty runtime limit validation --- go/cloud-query/README.md | 2 +- go/cloud-query/docs/api-reference.md | 2 +- go/cloud-query/internal/tools/python/README.md | 8 ++++---- .../internal/tools/python/internal/pool/pool.go | 2 +- .../internal/tools/python/internal/pool/pool_test.go | 6 ++++++ .../internal/tools/python/internal/worker/runtime.go | 6 +++--- .../internal/tools/python/internal/worker/server_test.go | 7 +++++++ 7 files changed, 23 insertions(+), 10 deletions(-) diff --git a/go/cloud-query/README.md b/go/cloud-query/README.md index 412b842534..5fc1308704 100644 --- a/go/cloud-query/README.md +++ b/go/cloud-query/README.md @@ -97,7 +97,7 @@ ToolQuery also supports cloud function invocation via `InvokeLambda` for AWS Lam `RunPython` executes Monty's limited Python subset in a crash-isolated `cloud-query python-worker` subprocess. The request's optional JSON object is available as `input`; `output` starts as an empty dictionary and must remain a JSON-serializable dictionary. The response returns that dictionary as `result_json` and standard-output text from `print()` separately as `stdout`. -This is not CPython. The sandbox has no host filesystem, environment, network, subprocess, shell, `pip`, third-party packages, or callback access. Two workers start with only `TMPDIR=/tmp`; each request gets a fresh gomonty REPL. Failed or canceled workers are killed and replaced, and healthy workers are periodically recycled. Source is limited to 64 KiB, input and result JSON to 1 MiB, stdout to 64 KiB, interpreter execution to 10 seconds, interpreter-managed memory to 64 MiB, and recursion to 200 frames. Two runs execute concurrently and up to 16 more wait in a FIFO queue. A full queue is rejected. A parent watchdog ends a run after 15 seconds or the caller's earlier deadline. +This is not CPython. The sandbox has no host filesystem, environment, network, subprocess, shell, `pip`, third-party packages, or callback access. Two workers start with only `TMPDIR=/tmp`; each request gets a fresh gomonty REPL. Failed or canceled workers are killed and replaced, and healthy workers are periodically recycled. Source is limited to 64 KiB, input and result JSON to 1 MiB, stdout to 64 KiB, interpreter execution to 60 seconds, interpreter-managed memory to 100 MiB, and recursion to 100 frames. Two runs execute concurrently and up to 16 more wait in a FIFO queue. A full queue is rejected. A parent watchdog ends a run after 65 seconds or the caller's earlier deadline. The image embeds gomonty `v0.0.14` and its platform-specific glibc library, built against official Monty commit `c9802b5f30d11fecf9f153feb1dfdab3abda070e`. It contains no separate `monty` executable. Both pins are recorded in OCI labels. Monty's memory limit covers interpreter-managed allocations rather than total pod RSS; operators should measure the workload before reducing the default cloud-query memory allocation. diff --git a/go/cloud-query/docs/api-reference.md b/go/cloud-query/docs/api-reference.md index 73cbda1b40..ac0616d967 100644 --- a/go/cloud-query/docs/api-reference.md +++ b/go/cloud-query/docs/api-reference.md @@ -1525,7 +1525,7 @@ message RunPythonOutput { } ``` -The runtime exposes no host filesystem, environment, network, subprocess, shell, package installation, third-party package, or host-tool callback. It limits source to 64 KiB, input and result JSON to 1 MiB, stdout to 64 KiB, execution to 10 seconds, memory to 64 MiB, recursion to 200 frames, wall time to 15 seconds, and concurrency to two runs per process. Up to 16 additional requests wait in a bounded FIFO queue. It uses gomonty `v0.0.14`, built against official Monty commit `c9802b5f30d11fecf9f153feb1dfdab3abda070e`; it is not CPython. +The runtime exposes no host filesystem, environment, network, subprocess, shell, package installation, third-party package, or host-tool callback. It limits source to 64 KiB, input and result JSON to 1 MiB, stdout to 64 KiB, execution to 60 seconds, memory to 100 MiB, recursion to 100 frames, wall time to 65 seconds, and concurrency to two runs per process. Up to 16 additional requests wait in a bounded FIFO queue. It uses gomonty `v0.0.14`, built against official Monty commit `c9802b5f30d11fecf9f153feb1dfdab3abda070e`; it is not CPython. ## Invoke Lambda diff --git a/go/cloud-query/internal/tools/python/README.md b/go/cloud-query/internal/tools/python/README.md index 8be87ac69b..168738c000 100644 --- a/go/cloud-query/internal/tools/python/README.md +++ b/go/cloud-query/internal/tools/python/README.md @@ -44,10 +44,10 @@ Limits: - Input and result JSON: 1 MiB each - Captured stdout: 64 KiB - Private diagnostics: 64 KiB -- Monty execution: 10 seconds -- Parent wall clock: 15 seconds or the caller's earlier deadline -- Monty-managed memory: 64 MiB -- Recursion: 200 frames +- Monty execution: 60 seconds +- Parent wall clock: 65 seconds or the caller's earlier deadline +- Monty-managed memory: 100 MiB +- Recursion: 100 frames - Active workers: 4 - Waiting requests: 16 by default - Worker recycling: 10 successful requests diff --git a/go/cloud-query/internal/tools/python/internal/pool/pool.go b/go/cloud-query/internal/tools/python/internal/pool/pool.go index 53b22be54d..3981fd0f39 100644 --- a/go/cloud-query/internal/tools/python/internal/pool/pool.go +++ b/go/cloud-query/internal/tools/python/internal/pool/pool.go @@ -14,7 +14,7 @@ import ( "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/protocol" ) -const wallTimeout = 15 * time.Second +const wallTimeout = 65 * time.Second // Runner owns a bounded set of isolated worker processes. It replaces workers // after failures and recycles them after the configured number of successes. diff --git a/go/cloud-query/internal/tools/python/internal/pool/pool_test.go b/go/cloud-query/internal/tools/python/internal/pool/pool_test.go index 1abe62ddf6..869c3d1969 100644 --- a/go/cloud-query/internal/tools/python/internal/pool/pool_test.go +++ b/go/cloud-query/internal/tools/python/internal/pool/pool_test.go @@ -125,6 +125,12 @@ func testConfig() Config { } } +func TestWallTimeoutAllowsMontyExecutionLimit(t *testing.T) { + if wallTimeout != 65*time.Second { + t.Fatalf("wall timeout = %s, want 65s", wallTimeout) + } +} + func TestPoolStartupRecycleAndRemoteFailure(t *testing.T) { factory := &fakeFactory{} runner, err := newRunner(context.Background(), testConfig(), factory.new) diff --git a/go/cloud-query/internal/tools/python/internal/worker/runtime.go b/go/cloud-query/internal/tools/python/internal/worker/runtime.go index a57f6b2855..e04bc17b8a 100644 --- a/go/cloud-query/internal/tools/python/internal/worker/runtime.go +++ b/go/cloud-query/internal/tools/python/internal/worker/runtime.go @@ -12,9 +12,9 @@ import ( "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" ) -const executionTimeout = 10 * time.Second -const maxMemoryBytes = 64 << 20 -const maxRecursionDepth = 200 +const executionTimeout = 60 * time.Second +const maxMemoryBytes = 100 << 20 +const maxRecursionDepth = 100 type montyRuntime struct { now func() time.Time diff --git a/go/cloud-query/internal/tools/python/internal/worker/server_test.go b/go/cloud-query/internal/tools/python/internal/worker/server_test.go index c739cfec98..b80b9f0826 100644 --- a/go/cloud-query/internal/tools/python/internal/worker/server_test.go +++ b/go/cloud-query/internal/tools/python/internal/worker/server_test.go @@ -66,6 +66,13 @@ func TestMontyRuntimeFreshStateAndStdout(t *testing.T) { } } +func TestMontyRuntimeLimitsMatchWorkbenchDefaults(t *testing.T) { + limits := newMontyRuntime().limits() + if limits.MaxDuration != 60*time.Second || limits.MaxMemory != 100<<20 || limits.MaxRecursionDepth != 100 { + t.Fatalf("limits = %#v", limits) + } +} + func TestMontyRuntimeUsesUTCClockOnly(t *testing.T) { runtime := &montyRuntime{ now: func() time.Time { From a47856414cae835eb981d214afb7b140409aa65b Mon Sep 17 00:00:00 2001 From: Sebastian Florek Date: Thu, 20 Aug 2026 17:01:06 +0200 Subject: [PATCH 04/15] chore(dependencies): update Go module dependencies - Bumped Go version in `console/go/helm-test` from 1.26.5 to 1.26.6 - Added various indirect dependencies to `console/go/ai-proxy` - Updated `google.golang.org/genproto/googleapis/rpc` to new version in multiple modules - Aligned proto dependencies across projects --- go/ai-proxy/go.mod | 14 +++++++++++++- go/ai-proxy/go.sum | 22 +++++++++++++++++----- go/client/go.mod | 2 +- go/client/go.sum | 3 +-- go/deployment-operator/terratest/go.mod | 2 +- go/deployment-operator/terratest/go.sum | 3 +-- go/helm-test/go.mod | 2 +- 7 files changed, 35 insertions(+), 13 deletions(-) diff --git a/go/ai-proxy/go.mod b/go/ai-proxy/go.mod index 5d6094334d..94169efefe 100644 --- a/go/ai-proxy/go.mod +++ b/go/ai-proxy/go.mod @@ -36,14 +36,26 @@ require ( github.com/aws/aws-sdk-go-v2/service/sts v1.42.2 // indirect github.com/aws/smithy-go v1.27.1 // indirect github.com/bahlo/generic-list-go v0.2.0 // indirect + github.com/beorn7/perks v1.0.1 // indirect github.com/buger/jsonparser v1.1.2 // indirect github.com/cenkalti/backoff v2.2.1+incompatible // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/go-logr/logr v1.4.3 // indirect - github.com/kr/pretty v0.3.1 // indirect + github.com/golang-jwt/jwt/v5 v5.3.1 // indirect + github.com/jpillora/backoff v1.0.0 // indirect github.com/mailru/easyjson v0.9.1 // indirect + github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect + github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f // indirect + github.com/prometheus/client_golang v1.23.2 // indirect + github.com/prometheus/client_model v0.6.2 // indirect + github.com/prometheus/common v0.67.5 // indirect + github.com/prometheus/procfs v0.20.1 // indirect github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect + go.yaml.in/yaml/v2 v2.4.4 // indirect golang.org/x/crypto v0.53.0 // indirect + golang.org/x/net v0.56.0 // indirect golang.org/x/sys v0.46.0 // indirect golang.org/x/text v0.39.0 // indirect + google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go/ai-proxy/go.sum b/go/ai-proxy/go.sum index d2aaae9143..d74c0e622a 100644 --- a/go/ai-proxy/go.sum +++ b/go/ai-proxy/go.sum @@ -36,37 +36,44 @@ github.com/aws/smithy-go v1.27.1 h1:4T340VFndXtADGF52gYa1POyL7s9E4Z1OeZ1hCscIw8= github.com/aws/smithy-go v1.27.1/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= +github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= github.com/buger/jsonparser v1.1.2 h1:frqHqw7otoVbk5M8LlE/L7HTnIq2v9RX6EJ48i9AxJk= github.com/buger/jsonparser v1.1.2/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= github.com/cenkalti/backoff v2.2.1+incompatible h1:tNowT99t7UNflLxfYYSlKYsBpXdEet03Pg2g16Swow4= github.com/cenkalti/backoff v2.2.1+incompatible/go.mod h1:90ReRw6GdpyfrHakVjL/QHaoyV4aDUVVkXQJJJ3NXXM= -github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= +github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63YCY= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/gorilla/mux v1.8.1 h1:TuBL49tXwgrFYWhqrNgrUNEY92u81SPhu7sTdzQEiWY= github.com/gorilla/mux v1.8.1/go.mod h1:AKf9I4AEqPTmMytcMc0KkNouC66V3BtZ4qD5fmWSiMQ= +github.com/jpillora/backoff v1.0.0 h1:uvFg412JmmHBHw7iwprIxkPMI+sGQ4kzOWsMeHnm2EA= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/mailru/easyjson v0.9.1 h1:LbtsOm5WAswyWbvTEOqhypdPeZzHavpZx96/n553mR8= github.com/mailru/easyjson v0.9.1/go.mod h1:1+xMtQp2MRNVL/V1bOzuP3aP8VNwRW55fQUto+XFtTU= +github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= +github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f h1:KUppIJq7/+SVif2QVs3tOP0zanoHgBEVAwHxUSIzRqU= github.com/ollama/ollama v0.21.1 h1:kF0PLEBucnoRDYADS1NmcQXGqC/B7M9WJy8ZyNIr/NA= github.com/ollama/ollama v0.21.1/go.mod h1:274niu48upWz/M7vL53i1WFe+TJRRw5oo4GiacbIYrA= -github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= -github.com/rogpeppe/go-internal v1.15.0 h1:D0RCU5rMAp+SpgkiNdrjfJ+LX4J1M32V2NeCY7EJ6hc= -github.com/rogpeppe/go-internal v1.15.0/go.mod h1:DrUVZyrJU+txYW5/1kwtXQSMFio52ZOxX7yM1VHvnxs= +github.com/prometheus/client_golang v1.23.2 h1:Je96obch5RDVy3FDMndoUsjAhG5Edi49h0RJWRi/o0o= +github.com/prometheus/client_model v0.6.2 h1:oBsgwpGs7iVziMvrGhE53c/GrLUsZdHnqNwqPLxwZyk= +github.com/prometheus/common v0.67.5 h1:pIgK94WWlQt1WLwAC5j2ynLaBRDiinoAb86HZHTUGI4= +github.com/prometheus/procfs v0.20.1 h1:XwbrGOIplXW/AU3YhIhLODXMJYyC1isLFfYCsTEycfc= github.com/prometheus/sigv4 v0.4.1 h1:EIc3j+8NBea9u1iV6O5ZAN8uvPq2xOIUPcqCTivHuXs= github.com/prometheus/sigv4 v0.4.1/go.mod h1:eu+ZbRvsc5TPiHwqh77OWuCnWK73IdkETYY46P4dXOU= +github.com/rogpeppe/go-internal v1.15.0 h1:D0RCU5rMAp+SpgkiNdrjfJ+LX4J1M32V2NeCY7EJ6hc= +github.com/rogpeppe/go-internal v1.15.0/go.mod h1:DrUVZyrJU+txYW5/1kwtXQSMFio52ZOxX7yM1VHvnxs= github.com/samber/lo v1.53.0 h1:t975lj2py4kJPQ6haz1QMgtId2gtmfktACxIXArw3HM= github.com/samber/lo v1.53.0/go.mod h1:4+MXEGsJzbKGaUEQFKBq2xtfuznW9oz/WrgyzMzRoM0= github.com/spf13/pflag v1.0.10 h1:4EBh2KAYBwaONj6b2Ye1GiHfwjqyROoF4RwYO+vPwFk= @@ -77,8 +84,11 @@ github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/ github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= +go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.yaml.in/yaml/v2 v2.4.4 h1:tuyd0P+2Ont/d6e2rl3be67goVK4R6deVxCUX5vyPaQ= golang.org/x/crypto v0.53.0 h1:QZ4Muo8THX6CizN2vPPd5fBGHyogrdK9fG4wLPFUsto= golang.org/x/crypto v0.53.0/go.mod h1:DNLU434OwVakk9PzuwV8w62mAJpRJL3vsgcfp4Qnsio= +golang.org/x/net v0.56.0 h1:Rw8j/hFzGvJUZwNBXnAtf5sVDVt+65SK2C7IxCxZt5o= 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/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw= @@ -87,9 +97,11 @@ golang.org/x/term v0.44.0 h1:0rLvDRCtNj0gZkyIXhCyOb2OAzEhLVqc4B+hrsBhrmc= golang.org/x/term v0.44.0/go.mod h1:7ze4MdzUzLXpSAoFP1H0bOI9aXDqveSvatT5vKcFh2Y= golang.org/x/text v0.39.0 h1:UbZz4pLOvn600D6Oh6GGEI6VAmndrEBLv8/6BEXzyus= golang.org/x/text v0.39.0/go.mod h1:3UwRclnC2g0TU9x8PZiyfOajCd1zaUNHF9cvqcQZ+ZM= +google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af h1:+5/Sw3GsDNlEmu7TfklWKPdQ0Ykja5VEmq2i817+jbI= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= +gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= k8s.io/klog/v2 v2.140.0 h1:Tf+J3AH7xnUzZyVVXhTgGhEKnFqye14aadWv7bzXdzc= diff --git a/go/client/go.mod b/go/client/go.mod index 3b0b1f4b99..3a455a64fb 100644 --- a/go/client/go.mod +++ b/go/client/go.mod @@ -86,7 +86,7 @@ require ( golang.org/x/sys v0.46.0 // indirect golang.org/x/time v0.15.0 // indirect golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect - google.golang.org/genproto/googleapis/rpc v0.0.0-20260406210006-6f92a3bedf2d // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 // indirect google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af // indirect gopkg.in/ini.v1 v1.67.2 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect diff --git a/go/client/go.sum b/go/client/go.sum index a2dfe3553e..de06e0c4c6 100644 --- a/go/client/go.sum +++ b/go/client/go.sum @@ -295,8 +295,7 @@ golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da h1:noIWHXmPHxILtqtCOPIhS golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da/go.mod h1:NDW/Ps6MPRej6fsCIbMTohpP40sJ/P/vI1MoTEGwX90= google.golang.org/appengine v1.6.8 h1:IhEN5q69dyKagZPYMSdIjS2HqprW324FRQZJcGqPAsM= google.golang.org/appengine v1.6.8/go.mod h1:1jJ3jBArFh5pcgW8gCtRJnepW8FzD1V44FJffLiz/Ds= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260406210006-6f92a3bedf2d h1:wT2n40TBqFY6wiwazVK9/iTWbsQrgk5ZfCSVFLO9LQA= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260406210006-6f92a3bedf2d/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 h1:RmoJA1ujG+/lRGNfUnOMfhCy5EipVMyvUE+KNbPbTlw= google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af h1:+5/Sw3GsDNlEmu7TfklWKPdQ0Ykja5VEmq2i817+jbI= google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/go/deployment-operator/terratest/go.mod b/go/deployment-operator/terratest/go.mod index 3d3c9c6ae3..3bfb918564 100644 --- a/go/deployment-operator/terratest/go.mod +++ b/go/deployment-operator/terratest/go.mod @@ -181,7 +181,7 @@ require ( golang.org/x/text v0.39.0 // indirect golang.org/x/time v0.15.0 // indirect golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect - google.golang.org/genproto/googleapis/rpc v0.0.0-20260406210006-6f92a3bedf2d // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 // indirect google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af // indirect gopkg.in/evanphx/json-patch.v4 v4.13.0 // indirect gopkg.in/inf.v0 v0.9.1 // indirect diff --git a/go/deployment-operator/terratest/go.sum b/go/deployment-operator/terratest/go.sum index feafeaf888..af9d3a03a3 100644 --- a/go/deployment-operator/terratest/go.sum +++ b/go/deployment-operator/terratest/go.sum @@ -513,8 +513,7 @@ golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da h1:noIWHXmPHxILtqtCOPIhS golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da/go.mod h1:NDW/Ps6MPRej6fsCIbMTohpP40sJ/P/vI1MoTEGwX90= google.golang.org/appengine v1.6.8 h1:IhEN5q69dyKagZPYMSdIjS2HqprW324FRQZJcGqPAsM= google.golang.org/appengine v1.6.8/go.mod h1:1jJ3jBArFh5pcgW8gCtRJnepW8FzD1V44FJffLiz/Ds= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260406210006-6f92a3bedf2d h1:wT2n40TBqFY6wiwazVK9/iTWbsQrgk5ZfCSVFLO9LQA= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260406210006-6f92a3bedf2d/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 h1:RmoJA1ujG+/lRGNfUnOMfhCy5EipVMyvUE+KNbPbTlw= google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af h1:+5/Sw3GsDNlEmu7TfklWKPdQ0Ykja5VEmq2i817+jbI= google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/go/helm-test/go.mod b/go/helm-test/go.mod index 9ef3e7985a..eadbf4871d 100644 --- a/go/helm-test/go.mod +++ b/go/helm-test/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/helm-test -go 1.26.5 +go 1.26.6 require ( github.com/onsi/ginkgo/v2 v2.28.1 From c17503d1206800489c63bad45aa7c0e58350d818 Mon Sep 17 00:00:00 2001 From: Sebastian Florek Date: Thu, 20 Aug 2026 17:03:24 +0200 Subject: [PATCH 05/15] chore(dependencies): bump Go version across all modules - Updated Go version from 1.26.5 to 1.26.6 in various modules - Ensured consistent Go version across all sub-projects --- go/ai-proxy/go.mod | 2 +- go/client/go.mod | 2 +- go/controller/go.mod | 2 +- go/datastore/go.mod | 2 +- go/demo/flaky-service/go.mod | 2 +- go/deployment-operator/go.mod | 2 +- go/deployment-operator/terratest/go.mod | 2 +- go/helmchartutils/go.mod | 2 +- go/kubernetes-agent/api/go.mod | 2 +- go/kubernetes-agent/kas/go.mod | 2 +- go/nexus/go.mod | 2 +- go/observability-proxy/go.mod | 2 +- go/oci-auth/go.mod | 2 +- go/polly/go.mod | 2 +- go/sign/go.mod | 2 +- go/tools/go.mod | 2 +- 16 files changed, 16 insertions(+), 16 deletions(-) diff --git a/go/ai-proxy/go.mod b/go/ai-proxy/go.mod index 94169efefe..44032a7463 100644 --- a/go/ai-proxy/go.mod +++ b/go/ai-proxy/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/ai-proxy -go 1.26.5 +go 1.26.6 replace github.com/pluralsh/console/go/polly => ../polly diff --git a/go/client/go.mod b/go/client/go.mod index 3a455a64fb..72d421453a 100644 --- a/go/client/go.mod +++ b/go/client/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/client -go 1.26.5 +go 1.26.6 require ( github.com/99designs/gqlgen v0.17.78 diff --git a/go/controller/go.mod b/go/controller/go.mod index 9a5e833d5c..4325e5d689 100644 --- a/go/controller/go.mod +++ b/go/controller/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/controller -go 1.26.5 +go 1.26.6 replace ( github.com/pluralsh/console/go/client => ../client diff --git a/go/datastore/go.mod b/go/datastore/go.mod index 58ea8edf37..6b903936cb 100644 --- a/go/datastore/go.mod +++ b/go/datastore/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/datastore -go 1.26.5 +go 1.26.6 replace ( github.com/pluralsh/console/go/client => ../client diff --git a/go/demo/flaky-service/go.mod b/go/demo/flaky-service/go.mod index 72aabb4930..5d4349cd47 100644 --- a/go/demo/flaky-service/go.mod +++ b/go/demo/flaky-service/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/demo/flaky-service -go 1.26.5 +go 1.26.6 require github.com/prometheus/client_golang v1.23.2 diff --git a/go/deployment-operator/go.mod b/go/deployment-operator/go.mod index 4da70cf306..3341879ff8 100644 --- a/go/deployment-operator/go.mod +++ b/go/deployment-operator/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/deployment-operator -go 1.26.5 +go 1.26.6 replace ( github.com/containerd/containerd => github.com/containerd/containerd v1.7.33 diff --git a/go/deployment-operator/terratest/go.mod b/go/deployment-operator/terratest/go.mod index 3bfb918564..77c5f42627 100644 --- a/go/deployment-operator/terratest/go.mod +++ b/go/deployment-operator/terratest/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/deployment-operator/terratest -go 1.26.5 +go 1.26.6 require ( github.com/gruntwork-io/terratest v1.0.1 diff --git a/go/helmchartutils/go.mod b/go/helmchartutils/go.mod index 45bcdf3118..a14fc5f943 100644 --- a/go/helmchartutils/go.mod +++ b/go/helmchartutils/go.mod @@ -1,3 +1,3 @@ module github.com/pluralsh/console/go/helmchartutils -go 1.26.5 +go 1.26.6 diff --git a/go/kubernetes-agent/api/go.mod b/go/kubernetes-agent/api/go.mod index 0bc7c7bce4..5bf74a10bf 100644 --- a/go/kubernetes-agent/api/go.mod +++ b/go/kubernetes-agent/api/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/kubernetes-agent/api -go 1.26.5 +go 1.26.6 require ( github.com/Yiling-J/theine-go v0.6.0 diff --git a/go/kubernetes-agent/kas/go.mod b/go/kubernetes-agent/kas/go.mod index aa0b4639ed..077fae9413 100644 --- a/go/kubernetes-agent/kas/go.mod +++ b/go/kubernetes-agent/kas/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/kubernetes-agent -go 1.26.5 +go 1.26.6 replace ( github.com/pluralsh/console/go/client => ../../client diff --git a/go/nexus/go.mod b/go/nexus/go.mod index d012eefe83..b478ab520a 100644 --- a/go/nexus/go.mod +++ b/go/nexus/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/nexus -go 1.26.5 +go 1.26.6 require ( github.com/bytedance/sonic v1.15.0 diff --git a/go/observability-proxy/go.mod b/go/observability-proxy/go.mod index e01023e260..74d18cc0d2 100644 --- a/go/observability-proxy/go.mod +++ b/go/observability-proxy/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/observability-proxy -go 1.26.5 +go 1.26.6 require ( golang.org/x/sync v0.22.0 diff --git a/go/oci-auth/go.mod b/go/oci-auth/go.mod index a535253269..ba29efd2ae 100644 --- a/go/oci-auth/go.mod +++ b/go/oci-auth/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/oci-auth -go 1.26.5 +go 1.26.6 require ( github.com/Azure/azure-sdk-for-go/sdk/azcore v1.21.1 diff --git a/go/polly/go.mod b/go/polly/go.mod index ef28240bff..7074063190 100644 --- a/go/polly/go.mod +++ b/go/polly/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/polly -go 1.26.5 +go 1.26.6 require ( dario.cat/mergo v1.0.2 diff --git a/go/sign/go.mod b/go/sign/go.mod index 33c828bac5..b569caafa7 100644 --- a/go/sign/go.mod +++ b/go/sign/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/sign -go 1.26.5 +go 1.26.6 require github.com/aws/aws-lambda-go v1.54.0 diff --git a/go/tools/go.mod b/go/tools/go.mod index 9f05ac013b..ac1ad8f725 100644 --- a/go/tools/go.mod +++ b/go/tools/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/tools -go 1.26.5 +go 1.26.6 require ( github.com/99designs/gqlgen v0.17.78 From 8ad5b0b2be7bd83ff37fc34bcd7d945bd0f28984 Mon Sep 17 00:00:00 2001 From: Sebastian Florek Date: Thu, 20 Aug 2026 17:10:15 +0200 Subject: [PATCH 06/15] chore(docker): bump Go version in Dockerfiles to 1.26.6 - Updated Go version from 1.26.5 to 1.26.6 for consistency across all Dockerfiles - Aligned version in various service and component Dockerfiles including datastore, deployment-operator, ai-proxy, and others --- go/ai-proxy/Dockerfile | 2 +- go/controller/Dockerfile | 2 +- go/datastore/Dockerfile | 2 +- go/demo/flaky-service/Dockerfile | 2 +- go/deployment-operator/Dockerfile | 2 +- .../dockerfiles/mcpserver/terraform-server/Dockerfile | 2 +- go/kubernetes-agent/api/Dockerfile | 2 +- go/kubernetes-agent/hack/docker/Dockerfile | 2 +- go/kubernetes-agent/hack/docker/Dockerfile.agentk | 2 +- go/nexus/Dockerfile | 2 +- go/observability-proxy/Dockerfile | 2 +- go/oci-auth/Dockerfile | 2 +- 12 files changed, 12 insertions(+), 12 deletions(-) diff --git a/go/ai-proxy/Dockerfile b/go/ai-proxy/Dockerfile index 9c0d20243e..3340d4993e 100644 --- a/go/ai-proxy/Dockerfile +++ b/go/ai-proxy/Dockerfile @@ -1,5 +1,5 @@ # Build the binary -FROM golang:1.26.5 as builder +FROM golang:1.26.6 as builder ARG TARGETOS ARG TARGETARCH ARG VERSION diff --git a/go/controller/Dockerfile b/go/controller/Dockerfile index 6876af21ee..c42372a1f6 100644 --- a/go/controller/Dockerfile +++ b/go/controller/Dockerfile @@ -1,5 +1,5 @@ # Build the manager binary -FROM golang:1.26.5 AS builder +FROM golang:1.26.6 AS builder ARG TARGETOS ARG TARGETARCH ARG GIT_COMMIT=unknown diff --git a/go/datastore/Dockerfile b/go/datastore/Dockerfile index ba96e3e7e9..998a0e9a96 100644 --- a/go/datastore/Dockerfile +++ b/go/datastore/Dockerfile @@ -1,5 +1,5 @@ # Build the manager binary -FROM golang:1.26.5 AS builder +FROM golang:1.26.6 AS builder ARG TARGETOS ARG TARGETARCH diff --git a/go/demo/flaky-service/Dockerfile b/go/demo/flaky-service/Dockerfile index 6d5bc9955a..d1da224db5 100644 --- a/go/demo/flaky-service/Dockerfile +++ b/go/demo/flaky-service/Dockerfile @@ -1,5 +1,5 @@ # Step 1: Build the Go binary -FROM golang:1.26.5 AS build +FROM golang:1.26.6 AS build # Set the Current Working Directory inside the container WORKDIR /app diff --git a/go/deployment-operator/Dockerfile b/go/deployment-operator/Dockerfile index c3060f1054..abc74dea9e 100644 --- a/go/deployment-operator/Dockerfile +++ b/go/deployment-operator/Dockerfile @@ -1,4 +1,4 @@ -FROM golang:1.26.5-alpine3.23 AS builder +FROM golang:1.26.6-alpine3.23 AS builder ARG HELM_VERSION=v3.21.2 ARG TARGETARCH diff --git a/go/deployment-operator/dockerfiles/mcpserver/terraform-server/Dockerfile b/go/deployment-operator/dockerfiles/mcpserver/terraform-server/Dockerfile index 9845de3f90..051a724fb3 100644 --- a/go/deployment-operator/dockerfiles/mcpserver/terraform-server/Dockerfile +++ b/go/deployment-operator/dockerfiles/mcpserver/terraform-server/Dockerfile @@ -1,4 +1,4 @@ -FROM golang:1.26.5-alpine AS builder +FROM golang:1.26.6-alpine AS builder WORKDIR /workspace diff --git a/go/kubernetes-agent/api/Dockerfile b/go/kubernetes-agent/api/Dockerfile index 195fd6bac0..cf715a248c 100644 --- a/go/kubernetes-agent/api/Dockerfile +++ b/go/kubernetes-agent/api/Dockerfile @@ -27,7 +27,7 @@ RUN adduser \ --uid "${UID}" \ "${USER}" -FROM golang:1.26.5-alpine AS builder +FROM golang:1.26.6-alpine AS builder ARG TARGETARCH ARG TARGETOS diff --git a/go/kubernetes-agent/hack/docker/Dockerfile b/go/kubernetes-agent/hack/docker/Dockerfile index 39ca83e6ed..5b6c53c147 100644 --- a/go/kubernetes-agent/hack/docker/Dockerfile +++ b/go/kubernetes-agent/hack/docker/Dockerfile @@ -3,7 +3,7 @@ # Usage: docker run [api|kas|agentk] [args...] # Builder stage for all binaries -FROM golang:1.26.5-alpine AS builder +FROM golang:1.26.6-alpine AS builder ARG TARGETARCH ARG TARGETOS diff --git a/go/kubernetes-agent/hack/docker/Dockerfile.agentk b/go/kubernetes-agent/hack/docker/Dockerfile.agentk index 61fcbcf9ae..df112223fc 100644 --- a/go/kubernetes-agent/hack/docker/Dockerfile.agentk +++ b/go/kubernetes-agent/hack/docker/Dockerfile.agentk @@ -3,7 +3,7 @@ # Usage: docker run [api|kas|agentk] [args...] # Builder stage for all binaries -FROM golang:1.26.5-alpine AS builder +FROM golang:1.26.6-alpine AS builder ARG TARGETARCH ARG TARGETOS diff --git a/go/nexus/Dockerfile b/go/nexus/Dockerfile index 05ee7c252d..05f4cb1f62 100644 --- a/go/nexus/Dockerfile +++ b/go/nexus/Dockerfile @@ -1,5 +1,5 @@ # Build stage -FROM golang:1.26.5-alpine AS builder +FROM golang:1.26.6-alpine AS builder WORKDIR /build diff --git a/go/observability-proxy/Dockerfile b/go/observability-proxy/Dockerfile index a3f1bca9ed..b5361c4abf 100644 --- a/go/observability-proxy/Dockerfile +++ b/go/observability-proxy/Dockerfile @@ -1,4 +1,4 @@ -FROM golang:1.26.5-alpine AS builder +FROM golang:1.26.6-alpine AS builder WORKDIR /build diff --git a/go/oci-auth/Dockerfile b/go/oci-auth/Dockerfile index cd2cf771aa..1c1a7ed09a 100644 --- a/go/oci-auth/Dockerfile +++ b/go/oci-auth/Dockerfile @@ -1,4 +1,4 @@ -FROM golang:1.26.5 as builder +FROM golang:1.26.6 as builder ARG TARGETOS ARG TARGETARCH From 93d827b8233a1f251d7007a70fee992c6a739657 Mon Sep 17 00:00:00 2001 From: Sebastian Florek Date: Thu, 20 Aug 2026 17:15:11 +0200 Subject: [PATCH 07/15] chore(docker): update Go version to 1.26.6 in Dockerfiles - Bumped Go version from 1.26.5 to 1.26.6 across all relevant Dockerfiles - Ensured version consistency in Kubernetes agent, cloud-query, and deployment-operator components - Modified GitHub workflow to align Go version with Dockerfile updates --- .../workflows/deployment-operator-cd-sentinel-harness.yaml | 4 ++-- go/build.Dockerfile | 2 +- go/cloud-query/db.Dockerfile | 2 +- .../dockerfiles/agent-harness/base.Dockerfile | 4 ++-- go/deployment-operator/dockerfiles/harness/base.Dockerfile | 2 +- .../dockerfiles/sentinel-harness/base.Dockerfile | 6 +++--- go/kubernetes-agent/api/dev.Dockerfile | 4 ++-- go/kubernetes-agent/hack/docker/dev.Dockerfile | 2 +- go/test.Dockerfile | 2 +- 9 files changed, 14 insertions(+), 14 deletions(-) diff --git a/.github/workflows/deployment-operator-cd-sentinel-harness.yaml b/.github/workflows/deployment-operator-cd-sentinel-harness.yaml index dc95878d3e..d95e9c2e41 100644 --- a/.github/workflows/deployment-operator-cd-sentinel-harness.yaml +++ b/.github/workflows/deployment-operator-cd-sentinel-harness.yaml @@ -202,6 +202,6 @@ jobs: cache-from: type=gha cache-to: type=gha,mode=max build-args: | - GO_VERSION=1.26.5 + GO_VERSION=1.26.6 SENTINEL_HARNESS_BASE_IMAGE_REPO=ghcr.io/pluralsh/sentinel-harness-base - SENTINEL_HARNESS_BASE_IMAGE_TAG=${{ needs.publish-base-image.outputs.version }} \ No newline at end of file + SENTINEL_HARNESS_BASE_IMAGE_TAG=${{ needs.publish-base-image.outputs.version }} diff --git a/go/build.Dockerfile b/go/build.Dockerfile index 128392af34..046d4b5c97 100644 --- a/go/build.Dockerfile +++ b/go/build.Dockerfile @@ -1,4 +1,4 @@ -FROM golang:1.26.5 +FROM golang:1.26.6 ARG MODULE_PATH diff --git a/go/cloud-query/db.Dockerfile b/go/cloud-query/db.Dockerfile index ed1ce3a24d..2f2f3a028a 100644 --- a/go/cloud-query/db.Dockerfile +++ b/go/cloud-query/db.Dockerfile @@ -1,7 +1,7 @@ ARG POSTGRES_MAJOR_VERSION=15 ARG POSTGRES_VERSION=${POSTGRES_MAJOR_VERSION}.18 -FROM golang:1.26.5 AS libraries +FROM golang:1.26.6 AS libraries # Configure versions for Steampipe extensions # Do not use latest versions here, as they may not be compatible diff --git a/go/deployment-operator/dockerfiles/agent-harness/base.Dockerfile b/go/deployment-operator/dockerfiles/agent-harness/base.Dockerfile index 30b77cf880..15ed04ec1f 100644 --- a/go/deployment-operator/dockerfiles/agent-harness/base.Dockerfile +++ b/go/deployment-operator/dockerfiles/agent-harness/base.Dockerfile @@ -1,4 +1,4 @@ -FROM golang:1.26.5-alpine AS builder +FROM golang:1.26.6-alpine AS builder ARG TARGETARCH ARG TARGETOS @@ -149,4 +149,4 @@ WORKDIR /plural HEALTHCHECK --interval=60s --timeout=10s --start-period=30s --retries=5 \ CMD kill -0 1 || exit 1 -ENTRYPOINT ["/entrypoint.sh", "/agent-harness", "--working-dir=/plural"] \ No newline at end of file +ENTRYPOINT ["/entrypoint.sh", "/agent-harness", "--working-dir=/plural"] diff --git a/go/deployment-operator/dockerfiles/harness/base.Dockerfile b/go/deployment-operator/dockerfiles/harness/base.Dockerfile index ae932b838e..729c6a5e3a 100644 --- a/go/deployment-operator/dockerfiles/harness/base.Dockerfile +++ b/go/deployment-operator/dockerfiles/harness/base.Dockerfile @@ -1,4 +1,4 @@ -FROM golang:1.26.5-alpine AS builder +FROM golang:1.26.6-alpine AS builder ARG TARGETARCH ARG TARGETOS diff --git a/go/deployment-operator/dockerfiles/sentinel-harness/base.Dockerfile b/go/deployment-operator/dockerfiles/sentinel-harness/base.Dockerfile index 6a256d445c..965b12a44d 100644 --- a/go/deployment-operator/dockerfiles/sentinel-harness/base.Dockerfile +++ b/go/deployment-operator/dockerfiles/sentinel-harness/base.Dockerfile @@ -1,4 +1,4 @@ -FROM golang:1.26.5-alpine AS builder +FROM golang:1.26.6-alpine AS builder ARG TARGETARCH ARG TARGETOS @@ -31,7 +31,7 @@ RUN CGO_ENABLED=0 \ -o /sentinel-harness \ cmd/sentinel-harness/main.go -FROM golang:1.26.5-alpine AS final +FROM golang:1.26.6-alpine AS final ARG TARGETARCH ARG TARGETOS @@ -55,4 +55,4 @@ WORKDIR /plural USER 65532:65532 HEALTHCHECK --interval=60s --timeout=10s --start-period=30s --retries=5 \ - CMD kill -0 1 || exit 1 \ No newline at end of file + CMD kill -0 1 || exit 1 diff --git a/go/kubernetes-agent/api/dev.Dockerfile b/go/kubernetes-agent/api/dev.Dockerfile index 38ad386373..e884849dd7 100644 --- a/go/kubernetes-agent/api/dev.Dockerfile +++ b/go/kubernetes-agent/api/dev.Dockerfile @@ -14,11 +14,11 @@ # ! Context expected to be set to "modules" dir ! -FROM golang:1.26.5-alpine AS AIR +FROM golang:1.26.6-alpine AS AIR RUN go install github.com/air-verse/air@latest -FROM golang:1.26.5-alpine +FROM golang:1.26.6-alpine # Copy air binary COPY --from=AIR $GOPATH/bin/air $GOPATH/bin/air diff --git a/go/kubernetes-agent/hack/docker/dev.Dockerfile b/go/kubernetes-agent/hack/docker/dev.Dockerfile index cb17f6e79b..aaa836a9a8 100644 --- a/go/kubernetes-agent/hack/docker/dev.Dockerfile +++ b/go/kubernetes-agent/hack/docker/dev.Dockerfile @@ -6,7 +6,7 @@ FROM busybox:uclibc AS busybox # Builder stage for all binaries with debug support -FROM golang:1.26.5-alpine AS builder +FROM golang:1.26.6-alpine AS builder ARG TARGETARCH ARG TARGETOS diff --git a/go/test.Dockerfile b/go/test.Dockerfile index d9f05d7bcb..1dd1edbdf7 100644 --- a/go/test.Dockerfile +++ b/go/test.Dockerfile @@ -1,4 +1,4 @@ -FROM golang:1.26.5 +FROM golang:1.26.6 ARG MODULE_PATH From f5ac395a277219ceb6e179a88530d311af784575 Mon Sep 17 00:00:00 2001 From: Sebastian Florek Date: Fri, 21 Aug 2026 11:09:30 +0200 Subject: [PATCH 08/15] refactor(errors): rename struct in Python contract error logic - Changed `truncatedCause` to `truncatedCauseError` for clarity and consistency - Simplified detail truncation logic in `truncateDetail` function --- .../tools/python/internal/contract/errors.go | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/go/cloud-query/internal/tools/python/internal/contract/errors.go b/go/cloud-query/internal/tools/python/internal/contract/errors.go index 0fef58a678..7561f2f626 100644 --- a/go/cloud-query/internal/tools/python/internal/contract/errors.go +++ b/go/cloud-query/internal/tools/python/internal/contract/errors.go @@ -117,7 +117,7 @@ func boundedCause(cause error) error { if len(cause.Error()) <= MaxDetailBytes { return cause } - return truncatedCause{detail: truncateDetail(cause.Error()), cause: cause} + return truncatedCauseError{detail: truncateDetail(cause.Error()), cause: cause} } func truncateDetail(detail string) string { @@ -125,17 +125,15 @@ func truncateDetail(detail string) string { return detail } limit := MaxDetailBytes - len(detailTruncationMarker) - if limit < 0 { - limit = 0 - } + return detail[:limit] + detailTruncationMarker } -type truncatedCause struct { +type truncatedCauseError struct { detail string cause error } -func (e truncatedCause) Error() string { return e.detail } +func (e truncatedCauseError) Error() string { return e.detail } -func (e truncatedCause) Unwrap() error { return e.cause } +func (e truncatedCauseError) Unwrap() error { return e.cause } From 10925d5272f17d1da844dae35b35784a5e97fae2 Mon Sep 17 00:00:00 2001 From: Sebastian Florek Date: Fri, 21 Aug 2026 11:21:26 +0200 Subject: [PATCH 09/15] chore(docker): set ownership for copied files in terratest Dockerfile - Added `--chown=65532:65532` to `COPY` command in `sentinel-harness/terratest.Dockerfile` to ensure correct file ownership within the container. --- .../dockerfiles/sentinel-harness/terratest.Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/go/deployment-operator/dockerfiles/sentinel-harness/terratest.Dockerfile b/go/deployment-operator/dockerfiles/sentinel-harness/terratest.Dockerfile index 398c1ba0b9..039063e78c 100644 --- a/go/deployment-operator/dockerfiles/sentinel-harness/terratest.Dockerfile +++ b/go/deployment-operator/dockerfiles/sentinel-harness/terratest.Dockerfile @@ -16,7 +16,7 @@ ENV CGO_ENABLED=0 \ WORKDIR /sentinel/terratest # Copy test files -COPY deployment-operator/terratest ./ +COPY --chown=65532:65532 deployment-operator/terratest ./ RUN go mod download From 5eb4203cabc0bfd3f78a0a6506c13fd2af827d94 Mon Sep 17 00:00:00 2001 From: Sebastian Florek Date: Fri, 21 Aug 2026 11:49:53 +0200 Subject: [PATCH 10/15] fix(tests): normalize newline character in Python test - Adjusted `stdout` newline representation in `python_test.exs` to use consistent newline character `\n` --- test/console/ai/tools/workbench/python_test.exs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/console/ai/tools/workbench/python_test.exs b/test/console/ai/tools/workbench/python_test.exs index be800a26f4..ed33300576 100644 --- a/test/console/ai/tools/workbench/python_test.exs +++ b/test/console/ai/tools/workbench/python_test.exs @@ -30,7 +30,7 @@ defmodule Console.AI.Tools.Workbench.PythonTest do assert request.input_json == ~s({"first":20,"second":22}) assert opts == Client.cloud_query_rpc_opts() - {:ok, %RunPythonOutput{result_json: ~s({"total":42}), stdout: "calculated total\\n"}} + {:ok, %RunPythonOutput{result_json: ~s({"total":42}), stdout: "calculated total\n"}} end) assert {:ok, %{result: %{"total" => 42}, stdout: "calculated total\n"}} = From c4e448490bdc3ea72344ab3a797ed68e33c5103e Mon Sep 17 00:00:00 2001 From: Sebastian Florek Date: Fri, 21 Aug 2026 12:48:34 +0200 Subject: [PATCH 11/15] test(memory_engine): add unit test for tool result serialization - Added `MemoryEngineTest` in `memory_engine_test.exs` to verify structured tool result serialization before next completion - Utilized `Mimic` for mocking provider calls - Included `MapTool` as an embedded schema for testing purposes - Enhanced `tool_msg` logic for proper result serialization in `memory_engine.ex` --- lib/console/ai/chat/memory_engine.ex | 9 +++- test/console/ai/chat/memory_engine_test.exs | 50 +++++++++++++++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) create mode 100644 test/console/ai/chat/memory_engine_test.exs diff --git a/lib/console/ai/chat/memory_engine.ex b/lib/console/ai/chat/memory_engine.ex index 96d7ffac93..1d0a6c33b8 100644 --- a/lib/console/ai/chat/memory_engine.ex +++ b/lib/console/ai/chat/memory_engine.ex @@ -7,6 +7,7 @@ defmodule Console.AI.Chat.MemoryEngine do alias Console.AI.{Provider, Tool} alias Console.AI.Chat.EnabledTools alias Console.AI.Tools.{EnableTools, ToolSearch} + alias Console.AI.Tools.Workbench.Output require Logger @type t :: %__MODULE__{} @@ -180,15 +181,19 @@ defmodule Console.AI.Chat.MemoryEngine do defp tool_msg(content, id, name, args, fun, attrs \\ %{}) defp tool_msg(content, id, name, args, _, attrs) when is_binary(content), do: {:tool, content, %{call_id: id, name: name, arguments: args, attributes: attrs}} - defp tool_msg(%{content: content} = msg, id, name, args, _, _), + defp tool_msg(%{content: content} = msg, id, name, args, _, _) when is_binary(content), do: {:tool, content, %{call_id: id, name: name, arguments: args, attributes: Map.delete(msg, :content)}} defp tool_msg(result, id, name, args, fmt, attrs) when is_function(fmt, 1) do case fmt.(result) do content when is_binary(content) -> {result, {:tool, content, %{call_id: id, name: name, arguments: args, attributes: attrs}}} - _ -> result + _ -> {result, tool_msg(tool_result_content(result), id, name, args, fmt, attrs)} end end + defp tool_result_content(result), do: tool_result_content(result, Output.json(result)) + defp tool_result_content(_, {:ok, content}), do: content + defp tool_result_content(result, {:error, _}), do: inspect(result) + defp msg({res, {:tool, _, _}}, :result), do: res defp msg({_, {:tool, _, _} = tool}, :tool), do: tool defp msg(pass, _), do: pass diff --git a/test/console/ai/chat/memory_engine_test.exs b/test/console/ai/chat/memory_engine_test.exs new file mode 100644 index 0000000000..04e534f5e9 --- /dev/null +++ b/test/console/ai/chat/memory_engine_test.exs @@ -0,0 +1,50 @@ +defmodule Console.AI.Chat.MemoryEngineTest do + use Console.DataCase, async: false + use Mimic + + alias Console.AI.{Provider, Tool} + alias Console.AI.Chat.MemoryEngine + alias Console.AI.Provider.Base + + defmodule MapTool do + use Ecto.Schema + import Ecto.Changeset + + embedded_schema do + end + + def name(), do: "map_tool" + def description(), do: "returns a structured result" + def json_schema(), do: %{"type" => "object", "properties" => %{}} + def changeset(model, attrs), do: cast(model, attrs, []) + def implement(_), do: {:ok, %{result: "ok", stdout: "done"}} + end + + setup :set_mimic_global + + test "serializes structured tool results before the next completion" do + deployment_settings( + ai: %{ + enabled: true, + provider: :openai, + openai: %{access_token: "test-key", model: "gpt-5.4-mini"} + } + ) + + expect(Provider, :completion, fn _, _ -> + {:ok, "", [%Tool{id: "call-1", name: MapTool.name(), arguments: %{}}]} + end) + + expect(Provider, :completion, fn messages, _ -> + assert {:tool, content, %{call_id: "call-1", name: "map_tool", arguments: %{}}} = List.last(messages) + assert {:ok, %{"result" => "ok", "stdout" => "done"}} = Jason.decode(content) + assert %ReqLLM.Context{} = Base.reqllm_messages(messages) + + {:ok, "done"} + end) + + assert {:ok, "done"} = + MemoryEngine.new([MapTool], 2, system_prompt: "") + |> MemoryEngine.reduce([{:user, "run the tool"}], fn messages, _ -> MemoryEngine.last_message(messages) end) + end +end From c5e67e536dbb1ec5b35a443e0fc7e9974edcaf84 Mon Sep 17 00:00:00 2001 From: Sebastian Florek Date: Fri, 21 Aug 2026 13:27:20 +0200 Subject: [PATCH 12/15] refactor(memory_engine): simplify `tool_msg` function logic - Removed unnecessary recursion in `tool_msg` function within `memory_engine.ex` for improved clarity and efficiency --- lib/console/ai/chat/memory_engine.ex | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/console/ai/chat/memory_engine.ex b/lib/console/ai/chat/memory_engine.ex index 1d0a6c33b8..fef76533d9 100644 --- a/lib/console/ai/chat/memory_engine.ex +++ b/lib/console/ai/chat/memory_engine.ex @@ -186,7 +186,7 @@ defmodule Console.AI.Chat.MemoryEngine do defp tool_msg(result, id, name, args, fmt, attrs) when is_function(fmt, 1) do case fmt.(result) do content when is_binary(content) -> {result, {:tool, content, %{call_id: id, name: name, arguments: args, attributes: attrs}}} - _ -> {result, tool_msg(tool_result_content(result), id, name, args, fmt, attrs)} + _ -> tool_msg(tool_result_content(result), id, name, args, fmt, attrs) end end From 4a658517f5fcbe73c4cd705702b94fc4f16cb175 Mon Sep 17 00:00:00 2001 From: Sebastian Florek Date: Fri, 21 Aug 2026 14:58:52 +0200 Subject: [PATCH 13/15] refactor(tests): remove MemoryEngineTest and streamline Python test outputs - Deleted `MemoryEngineTest` module for test suite simplification - Removed unused function `tool_result_content` in `memory_engine.ex` for cleaner code - Updated Python test cases to handle serialized outputs using `Output.json` - Added test for bounding oversized Python output in `python_test.exs` to ensure result size compliance --- lib/console/ai/chat/memory_engine.ex | 9 +--- lib/console/ai/tools/workbench/python.ex | 3 +- test/console/ai/chat/memory_engine_test.exs | 50 ------------------- .../ai/tools/workbench/python_test.exs | 24 +++++++-- 4 files changed, 24 insertions(+), 62 deletions(-) delete mode 100644 test/console/ai/chat/memory_engine_test.exs diff --git a/lib/console/ai/chat/memory_engine.ex b/lib/console/ai/chat/memory_engine.ex index fef76533d9..96d7ffac93 100644 --- a/lib/console/ai/chat/memory_engine.ex +++ b/lib/console/ai/chat/memory_engine.ex @@ -7,7 +7,6 @@ defmodule Console.AI.Chat.MemoryEngine do alias Console.AI.{Provider, Tool} alias Console.AI.Chat.EnabledTools alias Console.AI.Tools.{EnableTools, ToolSearch} - alias Console.AI.Tools.Workbench.Output require Logger @type t :: %__MODULE__{} @@ -181,19 +180,15 @@ defmodule Console.AI.Chat.MemoryEngine do defp tool_msg(content, id, name, args, fun, attrs \\ %{}) defp tool_msg(content, id, name, args, _, attrs) when is_binary(content), do: {:tool, content, %{call_id: id, name: name, arguments: args, attributes: attrs}} - defp tool_msg(%{content: content} = msg, id, name, args, _, _) when is_binary(content), + defp tool_msg(%{content: content} = msg, id, name, args, _, _), do: {:tool, content, %{call_id: id, name: name, arguments: args, attributes: Map.delete(msg, :content)}} defp tool_msg(result, id, name, args, fmt, attrs) when is_function(fmt, 1) do case fmt.(result) do content when is_binary(content) -> {result, {:tool, content, %{call_id: id, name: name, arguments: args, attributes: attrs}}} - _ -> tool_msg(tool_result_content(result), id, name, args, fmt, attrs) + _ -> result end end - defp tool_result_content(result), do: tool_result_content(result, Output.json(result)) - defp tool_result_content(_, {:ok, content}), do: content - defp tool_result_content(result, {:error, _}), do: inspect(result) - defp msg({res, {:tool, _, _}}, :result), do: res defp msg({_, {:tool, _, _} = tool}, :tool), do: tool defp msg(pass, _), do: pass diff --git a/lib/console/ai/tools/workbench/python.ex b/lib/console/ai/tools/workbench/python.ex index 6f1f3311b9..dbbfe07e58 100644 --- a/lib/console/ai/tools/workbench/python.ex +++ b/lib/console/ai/tools/workbench/python.ex @@ -1,6 +1,7 @@ defmodule Console.AI.Tools.Workbench.Python do use Console.AI.Tools.Workbench.Base alias CloudQuery.Client + alias Console.AI.Tools.Workbench.Output alias Toolquery.ToolQuery.Stub alias Toolquery.{RunPythonInput, RunPythonOutput} @@ -32,7 +33,7 @@ defmodule Console.AI.Tools.Workbench.Python do request = %RunPythonInput{script: code, input_json: input_json}, {:ok, %RunPythonOutput{result_json: result_json, stdout: stdout}} <- Stub.run_python(client, request, Client.cloud_query_rpc_opts()), {:ok, result} <- Jason.decode(result_json) do - {:ok, %{result: result, stdout: stdout}} + Output.json(%{result: result, stdout: stdout}) end end end diff --git a/test/console/ai/chat/memory_engine_test.exs b/test/console/ai/chat/memory_engine_test.exs deleted file mode 100644 index 04e534f5e9..0000000000 --- a/test/console/ai/chat/memory_engine_test.exs +++ /dev/null @@ -1,50 +0,0 @@ -defmodule Console.AI.Chat.MemoryEngineTest do - use Console.DataCase, async: false - use Mimic - - alias Console.AI.{Provider, Tool} - alias Console.AI.Chat.MemoryEngine - alias Console.AI.Provider.Base - - defmodule MapTool do - use Ecto.Schema - import Ecto.Changeset - - embedded_schema do - end - - def name(), do: "map_tool" - def description(), do: "returns a structured result" - def json_schema(), do: %{"type" => "object", "properties" => %{}} - def changeset(model, attrs), do: cast(model, attrs, []) - def implement(_), do: {:ok, %{result: "ok", stdout: "done"}} - end - - setup :set_mimic_global - - test "serializes structured tool results before the next completion" do - deployment_settings( - ai: %{ - enabled: true, - provider: :openai, - openai: %{access_token: "test-key", model: "gpt-5.4-mini"} - } - ) - - expect(Provider, :completion, fn _, _ -> - {:ok, "", [%Tool{id: "call-1", name: MapTool.name(), arguments: %{}}]} - end) - - expect(Provider, :completion, fn messages, _ -> - assert {:tool, content, %{call_id: "call-1", name: "map_tool", arguments: %{}}} = List.last(messages) - assert {:ok, %{"result" => "ok", "stdout" => "done"}} = Jason.decode(content) - assert %ReqLLM.Context{} = Base.reqllm_messages(messages) - - {:ok, "done"} - end) - - assert {:ok, "done"} = - MemoryEngine.new([MapTool], 2, system_prompt: "") - |> MemoryEngine.reduce([{:user, "run the tool"}], fn messages, _ -> MemoryEngine.last_message(messages) end) - end -end diff --git a/test/console/ai/tools/workbench/python_test.exs b/test/console/ai/tools/workbench/python_test.exs index ed33300576..cf17f1fcdc 100644 --- a/test/console/ai/tools/workbench/python_test.exs +++ b/test/console/ai/tools/workbench/python_test.exs @@ -3,7 +3,7 @@ defmodule Console.AI.Tools.Workbench.PythonTest do use Mimic alias CloudQuery.Client - alias Console.AI.Tools.Workbench.Python + alias Console.AI.Tools.Workbench.{Output, Python} alias Toolquery.{RunPythonOutput} alias Toolquery.ToolQuery.Stub @@ -22,7 +22,7 @@ defmodule Console.AI.Tools.Workbench.PythonTest do assert %{code: ["can't be blank"], explanation: ["can't be blank"]} = errors_on(changeset) end - test "runs Monty Python with JSON input and returns decoded output and stdout" do + test "runs Monty Python with JSON input and returns serialized output and stdout" do expect(Client, :connect, fn -> {:ok, :channel} end) expect(Stub, :run_python, fn :channel, request, opts -> @@ -33,11 +33,13 @@ defmodule Console.AI.Tools.Workbench.PythonTest do {:ok, %RunPythonOutput{result_json: ~s({"total":42}), stdout: "calculated total\n"}} end) - assert {:ok, %{result: %{"total" => 42}, stdout: "calculated total\n"}} = + assert {:ok, output} = Python.implement(%Python{ code: "output['total'] = input['first'] + input['second']", input: %{"first" => 20, "second" => 22} }) + + assert Jason.decode!(output) == %{"result" => %{"total" => 42}, "stdout" => "calculated total\n"} end test "defaults omitted input to an empty JSON object" do @@ -48,6 +50,20 @@ defmodule Console.AI.Tools.Workbench.PythonTest do {:ok, %RunPythonOutput{result_json: "{}", stdout: ""}} end) - assert {:ok, %{result: %{}, stdout: ""}} = Python.implement(%Python{code: "output = {}"}) + assert {:ok, output} = Python.implement(%Python{code: "output = {}"}) + assert Jason.decode!(output) == %{"result" => %{}, "stdout" => ""} + end + + test "bounds oversized Python output" do + expect(Client, :connect, fn -> {:ok, :channel} end) + + expect(Stub, :run_python, fn :channel, _request, _opts -> + result = Jason.encode!(%{"value" => String.duplicate("x", Output.max_bytes())}) + {:ok, %RunPythonOutput{result_json: result, stdout: ""}} + end) + + assert {:ok, output} = Python.implement(%Python{code: "output = {}"}) + assert byte_size(output) == Output.max_bytes() + assert output =~ "output truncated at 50 KiB" end end From f61b4d6a9e44256fb66ebe6e8ff0ba5dc0b7767b Mon Sep 17 00:00:00 2001 From: Sebastian Florek Date: Fri, 28 Aug 2026 12:24:36 +0200 Subject: [PATCH 14/15] refactor(python_tool): rename `workbench_python` to `python_sandbox` - Updated reference in `job.md.eex` to `python_sandbox` for consistent naming - Changed function `name()` in `workbench/python.ex` to return `python_sandbox` --- lib/console/ai/tools/workbench/python.ex | 2 +- priv/prompts/workbench/job.md.eex | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/lib/console/ai/tools/workbench/python.ex b/lib/console/ai/tools/workbench/python.ex index dbbfe07e58..96c7ff3af8 100644 --- a/lib/console/ai/tools/workbench/python.ex +++ b/lib/console/ai/tools/workbench/python.ex @@ -13,7 +13,7 @@ defmodule Console.AI.Tools.Workbench.Python do @json_schema Console.priv_file!("tools/workbench/python.json") |> Jason.decode!() - def name(), do: "workbench_python" + def name(), do: "python_sandbox" def description(), do: diff --git a/priv/prompts/workbench/job.md.eex b/priv/prompts/workbench/job.md.eex index 38bd6052b8..dcf32b4d58 100644 --- a/priv/prompts/workbench/job.md.eex +++ b/priv/prompts/workbench/job.md.eex @@ -155,7 +155,7 @@ do not run them simultaneously, since they won't be able to communicate with eac In general, you can call any of the subagents and tools as many times as you need to accomplish your task. -For any tool you delegate to, you should have them give as close to exact raw data as possible, rather than computing on it. When you need an exact, deterministic computation, always use either the `workbench_lua` or `workbench_python` tool with an appropriate supported script to get the result for you. +For any tool you delegate to, you should have them give as close to exact raw data as possible, rather than computing on it. When you need an exact, deterministic computation, always use `python_sandbox` tool with an appropriate supported script to get the result for you. For instance, From 6f442fe30e64eaa71a2ca4f0eb1edabe79e0d715 Mon Sep 17 00:00:00 2001 From: Sebastian Florek Date: Fri, 28 Aug 2026 12:29:37 +0200 Subject: [PATCH 15/15] refactor(python): update naming and remove redundant imports - Renamed `workbench_python` to `python_sandbox` in `python_test.exs` for consistency - Updated observability prompts to reflect tool name change - Removed redundant `Python` import in `infrastructure.ex` --- lib/console/ai/workbench/subagents/infrastructure.ex | 1 - priv/prompts/workbench/observability.md.eex | 2 +- test/console/ai/tools/workbench/python_test.exs | 2 +- 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/lib/console/ai/workbench/subagents/infrastructure.ex b/lib/console/ai/workbench/subagents/infrastructure.ex index fe7858cb67..36702fdc38 100644 --- a/lib/console/ai/workbench/subagents/infrastructure.ex +++ b/lib/console/ai/workbench/subagents/infrastructure.ex @@ -7,7 +7,6 @@ defmodule Console.AI.Workbench.Subagents.Infrastructure do Scratchpad, History, Codemode, - Python, Infrastructure.RawKubeGet, Infrastructure.RawKubeList, Infrastructure.Cluster, diff --git a/priv/prompts/workbench/observability.md.eex b/priv/prompts/workbench/observability.md.eex index 17bfad6387..84b8b115b4 100644 --- a/priv/prompts/workbench/observability.md.eex +++ b/priv/prompts/workbench/observability.md.eex @@ -21,7 +21,7 @@ data or use queries relative to the current time. 1. When querying observability tools, you should give at least some time range to the query. A 5 minute window minimum is good practice. 2. Feel free to query the same data source multiple times to probe. If you're investigating an alert, assume it's a real thing and search. 3. When searching logs, you can always search for contextual logs by filtering on specific facets (eg pod name that isolates to a specific running container), and using time ranges around the log in question. -4. You have `workbench_lua` and `workbench_python` tools available to do deterministic algorithmic calculations. You should be leveraging either instead of trying to infer based on prior tool calls and summaries, using exact raw data in the supported script used. +4. You have a `workbench_tool` tool available to you to do deterministic algorithmic calculations. You should be leveraging that instead of trying to infer based on prior tool calls and summaries, using exact raw data in the lua code used. 5. You are also given a workbench_history tool which can be used to search past work outside of this subagent. Use this to grab additional context that might not be present in your prompt, but don't rely on it if the prompt is sufficient. Also since many tools require a time anchor, the current time is <%= Timex.now() |> Timex.format!("{ISO:Extended}") %>. diff --git a/test/console/ai/tools/workbench/python_test.exs b/test/console/ai/tools/workbench/python_test.exs index cf17f1fcdc..d04955aefc 100644 --- a/test/console/ai/tools/workbench/python_test.exs +++ b/test/console/ai/tools/workbench/python_test.exs @@ -10,7 +10,7 @@ defmodule Console.AI.Tools.Workbench.PythonTest do setup :set_mimic_global test "exposes the Python tool contract" do - assert Python.name() == "workbench_python" + assert Python.name() == "python_sandbox" assert %{"input" => %{"type" => "object"}} = Python.json_schema()["properties"] assert Python.json_schema()["required"] == ["code", "explanation"] end