OpenShell/sdk/go/openshell/v1/sandbox_client.go

// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

package v1

import (
	"context"
	"fmt"
	"io"
	"time"

	"github.com/NVIDIA/OpenShell/sdk/go/openshell/v1/internal/converter"
	"github.com/NVIDIA/OpenShell/sdk/go/openshell/v1/types"
	pb "github.com/NVIDIA/OpenShell/sdk/go/proto/openshellv1"
	"google.golang.org/grpc"
)

const defaultPollInterval = 500 * time.Millisecond

type sandboxClient struct {
	client pb.OpenShellClient
}

var _ SandboxInterface = (*sandboxClient)(nil)
var _ SandboxTemplateCreateInterface = (*sandboxClient)(nil)

func newSandboxClient(conn grpc.ClientConnInterface) *sandboxClient {
	return &sandboxClient{client: pb.NewOpenShellClient(conn)}
}

func (s *sandboxClient) Create(ctx context.Context, workspace, name string, spec *SandboxSpec, labels map[string]string, opts ...CreateOptions) (*Sandbox, error) {
	protoSpec, err := converter.SandboxSpecToProtoChecked(spec)
	if err != nil {
		return nil, &StatusError{Code: ErrorInvalidArgument, Message: err.Error()}
	}
	req := &pb.CreateSandboxRequest{
		Name:           name,
		Spec:           protoSpec,
		Labels:         labels,
		WorkspaceScope: namedWorkspaceScope(workspace),
	}
	if len(opts) > 0 {
		req.Annotations = converter.CopyStringMap(opts[0].Annotations)
		req.ServiceExposures = serviceExposuresToProto(opts[0].ServiceExposures)
	}
	resp, err := s.client.CreateSandbox(ctx, req)
	if err != nil {
		return nil, converter.FromGRPCError(err)
	}
	sandbox := converter.SandboxFromProto(resp.GetSandbox())
	if sandbox != nil {
		sandbox.ServiceURLs = converter.CopyStringMap(resp.GetServiceUrls())
	}
	return sandbox, nil
}

func (s *sandboxClient) CreateFromTemplate(ctx context.Context, workspace, name, templateName string, spec *SandboxSpec, labels map[string]string, opts ...CreateOptions) (*Sandbox, error) {
	if templateName == "" {
		return nil, &StatusError{Code: ErrorInvalidArgument, Message: "template name is required"}
	}
	if err := validateTemplateCreateSpec(spec); err != nil {
		return nil, err
	}
	protoSpec, err := converter.SandboxSpecToProtoChecked(spec)
	if err != nil {
		return nil, &StatusError{Code: ErrorInvalidArgument, Message: err.Error()}
	}
	req := &pb.CreateSandboxRequest{
		Name:             name,
		Spec:             protoSpec,
		Labels:           labels,
		WorkspaceScope:   namedWorkspaceScope(workspace),
		WorkloadTemplate: templateName,
	}
	if len(opts) > 0 {
		req.Annotations = converter.CopyStringMap(opts[0].Annotations)
		req.ServiceExposures = serviceExposuresToProto(opts[0].ServiceExposures)
	}
	resp, err := s.client.CreateSandbox(ctx, req)
	if err != nil {
		return nil, converter.FromGRPCError(err)
	}
	sandbox := converter.SandboxFromProto(resp.GetSandbox())
	if sandbox != nil {
		sandbox.ServiceURLs = converter.CopyStringMap(resp.GetServiceUrls())
	}
	return sandbox, nil
}

func serviceExposuresToProto(exposures []types.ServiceExposure) []*pb.SandboxServiceExposure {
	result := make([]*pb.SandboxServiceExposure, 0, len(exposures))
	for _, exposure := range exposures {
		result = append(result, &pb.SandboxServiceExposure{
			Service:           exposure.Service,
			TargetPort:        exposure.TargetPort,
			AuthorizationMode: serviceAuthorizationModeToProto(exposure.AuthorizationMode),
		})
	}
	return result
}

func serviceAuthorizationModeToProto(mode types.ServiceAuthorizationMode) pb.ServiceAuthorizationMode {
	if mode == 0 {
		return pb.ServiceAuthorizationMode_SERVICE_AUTHORIZATION_MODE_STRIP
	}
	return pb.ServiceAuthorizationMode(mode)
}

func validateTemplateCreateSpec(spec *SandboxSpec) error {
	if spec == nil {
		return nil
	}
	if spec.LogLevel != "" || len(spec.Environment) > 0 || spec.Template != nil || spec.GPU || spec.GPUCount != nil {
		return &StatusError{Code: ErrorInvalidArgument, Message: "template creates only allow policy, providers, command, and tty in spec"}
	}
	return nil
}

func (s *sandboxClient) Get(ctx context.Context, workspace, name string) (*Sandbox, error) {
	resp, err := s.client.GetSandbox(ctx, &pb.GetSandboxRequest{
		Name:           name,
		WorkspaceScope: namedWorkspaceScope(workspace),
	})
	if err != nil {
		return nil, converter.FromGRPCError(err)
	}
	return converter.SandboxFromProto(resp.GetSandbox()), nil
}

func (s *sandboxClient) List(workspace string, opts ...ListOptions) (*Pager[*Sandbox], error) {
	pageSize, err := listPageSize(opts)
	if err != nil {
		return nil, err
	}
	var pageToken, labelSelector string
	var allWorkspaces bool
	if len(opts) > 0 {
		pageToken = opts[0].PageToken
		labelSelector = opts[0].LabelSelector
		allWorkspaces = opts[0].AllWorkspaces
	}
	workspaceScope := namedWorkspaceScope(workspace)
	if allWorkspaces {
		workspaceScope = allWorkspacesScope()
	}
	return newPager(pageToken, func(ctx context.Context, pageToken string) (*Page[*Sandbox], error) {
		req := &pb.ListSandboxesRequest{
			WorkspaceScope: workspaceScope,
			PageSize:       pageSize,
			PageToken:      pageToken,
			LabelSelector:  labelSelector,
		}
		resp, err := s.client.ListSandboxes(ctx, req)
		if err != nil {
			return nil, converter.FromGRPCError(err)
		}
		sandboxes := make([]*Sandbox, 0, len(resp.GetSandboxes()))
		for _, proto := range resp.GetSandboxes() {
			sandboxes = append(sandboxes, converter.SandboxFromProto(proto))
		}
		return &Page[*Sandbox]{Items: sandboxes, NextPageToken: resp.GetNextPageToken()}, nil
	}), nil
}

func (s *sandboxClient) ListAll(ctx context.Context, workspace string, opts ...ListOptions) ([]*Sandbox, error) {
	pager, err := s.List(workspace, opts...)
	if err != nil {
		return nil, err
	}
	return pager.All(ctx)
}

func (s *sandboxClient) Delete(ctx context.Context, workspace, name string, opts ...DeleteOptions) (*DeletionResult, error) {
	resp, err := s.client.DeleteSandbox(ctx, &pb.DeleteSandboxRequest{
		AllowMissing:   allowMissing(opts),
		Name:           name,
		WorkspaceScope: namedWorkspaceScope(workspace),
	})
	if err != nil {
		return nil, converter.FromGRPCError(err)
	}
	return &DeletionResult{Outcome: DeletionOutcome(resp.GetOutcome()), SandboxID: resp.GetSandboxId()}, nil
}

func (s *sandboxClient) Stop(ctx context.Context, workspace, name string) (*Sandbox, error) {
	resp, err := s.client.StopSandbox(ctx, &pb.StopSandboxRequest{
		Name:           name,
		WorkspaceScope: namedWorkspaceScope(workspace),
	})
	if err != nil {
		return nil, converter.FromGRPCError(err)
	}
	return converter.SandboxFromProto(resp.GetSandbox()), nil
}

func (s *sandboxClient) Start(ctx context.Context, workspace, name string) (*Sandbox, error) {
	resp, err := s.client.StartSandbox(ctx, &pb.StartSandboxRequest{
		Name:           name,
		WorkspaceScope: namedWorkspaceScope(workspace),
	})
	if err != nil {
		return nil, converter.FromGRPCError(err)
	}
	return converter.SandboxFromProto(resp.GetSandbox()), nil
}

func (s *sandboxClient) AttachProvider(ctx context.Context, workspace, sandboxName, providerName string, expectedResourceVersion uint64) (*AttachProviderResult, error) {
	resp, err := s.client.AttachSandboxProvider(ctx, &pb.AttachSandboxProviderRequest{
		Provider:                providerName,
		ExpectedResourceVersion: expectedResourceVersion,
		Sandbox:                 sandboxName,
		WorkspaceScope:          namedWorkspaceScope(workspace),
	})
	if err != nil {
		return nil, converter.FromGRPCError(err)
	}
	return &AttachProviderResult{
		Sandbox:  converter.SandboxFromProto(resp.GetSandbox()),
		Attached: resp.GetAttached(),
	}, nil
}

func (s *sandboxClient) DetachProvider(ctx context.Context, workspace, sandboxName, providerName string, expectedResourceVersion uint64) (*DetachProviderResult, error) {
	resp, err := s.client.DetachSandboxProvider(ctx, &pb.DetachSandboxProviderRequest{
		Provider:                providerName,
		ExpectedResourceVersion: expectedResourceVersion,
		Sandbox:                 sandboxName,
		WorkspaceScope:          namedWorkspaceScope(workspace),
	})
	if err != nil {
		return nil, converter.FromGRPCError(err)
	}
	return &DetachProviderResult{
		Sandbox:  converter.SandboxFromProto(resp.GetSandbox()),
		Detached: resp.GetDetached(),
	}, nil
}

func (s *sandboxClient) ListProviders(workspace, sandboxName string, opts ...ListOptions) (*Pager[*Provider], error) {
	pageSize, err := listPageSize(opts)
	if err != nil {
		return nil, err
	}
	pageToken := ""
	if len(opts) > 0 {
		pageToken = opts[0].PageToken
	}
	return newPager(pageToken, func(ctx context.Context, pageToken string) (*Page[*Provider], error) {
		resp, err := s.client.ListSandboxProviders(ctx, &pb.ListSandboxProvidersRequest{
			Sandbox:        sandboxName,
			WorkspaceScope: namedWorkspaceScope(workspace),
			PageSize:       pageSize,
			PageToken:      pageToken,
		})
		if err != nil {
			return nil, converter.FromGRPCError(err)
		}
		providers := make([]*Provider, 0, len(resp.GetProviders()))
		for _, proto := range resp.GetProviders() {
			providers = append(providers, converter.ProviderFromProto(proto))
		}
		return &Page[*Provider]{Items: providers, NextPageToken: resp.GetNextPageToken()}, nil
	}), nil
}

func (s *sandboxClient) ListAllProviders(ctx context.Context, workspace, sandboxName string, opts ...ListOptions) ([]*Provider, error) {
	pager, err := s.ListProviders(workspace, sandboxName, opts...)
	if err != nil {
		return nil, err
	}
	return pager.All(ctx)
}

func (s *sandboxClient) WaitReady(ctx context.Context, workspace, name string, opts ...WaitOptions) (*Sandbox, error) {
	return s.waitForPhase(ctx, workspace, name, SandboxReady, opts...)
}

func (s *sandboxClient) WaitStopped(ctx context.Context, workspace, name string, opts ...WaitOptions) (*Sandbox, error) {
	return s.waitForPhase(ctx, workspace, name, SandboxStopped, opts...)
}

func (s *sandboxClient) waitForPhase(ctx context.Context, workspace, name string, target SandboxPhase, opts ...WaitOptions) (*Sandbox, error) {
	interval := defaultPollInterval
	if len(opts) > 0 && opts[0].PollInterval > 0 {
		interval = opts[0].PollInterval
	}

	sb, err := s.Get(ctx, workspace, name)
	if err != nil {
		return nil, err
	}

	if result, termErr := checkTerminalPhase(sb, name, target); result != nil || termErr != nil {
		return result, termErr
	}

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return nil, contextError(ctx.Err())
		case <-ticker.C:
			sb, err = s.Get(ctx, workspace, name)
			if err != nil {
				return nil, err
			}
			if result, termErr := checkTerminalPhase(sb, name, target); result != nil || termErr != nil {
				return result, termErr
			}
		}
	}
}

func checkTerminalPhase(sb *Sandbox, name string, target SandboxPhase) (*Sandbox, error) {
	if sb.Status.Phase == target {
		return sb, nil
	}
	switch sb.Status.Phase {
	case SandboxCompleted:
		if target == SandboxReady {
			return sb, nil
		}
		return nil, &StatusError{Code: ErrorInternal, Message: fmt.Sprintf("sandbox %q completed before reaching %s", name, target)}
	case SandboxStopped:
		return nil, &StatusError{Code: ErrorInternal, Message: fmt.Sprintf("sandbox %q stopped before reaching %s", name, target)}
	case SandboxError:
		return nil, &StatusError{Code: ErrorInternal, Message: fmt.Sprintf("sandbox %q is in error state", name)}
	case SandboxDeleting:
		return nil, &StatusError{Code: ErrorInternal, Message: fmt.Sprintf("sandbox %q is being deleted", name)}
	default:
		return nil, nil
	}
}

func (s *sandboxClient) Watch(ctx context.Context, workspace, name string, opts ...WatchOptions) (WatchInterface[*Sandbox], error) {
	if name == "" {
		return nil, &StatusError{Code: ErrorInvalidArgument, Message: "sandbox name must not be empty"}
	}

	var watchOpts WatchOptions
	if len(opts) > 0 {
		watchOpts = opts[0]
	}
	if _, err := s.Get(ctx, workspace, name); err != nil {
		return nil, err
	}

	streamCtx, streamCancel := context.WithCancel(ctx)
	stream, err := s.client.WatchSandbox(streamCtx, &pb.WatchSandboxRequest{
		Sandbox:        name,
		WorkspaceScope: namedWorkspaceScope(workspace),
		FollowStatus:   true,
		StopOnTerminal: watchOpts.StopOnTerminal,
	})
	if err != nil {
		streamCancel()
		return nil, converter.FromGRPCError(err)
	}

	first, err := stream.Recv()
	if err != nil {
		streamCancel()
		return nil, converter.FromGRPCError(err)
	}

	ch := make(chan Event[*Sandbox], 64)
	w := newWatcher(ch, streamCancel)

	go func() {
		defer close(ch)
		defer streamCancel()
		ev := first
		isFirst := true
		for {
			if sbPayload, ok := ev.Payload.(*pb.SandboxStreamEvent_Sandbox); ok && sbPayload.Sandbox != nil {
				sandbox := converter.SandboxFromProto(sbPayload.Sandbox)
				eventType := EventModified
				if isFirst {
					eventType = EventAdded
					isFirst = false
				} else if sandbox.Status.Phase == SandboxDeleting {
					eventType = EventDeleted
				}
				select {
				case ch <- Event[*Sandbox]{Type: eventType, Object: sandbox}:
				case <-w.done:
					return
				}
				if watchOpts.StopOnTerminal && (sandbox.Status.Phase == SandboxReady || sandbox.Status.Phase == SandboxCompleted || sandbox.Status.Phase == SandboxStopped || sandbox.Status.Phase == SandboxError) {
					w.Stop()
					return
				}
			}
			var recvErr error
			ev, recvErr = stream.Recv()
			if recvErr != nil {
				if recvErr != io.EOF {
					select {
					case <-w.done:
						return
					default:
					}
					select {
					case ch <- Event[*Sandbox]{Type: EventError, Err: converter.FromGRPCError(recvErr)}:
					case <-w.done:
					}
				}
				return
			}
		}
	}()

	return w, nil
}

func (s *sandboxClient) GetLogs(ctx context.Context, workspace, sandboxName string, opts ...LogOption) (*LogResult, error) {
	if _, err := s.Get(ctx, workspace, sandboxName); err != nil {
		return nil, err
	}
	cfg := types.ApplyLogOptions(opts)
	req := &pb.GetSandboxLogsRequest{
		Sandbox:        sandboxName,
		WorkspaceScope: namedWorkspaceScope(workspace),
		Lines:          cfg.Lines(),
		Sources:        cfg.Sources(),
		MinLevel:       cfg.MinLevel(),
	}
	if !cfg.Since().IsZero() {
		req.SinceTime = converter.TimestampFromTime(cfg.Since())
	}

	resp, err := s.client.GetSandboxLogs(ctx, req)
	if err != nil {
		return nil, converter.FromGRPCError(err)
	}
	return converter.LogResultFromProto(resp), nil
}