review: sandbox on encoding/json/v2

e2e/cmd/browsersmoke/main.go

@@ -5,7 +5,7 @@ package main
 
 import (
 	"context"
-	"encoding/json"
+	"encoding/json/v2"
 	"flag"
 	"fmt"
 	"os"

e2e/cmd/envdsmoke/main.go

@@ -13,7 +13,8 @@ import (
 	"context"
 	"encoding/base64"
 	"encoding/binary"
-	"encoding/json"
+	"encoding/json/jsontext"
+	"encoding/json/v2"
 	"flag"
 	"fmt"
 	"io"
@@ -274,7 +275,7 @@ func readEnvelopes(r io.Reader) ([]string, error) {
 			continue
 		}
 		var end struct {
-			Error json.RawMessage `json:"error"`
+			Error jsontext.Value `json:"error"`
 		}
 		if err := json.Unmarshal(body, &end); err != nil {
 			return out, fmt.Errorf("parse end-of-stream envelope %q: %w", body, err)

e2e/cmd/qaab/main.go

@@ -4,7 +4,7 @@ package main
 
 import (
 	"context"
-	"encoding/json"
+	"encoding/json/v2"
 	"flag"
 	"fmt"
 	"os"

e2e/cmd/smoke/main.go

@@ -9,7 +9,8 @@ import (
 	"bufio"
 	"bytes"
 	"context"
-	"encoding/json"
+	"encoding/json/jsontext"
+	"encoding/json/v2"
 	"errors"
 	"flag"
 	"fmt"
@@ -709,7 +710,7 @@ func lspWrite(w io.Writer, body string) error {
 // lspReadResponse reads framed server messages until the response carrying
 // id arrives (server-initiated requests also carry ids but have a method;
 // notifications have no id — both are skipped), returning its result.
-func lspReadResponse(r *bufio.Reader, id int) (json.RawMessage, error) {
+func lspReadResponse(r *bufio.Reader, id int) (jsontext.Value, error) {
 	tp := textproto.NewReader(r)
 	for {
 		hdr, err := tp.ReadMIMEHeader()
@@ -725,10 +726,10 @@ func lspReadResponse(r *bufio.Reader, id int) (json.RawMessage, error) {
 			return nil, err
 		}
 		var msg struct {
-			ID     *int            `json:"id"`
-			Method string          `json:"method"`
-			Result json.RawMessage `json:"result"`
-			Error  json.RawMessage `json:"error"`
+			ID     *int           `json:"id"`
+			Method string         `json:"method"`
+			Result jsontext.Value `json:"result"`
+			Error  jsontext.Value `json:"error"`
 		}
 		if err := json.Unmarshal(body, &msg); err != nil {
 			return nil, fmt.Errorf("lsp frame %q: %w", body, err)

e2e/e2e_test.go

@@ -5,7 +5,7 @@
 package e2e
 
 import (
-	"encoding/json"
+	"encoding/json/v2"
 	"errors"
 	"fmt"
 	"io"

mcp/node_info_test.go

@@ -1,7 +1,7 @@
 package main
 
 import (
-	"encoding/json"
+	"encoding/json/v2"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -10,7 +10,7 @@ import (
 
 func TestNodeInfoReportsCapacityState(t *testing.T) {
 	node := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
-		_ = json.NewEncoder(w).Encode(map[string]any{
+		_ = json.MarshalWrite(w, map[string]any{
 			"pools":              []any{},
 			"claimed":            2,
 			"hibernated":         1,

mcp/server.go

@@ -3,7 +3,8 @@ package main
 import (
 	"bufio"
 	"context"
-	"encoding/json"
+	"encoding/json/jsontext"
+	"encoding/json/v2"
 	"errors"
 	"fmt"
 	"io"
@@ -23,16 +24,16 @@ const (
 )
 
 type rpcRequest struct {
-	ID     json.RawMessage `json:"id"`
-	Method string          `json:"method"`
-	Params json.RawMessage `json:"params"`
+	ID     jsontext.Value `json:"id"`
+	Method string         `json:"method"`
+	Params jsontext.Value `json:"params"`
 }
 
 type rpcResponse struct {
-	JSONRPC string          `json:"jsonrpc"`
-	ID      json.RawMessage `json:"id"`
-	Result  any             `json:"result,omitempty"`
-	Error   *rpcErrorBody   `json:"error,omitempty"`
+	JSONRPC string         `json:"jsonrpc"`
+	ID      jsontext.Value `json:"id"`
+	Result  any            `json:"result,omitzero"`
+	Error   *rpcErrorBody  `json:"error,omitempty"`
 }
 
 type rpcErrorBody struct {
@@ -85,14 +86,14 @@ func (s *server) serve(ctx context.Context, r *bufio.Reader, w io.Writer) error
 			return err
 		}
 		var req rpcRequest
-		if unmarshalErr := json.Unmarshal(line, &req); unmarshalErr != nil || req.Method == "" {
+		if unmarshalErr := unmarshalInput(line, &req); unmarshalErr != nil || req.Method == "" {
 			continue
 		}
 		if req.ID == nil {
 			continue // notification (e.g. notifications/initialized)
 		}
 		resp := s.dispatch(ctx, &req)
-		out, err := json.Marshal(resp)
+		out, err := json.Marshal(resp, jsontext.AllowInvalidUTF8(true), json.Deterministic(true))
 		if err != nil {
 			return err
 		}
@@ -114,10 +115,10 @@ func (s *server) dispatch(ctx context.Context, req *rpcRequest) rpcResponse {
 		return result(req.ID, map[string]any{"tools": toolSpecs()})
 	case "tools/call":
 		var call struct {
-			Name      string          `json:"name"`
-			Arguments json.RawMessage `json:"arguments"`
+			Name      string         `json:"name"`
+			Arguments jsontext.Value `json:"arguments"`
 		}
-		if err := json.Unmarshal(req.Params, &call); err != nil {
+		if err := unmarshalInput(req.Params, &call); err != nil {
 			return rpcError(req.ID, -32602, "invalid params")
 		}
 		text, err := s.callTool(ctx, call.Name, call.Arguments)
@@ -140,7 +141,7 @@ func (s *server) dispatch(ctx context.Context, req *rpcRequest) rpcResponse {
 	}
 }
 
-func (s *server) callTool(ctx context.Context, name string, args json.RawMessage) (string, error) {
+func (s *server) callTool(ctx context.Context, name string, args jsontext.Value) (string, error) {
 	ctx, cancel := context.WithTimeout(ctx, execTimeout)
 	defer cancel()
 	i := slices.IndexFunc(tools, func(t tool) bool { return t.name == name })
@@ -215,10 +216,10 @@ func (s *server) dropCkpt(id string) {
 	delete(s.ckpts, id)
 }
 
-func result(id json.RawMessage, v any) rpcResponse {
+func result(id jsontext.Value, v any) rpcResponse {
 	return rpcResponse{JSONRPC: "2.0", ID: id, Result: v}
 }
 
-func rpcError(id json.RawMessage, code int, msg string) rpcResponse {
+func rpcError(id jsontext.Value, code int, msg string) rpcResponse {
 	return rpcResponse{JSONRPC: "2.0", ID: id, Error: &rpcErrorBody{Code: code, Message: msg}}
 }

mcp/server_test.go

@@ -4,11 +4,12 @@ import (
 	"bufio"
 	"bytes"
 	"encoding/base64"
-	"encoding/json"
+	"encoding/json/v2"
 	"fmt"
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"slices"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -32,7 +33,7 @@ func TestServeSpeaksMCP(t *testing.T) {
 		if !strings.HasSuffix(r.URL.Path, "/checkpoint") {
 			return false
 		}
-		_ = json.NewEncoder(w).Encode(map[string]any{
+		_ = json.MarshalWrite(w, map[string]any{
 			"checkpoint": map[string]any{"id": "ck_0011223344556677", "name": "s1", "sandbox_id": "sb_1"},
 		})
 		return true
@@ -45,9 +46,11 @@ func TestServeSpeaksMCP(t *testing.T) {
 		`{"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"name":"checkpoint","arguments":{"sandbox_id":"sb_1","name":"s1"}}}`,
 		`{"jsonrpc":"2.0","id":5,"method":"tools/call","params":{"name":"release","arguments":{"sandbox_id":"sb_1"}}}`,
 		`{"jsonrpc":"2.0","id":6,"method":"tools/call","params":{"name":"exec","arguments":{"sandbox_id":"sb_1","command":"true"}}}`,
+		`{"jsonrpc":"2.0","id":7,"method":"ping"}`,
+		`{"jsonrpc":"2.0","id":8,"method":"ping","params":{"note":"\ud800"}}`,
 	)
-	if len(replies) != 6 {
-		t.Fatalf("got %d replies, want 6 (notification unanswered)", len(replies))
+	if len(replies) != 8 {
+		t.Fatalf("got %d replies, want 8 (notification unanswered)", len(replies))
 	}
 
 	initResult := replies[1]["result"].(map[string]any)
@@ -77,6 +80,27 @@ func TestServeSpeaksMCP(t *testing.T) {
 	if !strings.Contains(toolText(t, replies[6]), "sb_1") {
 		t.Errorf("exec error should name the released id: %q", toolText(t, replies[6]))
 	}
+	if pong, ok := replies[7]["result"].(map[string]any); !ok || len(pong) != 0 {
+		t.Errorf("ping reply %+v, want an empty result object", replies[7])
+	}
+	if _, ok := replies[8]["result"].(map[string]any); !ok {
+		t.Errorf("ping with a lone surrogate %+v, want a reply", replies[8])
+	}
+}
+
+func TestRepliesEncodeMapsInSortedOrder(t *testing.T) {
+	var out bytes.Buffer
+	in := bufio.NewReader(strings.NewReader(strings.Repeat(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`+"\n", 20)))
+	if err := newTestServer(t, nil).serve(t.Context(), in, &out); err != nil {
+		t.Fatalf("serve: %v", err)
+	}
+	lines := slices.Collect(bytes.Lines(out.Bytes()))
+	if len(lines) != 20 || slices.ContainsFunc(lines, func(l []byte) bool { return !bytes.Equal(l, lines[0]) }) {
+		t.Errorf("20 tools/list replies are not byte-identical (%d lines)", len(lines))
+	}
+	if got := jsonText(map[string]int{"b": 2, "a": 1, "d": 4, "c": 3, "f": 6, "e": 5, "h": 8, "g": 7}); got != `{"a":1,"b":2,"c":3,"d":4,"e":5,"f":6,"g":7,"h":8}` {
+		t.Errorf("jsonText = %s, want sorted keys", got)
+	}
 }
 
 func TestExecKeepsOutputWhenTheGuestDrops(t *testing.T) {
@@ -125,6 +149,31 @@ func TestLogsKeepsAtMostTheCap(t *testing.T) {
 	}
 }
 
+func TestToolRepliesReplaceInvalidUTF8(t *testing.T) {
+	lossy := base64.StdEncoding.EncodeToString([]byte("a\xffb"))
+	srv := newTestServer(t, agentRoute(t, func(req string) (string, bool) {
+		switch {
+		case strings.Contains(req, `"op":"fs_stat"`):
+			return `{"type":"stat","info":{"kind":"file","size":3}}` + "\n", false
+		case strings.Contains(req, `"op":"logs"`):
+			return `{"type":"stdout","data":"` + lossy + `"}` + "\n" + `{"type":"done"}` + "\n", false
+		}
+		return `{"type":"data","data":"` + lossy + `"}` + "\n" + `{"type":"done"}` + "\n", false
+	}))
+	replies := serveLines(t, srv,
+		`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"create_sandbox","arguments":{}}}`,
+		`{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"read_file","arguments":{"sandbox_id":"sb_1","path":"/bin/x"}}}`,
+		`{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"logs","arguments":{"sandbox_id":"sb_1","pid":7}}}`,
+	)
+	if text := toolText(t, replies[2]); text != "a\uFFFDb" {
+		t.Errorf("read_file answered %q, want the invalid byte replaced", text)
+	}
+	var logs map[string]any
+	if err := json.Unmarshal([]byte(toolText(t, replies[3])), &logs); err != nil || logs["stdout"] != "a\uFFFDb" {
+		t.Errorf("logs answered %v (%v), want JSON with the invalid byte replaced", logs, err)
+	}
+}
+
 func TestTrackBoxAfterCloseReleasesTheClaim(t *testing.T) {
 	released := make(chan string, 1)
 	node := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -230,7 +279,7 @@ func newTestServer(t *testing.T, extra func(http.ResponseWriter, *http.Request)
 	node := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		switch {
 		case r.URL.Path == "/v1/claim":
-			_ = json.NewEncoder(w).Encode(map[string]any{"id": "sb_1", "token": "tok"})
+			_ = json.MarshalWrite(w, map[string]any{"id": "sb_1", "token": "tok"})
 		case strings.HasSuffix(r.URL.Path, "/release"):
 			w.WriteHeader(http.StatusNoContent)
 		case extra != nil && extra(w, r):
@@ -283,11 +332,10 @@ func serveLines(t *testing.T, srv *server, lines ...string) map[int]map[string]a
 		t.Fatalf("serve: %v", err)
 	}
 	replies := map[int]map[string]any{}
-	dec := json.NewDecoder(&out)
-	for dec.More() {
+	for line := range bytes.Lines(out.Bytes()) {
 		var resp map[string]any
-		if err := dec.Decode(&resp); err != nil {
-			t.Fatalf("decode reply: %v", err)
+		if err := json.Unmarshal(line, &resp); err != nil {
+			t.Fatalf("decode reply %q: %v", line, err)
 		}
 		replies[int(resp["id"].(float64))] = resp
 	}