Merge pull request #76 from dispatchrun/pickled-value-container

Pickled value container support

Author: chriso

Makefile

@@ -37,6 +37,5 @@ push: image
 	$(DOCKER) push $(IMAGE)
 
 update:
-	buf mod update ./proto
 	for ref in $$(yq -r '.deps[] | .remote + "/gen/go/" + .owner + "/" + .repository + "/protocolbuffers/go@" + .commit' proto/buf.lock); do go get $$ref; done
 	go mod tidy

cli/any.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"log/slog"
 
+	pythonv1 "buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go/dispatch/sdk/python/v1"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/anypb"
 	"google.golang.org/protobuf/types/known/wrapperspb"
@@ -13,17 +14,31 @@ func anyString(any *anypb.Any) string {
 	if any == nil {
 		return "nil"
 	}
+
+	var s string
+	var err error
 	switch any.TypeUrl {
-	case "type.googleapis.com/google.protobuf.BytesValue":
-		s, err := anyBytesString(any)
-		if err != nil {
-			slog.Debug("cannot parse input/output value", "error", err)
-			// fallthrough
-		} else {
-			return s
+	case "buf.build/stealthrocket/dispatch-proto/dispatch.sdk.python.v1.Pickled":
+		var pickled proto.Message
+		pickled, err = any.UnmarshalNew()
+		if err == nil {
+			if p, ok := pickled.(*pythonv1.Pickled); ok {
+				s, err = pythonPickleString(p.PickledValue)
+			} else {
+				err = fmt.Errorf("invalid pickled message: %T", p)
+			}
 		}
+	case "type.googleapis.com/google.protobuf.BytesValue":
+		s, err = anyBytesString(any)
+	default:
+		// TODO: support unpacking other types of serialized values
+		err = fmt.Errorf("not implemented: %s", any.TypeUrl)
+	}
+	if err != nil {
+		slog.Debug("cannot parse input/output value", "error", err)
+		return fmt.Sprintf("%s(?)", any.TypeUrl)
 	}
-	return fmt.Sprintf("%s(?)", any.TypeUrl)
+	return s
 }
 
 func anyBytesString(any *anypb.Any) (string, error) {
@@ -37,6 +52,20 @@ func anyBytesString(any *anypb.Any) (string, error) {
 	}
 	b := bv.Value
 
-	// TODO: support unpacking other types of serialized values
-	return pythonPickleString(b)
+	// The Python SDK originally wrapped pickled values in a
+	// wrapperspb.BytesValue. Try to unpickle the bytes first,
+	// and return literal bytes if they cannot be unpickled.
+	s, err := pythonPickleString(b)
+	if err != nil {
+		s = string(truncateBytes(b))
+	}
+	return s, nil
+}
+
+func truncateBytes(b []byte) []byte {
+	const n = 4
+	if len(b) < n {
+		return b
+	}
+	return append(b[:n:n], "..."...)
 }

cli/tui.go

@@ -552,7 +552,20 @@ func (t *TUI) detailView(id DispatchID) string {
 			add("Input", rt.request.input)
 
 		case *sdkv1.RunRequest_PollResult:
-			add("Input", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of state>", len(d.PollResult.CoroutineState))))
+			switch s := d.PollResult.State.(type) {
+			case *sdkv1.PollResult_CoroutineState:
+				add("Input", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of opaque state>", len(s.CoroutineState))))
+			case *sdkv1.PollResult_TypedCoroutineState:
+				if any := s.TypedCoroutineState; any != nil {
+					add("Input", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of %s state>", len(any.Value), typeName(any.TypeUrl))))
+				} else {
+					add("Input", detailLowPriorityStyle.Render("<no state>"))
+				}
+			case nil:
+				add("Input", detailLowPriorityStyle.Render("<no state>"))
+			default:
+				add("Input", detailLowPriorityStyle.Render("<unknown state>"))
+			}
 			// TODO: show call results
 			// TODO: show poll error
 		}
@@ -593,7 +606,21 @@ func (t *TUI) detailView(id DispatchID) string {
 
 				case *sdkv1.RunResponse_Poll:
 					add("Status", suspendedStyle.Render("Suspended"))
-					add("Output", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of state>", len(d.Poll.CoroutineState))))
+
+					switch s := d.Poll.State.(type) {
+					case *sdkv1.Poll_CoroutineState:
+						add("Output", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of opaque state>", len(s.CoroutineState))))
+					case *sdkv1.Poll_TypedCoroutineState:
+						if any := s.TypedCoroutineState; any != nil {
+							add("Output", detailLowPriorityStyle.Render(fmt.Sprintf("<%d bytes of %s state>", len(any.Value), typeName(any.TypeUrl))))
+						} else {
+							add("Output", detailLowPriorityStyle.Render("<no state>"))
+						}
+					case nil:
+						add("Output", detailLowPriorityStyle.Render("<no state>"))
+					default:
+						add("Output", detailLowPriorityStyle.Render("<unknown state>"))
+					}
 
 					if len(d.Poll.Calls) > 0 {
 						var calls strings.Builder
@@ -1033,3 +1060,11 @@ func terminalHTTPStatusCode(code int) bool {
 		return true
 	}
 }
+
+func typeName(typeUrl string) string {
+	i := strings.LastIndexByte(typeUrl, '/')
+	if i < 0 {
+		return typeUrl
+	}
+	return typeUrl[i+1:]
+}

go.mod

@@ -3,7 +3,7 @@ module github.com/dispatchrun/dispatch
 go 1.22.0
 
 require (
-	buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go v1.33.0-20240429010127-639d52c5db75.1
+	buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go v1.34.2-20240612225639-f8a6c0a10402.2
 	github.com/charmbracelet/bubbles v0.18.0
 	github.com/charmbracelet/bubbletea v0.25.0
 	github.com/charmbracelet/lipgloss v0.9.1
@@ -14,11 +14,11 @@ require (
 	github.com/spf13/cobra v1.8.0
 	github.com/stretchr/testify v1.9.0
 	golang.org/x/term v0.19.0
-	google.golang.org/protobuf v1.33.0
+	google.golang.org/protobuf v1.34.2
 )
 
 require (
-	buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.33.0-20231115204500-e097f827e652.1 // indirect
+	buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.34.2-20231115204500-e097f827e652.2 // indirect
 	github.com/atotto/clipboard v0.1.4 // indirect
 	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
 	github.com/containerd/console v1.0.4-0.20230313162750-1ae8d489ac81 // indirect

go.sum

@@ -1,7 +1,7 @@
-buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.33.0-20231115204500-e097f827e652.1 h1:SbKXkoZduR4jQ+aCrX71I+dAcHKZPk1SiriRs/XYX9s=
-buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.33.0-20231115204500-e097f827e652.1/go.mod h1:Tgn5bgL220vkFOI0KPStlcClPeOJzAv4uT+V8JXGUnw=
-buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go v1.33.0-20240429010127-639d52c5db75.1 h1:C8qLXgA5Y2iK8VkXsNHi8E4qU3aUsZQ0mFqwgxIH7KI=
-buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go v1.33.0-20240429010127-639d52c5db75.1/go.mod h1:yJei/TBJwZBJ8ZUWCKVKceUHk/8gniSGs812SZA9TEE=
+buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.34.2-20231115204500-e097f827e652.2 h1:ilkjxsfdhrdeLL582R4XH+IQLrV2Y31lHB8jb3AQSJA=
+buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.34.2-20231115204500-e097f827e652.2/go.mod h1:ylS4c28ACSI59oJrOdW4pHS4n0Hw4TgSPHn8rpHl4Yw=
+buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go v1.34.2-20240612225639-f8a6c0a10402.2 h1:GfZ4I61ZZiXk6zXLZaMNz1rSW0MxlBRRc5n5GOUbIg4=
+buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go v1.34.2-20240612225639-f8a6c0a10402.2/go.mod h1:Z1Y+G5dzXnHxfdt7l3S9kLBWV6iO2Zw3z2DvY9yKpaw=
 github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4=
 github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI=
 github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=
@@ -18,7 +18,6 @@ github.com/cpuguy83/go-md2man/v2 v2.0.3/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46t
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
 github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
@@ -78,9 +77,8 @@ golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
 golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
-google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI=
-google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
+google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg=
+google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=