tutti/packages/agent/daemon/runtime/codexproto/internal/codegen/main.go

//revive:disable:file-length-limit
package main

import (
	"encoding/json"
	"fmt"
	"go/format"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/atombender/go-jsonschema/pkg/generator"
)

const (
	schemaBundleFile    = "codex_app_server_protocol.schemas.json"
	generatedHeaderBase = "// Code generated by runtime/codexproto/internal/codegen. DO NOT EDIT.\n"
	codexSourceCommit   = "6d2168f06ae275d5e1f73cabf935d2bcc8549998"
	codexSourceVersion  = "0.0.0"
)

func main() {
	packageRoot, err := repoRoot()
	if err != nil {
		fatal(err)
	}

	schemaDir := filepath.Join(packageRoot, "schema", "json")
	if err := generateProtocolTypes(schemaDir, packageRoot, codexSourceCommit, codexSourceVersion); err != nil {
		fatal(err)
	}

	if err := generateRPCStubs(schemaDir, packageRoot, codexSourceCommit); err != nil {
		fatal(err)
	}
}

func generateProtocolTypes(schemaDir, repoRoot, codexCommit, codexVersion string) error {
	cfg := generator.Config{
		DefaultPackageName: "codexproto",
		DefaultOutputName:  "types_gen.go",
		Tags:               []string{"json"},
		Capitalizations:    []string{"ID", "JSON", "RPC", "MCP", "HTTP", "URL", "SSE", "UUID"},
		OnlyModels:         true,
	}
	gen, err := generator.New(cfg)
	if err != nil {
		return err
	}

	schemaFiles, err := findSchemaFiles(schemaDir)
	if err != nil {
		return err
	}

	skip := skipSchemaFiles()
	var fallbacks []string
	var titles []string
	for _, file := range schemaFiles {
		if filepath.Base(file) == schemaBundleFile {
			continue
		}
		if skip[filepath.Base(file)] {
			continue
		}
		title, titleErr := schemaTitle(file)
		if titleErr != nil {
			return titleErr
		}
		if title != "" {
			titles = append(titles, title)
		}
		if err := safeDoFile(gen, file); err != nil {
			if title == "" {
				return fmt.Errorf("generate types from %s: %w", file, err)
			}
			fallbacks = append(fallbacks, title)
		}
	}

	sources, err := gen.Sources()
	if err != nil {
		return err
	}

	manualTypes := manualProtocolTypes()
	for name, src := range sources {
		sources[name] = filterManualTypeDeclarations(src, manualTypes)
	}

	outDir := repoRoot

	generated := collectGeneratedTypes(sources)
	for name := range manualTypes {
		generated[name] = struct{}{}
	}
	for _, title := range titles {
		if title == "" {
			continue
		}
		if _, ok := generated[title]; ok {
			continue
		}
		if _, ok := generated[title+"JSON"]; ok {
			continue
		}
		fallbacks = append(fallbacks, title)
	}

	for name, src := range sources {
		normalized := normalizeGeneratedHeader(src, codexCommit)
		if err := os.WriteFile(filepath.Join(outDir, name), normalized, 0o644); err != nil {
			return err
		}
	}

	if err := writeFallbackTypes(outDir, fallbacks, generated, codexCommit); err != nil {
		return err
	}

	if err := writeProtocolAliases(outDir, generated, fallbacks, codexCommit); err != nil {
		return err
	}

	return writeProtocolMetadata(outDir, codexCommit, codexVersion)
}

func filterManualTypeDeclarations(src []byte, manualTypes map[string]struct{}) []byte {
	lines := strings.SplitAfter(string(src), "\n")
	out := make([]string, 0, len(lines))
	for i := 0; i < len(lines); {
		line := lines[i]
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "type ") {
			fields := strings.Fields(trimmed)
			if len(fields) >= 2 {
				if _, ok := manualTypes[fields[1]]; ok {
					i++
					balance := strings.Count(line, "{") - strings.Count(line, "}")
					for i < len(lines) && balance > 0 {
						balance += strings.Count(lines[i], "{") - strings.Count(lines[i], "}")
						i++
					}
					continue
				}
			}
		}
		out = append(out, line)
		i++
	}
	return []byte(strings.Join(out, ""))
}

type schemaVariant struct {
	Properties map[string]schemaProperty `json:"properties"`
}

type schemaProperty struct {
	Enum  []string `json:"enum"`
	Ref   string   `json:"$ref"`
	Type  string   `json:"type"`
	Items *struct {
		Ref string `json:"$ref"`
	} `json:"items"`
}

type oneOfSchema struct {
	OneOf []schemaVariant `json:"oneOf"`
}

func generateRPCStubs(schemaDir, repoRoot, codexCommit string) error {
	defs, err := parseDefinitionNames(schemaDir)
	if err != nil {
		return err
	}

	clientMethods, err := loadMethods(filepath.Join(schemaDir, "ClientRequest.json"))
	if err != nil {
		return err
	}
	clientMethods = append(clientMethods, manualClientMethods()...)
	sort.Slice(clientMethods, func(i, j int) bool { return clientMethods[i].Method < clientMethods[j].Method })

	serverMethods, err := loadMethods(filepath.Join(schemaDir, "ServerRequest.json"))
	if err != nil {
		return err
	}

	notifications, err := loadNotifications(filepath.Join(schemaDir, "ServerNotification.json"))
	if err != nil {
		return err
	}

	clientRPC, err := buildRPCMethods(clientMethods, defs, clientResponseOverrides())
	if err != nil {
		return err
	}

	serverRPC, err := buildRPCMethods(serverMethods, defs, nil)
	if err != nil {
		return err
	}

	outDir := repoRoot

	if err := writeFile(outDir, "client_requests_gen.go", renderClientRequests(clientRPC, codexCommit)); err != nil {
		return err
	}

	if err := writeFile(outDir, "server_requests_gen.go", renderServerRequests(serverRPC, codexCommit)); err != nil {
		return err
	}

	return writeFile(outDir, "notifications_gen.go", renderNotifications(notifications, codexCommit))
}

func parseDefinitionNames(schemaDir string) (map[string]struct{}, error) {
	files, err := findSchemaFiles(schemaDir)
	if err != nil {
		return nil, err
	}
	defs := make(map[string]struct{})
	for _, file := range files {
		data, err := os.ReadFile(file)
		if err != nil {
			return nil, err
		}
		var doc struct {
			Title       string                     `json:"title"`
			Definitions map[string]json.RawMessage `json:"definitions"`
		}
		if err := json.Unmarshal(data, &doc); err != nil {
			return nil, err
		}
		if doc.Title != "" {
			defs[doc.Title] = struct{}{}
		}
		for name := range doc.Definitions {
			defs[name] = struct{}{}
		}
	}
	return defs, nil
}

func schemaTitle(path string) (string, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return "", err
	}
	var doc struct {
		Title string `json:"title"`
	}
	if err := json.Unmarshal(data, &doc); err != nil {
		return "", err
	}
	return doc.Title, nil
}

func writeFallbackTypes(outDir string, titles []string, generated map[string]struct{}, codexCommit string) error {
	fallbackPath := filepath.Join(outDir, "fallback_gen.go")
	if len(titles) == 0 {
		_ = os.Remove(fallbackPath)
		return nil
	}

	unique := map[string]struct{}{}
	for _, title := range titles {
		if title != "" {
			unique[title] = struct{}{}
		}
	}

	aliases := map[string]string{}
	var names []string
	for name := range unique {
		if _, exists := generated[name]; exists {
			continue
		}
		if target, ok := sanitizedFallbackTarget(name, generated); ok {
			aliases[name] = target
			continue
		}
		names = append(names, name)
	}
	if len(names) == 0 && len(aliases) == 0 {
		_ = os.Remove(fallbackPath)
		return nil
	}
	sort.Strings(names)
	aliasNames := make([]string, 0, len(aliases))
	for name := range aliases {
		aliasNames = append(aliasNames, name)
	}
	sort.Strings(aliasNames)

	var b strings.Builder
	b.WriteString(generatedHeader(codexCommit))
	b.WriteString("package codexproto\n\n")
	b.WriteString("// Fallback types for schemas that failed to generate.\n")
	for _, name := range aliasNames {
		b.WriteString("type ")
		b.WriteString(name)
		b.WriteString(" = ")
		b.WriteString(aliases[name])
		b.WriteString("\n")
	}
	for _, name := range names {
		b.WriteString("type ")
		b.WriteString(name)
		b.WriteString(" interface{}\n")
	}

	return os.WriteFile(fallbackPath, []byte(b.String()), 0o644)
}

func sanitizedFallbackTarget(name string, generated map[string]struct{}) (string, bool) {
	target := "Sanitized" + name + "JSON"
	if _, ok := generated[target]; ok {
		return target, true
	}
	return "", false
}

func collectGeneratedTypes(sources map[string][]byte) map[string]struct{} {
	generated := map[string]struct{}{}
	for _, src := range sources {
		lines := strings.Split(string(src), "\n")
		for _, line := range lines {
			line = strings.TrimSpace(line)
			if strings.HasPrefix(line, "type ") {
				fields := strings.Fields(line)
				if len(fields) >= 2 {
					generated[fields[1]] = struct{}{}
				}
			}
		}
	}
	return generated
}

func writeProtocolAliases(outDir string, generated map[string]struct{}, fallbacks []string, codexCommit string) error {
	existing := map[string]struct{}{}
	for name := range generated {
		existing[name] = struct{}{}
	}
	for _, name := range fallbacks {
		if name != "" {
			existing[name] = struct{}{}
		}
	}

	aliases := map[string]string{}
	for name := range generated {
		if !strings.HasSuffix(name, "JSON") {
			continue
		}
		base := strings.TrimSuffix(name, "JSON")
		if base == "" {
			continue
		}
		if _, ok := existing[base]; ok {
			continue
		}
		aliases[base] = name
	}

	aliasPath := filepath.Join(outDir, "aliases_gen.go")
	if len(aliases) == 0 {
		_ = os.Remove(aliasPath)
		return nil
	}

	names := make([]string, 0, len(aliases))
	for name := range aliases {
		names = append(names, name)
	}
	sort.Strings(names)

	var b strings.Builder
	b.WriteString(generatedHeader(codexCommit))
	b.WriteString("package codexproto\n\n")
	b.WriteString("// Aliases for JSON-suffixed schema types.\n")
	for _, name := range names {
		b.WriteString("type ")
		b.WriteString(name)
		b.WriteString(" = ")
		b.WriteString(aliases[name])
		b.WriteString("\n")
	}

	return os.WriteFile(aliasPath, []byte(b.String()), 0o644)
}

func skipSchemaFiles() map[string]bool {
	return map[string]bool{
		"JSONRPCMessage.json":      true,
		"JSONRPCRequest.json":      true,
		"JSONRPCResponse.json":     true,
		"JSONRPCNotification.json": true,
		"JSONRPCError.json":        true,
		"JSONRPCErrorError.json":   true,
		"EventMsg.json":            true,
		"RequestId.json":           true,
	}
}

func loadMethods(path string) ([]rpcMethod, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var doc oneOfSchema
	if err := json.Unmarshal(data, &doc); err != nil {
		return nil, err
	}

	methods := make([]rpcMethod, 0, len(doc.OneOf))
	for _, variant := range doc.OneOf {
		method := ""
		if prop, ok := variant.Properties["method"]; ok && len(prop.Enum) == 1 {
			method = prop.Enum[0]
		}
		if method == "" {
			continue
		}

		paramsType := ""
		if prop, ok := variant.Properties["params"]; ok {
			if prop.Ref != "" {
				paramsType = refName(prop.Ref)
			} else if prop.Type == "null" {
				paramsType = ""
			}
		}

		methods = append(methods, rpcMethod{
			Method:     method,
			ParamsType: paramsType,
		})
	}

	sort.Slice(methods, func(i, j int) bool { return methods[i].Method < methods[j].Method })

	return methods, nil
}

func loadNotifications(path string) ([]rpcNotification, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var doc oneOfSchema
	if err := json.Unmarshal(data, &doc); err != nil {
		return nil, err
	}

	notifications := make([]rpcNotification, 0, len(doc.OneOf))
	for _, variant := range doc.OneOf {
		method := ""
		if prop, ok := variant.Properties["method"]; ok && len(prop.Enum) == 1 {
			method = prop.Enum[0]
		}
		if method == "" {
			continue
		}

		paramsType := ""
		if prop, ok := variant.Properties["params"]; ok {
			if prop.Ref != "" {
				paramsType = refName(prop.Ref)
			} else if prop.Type == "null" {
				paramsType = ""
			}
		}

		notifications = append(notifications, rpcNotification{
			Method:     method,
			ParamsType: paramsType,
		})
	}

	sort.Slice(notifications, func(i, j int) bool { return notifications[i].Method < notifications[j].Method })

	return notifications, nil
}

func refName(ref string) string {
	const prefix = "#/definitions/"
	if strings.HasPrefix(ref, prefix) {
		return strings.TrimPrefix(ref, prefix)
	}
	return ref
}

type rpcMethod struct {
	Method       string
	ParamsType   string
	ResponseType string
}

type rpcNotification struct {
	Method     string
	ParamsType string
}

func buildRPCMethods(methods []rpcMethod, defs map[string]struct{}, overrides map[string]string) ([]rpcMethod, error) {
	resolved := make([]rpcMethod, 0, len(methods))
	for _, method := range methods {
		if method.ResponseType != "" {
			resolved = append(resolved, method)
			continue
		}
		responseType, err := resolveResponseType(method, defs, overrides)
		if err != nil {
			return nil, err
		}
		method.ResponseType = responseType
		resolved = append(resolved, method)
	}
	return resolved, nil
}

func resolveResponseType(method rpcMethod, defs map[string]struct{}, overrides map[string]string) (string, error) {
	if overrides != nil {
		if override, ok := overrides[method.Method]; ok {
			return override, nil
		}
	}

	if method.ParamsType != "" {
		base := strings.TrimSuffix(method.ParamsType, "Params")
		candidate := base + "Response"
		if _, ok := defs[candidate]; ok {
			return candidate, nil
		}
	}

	candidate := methodBaseName(method.Method) + "Response"
	if _, ok := defs[candidate]; ok {
		return candidate, nil
	}

	return "", fmt.Errorf("unable to resolve response type for method %q", method.Method)
}

func methodBaseName(method string) string {
	parts := strings.FieldsFunc(method, func(r rune) bool {
		return (r < 'a' || r > 'z') && (r < 'A' || r > 'Z') && (r < '0' || r > '9')
	})
	var b strings.Builder
	for _, part := range parts {
		if part == "" {
			continue
		}
		b.WriteString(strings.ToUpper(part[:1]))
		if len(part) > 1 {
			b.WriteString(part[1:])
		}
	}
	return b.String()
}

func clientResponseOverrides() map[string]string {
	return map[string]string{
		"account/logout":                           "LogoutAccountResponse",
		"account/rateLimits/read":                  "GetAccountRateLimitsResponse",
		"account/usage/read":                       "GetAccountTokenUsageResponse",
		"account/workspaceMessages/read":           "GetWorkspaceMessagesResponse",
		"addConversationListener":                  "AddConversationSubscriptionResponse",
		"config/value/write":                       "ConfigWriteResponse",
		"config/batchWrite":                        "ConfigWriteResponse",
		"config/mcpServer/reload":                  "McpServerRefreshResponse",
		"externalAgentConfig/import/readHistories": "ExternalAgentConfigImportHistoriesReadResponse",
		"removeConversationListener":               "RemoveConversationSubscriptionResponse",
	}
}

func manualClientMethods() []rpcMethod {
	return []rpcMethod{
		{
			Method:       "collaborationMode/list",
			ParamsType:   "CollaborationModeListParams",
			ResponseType: "CollaborationModeListResponse",
		},
	}
}

func renderClientRequests(methods []rpcMethod, codexCommit string) []byte {
	var b strings.Builder
	b.WriteString(generatedHeader(codexCommit))
	b.WriteString("package codexproto\n\n")
	b.WriteString("import (\n\t\"context\"\n\t\"encoding/json\"\n)\n\n")

	b.WriteString("// ClientRequests exposes typed JSON-RPC calls supported by the app-server.\n")
	b.WriteString("type ClientRequests interface {\n")
	for _, method := range methods {
		if method.ParamsType == "" {
			fmt.Fprintf(&b, "\t%s(ctx context.Context) (*%s, error)\n", methodName(method.Method), method.ResponseType)
		} else {
			fmt.Fprintf(&b, "\t%s(ctx context.Context, params %s) (*%s, error)\n", methodName(method.Method), paramsType(method.ParamsType), method.ResponseType)
		}
	}
	b.WriteString("}\n\n")

	b.WriteString("// Client implements ClientRequests using JSON-RPC.\n")
	for _, method := range methods {
		if method.ParamsType == "" {
			fmt.Fprintf(&b, "func (c *Client) %s(ctx context.Context) (*%s, error) {\n", methodName(method.Method), method.ResponseType)
		} else {
			fmt.Fprintf(&b, "func (c *Client) %s(ctx context.Context, params %s) (*%s, error) {\n", methodName(method.Method), paramsType(method.ParamsType), method.ResponseType)
		}
		b.WriteString("\tvar result ")
		b.WriteString(method.ResponseType)
		b.WriteString("\n")
		if method.ParamsType == "" {
			fmt.Fprintf(&b, "\tif err := c.Call(ctx, %q, nil, &result); err != nil {\n", method.Method)
		} else {
			fmt.Fprintf(&b, "\tif err := c.Call(ctx, %q, params, &result); err != nil {\n", method.Method)
		}
		b.WriteString("\t\treturn nil, err\n\t}\n\treturn &result, nil\n}\n\n")
	}

	b.WriteString("// BuildClientRequest builds a JSON-RPC request for a client method.\n")
	b.WriteString("func BuildClientRequest(method string, params any, id RequestID) (JSONRPCRequest, error) {\n")
	b.WriteString("\tvar raw json.RawMessage\n\tif params != nil {\n\t\tdata, err := json.Marshal(params)\n\t\tif err != nil {\n\t\t\treturn JSONRPCRequest{}, err\n\t\t}\n\t\traw = data\n\t}\n\treturn JSONRPCRequest{ID: id, Method: method, Params: raw}, nil\n}\n")

	return []byte(b.String())
}

func renderServerRequests(methods []rpcMethod, codexCommit string) []byte {
	var b strings.Builder
	b.WriteString(generatedHeader(codexCommit))
	b.WriteString("package codexproto\n\n")
	b.WriteString("import (\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n)\n\n")

	b.WriteString("// ServerRequestHandler handles requests initiated by the app-server.\n")
	b.WriteString("type ServerRequestHandler interface {\n")
	for _, method := range methods {
		if method.ParamsType == "" {
			fmt.Fprintf(&b, "\t%s(ctx context.Context) (*%s, error)\n", methodName(method.Method), method.ResponseType)
		} else {
			fmt.Fprintf(&b, "\t%s(ctx context.Context, params %s) (*%s, error)\n", methodName(method.Method), paramsType(method.ParamsType), method.ResponseType)
		}
	}
	b.WriteString("}\n\n")

	b.WriteString("// ParsedServerRequest represents a typed server request envelope.\n")
	b.WriteString("type ParsedServerRequest struct {\n\tMethod string\n\tParams any\n\tRaw json.RawMessage\n}\n\n")

	b.WriteString("type serverRequestParser func(json.RawMessage) (ParsedServerRequest, error)\n\n")
	b.WriteString("var serverRequestParsers = map[string]serverRequestParser{\n")
	for _, method := range methods {
		fmt.Fprintf(&b, "\t%q: func(params json.RawMessage) (ParsedServerRequest, error) {\n", method.Method)
		if method.ParamsType == "" {
			fmt.Fprintf(&b, "\t\treturn ParsedServerRequest{Method: %q, Params: nil, Raw: params}, nil\n", method.Method)
		} else {
			b.WriteString("\t\tvar payload ")
			b.WriteString(method.ParamsType)
			b.WriteString("\n")
			b.WriteString("\t\tif len(params) > 0 {\n\t\t\tif err := json.Unmarshal(params, &payload); err != nil {\n\t\t\t\treturn ParsedServerRequest{Method: ")
			fmt.Fprintf(&b, "%q", method.Method)
			b.WriteString(", Raw: params}, err\n\t\t\t}\n\t\t}\n")
			b.WriteString("\t\treturn ParsedServerRequest{Method: ")
			fmt.Fprintf(&b, "%q", method.Method)
			b.WriteString(", Params: payload, Raw: params}, nil\n")
		}
		b.WriteString("\t},\n")
	}
	b.WriteString("}\n\n")

	b.WriteString("// ParseServerRequest parses a server request by method.\n")
	b.WriteString("func ParseServerRequest(method string, params json.RawMessage) (ParsedServerRequest, error) {\n")
	b.WriteString("\tif parser, ok := serverRequestParsers[method]; ok {\n\t\treturn parser(params)\n\t}\n")
	b.WriteString("\treturn ParsedServerRequest{Method: method, Raw: params}, nil\n}\n\n")

	b.WriteString("// IsKnownServerRequestMethod reports whether method is in the generated server request surface.\n")
	b.WriteString("func IsKnownServerRequestMethod(method string) bool {\n")
	b.WriteString("\t_, ok := serverRequestParsers[method]\n\treturn ok\n}\n\n")

	b.WriteString("// ServerRequestMethods returns the generated server request surface in stable order.\n")
	b.WriteString("func ServerRequestMethods() []string {\n")
	b.WriteString("\treturn []string{\n")
	for _, method := range methods {
		fmt.Fprintf(&b, "\t\t%q,\n", method.Method)
	}
	b.WriteString("\t}\n}\n\n")

	b.WriteString("func dispatchServerRequest(ctx context.Context, handler ServerRequestHandler, req JSONRPCRequest) (any, error) {\n")
	b.WriteString("\tswitch req.Method {\n")
	for _, method := range methods {
		fmt.Fprintf(&b, "\tcase %q:\n", method.Method)
		if method.ParamsType == "" {
			b.WriteString("\t\treturn handler.")
			b.WriteString(methodName(method.Method))
			b.WriteString("(ctx)\n")
		} else {
			b.WriteString("\t\tvar params ")
			b.WriteString(method.ParamsType)
			b.WriteString("\n")
			b.WriteString("\t\tif len(req.Params) > 0 {\n\t\t\tif err := json.Unmarshal(req.Params, &params); err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t}\n")
			b.WriteString("\t\treturn handler.")
			b.WriteString(methodName(method.Method))
			b.WriteString("(ctx, params)\n")
		}
	}
	b.WriteString("\tdefault:\n\t\treturn nil, fmt.Errorf(\"unsupported server request %q\", req.Method)\n\t}\n}\n")

	return []byte(b.String())
}

func renderNotifications(notifications []rpcNotification, codexCommit string) []byte {
	var b strings.Builder
	b.WriteString(generatedHeader(codexCommit))
	b.WriteString("package codexproto\n\n")
	b.WriteString("import \"encoding/json\"\n\n")

	b.WriteString("// Notification represents a typed server notification.\n")
	b.WriteString("type Notification struct {\n\tMethod string\n\tParams any\n\tRaw json.RawMessage\n}\n\n")

	b.WriteString("type notificationParser func(json.RawMessage) (Notification, error)\n\n")
	b.WriteString("var notificationParsers = map[string]notificationParser{\n")
	for _, notification := range notifications {
		fmt.Fprintf(&b, "\t%q: func(params json.RawMessage) (Notification, error) {\n", notification.Method)
		if notification.ParamsType == "" {
			fmt.Fprintf(&b, "\t\treturn Notification{Method: %q, Params: nil, Raw: params}, nil\n", notification.Method)
		} else {
			b.WriteString("\t\tvar payload ")
			b.WriteString(notification.ParamsType)
			b.WriteString("\n")
			b.WriteString("\t\tif len(params) > 0 {\n\t\t\tif err := json.Unmarshal(params, &payload); err != nil {\n\t\t\t\treturn Notification{Method: ")
			fmt.Fprintf(&b, "%q", notification.Method)
			b.WriteString(", Raw: params}, err\n\t\t\t}\n\t\t}\n")
			b.WriteString("\t\treturn Notification{Method: ")
			fmt.Fprintf(&b, "%q", notification.Method)
			b.WriteString(", Params: payload, Raw: params}, nil\n")
		}
		b.WriteString("\t},\n")
	}
	b.WriteString("}\n\n")

	b.WriteString("// ParseServerNotification parses a server notification by method.\n")
	b.WriteString("func ParseServerNotification(method string, params json.RawMessage) (Notification, error) {\n")
	b.WriteString("\tif handler, ok := notificationParsers[method]; ok {\n\t\treturn handler(params)\n\t}\n")
	b.WriteString("\treturn Notification{Method: method, Raw: params}, nil\n}\n")

	b.WriteString("\n// IsKnownServerNotificationMethod reports whether method is in the generated server notification surface.\n")
	b.WriteString("func IsKnownServerNotificationMethod(method string) bool {\n")
	b.WriteString("\t_, ok := notificationParsers[method]\n\treturn ok\n}\n")

	return []byte(b.String())
}

func paramsType(name string) string {
	if name == "" {
		return "struct{}"
	}
	return name
}

func manualProtocolTypes() map[string]struct{} {
	return map[string]struct{}{
		"ApplyPatchApprovalParams":                {},
		"ApplyPatchApprovalResponse":              {},
		"CommandExecutionRequestApprovalParams":   {},
		"CommandExecutionRequestApprovalResponse": {},
		"ErrorNotification":                       {},
		"ExecCommandApprovalParams":               {},
		"ExecCommandApprovalResponse":             {},
		"FileChangeRequestApprovalParams":         {},
		"FileChangeRequestApprovalResponse":       {},
		"ItemCompletedNotification":               {},
		"PermissionsRequestApprovalParams":        {},
		"PermissionsRequestApprovalResponse":      {},
		"ReviewStartParams":                       {},
		"ThreadResumeResponse":                    {},
		"ThreadGoal":                              {},
		"ThreadGoalUpdatedNotification":           {},
		"ThreadForkParams":                        {},
		"ThreadForkResponse":                      {},
		"ThreadResumeParams":                      {},
		"ThreadStartParams":                       {},
		"ThreadStartResponse":                     {},
		"ToolRequestUserInputParams":              {},
		"ToolRequestUserInputResponse":            {},
		"TurnStartParams":                         {},
		"TurnStartParamsInputElem":                {},
		"TurnSteerParamsInputElem":                {},
		"TurnCompletedNotification":               {},
		"TurnStartedNotification":                 {},
	}
}

func methodName(method string) string {
	base := methodBaseName(method)
	if base == "" {
		return "Unknown"
	}
	return base
}

func normalizeGeneratedHeader(src []byte, codexCommit string) []byte {
	if len(src) == 0 {
		return src
	}
	parts := strings.SplitN(string(src), "\n", 2)
	if len(parts) == 0 {
		return src
	}
	header := strings.TrimSpace(parts[0])
	headerLower := strings.ToLower(header)
	if !strings.HasPrefix(headerLower, "// code generated") && !strings.HasPrefix(headerLower, "// generated") {
		return src
	}
	rest := ""
	if len(parts) > 1 {
		rest = strings.TrimLeft(parts[1], "\n")
	}
	for strings.Contains(rest, "\n\n\n") {
		rest = strings.ReplaceAll(rest, "\n\n\n", "\n\n")
	}
	return []byte(generatedHeader(codexCommit) + rest)
}

func generatedHeader(codexCommit string) string {
	return generatedHeaderBase + "// Source codex commit: " + codexCommit + "\n\n"
}

func writeProtocolMetadata(outDir, codexCommit, codexVersion string) error {
	var b strings.Builder
	b.WriteString(generatedHeader(codexCommit))
	b.WriteString("package codexproto\n\n")
	b.WriteString("// GeneratedCodexCommit is the Codex source commit used to generate this protocol package.\n")
	fmt.Fprintf(&b, "const GeneratedCodexCommit = %q\n\n", codexCommit)
	b.WriteString("// GeneratedCodexVersion is the Codex package version used to generate this protocol package.\n")
	fmt.Fprintf(&b, "const GeneratedCodexVersion = %q\n", codexVersion)
	return writeFile(outDir, "metadata_gen.go", []byte(b.String()))
}

func writeFile(dir, name string, contents []byte) error {
	path := filepath.Join(dir, name)
	if filepath.Ext(name) == ".go" {
		formatted, err := format.Source(contents)
		if err != nil {
			return fmt.Errorf("format %s: %w", name, err)
		}
		contents = formatted
	}
	return os.WriteFile(path, contents, 0o644)
}

func findSchemaFiles(schemaDir string) ([]string, error) {
	var files []string
	err := filepath.WalkDir(schemaDir, func(path string, d os.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if d.IsDir() {
			return nil
		}
		if filepath.Ext(path) != ".json" {
			return nil
		}
		files = append(files, path)
		return nil
	})
	if err != nil {
		return nil, err
	}
	sort.Strings(files)
	return files, nil
}

func safeDoFile(gen *generator.Generator, file string) (err error) {
	sanitized, cleanup, err := sanitizeSchemaFile(file)
	if err != nil {
		return err
	}
	defer cleanup()
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("panic: %v", r)
		}
	}()
	return gen.DoFile(sanitized)
}

func sanitizeSchemaFile(path string) (string, func(), error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return "", nil, err
	}
	var node any
	if err := json.Unmarshal(data, &node); err != nil {
		return "", nil, err
	}
	stripSubschemas(node)
	out, err := json.Marshal(node)
	if err != nil {
		return "", nil, err
	}
	tempDir, err := os.MkdirTemp("", "tutti-codexproto-sanitized-*")
	if err != nil {
		return "", nil, err
	}
	sanitizedName := "sanitized-" + filepath.Base(path)
	sanitizedPath := filepath.Join(tempDir, sanitizedName)
	if err := os.WriteFile(sanitizedPath, out, 0o644); err != nil {
		_ = os.RemoveAll(tempDir)
		return "", nil, err
	}
	cleanup := func() {
		_ = os.RemoveAll(tempDir)
	}
	return sanitizedPath, cleanup, nil
}

func stripSubschemas(node any) {
	switch value := node.(type) {
	case map[string]any:
		delete(value, "oneOf")
		delete(value, "anyOf")
		for _, child := range value {
			stripSubschemas(child)
		}
	case []any:
		for _, child := range value {
			stripSubschemas(child)
		}
	}
}

func repoRoot() (string, error) {
	cwd, err := os.Getwd()
	if err != nil {
		return "", err
	}
	if packageRoot := filepath.Join(cwd, "runtime", "codexproto"); isPackageRoot(packageRoot) {
		return packageRoot, nil
	}

	// Walk up until we find the codexproto package root.
	current := cwd
	for {
		if isPackageRoot(current) {
			return current, nil
		}
		parent := filepath.Dir(current)
		if parent == current {
			break
		}
		current = parent
	}

	return "", fmt.Errorf("unable to locate codexproto package root from %s", cwd)
}

func isPackageRoot(path string) bool {
	return exists(filepath.Join(path, "schema", "json", schemaBundleFile)) &&
		exists(filepath.Join(path, "internal", "codegen", "main.go"))
}

func exists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

func fatal(err error) {
	fmt.Fprintln(os.Stderr, err)
	os.Exit(1)
}