diff --git a/internal/cli/custom/environments_runtime.go b/internal/cli/custom/environments_runtime.go index 3fa3e83..ed5f46a 100644 --- a/internal/cli/custom/environments_runtime.go +++ b/internal/cli/custom/environments_runtime.go @@ -15,6 +15,7 @@ package custom import ( + "cmp" "context" "errors" "fmt" @@ -25,6 +26,7 @@ import ( "strings" "time" + "github.com/google-gemini/gemini-api-cli/internal/cli/custom/envruntime" "github.com/google-gemini/gemini-api-cli/internal/config" "github.com/google-gemini/gemini-api-cli/internal/flagutil" "github.com/google-gemini/gemini-api-cli/internal/output" @@ -56,7 +58,7 @@ func resolveNoRetries(cmd *cobra.Command) bool { } func runtimeSignalContext(parent context.Context) (context.Context, context.CancelFunc) { - ctx, stop := signal.NotifyContext(parent, runtimeShutdownSignals()...) + ctx, stop := signal.NotifyContext(parent, envruntime.ShutdownSignals()...) go func() { <-ctx.Done() stop() @@ -88,15 +90,15 @@ func (f *workerCommonFlags) register(cmd *cobra.Command) { cmd.Flags().StringVar(&f.workerID, "worker-id", "", "Unique worker identifier (or GEMINI_WORKER_ID; default: worker--)") cmd.Flags().StringVar(&f.rootDir, "root-dir", ".", "Root directory confining built-in file tools (execute_bash is not sandboxed)") cmd.Flags().StringVar(&f.workDir, "workdir", ".", "Working directory for bash and relative tool paths (must reside within --root-dir)") - cmd.Flags().DurationVar(&f.heartbeatInterval, "heartbeat-interval", defaultHeartbeatInterval, "In-band stream heartbeat interval") - cmd.Flags().DurationVar(&f.commandTimeout, "command-timeout", 2*time.Minute, "Maximum timeout for execute_bash commands") + cmd.Flags().DurationVar(&f.heartbeatInterval, "heartbeat-interval", envruntime.DefaultHeartbeatInterval, "In-band stream heartbeat interval") + cmd.Flags().DurationVar(&f.commandTimeout, "command-timeout", envruntime.DefaultCommandTimeout, "Maximum timeout for execute_bash commands") cmd.Flags().Lookup("command-timeout").DefValue = "2m" cmd.Flags().BoolVar(&f.disableBash, "disable-bash", false, "Disable the execute_bash tool completely") cmd.Flags().BoolVar(&f.quiet, "quiet", false, "Suppress operational progress logs on stderr") } func (f *workerCommonFlags) validate() (string, string, error) { - if _, _, err := resolveSandboxDirs(f.rootDir, f.workDir, false); err != nil { + if _, _, err := envruntime.ResolveSandboxDirs(f.rootDir, f.workDir, false); err != nil { return "", "", usageError(err.Error()) } if f.heartbeatInterval <= 0 { @@ -105,14 +107,40 @@ func (f *workerCommonFlags) validate() (string, string, error) { if f.commandTimeout <= 0 { return "", "", usageError("--command-timeout must be > 0") } - resolvedEnvKey := firstNonEmpty(f.envKey, os.Getenv("GEMINI_ENVIRONMENT_KEY")) - if err := validateHeaderSecret("--environment-key", resolvedEnvKey); err != nil { + resolvedEnvKey := cmp.Or(f.envKey, os.Getenv("GEMINI_ENVIRONMENT_KEY")) + if err := envruntime.ValidateHeaderSecret("--environment-key", resolvedEnvKey); err != nil { return "", "", usageError(err.Error()) } - resolvedWorkerID := firstNonEmpty(f.workerID, os.Getenv("GEMINI_WORKER_ID")) + resolvedWorkerID := cmp.Or(f.workerID, os.Getenv("GEMINI_WORKER_ID")) return resolvedEnvKey, resolvedWorkerID, nil } +func newCLIPullWorker(cmd *cobra.Command, envClient envruntime.EnvironmentClient, common workerCommonFlags, absRoot, absWork, envKey, workerID string, maxConcurrent int, pollInterval time.Duration) *envruntime.PullWorker { + cfg := envruntime.CLIPullWorkerDefaults() + cfg.Client = envClient + cfg.Tools = envruntime.NewToolDispatcher(envruntime.ToolDispatcherConfig{ + FS: envruntime.NewSandboxFS(absRoot, absWork), + Shell: envruntime.NewSandboxShell(absWork), + DisableBash: common.disableBash, + CommandTimeout: common.commandTimeout, + DefaultWriteToFileOverwrite: true, + }) + cfg.EnvironmentKey = envKey + cfg.WorkerID = workerID + if maxConcurrent > 0 { + cfg.MaxConcurrent = maxConcurrent + } + if pollInterval > 0 { + cfg.PollInterval = pollInterval + } + cfg.HeartbeatInterval = common.heartbeatInterval + cfg.CommandTimeout = common.commandTimeout + cfg.NoRetries = resolveNoRetries(cmd) + cfg.NonRetryableCheck = isNonRetryableControlPlaneError + cfg.Logger = newRuntimeLogger(cmd, common.quiet) + return envruntime.NewPullWorker(cfg) +} + func newEnvironmentsConnectCmd() *cobra.Command { var ( common workerCommonFlags @@ -156,7 +184,7 @@ func newEnvironmentsConnectCmd() *cobra.Command { }) if isDryRun(cmd) { - _, err := envClient.LeaseInteractions(cmd.Context(), firstNonEmpty(resolvedWorkerID, "worker-dry-run"), int32(maxConcurrent)) + _, err := envClient.LeaseSimple(cmd.Context(), cmp.Or(resolvedWorkerID, "worker-dry-run"), int32(maxConcurrent)) if err != nil { return output.Error(cmd, err) } @@ -169,7 +197,7 @@ func newEnvironmentsConnectCmd() *cobra.Command { )) } - absRoot, absWork, err := resolveSandboxDirs(common.rootDir, common.workDir, true) + absRoot, absWork, err := envruntime.ResolveSandboxDirs(common.rootDir, common.workDir, true) if err != nil { return usageError(err.Error()) } @@ -177,16 +205,7 @@ func newEnvironmentsConnectCmd() *cobra.Command { ctx, stop := runtimeSignalContext(cmd.Context()) defer stop() - worker := newPullWorker(envClient, newFSServer(absRoot, absWork), newShellServer(absWork), pullWorkerOptions{ - WorkerID: resolvedWorkerID, - MaxConcurrent: maxConcurrent, - PollInterval: pollInterval, - HeartbeatInterval: common.heartbeatInterval, - CommandTimeout: common.commandTimeout, - DisableBash: common.disableBash, - NoRetries: resolveNoRetries(cmd), - Logger: newRuntimeLogger(cmd, common.quiet), - }) + worker := newCLIPullWorker(cmd, envClient, common, absRoot, absWork, resolvedEnvKey, resolvedWorkerID, maxConcurrent, pollInterval) if err := worker.Run(ctx); err != nil && !errors.Is(err, context.Canceled) { return output.Error(cmd, err) @@ -250,11 +269,11 @@ func newEnvironmentsPollCmd() *cobra.Command { return usageError("--handoff-heartbeat-interval must be > 0") } - resolvedEnvKey := firstNonEmpty(envKey, os.Getenv("GEMINI_ENVIRONMENT_KEY")) - if err := validateHeaderSecret("--environment-key", resolvedEnvKey); err != nil { + resolvedEnvKey := cmp.Or(envKey, os.Getenv("GEMINI_ENVIRONMENT_KEY")) + if err := envruntime.ValidateHeaderSecret("--environment-key", resolvedEnvKey); err != nil { return usageError(err.Error()) } - rawEnvID := firstNonEmpty(environmentID, os.Getenv("GEMINI_ENVIRONMENT_ID")) + rawEnvID := cmp.Or(environmentID, os.Getenv("GEMINI_ENVIRONMENT_ID")) var resolvedEnvID string if rawEnvID != "" { normID, err := normalizeInteractionID(strings.TrimPrefix(strings.TrimSpace(rawEnvID), "environments/")) @@ -263,7 +282,7 @@ func newEnvironmentsPollCmd() *cobra.Command { } resolvedEnvID = normID } - resolvedPollerID := firstNonEmpty(pollerID, os.Getenv("GEMINI_WORKER_ID")) + resolvedPollerID := cmp.Or(pollerID, os.Getenv("GEMINI_WORKER_ID")) rt, err := newRuntimeTransport(cmd, resolvedEnvKey != "") if err != nil { @@ -277,7 +296,7 @@ func newEnvironmentsPollCmd() *cobra.Command { }) if isDryRun(cmd) { - _, err := envClient.LeaseInteractions(cmd.Context(), firstNonEmpty(resolvedPollerID, "poller-dry-run"), int32(maxInteractions)) + _, err := envClient.LeaseSimple(cmd.Context(), cmp.Or(resolvedPollerID, "poller-dry-run"), int32(maxInteractions)) if err != nil { return output.Error(cmd, err) } @@ -296,13 +315,20 @@ func newEnvironmentsPollCmd() *cobra.Command { ctx, stop := runtimeSignalContext(cmd.Context()) defer stop() - poller := NewPollerController(envClient, NewCommandWorkerProvisioner(provisionCmd), PollerOptions{ + poller := envruntime.NewStatelessPoller(envruntime.StatelessPollerConfig{ + Client: envClient, + Heartbeat: envClient.HeartbeatInteraction, + ServerURL: envClient.ServerURL(), + Provisioner: envruntime.NewShellCommandProvisioner(provisionCmd), PollerID: resolvedPollerID, + EnvironmentID: resolvedEnvID, + EnvironmentKey: resolvedEnvKey, MaxInteractions: maxInteractions, PollInterval: pollInterval, ProvisionTimeout: provisionTimeout, HandoffHeartbeatInterval: handoffHeartbeatInterval, NoRetries: resolveNoRetries(cmd), + NonRetryableCheck: isNonRetryableControlPlaneError, Logger: newRuntimeLogger(cmd, quiet), }) @@ -328,7 +354,7 @@ func newEnvironmentsPollCmd() *cobra.Command { cmd.Flags().StringVar(&provisionCmd, "provision-cmd", "", "Shell command executed via sh -c to provision a worker and dispatch a leased interaction (receives v1 JSON on stdin)") cmd.Flags().DurationVar(&provisionTimeout, "provision-timeout", 2*time.Minute, "Timeout for --provision-cmd handoff") cmd.Flags().Lookup("provision-timeout").DefValue = "2m" - cmd.Flags().DurationVar(&handoffHeartbeatInterval, "handoff-heartbeat-interval", defaultHeartbeatInterval, "Heartbeat interval while provisioning and handing off a lease to a cold-starting worker") + cmd.Flags().DurationVar(&handoffHeartbeatInterval, "handoff-heartbeat-interval", envruntime.DefaultHeartbeatInterval, "Heartbeat interval while provisioning and handing off a lease to a cold-starting worker") cmd.Flags().BoolVar(&quiet, "quiet", false, "Suppress operational progress logs on stderr") return cmd @@ -389,7 +415,7 @@ func newEnvironmentsServeCmd() *cobra.Command { return output.UsageHelpError(cmd, errors.New("missing lease: pass --interaction-id and --lease-token for 'environments serve'")) } - absRoot, absWork, err := resolveSandboxDirs(common.rootDir, common.workDir, true) + absRoot, absWork, err := envruntime.ResolveSandboxDirs(common.rootDir, common.workDir, true) if err != nil { return usageError(err.Error()) } @@ -397,16 +423,9 @@ func newEnvironmentsServeCmd() *cobra.Command { ctx, stop := runtimeSignalContext(cmd.Context()) defer stop() - worker := newPullWorker(envClient, newFSServer(absRoot, absWork), newShellServer(absWork), pullWorkerOptions{ - WorkerID: resolvedWorkerID, - HeartbeatInterval: common.heartbeatInterval, - CommandTimeout: common.commandTimeout, - DisableBash: common.disableBash, - NoRetries: resolveNoRetries(cmd), - Logger: newRuntimeLogger(cmd, common.quiet), - }) + worker := newCLIPullWorker(cmd, envClient, common, absRoot, absWork, resolvedEnvKey, resolvedWorkerID, 0, 0) - if err := runStartupLeases(ctx, worker, startupLeases); err != nil && !errors.Is(err, context.Canceled) { + if err := worker.ServeLeases(ctx, startupLeases); err != nil && !errors.Is(err, context.Canceled) { return output.Error(cmd, err) } return nil @@ -441,7 +460,7 @@ func parseStartupLease(interactionID, leaseToken string) ([]*LeasedInteraction, if err != nil { return nil, err } - if err := validateHeaderSecret("--lease-token", tok); err != nil { + if err := envruntime.ValidateHeaderSecret("--lease-token", tok); err != nil { return nil, err } return []*LeasedInteraction{{ diff --git a/internal/cli/custom/envruntime/client.go b/internal/cli/custom/envruntime/client.go new file mode 100644 index 0000000..938a141 --- /dev/null +++ b/internal/cli/custom/envruntime/client.go @@ -0,0 +1,1328 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package envruntime + +import ( + "bytes" + "cmp" + "context" + "crypto/tls" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/url" + "os" + "regexp" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" +) + +const ( + // EnvironmentKeyHeader is the HTTP header carrying a self-hosted environment key. + EnvironmentKeyHeader = "x-goog-environment-key" + + // EnvironmentIDHeader is the HTTP header carrying an explicit environment identifier. + EnvironmentIDHeader = "x-goog-environment-id" + + // APIKeyHeader is the HTTP header carrying a Gemini API key. + APIKeyHeader = "x-goog-api-key" + + // LeaseTokenHeader is the HTTP header carrying an interaction lease token. + LeaseTokenHeader = "x-goog-lease-token" + + // WorkerIDHeader is the HTTP header carrying the worker identifier. + WorkerIDHeader = "x-goog-worker-id" + + // DefaultEnvironmentClientID is the client_id sent in ReconnectInteraction for environment workers. + DefaultEnvironmentClientID = "environment" + + // DefaultHeartbeatInterval is the default in-band stream heartbeat interval (5s = 15s server TTL / 3). + DefaultHeartbeatInterval = 5 * time.Second + + // DefaultLeaseTTL is the default lease duration requested in LeaseInteractions. + DefaultLeaseTTL = 15 * time.Second + + // DefaultPollBackoff is the fallback backoff duration between LeaseInteractions requests. + DefaultPollBackoff = 5 * time.Second + + // MaxControlPlaneBodyBytes caps HTTP control-plane error and JSON response reads at 32 MiB. + MaxControlPlaneBodyBytes = 32 << 20 +) + +var ( + validIdentifierRe = regexp.MustCompile(`^[a-zA-Z0-9._:\-]+$`) + leaseTokenJSONPattern = regexp.MustCompile(`("(?:lease_token|environment_key)"\s*:\s*)"(?:\\.|[^"\\\r\n])*"?`) +) + +// IsValidIdentifier reports whether s is a safe resource identifier without path traversal segments. +func IsValidIdentifier(s string) bool { + return s != "" && s != "." && s != ".." && !strings.Contains(s, "..") && validIdentifierRe.MatchString(s) +} + +// ValidateHeaderSecret verifies that val contains no ASCII control characters. +func ValidateHeaderSecret(flagName, val string) error { + for _, r := range val { + if r < 0x20 || r == 0x7f { + return fmt.Errorf("%s contains control characters", flagName) + } + } + return nil +} + +// DefaultRuntimeID returns a default worker or poller identifier combining prefix, hostname, and PID. +func DefaultRuntimeID(prefix string) string { + host, _ := os.Hostname() + if host == "" { + host = "local" + } + return fmt.Sprintf("%s-%s-%d", prefix, host, os.Getpid()) +} + +// HTTPDoer executes an *http.Request and is satisfied by *http.Client and CLI diagnostic wrappers. +type HTTPDoer interface { + Do(req *http.Request) (*http.Response, error) +} + +// EnvironmentClient abstracts control-plane leasing and data-plane stream attachment for workers and pollers. +type EnvironmentClient interface { + LeaseInteractions(ctx context.Context, req *LeaseInteractionsRequest) (*LeaseInteractionsResult, error) + ConnectInteraction(ctx context.Context, req *ConnectInteractionRequest) (InteractionStream, error) +} + +// InteractionStream represents an active bidirectional event stream for a leased interaction. +type InteractionStream interface { + Send(ctx context.Context, batch *ClientEventBatch) error + Recv(ctx context.Context) (*ServerEventBatch, error) + CloseSend() error + Close() error +} + +// HTTPStreamDialer opens a custom bidirectional interaction stream for HTTPClient. +type HTTPStreamDialer func(ctx context.Context, endpoint string, headers http.Header, req *ConnectInteractionRequest) (InteractionStream, error) + +// TransportError represents an HTTP or gRPC status failure returned by HTTPClient. +type TransportError struct { + Op string + StatusCode int + GRPCCode int + Message string + Body string + Resp *http.Response +} + +// Error implements error. +func (e *TransportError) Error() string { + detail := cmp.Or(strings.TrimSpace(e.Body), e.Message, http.StatusText(e.StatusCode)) + if e.Op != "" { + return fmt.Sprintf("%s: HTTP %d: %s", e.Op, e.StatusCode, detail) + } + if e.StatusCode > 0 { + return fmt.Sprintf("HTTP %d: %s", e.StatusCode, detail) + } + return detail +} + +// LeaseEntry represents an active interaction lease held by a worker or poller for renewal. +type LeaseEntry struct { + InteractionID string `json:"interaction_id"` + LeaseToken string `json:"lease_token,omitempty"` +} + +// LeaseInteractionsRequest contains parameters for heartbeating active leases and claiming new ones. +type LeaseInteractionsRequest struct { + EnvironmentID string + EnvironmentKey string + WorkerID string + ActiveLeases []LeaseEntry + MaxNewInteractions int32 + SendZeroMaxNew bool + // RequestedLeaseTTL and Labels are forwarded by custom or gRPC EnvironmentClient + // implementations and omitted by HTTPClient over REST. + RequestedLeaseTTL time.Duration + Labels map[string]string +} + +// EnvironmentSource represents a source mounted into an environment. +type EnvironmentSource struct { + Type string `json:"type,omitempty"` + Source string `json:"source,omitempty"` + Target string `json:"target,omitempty"` + Content string `json:"content,omitempty"` + Encoding string `json:"encoding,omitempty"` +} + +// EnvironmentEgressRule represents a domain egress rule in an EnvironmentConfig network allowlist. +type EnvironmentEgressRule struct { + Domain string `json:"domain,omitempty"` + Transform map[string]string `json:"transform,omitempty"` + Credential string `json:"credential,omitempty"` + Mode string `json:"mode,omitempty"` +} + +// EnvironmentEnvVar represents an environment variable value or credential reference. +type EnvironmentEnvVar struct { + Value string `json:"value,omitempty"` + Credential string `json:"credential,omitempty"` +} + +// EnvironmentConfig represents the sandbox environment configuration attached to a LeasedInteraction. +type EnvironmentConfig struct { + NetworkMode string `json:"network_mode,omitempty"` + NetworkAllowlist []EnvironmentEgressRule `json:"network_allowlist,omitempty"` + AllowlistDomains []string `json:"allowlist_domains,omitempty"` + DenylistDomains []string `json:"denylist_domains,omitempty"` + Sources []EnvironmentSource `json:"sources,omitempty"` + Env map[string]string `json:"env,omitempty"` + EnvVars map[string]EnvironmentEnvVar `json:"env_vars,omitempty"` + RawJSON json.RawMessage `json:"-"` +} + +// LeasedInteraction represents an interaction lease granted or renewed by LeaseInteractions. +type LeasedInteraction struct { + InteractionID string + EnvironmentID string + LeaseToken string + LeaseGeneration int64 + LeaseExpirationTime time.Time + RawLeaseExpirationTime string + EnvironmentKey string + EnvironmentConfig *EnvironmentConfig + RawEnvironmentConfig json.RawMessage + HeartbeatInterval time.Duration +} + +type leasedInteractionJSON struct { + InteractionID string `json:"interaction_id"` + EnvironmentID string `json:"environment_id,omitempty"` + LeaseToken string `json:"lease_token"` + LeaseGeneration json.RawMessage `json:"lease_generation,omitempty"` + LeaseExpirationTime string `json:"lease_expiration_time,omitempty"` + EnvironmentKey string `json:"environment_key,omitempty"` + EnvironmentConfig json.RawMessage `json:"environment_config,omitempty"` +} + +// UnmarshalJSON decodes LeasedInteraction, accepting ProtoJSON string or numeric lease_generation. +func (l *LeasedInteraction) UnmarshalJSON(data []byte) error { + var wire leasedInteractionJSON + if err := json.Unmarshal(data, &wire); err != nil { + return err + } + gen, err := ParseProtoInt64(wire.LeaseGeneration) + if err != nil { + return fmt.Errorf("invalid lease_generation: %w", err) + } + var expTime time.Time + if parsed, err := time.Parse(time.RFC3339Nano, wire.LeaseExpirationTime); err == nil { + expTime = parsed + } + var rawCfg json.RawMessage + if trimmed := bytes.TrimSpace(wire.EnvironmentConfig); string(trimmed) != "null" { + rawCfg = append(json.RawMessage(nil), trimmed...) + } + *l = LeasedInteraction{ + InteractionID: wire.InteractionID, + EnvironmentID: wire.EnvironmentID, + LeaseToken: wire.LeaseToken, + LeaseGeneration: gen, + LeaseExpirationTime: expTime, + RawLeaseExpirationTime: wire.LeaseExpirationTime, + EnvironmentKey: wire.EnvironmentKey, + EnvironmentConfig: ParseEnvironmentConfig(rawCfg), + RawEnvironmentConfig: rawCfg, + } + return nil +} + +// MarshalJSON encodes LeasedInteraction into JSON preserving raw environment_config bytes when present. +func (l LeasedInteraction) MarshalJSON() ([]byte, error) { + expStr := l.RawLeaseExpirationTime + if expStr == "" && !l.LeaseExpirationTime.IsZero() { + expStr = l.LeaseExpirationTime.UTC().Format(time.RFC3339Nano) + } + rawCfg := l.RawEnvironmentConfig + if len(rawCfg) == 0 && l.EnvironmentConfig != nil { + if len(l.EnvironmentConfig.RawJSON) > 0 { + rawCfg = l.EnvironmentConfig.RawJSON + } else { + encoded, err := json.Marshal(l.EnvironmentConfig) + if err != nil { + return nil, err + } + rawCfg = encoded + } + } + var rawGen json.RawMessage + if l.LeaseGeneration != 0 { + rawGen = json.RawMessage(strconv.FormatInt(l.LeaseGeneration, 10)) + } + return json.Marshal(leasedInteractionJSON{ + InteractionID: l.InteractionID, + EnvironmentID: l.EnvironmentID, + LeaseToken: l.LeaseToken, + LeaseGeneration: rawGen, + LeaseExpirationTime: expStr, + EnvironmentKey: l.EnvironmentKey, + EnvironmentConfig: rawCfg, + }) +} + +// LeaseInteractionsResult holds leases returned by LeaseInteractions and the recommended poll backoff. +type LeaseInteractionsResult struct { + Leases []LeasedInteraction + PollBackoff time.Duration +} + +// ConnectInteractionRequest contains parameters for attaching to an interaction stream as a worker. +type ConnectInteractionRequest struct { + InteractionID string + EnvironmentID string + EnvironmentKey string + LeaseToken string + WorkerID string + ClientID string + NextOutputSeqNum int64 + AckOutputSeqNum int64 + HeartbeatInterval time.Duration +} + +// HeartbeatResult summarizes the result of a lease renewal request during worker handoff. +type HeartbeatResult struct { + Renewed bool + ShouldCancel bool +} + +// NewLeaseTokenRedactor wraps w so any JSON "lease_token" string value written to it is replaced with "[REDACTED]". +func NewLeaseTokenRedactor(w io.Writer) io.Writer { + return &leaseTokenRedactor{inner: w} +} + +type leaseTokenRedactor struct { + mu sync.Mutex + inner io.Writer +} + +func (r *leaseTokenRedactor) Write(p []byte) (int, error) { + redacted := leaseTokenJSONPattern.ReplaceAll(p, []byte(`${1}"[REDACTED]"`)) + r.mu.Lock() + defer r.mu.Unlock() + if _, err := r.inner.Write(redacted); err != nil { + return 0, err + } + return len(p), nil +} + +// ParseProtoInt64 decodes a ProtoJSON int64 field encoded either as a JSON string or JSON number. +func ParseProtoInt64(raw json.RawMessage) (int64, error) { + trimmed := bytes.TrimSpace(raw) + if len(trimmed) == 0 || string(trimmed) == "null" { + return 0, nil + } + if trimmed[0] == '"' { + var s string + if err := json.Unmarshal(trimmed, &s); err != nil { + return 0, err + } + if s == "" { + return 0, nil + } + return strconv.ParseInt(strings.TrimSpace(s), 10, 64) + } + var n int64 + if err := json.Unmarshal(trimmed, &n); err != nil { + return 0, err + } + return n, nil +} + +// ParseDurationRaw decodes a ProtoJSON duration string ("5s"), object ({seconds,nanos}), or numeric seconds. +func ParseDurationRaw(raw json.RawMessage, fallback time.Duration) time.Duration { + trimmed := bytes.TrimSpace(raw) + var s string + if err := json.Unmarshal(trimmed, &s); err == nil { + if d, err := time.ParseDuration(s); err == nil && d > 0 { + return d + } + } + var obj struct { + Seconds int64 `json:"seconds"` + Nanos int64 `json:"nanos"` + } + if err := json.Unmarshal(trimmed, &obj); err == nil && (obj.Seconds > 0 || obj.Nanos > 0) { + return time.Duration(obj.Seconds)*time.Second + time.Duration(obj.Nanos)*time.Nanosecond + } + var sec float64 + if err := json.Unmarshal(trimmed, &sec); err == nil && sec > 0 { + return time.Duration(sec * float64(time.Second)) + } + return fallback +} + +type rawEnvironmentConfigWire struct { + NetworkMode string `json:"network_mode"` + NetworkAllowlist json.RawMessage `json:"network_allowlist"` + NetworkDenylist json.RawMessage `json:"network_denylist"` + AllowlistDomains []string `json:"allowlist_domains"` + DenylistDomains []string `json:"denylist_domains"` + Sources []EnvironmentSource `json:"sources"` + Env map[string]json.RawMessage `json:"env"` + EnvVars map[string]json.RawMessage `json:"env_vars"` +} + +// ParseEnvironmentConfig parses rawCfg into an *EnvironmentConfig while retaining RawJSON. +func ParseEnvironmentConfig(rawCfg json.RawMessage) *EnvironmentConfig { + raw := bytes.TrimSpace(rawCfg) + if len(raw) == 0 || string(raw) == "null" { + return nil + } + var parsed rawEnvironmentConfigWire + if err := json.Unmarshal(raw, &parsed); err != nil { + return &EnvironmentConfig{ + RawJSON: append(json.RawMessage(nil), raw...), + } + } + allowlist := parseNetworkAllowlist(parsed.NetworkAllowlist) + envMap, envVars := parseEnvironmentEnvVars(parsed.Env, parsed.EnvVars) + return &EnvironmentConfig{ + NetworkMode: parsed.NetworkMode, + NetworkAllowlist: allowlist, + AllowlistDomains: parsed.AllowlistDomains, + DenylistDomains: parsed.DenylistDomains, + Sources: parsed.Sources, + Env: envMap, + EnvVars: envVars, + RawJSON: append(json.RawMessage(nil), raw...), + } +} + +func parseNetworkAllowlist(raw json.RawMessage) []EnvironmentEgressRule { + trimmed := bytes.TrimSpace(raw) + var wrapper struct { + Allowlist []EnvironmentEgressRule `json:"allowlist"` + } + if err := json.Unmarshal(trimmed, &wrapper); err == nil { + return wrapper.Allowlist + } + var directRules []EnvironmentEgressRule + if err := json.Unmarshal(trimmed, &directRules); err == nil { + return directRules + } + var domains []string + if err := json.Unmarshal(trimmed, &domains); err == nil { + var rules []EnvironmentEgressRule + for _, d := range domains { + if d != "" { + rules = append(rules, EnvironmentEgressRule{Domain: d}) + } + } + return rules + } + return nil +} + +func parseEnvironmentEnvVars(maps ...map[string]json.RawMessage) (map[string]string, map[string]EnvironmentEnvVar) { + var envMap map[string]string + var envVars map[string]EnvironmentEnvVar + for _, m := range maps { + for k, rawVal := range m { + trimmed := bytes.TrimSpace(rawVal) + if len(trimmed) == 0 || string(trimmed) == "null" { + continue + } + var strVal string + if err := json.Unmarshal(trimmed, &strVal); err == nil { + if envMap == nil { + envMap = make(map[string]string, len(m)) + envVars = make(map[string]EnvironmentEnvVar, len(m)) + } + envMap[k] = strVal + envVars[k] = EnvironmentEnvVar{Value: strVal} + continue + } + var objVal EnvironmentEnvVar + if err := json.Unmarshal(trimmed, &objVal); err == nil { + if envMap == nil { + envMap = make(map[string]string, len(m)) + envVars = make(map[string]EnvironmentEnvVar, len(m)) + } + envVars[k] = objVal + if objVal.Value != "" { + envMap[k] = objVal.Value + } else if objVal.Credential != "" { + envMap[k] = objVal.Credential + } + } + } + } + return envMap, envVars +} + +// HTTPClientConfig configures an HTTPClient for control-plane REST and gRPC/HTTP2 streams. +type HTTPClientConfig struct { + BaseURL string + APIVersion string + StreamAPIVersion string + APIKey string + AccessToken string + EnvironmentID string + EnvironmentKey string + WorkerID string + UserProject string + APIRevision string + UserAgent string + Headers map[string]string + Timeout time.Duration + HTTPDoer HTTPDoer + StreamDoer HTTPDoer + TLSConfig *tls.Config + StreamDialer HTTPStreamDialer + Debug bool + DebugOut io.Writer + ProbePlainHTTPOnStreamError bool + RequireCredentialsOnLease bool +} + +// HTTPClient implements EnvironmentClient over Gemini API REST and HTTP/2 gRPC bidi endpoints. +type HTTPClient struct { + baseURL string + scheme string + apiVersion string + streamAPIVersion string + apiKey string + accessToken string + environmentID string + environmentKey string + workerID string + userProject string + apiRevision string + userAgent string + headers map[string]string + timeout time.Duration + httpDoer HTTPDoer + streamDoer HTTPDoer + streamDialer HTTPStreamDialer + debug bool + debugOut io.Writer + probePlainHTTP bool + requireLeaseCred bool +} + +func newDefaultHTTPTransport(tlsCfg *tls.Config) *http.Transport { + if dt, ok := http.DefaultTransport.(*http.Transport); ok { + tr := dt.Clone() + tr.Proxy = http.ProxyFromEnvironment + tr.ForceAttemptHTTP2 = true + tr.TLSClientConfig = tlsCfg + return tr + } + return &http.Transport{ + Proxy: http.ProxyFromEnvironment, + ForceAttemptHTTP2: true, + TLSClientConfig: tlsCfg, + } +} + +// NewHTTPClient creates an HTTPClient from cfg. +func NewHTTPClient(cfg HTTPClientConfig) *HTTPClient { + baseURL := strings.TrimRight(strings.TrimSpace(cfg.BaseURL), "/") + var scheme string + if u, err := url.Parse(baseURL); err == nil { + scheme = strings.ToLower(u.Scheme) + } + apiVersion := strings.Trim(strings.TrimSpace(cfg.APIVersion), "/") + if apiVersion == "" { + apiVersion = "v1beta" + } + streamAPIVersion := strings.Trim(strings.TrimSpace(cfg.StreamAPIVersion), "/") + if streamAPIVersion == "" { + streamAPIVersion = apiVersion + } + + var tlsCfg *tls.Config + if cfg.TLSConfig != nil { + tlsCfg = cfg.TLSConfig.Clone() + } else { + tlsCfg = &tls.Config{MinVersion: tls.VersionTLS12} + } + + httpDoer := cfg.HTTPDoer + if httpDoer == nil { + httpDoer = &http.Client{ + CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }, + Transport: newDefaultHTTPTransport(tlsCfg), + } + } + + streamDoer := cmp.Or(cfg.StreamDoer, httpDoer) + if cfg.StreamDoer == nil && scheme == "http" { + var h2cProtos http.Protocols + h2cProtos.SetUnencryptedHTTP2(true) + streamDoer = &http.Client{ + CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }, + Transport: &http.Transport{ + Protocols: &h2cProtos, + DialContext: (&net.Dialer{ + Timeout: 15 * time.Second, + KeepAlive: 30 * time.Second, + }).DialContext, + }, + } + } + + copiedHeaders := make(map[string]string, len(cfg.Headers)) + for k, v := range cfg.Headers { + copiedHeaders[k] = v + } + + return &HTTPClient{ + baseURL: baseURL, + scheme: scheme, + apiVersion: apiVersion, + streamAPIVersion: streamAPIVersion, + apiKey: cfg.APIKey, + accessToken: cfg.AccessToken, + environmentID: strings.TrimSpace(cfg.EnvironmentID), + environmentKey: strings.TrimSpace(cfg.EnvironmentKey), + workerID: strings.TrimSpace(cfg.WorkerID), + userProject: cfg.UserProject, + apiRevision: cfg.APIRevision, + userAgent: cfg.UserAgent, + headers: copiedHeaders, + timeout: cfg.Timeout, + httpDoer: httpDoer, + streamDoer: streamDoer, + streamDialer: cfg.StreamDialer, + debug: cfg.Debug, + debugOut: cmp.Or[io.Writer](cfg.DebugOut, os.Stderr), + probePlainHTTP: cfg.ProbePlainHTTPOnStreamError, + requireLeaseCred: cfg.RequireCredentialsOnLease, + } +} + +// ServerURL returns the normalized base URL configured on the client. +func (c *HTTPClient) ServerURL() string { + return c.baseURL +} + +// EndpointURL joins path onto BaseURL/. +func (c *HTTPClient) EndpointURL(path string) string { + return c.baseURL + "/" + c.apiVersion + "/" + strings.TrimLeft(path, "/") +} + +// BidiStreamMethodPath returns the gRPC HTTP/2 path for CreateInteractionBidiStreamByEnvironmentKey. +func (c *HTTPClient) BidiStreamMethodPath() string { + return fmt.Sprintf("google.learning.gemini.api.interactions.%s.InteractionsService/CreateInteractionBidiStreamByEnvironmentKey", c.streamAPIVersion) +} + +// BuildHeaders constructs standard authentication and diagnostic headers for an outbound request. +func (c *HTTPClient) BuildHeaders(envID, envKey, leaseToken string) http.Header { + h := make(http.Header) + if c.userAgent != "" { + h.Set("User-Agent", c.userAgent) + } + switch { + case c.apiKey != "": + h.Set(APIKeyHeader, c.apiKey) + case c.accessToken != "": + h.Set("Authorization", "Bearer "+c.accessToken) + } + if c.userProject != "" { + h.Set("x-goog-user-project", c.userProject) + } + if c.apiRevision != "" { + h.Set("Api-Revision", c.apiRevision) + } + hasCustomAuth := false + hasCustomAPIKey := false + for k, v := range c.headers { + if http.CanonicalHeaderKey(k) == "Authorization" { + hasCustomAuth = true + } + if strings.EqualFold(k, APIKeyHeader) { + hasCustomAPIKey = true + } + h.Set(k, v) + } + if hasCustomAuth && !hasCustomAPIKey { + h.Del(APIKeyHeader) + } + if envKey != "" { + h.Set(EnvironmentKeyHeader, envKey) + } + if envID != "" { + h.Set(EnvironmentIDHeader, envID) + } + if leaseToken != "" { + h.Set(LeaseTokenHeader, leaseToken) + } + return h +} + +func (c *HTTPClient) newRequest(ctx context.Context, method, path, envID, envKey, leaseToken string, body io.Reader) (*http.Request, error) { + fullURL := c.baseURL + "/" + strings.TrimLeft(path, "/") + req, err := http.NewRequestWithContext(ctx, method, fullURL, body) + if err != nil { + return nil, err + } + for k, vals := range c.BuildHeaders(envID, envKey, leaseToken) { + for _, v := range vals { + req.Header.Add(k, v) + } + } + return req, nil +} + +// DoJSON sends a JSON HTTP request to BaseURL/ and decodes any 2xx JSON body into respOut. +func (c *HTTPClient) DoJSON(ctx context.Context, method, relPath, envID, envKey, leaseToken, errPrefix string, reqBody, respOut any) error { + var bodyReader io.Reader + if reqBody != nil { + payload, err := json.Marshal(reqBody) + if err != nil { + return fmt.Errorf("marshaling request JSON: %w", err) + } + bodyReader = bytes.NewReader(payload) + } + + reqCtx := ctx + var cancel context.CancelFunc + if c.timeout > 0 { + reqCtx, cancel = context.WithTimeout(ctx, c.timeout) + defer cancel() + } + + req, err := c.newRequest(reqCtx, method, relPath, envID, envKey, leaseToken, bodyReader) + if err != nil { + return err + } + if reqBody != nil { + req.Header.Set("Content-Type", "application/json") + } + + resp, err := c.httpDoer.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + + body, err := io.ReadAll(io.LimitReader(resp.Body, MaxControlPlaneBodyBytes+1)) + if err != nil { + return err + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + if len(body) > MaxControlPlaneBodyBytes { + body = body[:MaxControlPlaneBodyBytes] + } + return &TransportError{ + Op: errPrefix, + StatusCode: resp.StatusCode, + GRPCCode: HTTPStatusToGRPCCode(resp.StatusCode), + Message: errPrefix, + Body: string(body), + Resp: resp, + } + } + if len(body) > MaxControlPlaneBodyBytes { + return fmt.Errorf("%s response exceeds %d-byte limit", cmp.Or(errPrefix, "HTTP JSON"), MaxControlPlaneBodyBytes) + } + if respOut != nil && len(bytes.TrimSpace(body)) > 0 { + if err := json.Unmarshal(body, respOut); err != nil { + return fmt.Errorf("decoding %s response: %w", cmp.Or(errPrefix, "HTTP JSON"), err) + } + } + return nil +} + +type httpLeaseInteractionsBody struct { + WorkerID string `json:"worker_id,omitempty"` + ActiveLeases []LeaseEntry `json:"active_leases,omitempty"` + MaxNewInteractions *int32 `json:"max_new_interactions,omitempty"` + EnvironmentID string `json:"environment_id,omitempty"` + // NOTE: requested_lease_ttl and labels are not part of the public REST schema + // and are rejected by the HTTP endpoint, so HTTPClient intentionally omits them. +} + +type httpLeaseInteractionsResponse struct { + Leases []LeasedInteraction `json:"leases"` + PollBackoff json.RawMessage `json:"poll_backoff,omitempty"` +} + +// LeaseInteractions heartbeats active leases and claims new queued interactions via POST //environments:leaseInteractions. +func (c *HTTPClient) LeaseInteractions(ctx context.Context, req *LeaseInteractionsRequest) (*LeaseInteractionsResult, error) { + return c.doLeaseRPC(ctx, req, "LeaseInteractions failed") +} + +func (c *HTTPClient) doLeaseRPC(ctx context.Context, req *LeaseInteractionsRequest, errPrefix string) (*LeaseInteractionsResult, error) { + if req == nil { + req = &LeaseInteractionsRequest{} + } + envID := cmp.Or(req.EnvironmentID, c.environmentID) + envKey := cmp.Or(req.EnvironmentKey, c.environmentKey) + if c.requireLeaseCred && envKey == "" && c.apiKey == "" && c.accessToken == "" { + return nil, &TransportError{ + Op: errPrefix, + StatusCode: http.StatusBadRequest, + GRPCCode: 3, + Message: "environment key or API key is required for LeaseInteractions", + } + } + + var maxNewPtr *int32 + if req.MaxNewInteractions > 0 || req.SendZeroMaxNew { + v := req.MaxNewInteractions + maxNewPtr = &v + } + + wireReq := httpLeaseInteractionsBody{ + WorkerID: cmp.Or(req.WorkerID, c.workerID), + ActiveLeases: req.ActiveLeases, + MaxNewInteractions: maxNewPtr, + EnvironmentID: envID, + } + + var parsed httpLeaseInteractionsResponse + if err := c.DoJSON(ctx, http.MethodPost, c.apiVersion+"/environments:leaseInteractions", envID, envKey, "", errPrefix, wireReq, &parsed); err != nil { + return nil, err + } + + leases := make([]LeasedInteraction, 0, len(parsed.Leases)) + for _, item := range parsed.Leases { + if item.InteractionID == "" { + continue + } + if item.EnvironmentID == "" && envID != "" { + item.EnvironmentID = envID + } + leases = append(leases, item) + } + + return &LeaseInteractionsResult{ + Leases: leases, + PollBackoff: ParseDurationRaw(parsed.PollBackoff, DefaultPollBackoff), + }, nil +} + +// LeaseSimple claims up to maxNew queued interactions for workerID. +func (c *HTTPClient) LeaseSimple(ctx context.Context, workerID string, maxNew int32) ([]*LeasedInteraction, error) { + res, err := c.LeaseInteractions(ctx, &LeaseInteractionsRequest{ + WorkerID: workerID, + MaxNewInteractions: maxNew, + SendZeroMaxNew: true, + EnvironmentID: c.environmentID, + EnvironmentKey: c.environmentKey, + }) + if err != nil { + return nil, err + } + out := make([]*LeasedInteraction, 0, len(res.Leases)) + for i := range res.Leases { + cp := res.Leases[i] + out = append(out, &cp) + } + return out, nil +} + +// HeartbeatInteraction renews a single interaction lease over HTTP with max_new_interactions=0. +func (c *HTTPClient) HeartbeatInteraction(ctx context.Context, pollerID, interactionID, envID, leaseToken string) (*HeartbeatResult, error) { + resolvedEnvID := cmp.Or(envID, c.environmentID) + res, err := c.doLeaseRPC(ctx, &LeaseInteractionsRequest{ + WorkerID: pollerID, + EnvironmentID: resolvedEnvID, + EnvironmentKey: c.environmentKey, + ActiveLeases: []LeaseEntry{{ + InteractionID: interactionID, + LeaseToken: leaseToken, + }}, + MaxNewInteractions: 0, + SendZeroMaxNew: true, + }, "LeaseInteractions heartbeat failed") + if err != nil { + var te *TransportError + if errors.As(err, &te) { + switch te.StatusCode { + case http.StatusNotFound, http.StatusConflict, http.StatusPreconditionFailed: + return &HeartbeatResult{ShouldCancel: true}, nil + } + } + return nil, err + } + for _, l := range res.Leases { + if l.InteractionID == interactionID { + return &HeartbeatResult{Renewed: true}, nil + } + } + return &HeartbeatResult{ShouldCancel: true}, nil +} + +// PreviewConnect sends a dry-run preview request for CreateInteractionBidiStreamByEnvironmentKey without opening a stream pipe. +func (c *HTTPClient) PreviewConnect(ctx context.Context, leaseToken string) error { + req, err := c.newRequest(ctx, http.MethodPost, c.BidiStreamMethodPath(), c.environmentID, c.environmentKey, leaseToken, http.NoBody) + if err != nil { + return err + } + req.Header.Set("Content-Type", "application/grpc") + req.Header.Set("TE", "trailers") + resp, err := c.httpDoer.Do(req) + if err != nil { + return err + } + _ = resp.Body.Close() + return nil +} + +// ConnectInteraction opens a bidirectional interaction event stream and sends the initial ReconnectInteraction handshake. +func (c *HTTPClient) ConnectInteraction(ctx context.Context, req *ConnectInteractionRequest) (InteractionStream, error) { + if req == nil || strings.TrimSpace(req.InteractionID) == "" { + return nil, &TransportError{ + Op: "ConnectInteraction failed", + StatusCode: http.StatusBadRequest, + GRPCCode: 3, + Message: "interaction ID is required", + } + } + envID := cmp.Or(req.EnvironmentID, c.environmentID) + envKey := cmp.Or(req.EnvironmentKey, c.environmentKey) + workerID := cmp.Or(req.WorkerID, c.workerID) + clientID := cmp.Or(req.ClientID, DefaultEnvironmentClientID) + nextOut := req.NextOutputSeqNum + if nextOut <= 0 { + nextOut = 1 + } + + if c.streamDialer != nil { + headers := c.BuildHeaders(envID, envKey, req.LeaseToken) + if workerID != "" { + headers.Set(WorkerIDHeader, workerID) + } + endpoint := c.EndpointURL("/interactions/" + url.PathEscape(req.InteractionID) + ":stream") + stream, err := c.streamDialer(ctx, endpoint, headers, &ConnectInteractionRequest{ + InteractionID: req.InteractionID, + EnvironmentID: envID, + EnvironmentKey: envKey, + LeaseToken: req.LeaseToken, + WorkerID: workerID, + ClientID: clientID, + NextOutputSeqNum: nextOut, + AckOutputSeqNum: req.AckOutputSeqNum, + HeartbeatInterval: req.HeartbeatInterval, + }) + if err != nil { + return nil, err + } + if req.HeartbeatInterval > 0 { + return NewHeartbeatStream(ctx, stream, req.HeartbeatInterval, req.AckOutputSeqNum), nil + } + return stream, nil + } + + pr, pw := io.Pipe() + methodPath := c.BidiStreamMethodPath() + httpReq, err := c.newRequest(ctx, http.MethodPost, methodPath, envID, envKey, req.LeaseToken, pr) + if err != nil { + _ = pw.Close() + _ = pr.Close() + return nil, err + } + httpReq.Header.Set("Content-Type", "application/grpc") + httpReq.Header.Set("TE", "trailers") + if workerID != "" { + httpReq.Header.Set(WorkerIDHeader, workerID) + } + + if c.debug { + _, _ = fmt.Fprintf(c.debugOut, "=== DEBUG STREAM ===\n%s %s\n", httpReq.Method, httpReq.URL.String()) + } + + handshake := MarshalClientEventBatch(&ClientEventBatch{ + AckOutputSeqNum: req.AckOutputSeqNum, + ReconnectInteraction: &ReconnectInteractionEvent{ + InteractionID: req.InteractionID, + NextOutputSeqNum: nextOut, + ClientID: clientID, + }, + }) + + writeErrCh := make(chan error, 1) + go func() { + writeErrCh <- GRPCWriteFrame(pw, handshake) + }() + + resp, err := c.streamDoer.Do(httpReq) + if err != nil { + _ = pw.CloseWithError(err) + _ = pr.Close() + if c.probePlainHTTP && c.scheme == "http" && ctx.Err() == nil { + if httpErr := c.probePlainHTTPError(ctx, methodPath, envID, envKey, req.LeaseToken); httpErr != nil { + return nil, httpErr + } + } + return nil, err + } + + if c.debug { + _, _ = fmt.Fprintf(c.debugOut, "=== DEBUG STREAM RESPONSE ===\n%s\n", resp.Status) + } + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + _ = pr.Close() + _ = pw.Close() + body, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<20)) + _ = resp.Body.Close() + return nil, &TransportError{ + Op: "ConnectInteraction failed", + StatusCode: resp.StatusCode, + GRPCCode: HTTPStatusToGRPCCode(resp.StatusCode), + Message: "ConnectInteraction failed", + Body: string(body), + Resp: resp, + } + } + if stErr := checkGRPCHeaders(resp); stErr != nil { + _ = pr.Close() + _ = pw.Close() + _ = resp.Body.Close() + return nil, stErr + } + select { + case werr := <-writeErrCh: + if werr != nil { + _ = pw.Close() + if !errors.Is(werr, io.ErrClosedPipe) { + _ = resp.Body.Close() + return nil, werr + } + // The server closed its read side early (for example, returning gRPC + // trailers without draining the request stream); keep resp.Body open so + // Recv surfaces buffered frames and trailer status codes. + } + case <-ctx.Done(): + _ = pw.CloseWithError(ctx.Err()) + _ = pr.Close() + _ = resp.Body.Close() + return nil, ctx.Err() + } + + var stream InteractionStream = &grpcInteractionStream{ + streamCtx: ctx, + resp: resp, + pw: pw, + } + if req.HeartbeatInterval > 0 { + stream = NewHeartbeatStream(ctx, stream, req.HeartbeatInterval, req.AckOutputSeqNum) + } + return stream, nil +} + +func (c *HTTPClient) probePlainHTTPError(ctx context.Context, methodPath, envID, envKey, leaseToken string) error { + probeReq, err := c.newRequest(ctx, http.MethodPost, methodPath, envID, envKey, leaseToken, http.NoBody) + if err != nil { + return nil + } + probeReq.Header.Set("Content-Type", "application/grpc") + resp, err := c.httpDoer.Do(probeReq) + if err != nil { + return nil + } + defer resp.Body.Close() + if resp.StatusCode >= 200 && resp.StatusCode < 300 { + return nil + } + body, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<20)) + return &TransportError{ + Op: "ConnectInteraction failed", + StatusCode: resp.StatusCode, + GRPCCode: HTTPStatusToGRPCCode(resp.StatusCode), + Message: "ConnectInteraction failed", + Body: string(body), + Resp: resp, + } +} + +type grpcInteractionStream struct { + streamCtx context.Context + resp *http.Response + pw *io.PipeWriter + sendMu sync.Mutex +} + +func (s *grpcInteractionStream) Send(ctx context.Context, batch *ClientEventBatch) error { + if ctx != nil && ctx.Err() != nil { + return ctx.Err() + } + payload := MarshalClientEventBatch(batch) + s.sendMu.Lock() + defer s.sendMu.Unlock() + if ctx != nil { + stop := context.AfterFunc(ctx, func() { + _ = s.pw.CloseWithError(ctx.Err()) + }) + defer stop() + } + err := GRPCWriteFrame(s.pw, payload) + if err != nil && ctx != nil && ctx.Err() != nil { + return ctx.Err() + } + return err +} + +func (s *grpcInteractionStream) Recv(ctx context.Context) (*ServerEventBatch, error) { + if ctx != nil && ctx.Err() != nil { + return nil, ctx.Err() + } + if ctx != nil && ctx != s.streamCtx { + stop := context.AfterFunc(ctx, func() { + _ = s.Close() + }) + defer stop() + } + frame, err := GRPCReadFrame(s.resp.Body) + if err != nil { + if ctx != nil && ctx.Err() != nil { + return nil, ctx.Err() + } + if errors.Is(err, io.EOF) { + if stErr := checkGRPCTrailers(s.resp); stErr != nil { + return nil, stErr + } + } + return nil, err + } + return UnmarshalServerEventBatch(frame) +} + +func (s *grpcInteractionStream) CloseSend() error { + s.sendMu.Lock() + defer s.sendMu.Unlock() + return s.pw.Close() +} + +func (s *grpcInteractionStream) Close() error { + _ = s.pw.CloseWithError(io.ErrClosedPipe) + return s.resp.Body.Close() +} + +// SendHeartbeat transmits an AckOutputSeqNum-only ClientEventBatch over stream to renew the server-side lease. +func SendHeartbeat(ctx context.Context, stream InteractionStream, ackOutputSeqNum int64) error { + if stream == nil { + return fmt.Errorf("interaction stream is nil") + } + return stream.Send(ctx, &ClientEventBatch{ + AckOutputSeqNum: ackOutputSeqNum, + }) +} + +// HeartbeatStream wraps an InteractionStream to send periodic in-band AckOutputSeqNum frames. +type HeartbeatStream struct { + inner InteractionStream + hbCtx context.Context + hbCancel context.CancelFunc + sendMu sync.Mutex + ackSeq atomic.Int64 + stoppedWg sync.WaitGroup +} + +// NewHeartbeatStream wraps inner with a background ticker sending SendHeartbeat every interval. +func NewHeartbeatStream(ctx context.Context, inner InteractionStream, interval time.Duration, initialAckSeqNum int64) *HeartbeatStream { + hbCtx, hbCancel := context.WithCancel(ctx) + hs := &HeartbeatStream{ + inner: inner, + hbCtx: hbCtx, + hbCancel: hbCancel, + } + hs.ackSeq.Store(initialAckSeqNum) + if interval <= 0 { + return hs + } + hs.stoppedWg.Add(1) + go func() { + defer hs.stoppedWg.Done() + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-hbCtx.Done(): + return + case <-ticker.C: + hs.sendMu.Lock() + err := SendHeartbeat(hbCtx, hs.inner, hs.ackSeq.Load()) + hs.sendMu.Unlock() + if err != nil { + return + } + } + } + }() + return hs +} + +// Send updates the tracked AckOutputSeqNum and forwards batch to the underlying stream. +func (s *HeartbeatStream) Send(ctx context.Context, batch *ClientEventBatch) error { + if batch != nil && batch.AckOutputSeqNum > 0 { + for { + cur := s.ackSeq.Load() + if batch.AckOutputSeqNum <= cur || s.ackSeq.CompareAndSwap(cur, batch.AckOutputSeqNum) { + break + } + } + } + s.sendMu.Lock() + defer s.sendMu.Unlock() + return s.inner.Send(ctx, batch) +} + +// Recv reads the next ServerEventBatch from the underlying stream. +func (s *HeartbeatStream) Recv(ctx context.Context) (*ServerEventBatch, error) { + batch, err := s.inner.Recv(ctx) + if err != nil { + s.hbCancel() + } + return batch, err +} + +// CloseSend stops the background heartbeat ticker and half-closes the underlying stream. +func (s *HeartbeatStream) CloseSend() error { + s.hbCancel() + s.stoppedWg.Wait() + s.sendMu.Lock() + defer s.sendMu.Unlock() + return s.inner.CloseSend() +} + +// Close stops the background heartbeat ticker and closes the underlying stream. +func (s *HeartbeatStream) Close() error { + s.hbCancel() + err := s.inner.Close() + s.stoppedWg.Wait() + return err +} + +func checkGRPCHeaders(resp *http.Response) error { + return parseGRPCStatus(resp.Header.Get("Grpc-Status"), resp.Header.Get("Grpc-Message"), resp) +} + +func checkGRPCTrailers(resp *http.Response) error { + return parseGRPCStatus(resp.Trailer.Get("Grpc-Status"), resp.Trailer.Get("Grpc-Message"), resp) +} + +func parseGRPCStatus(rawCode, rawMsg string, resp *http.Response) error { + trimmed := strings.TrimSpace(rawCode) + if trimmed == "" || trimmed == "0" { + return nil + } + code, err := strconv.Atoi(trimmed) + if err != nil { + code = 2 + } + msg := DecodeGRPCMessage(rawMsg) + if msg == "" { + msg = fmt.Sprintf("gRPC status %d", code) + } + httpStatus := GRPCCodeToHTTPStatus(code) + errBody, _ := json.Marshal(map[string]any{ + "error": map[string]any{ + "code": httpStatus, + "message": msg, + }, + }) + return &TransportError{ + Op: msg, + StatusCode: httpStatus, + GRPCCode: code, + Message: msg, + Body: string(errBody), + Resp: resp, + } +} + +// GRPCCodeToHTTPStatus maps a canonical gRPC status code to its HTTP status equivalent. +func GRPCCodeToHTTPStatus(code int) int { + switch code { + case 3, 9, 11: + return http.StatusBadRequest + case 4: + return http.StatusGatewayTimeout + case 5: + return http.StatusNotFound + case 6, 10: + return http.StatusConflict + case 7: + return http.StatusForbidden + case 8: + return http.StatusTooManyRequests + case 12: + return http.StatusNotImplemented + case 14: + return http.StatusServiceUnavailable + case 16: + return http.StatusUnauthorized + default: + return http.StatusInternalServerError + } +} + +// HTTPStatusToGRPCCode maps an HTTP status code to its canonical gRPC status code integer. +func HTTPStatusToGRPCCode(statusCode int) int { + switch statusCode { + case http.StatusBadRequest: + return 3 + case http.StatusUnauthorized: + return 16 + case http.StatusForbidden: + return 7 + case http.StatusNotFound: + return 5 + case http.StatusConflict: + return 10 + case http.StatusTooManyRequests: + return 8 + case http.StatusNotImplemented: + return 12 + case http.StatusServiceUnavailable, http.StatusBadGateway, http.StatusGatewayTimeout: + return 14 + default: + return 13 + } +} + +// IsNonRetryableStatus reports whether statusCode represents a non-retryable 400/401/403/404 control-plane HTTP status. +func IsNonRetryableStatus(statusCode int) bool { + switch statusCode { + case http.StatusBadRequest, http.StatusUnauthorized, http.StatusForbidden, http.StatusNotFound: + return true + } + return false +} + +// IsNonRetryableError reports whether err represents a non-retryable 400/401/403/404 control-plane error. +func IsNonRetryableError(err error) bool { + var te *TransportError + return errors.As(err, &te) && IsNonRetryableStatus(te.StatusCode) +} + +// RetryBackoff computes bounded exponential backoff capped at 30 seconds. +func RetryBackoff(base time.Duration, consecutive int) time.Duration { + if base <= 0 { + base = time.Second + } + var shift int + if consecutive > 1 { + shift = min(consecutive-1, 5) + } + return min(base*time.Duration(1<= 1 { + _ = GRPCWriteFrame(w, MarshalServerEventBatch(&ServerEventBatch{ + InteractionCompleted: &InteractionCompletedEvent{}, + })) + flusher.Flush() + return + } + } + } + default: + http.NotFound(w, r) + } + })) + + // Verify BuildHeaders, StreamAPIVersion, and custom StreamDoer. + hdrClient := NewHTTPClient(HTTPClientConfig{ + BaseURL: srvURL, + APIVersion: "v1beta", + StreamAPIVersion: "v1alpha", + UserAgent: "ua-custom", + AccessToken: "tok-bearer", + }) + if !strings.Contains(hdrClient.BidiStreamMethodPath(), ".v1alpha.") { + t.Errorf("BidiStreamMethodPath() = %q, want .v1alpha.", hdrClient.BidiStreamMethodPath()) + } + hFull := hdrClient.BuildHeaders("env-id-1", "ek-1", "lt-1") + if hFull.Get("User-Agent") != "ua-custom" || hFull.Get("Authorization") != "Bearer tok-bearer" || hFull.Get(EnvironmentIDHeader) != "env-id-1" || hFull.Get(EnvironmentKeyHeader) != "ek-1" || hFull.Get(LeaseTokenHeader) != "lt-1" { + t.Errorf("BuildHeaders(full) = %v", hFull) + } + hEmpty := hdrClient.BuildHeaders("", "", "") + if hEmpty.Values(EnvironmentIDHeader) != nil || hEmpty.Values(EnvironmentKeyHeader) != nil || hEmpty.Values(LeaseTokenHeader) != nil { + t.Errorf("BuildHeaders(empty) = %v, want no env/lease headers", hEmpty) + } + overrideHdrClient := NewHTTPClient(HTTPClientConfig{ + BaseURL: srvURL, + UserAgent: "ua-default", + APIKey: "default-api-key", + UserProject: "proj-default", + APIRevision: "rev-default", + Headers: map[string]string{ + "User-Agent": "ua-override", + "Authorization": "Bearer custom-override", + "x-goog-user-project": "proj-override", + "Api-Revision": "rev-override", + }, + }) + hOverride := overrideHdrClient.BuildHeaders("", "", "") + if hOverride.Get("User-Agent") != "ua-override" || hOverride.Get("Authorization") != "Bearer custom-override" || hOverride.Values("x-goog-api-key") != nil || hOverride.Get("x-goog-user-project") != "proj-override" || hOverride.Get("Api-Revision") != "rev-override" { + t.Errorf("BuildHeaders(override) = %v", hOverride) + } + explicitDualCredClient := NewHTTPClient(HTTPClientConfig{ + BaseURL: srvURL, + APIKey: "default-api-key", + Headers: map[string]string{ + "Authorization": "Bearer custom-override", + "x-goog-api-key": "explicit-hdr-key", + }, + }) + if hDual := explicitDualCredClient.BuildHeaders("", "", ""); hDual.Get("Authorization") != "Bearer custom-override" || hDual.Get("x-goog-api-key") != "explicit-hdr-key" { + t.Errorf("BuildHeaders(explicit dual headers) = %v", hDual) + } + + oversizedJSONClient := NewHTTPClient(HTTPClientConfig{ + BaseURL: srvURL, + HTTPDoer: httpDoerFunc(func(_ *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(strings.Repeat(" ", MaxControlPlaneBodyBytes+16))), + }, nil + }), + }) + if err := oversizedJSONClient.DoJSON(t.Context(), http.MethodGet, "v1beta/oversized", "", "", "", "oversized", nil, nil); err == nil || !strings.Contains(err.Error(), "response exceeds") { + t.Errorf("DoJSON(>MaxControlPlaneBodyBytes) = %v, want response exceeds limit error", err) + } + + var customStreamDoerCalled bool + customStreamClient := NewHTTPClient(HTTPClientConfig{ + BaseURL: srvURL, + StreamDoer: httpDoerFunc(func(_ *http.Request) (*http.Response, error) { + customStreamDoerCalled = true + return nil, errors.New("custom stream doer err") + }), + }) + if _, err := customStreamClient.ConnectInteraction(t.Context(), &ConnectInteractionRequest{InteractionID: "int_custom"}); err == nil || !customStreamDoerCalled { + t.Errorf("ConnectInteraction(custom StreamDoer) = %v, called=%v", err, customStreamDoerCalled) + } + + var dialReq *ConnectInteractionRequest + dialClient := NewHTTPClient(HTTPClientConfig{ + BaseURL: srvURL, + StreamDialer: func(_ context.Context, _ string, _ http.Header, req *ConnectInteractionRequest) (InteractionStream, error) { + cp := *req + dialReq = &cp + return nil, errors.New("dialer stop") + }, + }) + _, _ = dialClient.ConnectInteraction(t.Context(), &ConnectInteractionRequest{InteractionID: "int_d", NextOutputSeqNum: 5}) + if dialReq == nil || dialReq.NextOutputSeqNum != 5 { + t.Errorf("ConnectInteraction(NextOutputSeqNum=5) = %+v, want 5", dialReq) + } + + var debugBuf bytes.Buffer + client := NewHTTPClient(HTTPClientConfig{ + BaseURL: srvURL + "/", + EnvironmentKey: "ek_live", + UserProject: "proj-1", + APIRevision: "rev-1", + Headers: map[string]string{"X-Extra": "v"}, + Debug: true, + DebugOut: NewLeaseTokenRedactor(&debugBuf), + ProbePlainHTTPOnStreamError: true, + }) + if client.ServerURL() != srvURL { + t.Errorf("ServerURL() = %q, want %q", client.ServerURL(), srvURL) + } + if err := client.DoJSON(t.Context(), http.MethodGet, "v1beta/ping?body=json", "", "", "", "ping", nil, nil); err != nil { + t.Errorf("DoJSON(nil body, nil respOut with json body) failed: %v", err) + } + if _, hasUA := client.BuildHeaders("", "", "")["User-Agent"]; hasUA { + t.Errorf("BuildHeaders() set User-Agent when UserAgent config is empty: %v", client.BuildHeaders("", "", "")) + } + if gotPingUA, _ := lastPingUA.Load().(string); gotPingUA == "" { + t.Error("expected Go default User-Agent header on wire when UserAgent config is empty") + } + var emptyOut map[string]any + if err := client.DoJSON(t.Context(), http.MethodGet, "v1beta/ping", "", "", "", "ping", nil, &emptyOut); err != nil || emptyOut != nil { + t.Errorf("DoJSON(nil body, non-nil respOut with whitespace body) = (%v, %v)", emptyOut, err) + } + + // EnvironmentID inheritance + PollBackoff parsing + omit MaxNewInteractions when 0 and !SendZeroMaxNew. + res, err := client.LeaseInteractions(t.Context(), &LeaseInteractionsRequest{ + EnvironmentID: "env_inherited", + WorkerID: "backoff-parse", + MaxNewInteractions: 2, + }) + if err != nil || len(res.Leases) != 2 || res.Leases[0].InteractionID != "int_p1" || res.Leases[0].EnvironmentID != "env_inherited" || res.Leases[1].EnvironmentID != "env_explicit" || res.PollBackoff != 1500*time.Millisecond { + t.Fatalf("LeaseInteractions(backoff-parse) = (%+v, %v)", res, err) + } + if _, err := client.LeaseInteractions(t.Context(), &LeaseInteractionsRequest{ + WorkerID: "omit-max-new", + MaxNewInteractions: 0, + SendZeroMaxNew: false, + }); err != nil { + t.Errorf("LeaseInteractions(omit-max-new) failed: %v", err) + } + + // RequireCredentialsOnLease check. + strictNoCred := NewHTTPClient(HTTPClientConfig{BaseURL: srvURL, RequireCredentialsOnLease: true}) + if _, err := strictNoCred.LeaseSimple(t.Context(), "w1", 1); err == nil { + t.Error("LeaseSimple(RequireCredentialsOnLease without creds) = nil, want error") + } + + // Heartbeat outcomes: renewed, empty->cancel, [null]->neither, 409->cancel. + hbOK, err := client.HeartbeatInteraction(t.Context(), "hb-ok", "int_hb", "env_1", "tok_hb") + if err != nil || !hbOK.Renewed || hbOK.ShouldCancel { + t.Errorf("HeartbeatInteraction(ok) = (%+v, %v)", hbOK, err) + } + hbEmpty, err := client.HeartbeatInteraction(t.Context(), "hb-empty", "int_hb", "env_1", "tok_hb") + if err != nil || hbEmpty.Renewed || !hbEmpty.ShouldCancel { + t.Errorf("HeartbeatInteraction(empty) = (%+v, %v)", hbEmpty, err) + } + hbNull, err := client.HeartbeatInteraction(t.Context(), "hb-null", "int_hb", "env_1", "tok_hb") + if err != nil || hbNull.Renewed || !hbNull.ShouldCancel { + t.Errorf("HeartbeatInteraction(null) = (%+v, %v), want ShouldCancel=true", hbNull, err) + } + hbConflict, err := client.HeartbeatInteraction(t.Context(), "hb-conflict", "int_hb", "env_1", "tok_hb") + if err != nil || hbConflict.Renewed || !hbConflict.ShouldCancel { + t.Errorf("HeartbeatInteraction(409) = (%+v, %v)", hbConflict, err) + } + + // 403 TransportError on LeaseSimple. + deniedClient := NewHTTPClient(HTTPClientConfig{BaseURL: srvURL, EnvironmentKey: "ek_403"}) + if _, err := deniedClient.LeaseSimple(t.Context(), "w1", 1); err == nil || !IsNonRetryableError(err) { + t.Errorf("LeaseSimple(403) = %v, want non-retryable TransportError", err) + } + + // PreviewConnect invokes HTTPDoer with grpc headers and closes response body. + var previewPath, previewCT, previewTE string + var previewBodyClosed bool + previewClient := NewHTTPClient(HTTPClientConfig{ + BaseURL: srvURL, + EnvironmentKey: "ek_prev", + HTTPDoer: httpDoerFunc(func(req *http.Request) (*http.Response, error) { + previewPath = req.URL.Path + previewCT = req.Header.Get("Content-Type") + previewTE = req.Header.Get("TE") + return &http.Response{ + StatusCode: http.StatusOK, + Body: &trackingReadCloser{Reader: strings.NewReader(""), onClose: func() { previewBodyClosed = true }}, + }, nil + }), + }) + if err := previewClient.PreviewConnect(t.Context(), "tok_prev"); err != nil || !strings.Contains(previewPath, "CreateInteractionBidiStreamByEnvironmentKey") || previewCT != "application/grpc" || previewTE != "trailers" || !previewBodyClosed { + t.Errorf("PreviewConnect() = %v, path=%q ct=%q te=%q closed=%v", err, previewPath, previewCT, previewTE, previewBodyClosed) + } + + // ConnectInteraction with in-band HeartbeatStream, gRPC header error, and gRPC trailer error. + stream, err := client.ConnectInteraction(t.Context(), &ConnectInteractionRequest{ + InteractionID: "int_live", + LeaseToken: "tok_secret_stream", + WorkerID: "w-live", + AckOutputSeqNum: 9, + HeartbeatInterval: 20 * time.Millisecond, + }) + if err != nil { + t.Fatalf("ConnectInteraction() failed: %v", err) + } + if err := stream.Send(t.Context(), &ClientEventBatch{ + ReconnectInteraction: &ReconnectInteractionEvent{InteractionID: "int_live", ClientID: DefaultEnvironmentClientID}, + }); err != nil { + t.Fatalf("stream.Send() failed: %v", err) + } + evBatch, err := stream.Recv(t.Context()) + if err != nil || evBatch.InteractionCompleted == nil { + t.Fatalf("stream.Recv() = (%+v, %v), want InteractionCompleted", evBatch, err) + } + if err := stream.CloseSend(); err != nil { + t.Errorf("stream.CloseSend() = %v", err) + } + if err := stream.Send(t.Context(), &ClientEventBatch{}); !errors.Is(err, io.ErrClosedPipe) { + t.Errorf("stream.Send() after CloseSend() = %v, want io.ErrClosedPipe", err) + } + _ = stream.Close() + if heartbeatFrames.Load() < 1 { + t.Errorf("heartbeatFrames = %d, want >= 1", heartbeatFrames.Load()) + } + if !strings.Contains(debugBuf.String(), "=== DEBUG STREAM ===") || !strings.Contains(debugBuf.String(), "=== DEBUG STREAM RESPONSE ===") { + t.Errorf("expected debug stream output in debugBuf, got %q", debugBuf.String()) + } + if strings.Contains(debugBuf.String(), "tok_secret_stream") { + t.Errorf("debug output leaked lease_token: %q", debugBuf.String()) + } + + // Verify grpcInteractionStream.Close() closes underlying resp.Body without deadlocking an in-flight blocked Send, + // and HeartbeatStream CloseSend/Close wait for or unblock in-flight sends. + var rawBodyClosed bool + prBlocked, pwBlocked := io.Pipe() + defer prBlocked.Close() + rawStream := &grpcInteractionStream{ + resp: &http.Response{ + Body: &trackingReadCloser{Reader: strings.NewReader(""), onClose: func() { rawBodyClosed = true }}, + }, + pw: pwBlocked, + } + blockedSendErr := make(chan error, 1) + go func() { + blockedSendErr <- rawStream.Send(t.Context(), &ClientEventBatch{ + Events: []*ClientEvent{{SeqNum: 1, Payload: &EventPayload{SubInteractionID: strings.Repeat("x", 1024)}}}, + }) + }() + time.Sleep(10 * time.Millisecond) + if err := rawStream.Close(); err != nil || !rawBodyClosed { + t.Errorf("rawStream.Close() = %v, rawBodyClosed=%v", err, rawBodyClosed) + } + select { + case sErr := <-blockedSendErr: + if !errors.Is(sErr, io.ErrClosedPipe) { + t.Errorf("blocked rawStream.Send() after Close() = %v, want io.ErrClosedPipe", sErr) + } + case <-time.After(time.Second): + t.Fatal("rawStream.Close() deadlocked or failed to unblock in-flight Send") + } + + for _, useClose := range []bool{false, true} { + sendEntered := make(chan struct{}) + var sendCompleted atomic.Bool + var sendCalls atomic.Int32 + slowInner := &scriptedStream{ + sendFn: func(call int, _ *ClientEventBatch) error { + sendCalls.Add(1) + if call == 1 { + close(sendEntered) + time.Sleep(20 * time.Millisecond) + sendCompleted.Store(true) + } + return nil + }, + } + hs := NewHeartbeatStream(t.Context(), slowInner, 10*time.Millisecond, 3) + <-sendEntered + if useClose { + _ = hs.Close() + } else { + _ = hs.CloseSend() + } + if !sendCompleted.Load() { + t.Errorf("HeartbeatStream stop (useClose=%v) returned before in-flight Send completed", useClose) + } + afterStopCount := sendCalls.Load() + time.Sleep(25 * time.Millisecond) + if sendCalls.Load() != afterStopCount { + t.Errorf("HeartbeatStream (useClose=%v) continued sending after stop: before=%d after=%d", useClose, afterStopCount, sendCalls.Load()) + } + } + + var noDebugBuf bytes.Buffer + noDebugClient := NewHTTPClient(HTTPClientConfig{ + BaseURL: srvURL, + EnvironmentKey: "ek_live", + Debug: false, + DebugOut: &noDebugBuf, + }) + if _, err := noDebugClient.ConnectInteraction(t.Context(), &ConnectInteractionRequest{ + InteractionID: "int_hdr", + LeaseToken: "tok_grpc_hdr_err", + }); err == nil || !strings.Contains(err.Error(), "permission denied") { + t.Errorf("ConnectInteraction(Grpc-Status header 7) = %v, want permission denied", err) + } + if noDebugBuf.Len() != 0 { + t.Errorf("expected no debug output when Debug=false, got %q", noDebugBuf.String()) + } + + nilDebugOutClient := NewHTTPClient(HTTPClientConfig{ + BaseURL: srvURL, + EnvironmentKey: "ek_live", + Debug: true, + }) + if nilDebugOutClient.debugOut != os.Stderr { + t.Errorf("nilDebugOutClient.debugOut = %v, want os.Stderr", nilDebugOutClient.debugOut) + } + if _, err := nilDebugOutClient.ConnectInteraction(t.Context(), &ConnectInteractionRequest{ + InteractionID: "int_hdr_default_debug", + LeaseToken: "tok_grpc_hdr_err", + }); err == nil || !strings.Contains(err.Error(), "permission denied") { + t.Errorf("nilDebugOutClient.ConnectInteraction(Debug=true, DebugOut=nil) = %v, want permission denied", err) + } + + trStream, err := client.ConnectInteraction(t.Context(), &ConnectInteractionRequest{ + InteractionID: "int_tr", + LeaseToken: "tok_grpc_trailer_err", + }) + if err != nil { + t.Fatalf("ConnectInteraction(trailer setup) failed: %v", err) + } + if _, err := trStream.Recv(t.Context()); err == nil || !strings.Contains(err.Error(), "backend unavailable") { + t.Errorf("trStream.Recv(Grpc-Status trailer 14) = %v, want backend unavailable", err) + } + _ = trStream.Close() +} + +type trackingReadCloser struct { + io.Reader + onClose func() +} + +func (t *trackingReadCloser) Close() error { + if t.onClose != nil { + t.onClose() + } + return nil +} + +type httpDoerFunc func(*http.Request) (*http.Response, error) + +func (f httpDoerFunc) Do(req *http.Request) (*http.Response, error) { + return f(req) +} diff --git a/internal/cli/custom/envruntime/doc.go b/internal/cli/custom/envruntime/doc.go new file mode 100644 index 0000000..cedd474 --- /dev/null +++ b/internal/cli/custom/envruntime/doc.go @@ -0,0 +1,18 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package envruntime implements the shared self-hosted environment runtime +// transport, wire codec, filesystem and shell tool dispatchers, pull worker, +// and control-plane poller using only the Go standard library. +package envruntime diff --git a/internal/cli/custom/envruntime/events.go b/internal/cli/custom/envruntime/events.go new file mode 100644 index 0000000..d1fc1f7 --- /dev/null +++ b/internal/cli/custom/envruntime/events.go @@ -0,0 +1,242 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package envruntime + +// ClientEventBatch represents a batch of outbound client events and stream control signals. +type ClientEventBatch struct { + Events []*ClientEvent + AckOutputSeqNum int64 + ReconnectInteraction *ReconnectInteractionEvent + Disconnect *DisconnectEvent +} + +// ReconnectInteractionEvent attaches or resumes an interaction stream as a worker client. +type ReconnectInteractionEvent struct { + InteractionID string + NextOutputSeqNum int64 + Background bool + ClientID string +} + +// DisconnectEvent requests a graceful disconnection without terminating the interaction. +type DisconnectEvent struct{} + +// ClientEvent represents a sequenced outbound event from the worker to the server. +type ClientEvent struct { + Payload *EventPayload + SeqNum int64 + TimestampMicros int64 +} + +// EventPayload carries the typed event union shared by ClientEvent and ServerEvent. +type EventPayload struct { + SubInteractionID string + Step *Step + FunctionInvocation *FunctionInvocation + FunctionResult *FunctionResultEvent + CallHookRequest *CallHookRequest + CallHookResponse *CallHookResponse +} + +// Clone returns a deep copy of EventPayload suitable for caching across reconnects. +func (p *EventPayload) Clone() *EventPayload { + if p == nil { + return nil + } + cp := *p + if p.Step != nil { + st := *p.Step + if p.Step.ToolCall != nil { + tc := *p.Step.ToolCall + if p.Step.ToolCall.FunctionCall != nil { + fc := *p.Step.ToolCall.FunctionCall + fc.Arguments = cloneMap(p.Step.ToolCall.FunctionCall.Arguments) + tc.FunctionCall = &fc + } + st.ToolCall = &tc + } + if p.Step.ToolResult != nil { + tr := *p.Step.ToolResult + if p.Step.ToolResult.FunctionResult != nil { + fr := *p.Step.ToolResult.FunctionResult + fr.Result = cloneAny(p.Step.ToolResult.FunctionResult.Result) + tr.FunctionResult = &fr + } + st.ToolResult = &tr + } + cp.Step = &st + } + if p.FunctionInvocation != nil { + fi := *p.FunctionInvocation + fi.Arguments = cloneMap(p.FunctionInvocation.Arguments) + cp.FunctionInvocation = &fi + } + if p.FunctionResult != nil { + fr := *p.FunctionResult + fr.Result = cloneAny(p.FunctionResult.Result) + cp.FunctionResult = &fr + } + if p.CallHookRequest != nil { + hr := *p.CallHookRequest + if len(p.CallHookRequest.RawWire) > 0 { + hr.RawWire = append([]byte(nil), p.CallHookRequest.RawWire...) + } + cp.CallHookRequest = &hr + } + if p.CallHookResponse != nil { + resp := *p.CallHookResponse + if len(p.CallHookResponse.RawWire) > 0 { + resp.RawWire = append([]byte(nil), p.CallHookResponse.RawWire...) + } + cp.CallHookResponse = &resp + } + return &cp +} + +func cloneMap(m map[string]any) map[string]any { + if m == nil { + return nil + } + out := make(map[string]any, len(m)) + for k, v := range m { + out[k] = cloneAny(v) + } + return out +} + +func cloneAny(v any) any { + switch val := v.(type) { + case map[string]any: + return cloneMap(val) + case []any: + out := make([]any, len(val)) + for i, elem := range val { + out[i] = cloneAny(elem) + } + return out + case []string: + return append([]string(nil), val...) + case []byte: + return append([]byte(nil), val...) + default: + return val + } +} + +// Step represents a step payload containing a tool call or tool result. +type Step struct { + ToolCall *ToolCallStep + ToolResult *ToolResultStep +} + +// ToolCallStep represents an inbound tool call step. +type ToolCallStep struct { + ID string + FunctionCall *FunctionCallStep +} + +// FunctionCallStep represents a function name and structured argument map. +type FunctionCallStep struct { + Name string + Arguments map[string]any +} + +// ToolResultStep represents an outbound tool result step. +type ToolResultStep struct { + CallID string + FunctionResult *FunctionResultStep +} + +// FunctionResultStep represents a function execution result inside a ToolResultStep. +type FunctionResultStep struct { + Name string + IsError bool + Result any +} + +// FunctionInvocation represents a top-level server-to-client function call payload. +type FunctionInvocation struct { + CallID string + Name string + Arguments map[string]any +} + +// FunctionResultEvent represents a top-level client-to-server function result payload. +type FunctionResultEvent struct { + CallID string + Name string + IsError bool + Result any +} + +// CallHookRequest represents a server-to-client hook execution request. +// RawWire preserves the complete serialized protobuf bytes so callers with full hook descriptors +// can unmarshal nested hook-type payloads without data loss. +type CallHookRequest struct { + RequestID string + Name string + HookType int32 + RawWire []byte +} + +// CallHookResponse represents a client-to-server hook execution response. +// When RawWire is non-empty, MarshalClientEventBatch emits RawWire verbatim as the CallHookResponse message. +type CallHookResponse struct { + RequestID string + EmptyResult bool + ErrorMessage string + RawWire []byte +} + +// ServerEventBatch represents a batch of inbound server events and stream lifecycle signals. +type ServerEventBatch struct { + Events []*ServerEvent + AckInputSeqNum int64 + InteractionCompleted *InteractionCompletedEvent + InteractionReconnected *InteractionReconnectedEvent + ReconnectionRequired *ReconnectionRequiredEvent + Yield *YieldEvent +} + +// ServerEvent represents a sequenced inbound event from the server. +type ServerEvent struct { + Payload *EventPayload + SeqNum int64 + TimestampMicros int64 +} + +// RPCStatus represents a status code and message attached to terminal events. +type RPCStatus struct { + Code int32 + Message string +} + +// InteractionCompletedEvent signals that the interaction has reached a terminal state. +type InteractionCompletedEvent struct { + Status *RPCStatus +} + +// InteractionReconnectedEvent confirms attachment to a live interaction stream. +type InteractionReconnectedEvent struct { + NextInputSeqNum int64 +} + +// ReconnectionRequiredEvent instructs the worker to open a new stream and send ReconnectInteraction. +type ReconnectionRequiredEvent struct{} + +// YieldEvent signals that the server actor is currently idle and eligible for stream disconnection. +type YieldEvent struct { + NextInputSeqNum int64 +} diff --git a/internal/cli/custom/envruntime/hook.go b/internal/cli/custom/envruntime/hook.go new file mode 100644 index 0000000..ffb06b5 --- /dev/null +++ b/internal/cli/custom/envruntime/hook.go @@ -0,0 +1,49 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package envruntime + +import "context" + +// HookHandler serves CallHookRequest events over an interaction stream. +type HookHandler interface { + CallHook(ctx context.Context, req *CallHookRequest) (*CallHookResponse, error) +} + +// Observer receives lifecycle telemetry notifications from workers and pollers. +type Observer interface { + LeasePolled(leased int, err error) + WorkerProvisioned(reused bool, err error) + HandoffFinished(outcome string) + InteractionStarted(interactionID string) + InteractionFinished(interactionID string, err error) +} + +// NopObserver is an Observer whose methods are all no-ops. +type NopObserver struct{} + +// LeasePolled implements Observer. +func (NopObserver) LeasePolled(int, error) {} + +// WorkerProvisioned implements Observer. +func (NopObserver) WorkerProvisioned(bool, error) {} + +// HandoffFinished implements Observer. +func (NopObserver) HandoffFinished(string) {} + +// InteractionStarted implements Observer. +func (NopObserver) InteractionStarted(string) {} + +// InteractionFinished implements Observer. +func (NopObserver) InteractionFinished(string, error) {} diff --git a/internal/cli/custom/envruntime/hygiene_test.go b/internal/cli/custom/envruntime/hygiene_test.go new file mode 100644 index 0000000..83d58b4 --- /dev/null +++ b/internal/cli/custom/envruntime/hygiene_test.go @@ -0,0 +1,154 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package envruntime + +import ( + "go/ast" + "go/parser" + "go/token" + "os" + "path/filepath" + "regexp" + "strconv" + "strings" + "testing" +) + +func findPackageSourceDir(t *testing.T) string { + t.Helper() + const relPkg = "third_party/gemini_api_cli/internal/cli/custom/envruntime" + candidates := []string{ + ".", + relPkg, + } + if srcDir := os.Getenv("TEST_SRCDIR"); srcDir != "" { + ws := os.Getenv("TEST_WORKSPACE") + if ws == "" { + ws = "google3" + } + candidates = append(candidates, filepath.Join(srcDir, ws, relPkg)) + } + for _, dir := range candidates { + entries, err := os.ReadDir(dir) + if err != nil { + continue + } + for _, e := range entries { + if !e.IsDir() && strings.HasSuffix(e.Name(), ".go") { + return dir + } + } + } + t.Fatalf("could not locate package .go sources in any candidate directory: %v", candidates) + return "." +} + +func TestStdlibOnlyAndExportHygiene(t *testing.T) { + pkgDir := findPackageSourceDir(t) + entries, err := os.ReadDir(pkgDir) + if err != nil { + t.Fatalf("os.ReadDir(%q) failed: %v", pkgDir, err) + } + fset := token.NewFileSet() + bannedMutants := "NO" + "MUTANTS" + bannedNolint := "//" + "nolint" + bannedNolintUpper := "// " + "NOLINT" + internalTermPattern := regexp.MustCompile(`\b(` + strings.Join([]string{ + "One" + "Platform", + "GOOGLE_" + "INTERNAL", + "E" + "SF", + "Stub" + "by", + "stub" + "by", + "Pi" + "per", + "Criti" + "que", + "Bo" + "rg", + "Span" + "ner", + "Saw" + "mill", + "Sher" + "log", + }, "|") + `)\b`) + internalLinkPattern := regexp.MustCompile(`(^|[\s(])(go/[a-zA-Z0-9_-]+|(b|cl)/[0-9]+)`) + goFileCount := 0 + + for _, entry := range entries { + name := entry.Name() + if entry.IsDir() || !strings.HasSuffix(name, ".go") { + continue + } + if name == "hooks.go" { + t.Errorf("file %q is named hooks.go, which Speakeasy staging skips", name) + } + goFileCount++ + rawBytes, err := os.ReadFile(filepath.Join(pkgDir, name)) + if err != nil { + t.Fatalf("os.ReadFile(%q) failed: %v", name, err) + } + src := string(rawBytes) + if !strings.HasPrefix(src, "// Copyright 2026 Google LLC\n") { + t.Errorf("file %q missing Apache 2.0 Copyright 2026 Google LLC header", name) + } + + fileAST, err := parser.ParseFile(fset, name, rawBytes, parser.ParseComments) + if err != nil { + t.Fatalf("parser.ParseFile(%q) failed: %v", name, err) + } + importLiterals := make(map[*ast.BasicLit]bool, len(fileAST.Imports)) + for _, imp := range fileAST.Imports { + importLiterals[imp.Path] = true + impPath, err := strconv.Unquote(imp.Path.Value) + if err != nil { + t.Fatalf("strconv.Unquote(%s) failed: %v", imp.Path.Value, err) + } + firstElem, _, _ := strings.Cut(impPath, "/") + if strings.HasPrefix(impPath, "google3/") || strings.Contains(firstElem, ".") { + t.Errorf("file %q imports non-stdlib package %q, want Go standard library only", name, impPath) + } + } + for _, cg := range fileAST.Comments { + for _, c := range cg.List { + if strings.Contains(c.Text, bannedMutants) || strings.Contains(c.Text, bannedNolint) || strings.Contains(c.Text, bannedNolintUpper) { + t.Errorf("file %q contains disallowed directive comment %q", name, c.Text) + } + if match := internalTermPattern.FindString(c.Text); match != "" { + t.Errorf("file %q comment contains Google-internal term %q: %q", name, match, c.Text) + } + if match := internalLinkPattern.FindString(c.Text); match != "" { + t.Errorf("file %q comment contains Google-internal shortlink %q: %q", name, strings.TrimSpace(match), c.Text) + } + } + } + if name != "hygiene_test.go" { + ast.Inspect(fileAST, func(n ast.Node) bool { + lit, ok := n.(*ast.BasicLit) + if !ok || lit.Kind != token.STRING || importLiterals[lit] { + return true + } + val, err := strconv.Unquote(lit.Value) + if err != nil { + return true + } + if match := internalTermPattern.FindString(val); match != "" { + t.Errorf("file %q string literal contains Google-internal term %q: %q", name, match, val) + } + if match := internalLinkPattern.FindString(val); match != "" { + t.Errorf("file %q string literal contains Google-internal shortlink %q: %q", name, strings.TrimSpace(match), val) + } + return true + }) + } + } + if goFileCount == 0 { + t.Fatal("TestStdlibOnlyAndExportHygiene found 0 .go files") + } +} diff --git a/internal/cli/custom/envruntime/poller.go b/internal/cli/custom/envruntime/poller.go new file mode 100644 index 0000000..6d662dc --- /dev/null +++ b/internal/cli/custom/envruntime/poller.go @@ -0,0 +1,1777 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package envruntime + +import ( + "bytes" + "cmp" + "context" + "encoding/json" + "errors" + "fmt" + "log" + "log/slog" + "net/http" + "os" + "os/exec" + "runtime" + "slices" + "strings" + "sync" + "time" +) + +const ( + // WorkerProvisionSchemaVersionV1 is the v1 schema version for worker provision and suspend JSON contracts. + WorkerProvisionSchemaVersionV1 = "v1" + + // WorkerStatusAccepted is the status returned by StartWork when an interaction dispatch is accepted. + WorkerStatusAccepted = "accepted" + + // DefaultPollerMaxConcurrentInteractions is the default maximum concurrent in-flight handoffs for Tier2Poller. + DefaultPollerMaxConcurrentInteractions = 16 + + // DefaultMaxHandoffRetries is the default consecutive StartWork failure limit before poison-pill cooldown. + DefaultMaxHandoffRetries = 3 + + // DefaultHandoffRetryBackoff is the default initial backoff between handoff retries. + DefaultHandoffRetryBackoff = 50 * time.Millisecond + + // DefaultPoisonCooldown is the default duration during which an interaction that exhausted retries is skipped. + DefaultPoisonCooldown = 30 * time.Second + + // DefaultSweepInterval is the default interval between idle worker sweeps. + DefaultSweepInterval = 10 * time.Second + + // DefaultWorkerIdleTimeout is the default idle duration before a Tier-2 worker is suspended. + DefaultWorkerIdleTimeout = 60 * time.Second + + // DefaultWorkerResumeGracePeriod is the default cooldown window after provisioning before idle sweeps apply. + DefaultWorkerResumeGracePeriod = 3 * time.Second + + // DefaultDrainTimeout is the default graceful shutdown drain timeout for in-flight handoffs. + DefaultDrainTimeout = 15 * time.Second + + // DefaultLeaseCallTimeout bounds a single LeaseInteractions RPC call in Tier2Poller. + DefaultLeaseCallTimeout = 10 * time.Second + + // DefaultHandoffAttemptTimeout bounds a single worker provisioning + StartWork handoff cycle. + DefaultHandoffAttemptTimeout = 2 * time.Minute + + maxCommandErrorBytes = 256 +) + +// SingleflightResult holds the outcome delivered on a Singleflight.DoChan channel. +type SingleflightResult struct { + Val any + Err error + Shared bool +} + +type singleflightCall struct { + wg sync.WaitGroup + val any + err error + dups int + chans []chan<- SingleflightResult +} + +// Singleflight deduplicates concurrent in-flight function calls keyed by string using only the Go standard library. +type Singleflight struct { + mu sync.Mutex + m map[string]*singleflightCall +} + +// Do executes fn once for concurrent callers sharing key, returning (val, shared, err). +func (g *Singleflight) Do(key string, fn func() (any, error)) (any, bool, error) { + g.mu.Lock() + if g.m == nil { + g.m = make(map[string]*singleflightCall) + } + if c, ok := g.m[key]; ok { + c.dups++ + g.mu.Unlock() + c.wg.Wait() + return c.val, true, c.err + } + c := &singleflightCall{} + c.wg.Add(1) + g.m[key] = c + g.mu.Unlock() + + g.doCall(c, key, fn) + return c.val, c.dups > 0, c.err +} + +// DoChan executes fn once for concurrent callers sharing key and delivers the result on a 1-buffered channel. +func (g *Singleflight) DoChan(key string, fn func() (any, error)) <-chan SingleflightResult { + ch := make(chan SingleflightResult, 1) + g.mu.Lock() + if g.m == nil { + g.m = make(map[string]*singleflightCall) + } + if c, ok := g.m[key]; ok { + c.dups++ + c.chans = append(c.chans, ch) + g.mu.Unlock() + return ch + } + c := &singleflightCall{chans: []chan<- SingleflightResult{ch}} + c.wg.Add(1) + g.m[key] = c + g.mu.Unlock() + + go g.doCall(c, key, fn) + return ch +} + +func (g *Singleflight) doCall(c *singleflightCall, key string, fn func() (any, error)) { + normalReturn := false + defer func() { + if !normalReturn { + if r := recover(); r != nil { + c.err = fmt.Errorf("singleflight panic: %v", r) + } else if c.err == nil { + c.err = errors.New("singleflight goroutine exited unexpectedly") + } + } + g.mu.Lock() + delete(g.m, key) + chans := c.chans + shared := c.dups > 0 + g.mu.Unlock() + + c.wg.Done() + for _, ch := range chans { + ch <- SingleflightResult{Val: c.val, Err: c.err, Shared: shared} + } + }() + c.val, c.err = fn() + normalReturn = true +} + +// LeasePollLoopConfig configures RunLeasePollLoop for CLI workers and stateless pollers. +type LeasePollLoopConfig struct { + Client EnvironmentClient + CallerID string + EnvironmentID string + EnvironmentKey string + MaxSlots int + PollInterval time.Duration + NoRetries bool + NonRetryableCheck func(error) bool + Logger *log.Logger + Observer Observer + ErrLabel string + StartMsg string +} + +// RunLeasePollLoop polls LeaseInteractions up to MaxSlots concurrency and dispatches each claimed lease to handle. +func RunLeasePollLoop(ctx context.Context, cfg LeasePollLoopConfig, handle func(context.Context, *LeasedInteraction)) error { + maxSlots := cfg.MaxSlots + if maxSlots <= 0 { + maxSlots = 1 + } + pollInterval := cfg.PollInterval + if pollInterval <= 0 { + pollInterval = 2 * time.Second + } + logger := cfg.Logger + if logger == nil { + logger = log.New(os.Stderr, "", log.LstdFlags) + } + obs := cfg.Observer + if obs == nil { + obs = NopObserver{} + } + + sem := make(chan struct{}, maxSlots) + inFlight := make(map[string]bool) + var mu sync.Mutex + var wg sync.WaitGroup + defer wg.Wait() + + startedLogged := false + consecutiveErrors := 0 + + for { + select { + case <-ctx.Done(): + return ctx.Err() + default: + } + + availSlots := maxSlots - len(sem) + if availSlots <= 0 { + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(100 * time.Millisecond): + continue + } + } + + res, err := cfg.Client.LeaseInteractions(ctx, &LeaseInteractionsRequest{ + WorkerID: cfg.CallerID, + EnvironmentID: cfg.EnvironmentID, + EnvironmentKey: cfg.EnvironmentKey, + MaxNewInteractions: int32(availSlots), + SendZeroMaxNew: true, + }) + if err != nil { + obs.LeasePolled(0, err) + nonRetryable := IsNonRetryableError(err) || (cfg.NonRetryableCheck != nil && cfg.NonRetryableCheck(err)) + if nonRetryable || cfg.NoRetries { + return err + } + consecutiveErrors++ + delay := RetryBackoff(pollInterval, consecutiveErrors) + if ctx.Err() == nil { + logger.Printf("%s (retrying in %v): %v", cfg.ErrLabel, delay, err) + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(delay): + continue + } + } + obs.LeasePolled(len(res.Leases), nil) + consecutiveErrors = 0 + if !startedLogged { + startedLogged = true + logger.Print(cfg.StartMsg) + } + + if len(res.Leases) == 0 { + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(pollInterval): + continue + } + } + + for i := range res.Leases { + item := res.Leases[i] + if item.InteractionID == "" { + continue + } + mu.Lock() + if inFlight[item.InteractionID] { + mu.Unlock() + continue + } + inFlight[item.InteractionID] = true + mu.Unlock() + + select { + case sem <- struct{}{}: + case <-ctx.Done(): + mu.Lock() + delete(inFlight, item.InteractionID) + mu.Unlock() + return ctx.Err() + } + + wg.Add(1) + leaseCopy := item + go func(lease *LeasedInteraction) { + defer wg.Done() + defer func() { <-sem }() + defer func() { + mu.Lock() + delete(inFlight, lease.InteractionID) + mu.Unlock() + }() + handle(ctx, lease) + }(&leaseCopy) + } + } +} + +// ProvisionInput is the v1 JSON payload passed to --provision-cmd on stdin by StatelessPoller. +type ProvisionInput struct { + SchemaVersion string `json:"schema_version"` + InteractionID string `json:"interaction_id"` + EnvironmentID string `json:"environment_id"` + LeaseToken string `json:"lease_token"` + ServerURL string `json:"server_url,omitempty"` + EnvironmentConfig json.RawMessage `json:"environment_config,omitempty"` +} + +// ShellCommandProvisioner dispatches a leased interaction by invoking a shell command with ProvisionInput JSON on stdin. +type ShellCommandProvisioner struct { + provisionCmd string +} + +// NewShellCommandProvisioner creates a ShellCommandProvisioner for provisionCmd. +func NewShellCommandProvisioner(provisionCmd string) *ShellCommandProvisioner { + return &ShellCommandProvisioner{ + provisionCmd: strings.TrimSpace(provisionCmd), + } +} + +// ProvisionAndDispatch invokes provisionCmd with input serialized as JSON on stdin and LeaseToken redacted from stderr. +func (p *ShellCommandProvisioner) ProvisionAndDispatch(ctx context.Context, input ProvisionInput) error { + if p.provisionCmd == "" { + return fmt.Errorf("--provision-cmd is required to dispatch interaction %q (environment %q)", input.InteractionID, input.EnvironmentID) + } + if input.SchemaVersion == "" { + input.SchemaVersion = WorkerProvisionSchemaVersionV1 + } + + stdinBytes, err := json.Marshal(input) + if err != nil { + return fmt.Errorf("marshaling provision input: %w", err) + } + + shellBin, shellArgs := ResolveProvisionShell(runtime.GOOS, p.provisionCmd) + cmd := exec.CommandContext(ctx, shellBin, shellArgs...) + cmd.Stdin = bytes.NewReader(stdinBytes) + cmd.Env = FilterScrubbedEnv(os.Environ()) + ConfigureSubprocessGroup(cmd) + cmd.WaitDelay = 5 * time.Second + stderr := &CappedBuffer{MaxBytes: MaxBashOutputBytes} + cmd.Stderr = stderr + + if err := cmd.Run(); err != nil { + trimmed := strings.TrimSpace(stderr.String()) + if input.LeaseToken != "" { + trimmed = strings.ReplaceAll(trimmed, input.LeaseToken, "[REDACTED]") + } + if trimmed != "" { + return fmt.Errorf("provision command failed (%s): %w", trimmed, err) + } + return fmt.Errorf("provision command failed: %w", err) + } + return nil +} + +// ResolveProvisionShell selects sh or cmd.exe to execute a stateless --provision-cmd string. +func ResolveProvisionShell(goos, provisionCmd string) (string, []string) { + if pth, err := exec.LookPath("sh"); err == nil { + return pth, []string{"-c", provisionCmd} + } + if goos == "windows" { + return "cmd.exe", []string{"/C", provisionCmd} + } + return "sh", []string{"-c", provisionCmd} +} + +// HeartbeatFunc renews a lease while a stateless poller dispatches an interaction to a cold-starting worker. +type HeartbeatFunc func(ctx context.Context, pollerID, interactionID, envID, leaseToken string) (*HeartbeatResult, error) + +// StatelessPollerConfig configures a StatelessPoller (`gemini-api environments poll`). +type StatelessPollerConfig struct { + Client EnvironmentClient + Heartbeat HeartbeatFunc + ServerURL string + Provisioner *ShellCommandProvisioner + ProvisionFunc func(ctx context.Context, input ProvisionInput) error + PollerID string + EnvironmentID string + EnvironmentKey string + MaxInteractions int + PollInterval time.Duration + ProvisionTimeout time.Duration + HandoffHeartbeatInterval time.Duration + NoRetries bool + NonRetryableCheck func(error) bool + Logger *log.Logger + Observer Observer +} + +type statelessFailureState struct { + consecutive int + cooldownEnd time.Time +} + +// StatelessPoller polls LeaseInteractions and dispatches leases via ShellCommandProvisioner while heartbeating cold starts. +type StatelessPoller struct { + cfg StatelessPollerConfig + pollerID string + maxInteractions int + pollInterval time.Duration + provisionTimeout time.Duration + handoffHeartbeatInterval time.Duration + logger *log.Logger + observer Observer + + mu sync.Mutex + failures map[string]*statelessFailureState +} + +// NewStatelessPoller creates a new StatelessPoller from cfg. +func NewStatelessPoller(cfg StatelessPollerConfig) *StatelessPoller { + pollerID := strings.TrimSpace(cfg.PollerID) + if pollerID == "" { + pollerID = DefaultRuntimeID("poller") + } + maxInteractions := cfg.MaxInteractions + if maxInteractions <= 0 { + maxInteractions = 10 + } + pollInterval := cfg.PollInterval + if pollInterval <= 0 { + pollInterval = 2 * time.Second + } + provisionTimeout := cfg.ProvisionTimeout + if provisionTimeout <= 0 { + provisionTimeout = 2 * time.Minute + } + hbInterval := cfg.HandoffHeartbeatInterval + if hbInterval <= 0 { + hbInterval = DefaultHeartbeatInterval + } + logger := cfg.Logger + if logger == nil { + logger = log.New(os.Stderr, "", log.LstdFlags) + } + obs := cfg.Observer + if obs == nil { + obs = NopObserver{} + } + return &StatelessPoller{ + cfg: cfg, + pollerID: pollerID, + maxInteractions: maxInteractions, + pollInterval: pollInterval, + provisionTimeout: provisionTimeout, + handoffHeartbeatInterval: hbInterval, + logger: logger, + observer: obs, + failures: make(map[string]*statelessFailureState), + } +} + +// Run executes the stateless poller loop until ctx is canceled or a fatal non-retryable error occurs. +func (p *StatelessPoller) Run(ctx context.Context) error { + return RunLeasePollLoop(ctx, LeasePollLoopConfig{ + Client: p.cfg.Client, + CallerID: p.pollerID, + EnvironmentID: p.cfg.EnvironmentID, + EnvironmentKey: p.cfg.EnvironmentKey, + MaxSlots: p.maxInteractions, + PollInterval: p.pollInterval, + NoRetries: p.cfg.NoRetries, + NonRetryableCheck: p.cfg.NonRetryableCheck, + Logger: p.logger, + Observer: p.observer, + ErrLabel: "Poller LeaseInteractions error", + StartMsg: fmt.Sprintf("PollerController started (maxInteractions=%d)", p.maxInteractions), + }, p.dispatchLease) +} + +func (p *StatelessPoller) heartbeatOnce(ctx context.Context, interactionID, envID, leaseToken string) (*HeartbeatResult, error) { + if p.cfg.Heartbeat != nil { + return p.cfg.Heartbeat(ctx, p.pollerID, interactionID, envID, leaseToken) + } + if hc, ok := p.cfg.Client.(*HTTPClient); ok { + return hc.HeartbeatInteraction(ctx, p.pollerID, interactionID, envID, leaseToken) + } + resolvedEnvID := cmp.Or(envID, p.cfg.EnvironmentID) + res, err := p.cfg.Client.LeaseInteractions(ctx, &LeaseInteractionsRequest{ + WorkerID: p.pollerID, + EnvironmentID: resolvedEnvID, + EnvironmentKey: p.cfg.EnvironmentKey, + ActiveLeases: []LeaseEntry{{ + InteractionID: interactionID, + LeaseToken: leaseToken, + }}, + MaxNewInteractions: 0, + SendZeroMaxNew: true, + }) + if err != nil { + var te *TransportError + if errors.As(err, &te) { + switch te.StatusCode { + case http.StatusNotFound, http.StatusConflict, http.StatusPreconditionFailed: + return &HeartbeatResult{ShouldCancel: true}, nil + } + } + return nil, err + } + for _, l := range res.Leases { + if l.InteractionID == interactionID { + return &HeartbeatResult{Renewed: true}, nil + } + } + return &HeartbeatResult{ShouldCancel: true}, nil +} + +func (p *StatelessPoller) dispatchLease(ctx context.Context, lease *LeasedInteraction) { + envID := cmp.Or(lease.EnvironmentID, "default") + p.logger.Printf("[%s] Claimed interaction lease (env=%s)", lease.InteractionID, envID) + + p.mu.Lock() + if st, ok := p.failures[lease.InteractionID]; ok && time.Now().Before(st.cooldownEnd) { + waitDur := time.Until(st.cooldownEnd) + p.mu.Unlock() + select { + case <-ctx.Done(): + return + case <-time.After(waitDur): + } + } else { + p.mu.Unlock() + } + + dispatchCtx, cancel := context.WithTimeout(ctx, p.provisionTimeout) + defer cancel() + + hbDone := make(chan struct{}) + go func() { + defer close(hbDone) + ticker := time.NewTicker(p.handoffHeartbeatInterval) + defer ticker.Stop() + for { + select { + case <-dispatchCtx.Done(): + return + case <-ticker.C: + resp, err := p.heartbeatOnce(dispatchCtx, lease.InteractionID, lease.EnvironmentID, lease.LeaseToken) + if err == nil && resp != nil && resp.ShouldCancel { + cancel() + return + } + } + } + }() + defer func() { + cancel() + <-hbDone + }() + + serverURL := p.cfg.ServerURL + if serverURL == "" { + if hc, ok := p.cfg.Client.(*HTTPClient); ok { + serverURL = hc.ServerURL() + } + } + rawCfg := lease.RawEnvironmentConfig + if len(rawCfg) == 0 && lease.EnvironmentConfig != nil { + rawCfg = lease.EnvironmentConfig.RawJSON + } + input := ProvisionInput{ + SchemaVersion: WorkerProvisionSchemaVersionV1, + InteractionID: lease.InteractionID, + EnvironmentID: envID, + LeaseToken: lease.LeaseToken, + ServerURL: serverURL, + EnvironmentConfig: rawCfg, + } + + var err error + switch { + case p.cfg.ProvisionFunc != nil: + err = p.cfg.ProvisionFunc(dispatchCtx, input) + case p.cfg.Provisioner != nil: + err = p.cfg.Provisioner.ProvisionAndDispatch(dispatchCtx, input) + default: + err = fmt.Errorf("provisioner is not configured") + } + if err != nil { + p.recordDispatchFailure(lease.InteractionID) + p.observer.HandoffFinished("error") + p.logger.Printf("[%s] ProvisionAndDispatch failed for env %s: %v", lease.InteractionID, envID, err) + return + } + + p.mu.Lock() + delete(p.failures, lease.InteractionID) + p.mu.Unlock() + p.observer.HandoffFinished("accepted") + p.logger.Printf("[%s] Dispatched interaction lease for environment %s", lease.InteractionID, envID) +} + +func (p *StatelessPoller) recordDispatchFailure(interactionID string) { + p.mu.Lock() + defer p.mu.Unlock() + now := time.Now() + for id, entry := range p.failures { + if id != interactionID && now.Sub(entry.cooldownEnd) > 10*time.Minute { + delete(p.failures, id) + } + } + st, ok := p.failures[interactionID] + if !ok { + st = &statelessFailureState{} + p.failures[interactionID] = st + } + st.consecutive++ + shift := min(st.consecutive-1, 5) + st.cooldownEnd = now.Add(min(time.Duration(1< 0 { + suspBin = suspFields[0] + suspArgs = suspFields[1:] + } + return &DirectCommandProvisioner{ + provisionBinary: provFields[0], + provisionArgs: provFields[1:], + suspendBinary: suspBin, + suspendArgs: suspArgs, + }, nil +} + +func truncateCommandOutput(s string) string { + if trimmed := strings.TrimSpace(s); len(trimmed) > maxCommandErrorBytes { + return strings.ToValidUTF8(trimmed[:maxCommandErrorBytes], "") + "...(truncated)" + } + return strings.TrimSpace(s) +} + +// ProvisionWorker executes the provision binary with WorkerProvisionRequest JSON on stdin and parses WorkerProvisionResponse from stdout. +func (p *DirectCommandProvisioner) ProvisionWorker(ctx context.Context, req WorkerProvisionRequest) (*WorkerProvisionResponse, error) { + envID := strings.TrimSpace(req.EnvironmentID) + if envID == "" { + return nil, &PathError{Kind: PathErrInvalidArgument, Message: "environment_id is required"} + } + req.SchemaVersion = WorkerProvisionSchemaVersionV1 + req.EnvironmentID = envID + req.InteractionID = strings.TrimSpace(req.InteractionID) + + stdinBytes, err := json.Marshal(req) + if err != nil { + return nil, fmt.Errorf("marshaling WorkerProvisionRequest: %w", err) + } + + cmd := exec.CommandContext(ctx, p.provisionBinary, p.provisionArgs...) + cmd.Stdin = bytes.NewReader(stdinBytes) + cmd.Env = FilterScrubbedEnv(os.Environ()) + ConfigureSubprocessGroup(cmd) + cmd.WaitDelay = 5 * time.Second + + stdout := &CappedBuffer{MaxBytes: MaxBashOutputBytes} + stderr := &CappedBuffer{MaxBytes: MaxBashOutputBytes} + cmd.Stdout = stdout + cmd.Stderr = stderr + if err := cmd.Run(); err != nil { + return nil, fmt.Errorf("worker provision command failed (%v): %s", err, truncateCommandOutput(stderr.buf.String())) + } + if stdout.truncated { + return nil, &PathError{Kind: PathErrInternal, Message: "worker provision command output exceeded maximum size"} + } + + var resp WorkerProvisionResponse + if err := json.Unmarshal(bytes.TrimSpace(stdout.buf.Bytes()), &resp); err != nil { + return nil, fmt.Errorf("invalid WorkerProvisionResponse JSON from provision command: %w", err) + } + resp.URI = strings.TrimSpace(resp.URI) + if resp.URI == "" { + return nil, &PathError{Kind: PathErrInternal, Message: "worker provision command returned empty uri"} + } + if strings.TrimSpace(resp.SchemaVersion) == "" { + resp.SchemaVersion = WorkerProvisionSchemaVersionV1 + } + return &resp, nil +} + +// SuspendWorker executes the optional suspend binary passing WorkerSuspendRequest JSON on stdin. +func (p *DirectCommandProvisioner) SuspendWorker(ctx context.Context, req WorkerSuspendRequest) error { + if p.suspendBinary == "" { + return nil + } + envID := strings.TrimSpace(req.EnvironmentID) + workerID := strings.TrimSpace(req.WorkerID) + if envID == "" && workerID == "" { + return &PathError{Kind: PathErrInvalidArgument, Message: "environment_id or worker_id is required"} + } + req.SchemaVersion = WorkerProvisionSchemaVersionV1 + req.EnvironmentID = envID + req.WorkerID = workerID + req.IdleDuration = strings.TrimSpace(req.IdleDuration) + + stdinBytes, err := json.Marshal(req) + if err != nil { + return fmt.Errorf("marshaling WorkerSuspendRequest: %w", err) + } + + cmd := exec.CommandContext(ctx, p.suspendBinary, p.suspendArgs...) + cmd.Stdin = bytes.NewReader(stdinBytes) + cmd.Env = FilterScrubbedEnv(os.Environ()) + ConfigureSubprocessGroup(cmd) + cmd.WaitDelay = 5 * time.Second + + stderr := &CappedBuffer{MaxBytes: MaxBashOutputBytes} + cmd.Stderr = stderr + if err := cmd.Run(); err != nil { + return fmt.Errorf("worker suspend command failed (%v): %s", err, truncateCommandOutput(stderr.buf.String())) + } + return nil +} + +// ConvertToProvisionEnvironmentConfig converts a parsed *EnvironmentConfig into a *WorkerProvisionEnvironmentConfig. +func ConvertToProvisionEnvironmentConfig(cfg *EnvironmentConfig) *WorkerProvisionEnvironmentConfig { + if cfg == nil { + return nil + } + var envMap map[string]string + for k, v := range cfg.Env { + if envMap == nil { + envMap = make(map[string]string) + } + envMap[k] = v + } + for k, ev := range cfg.EnvVars { + if _, exists := envMap[k]; !exists { + if val := cmp.Or(ev.Value, ev.Credential); val != "" { + if envMap == nil { + envMap = make(map[string]string) + } + envMap[k] = val + } + } + } + + seenAllow := make(map[string]bool) + var allowlist []string + addAllow := func(d string) { + d = strings.TrimSpace(d) + if d != "" && !seenAllow[d] { + seenAllow[d] = true + allowlist = append(allowlist, d) + } + } + + seenDeny := make(map[string]bool) + var denylist []string + addDeny := func(d string) { + d = strings.TrimSpace(d) + if d != "" && !seenDeny[d] { + seenDeny[d] = true + denylist = append(denylist, d) + } + } + + for _, d := range cfg.AllowlistDomains { + addAllow(d) + } + for _, d := range cfg.DenylistDomains { + addDeny(d) + } + for _, rule := range cfg.NetworkAllowlist { + if strings.EqualFold(strings.TrimSpace(rule.Mode), "deny") { + addDeny(rule.Domain) + } else { + addAllow(rule.Domain) + } + } + + networkMode := strings.TrimSpace(cfg.NetworkMode) + if len(bytes.TrimSpace(cfg.RawJSON)) > 0 { + var raw rawEnvironmentConfigWire + if err := json.Unmarshal(cfg.RawJSON, &raw); err == nil { + for _, d := range raw.AllowlistDomains { + addAllow(d) + } + for _, d := range raw.DenylistDomains { + addDeny(d) + } + for _, rule := range parseNetworkAllowlist(raw.NetworkDenylist) { + addDeny(rule.Domain) + } + networkMode = cmp.Or(networkMode, strings.TrimSpace(raw.NetworkMode)) + } + } + + var sources []EnvironmentSource + if len(cfg.Sources) > 0 { + sources = append([]EnvironmentSource(nil), cfg.Sources...) + } + + if len(envMap) == 0 && len(allowlist) == 0 && len(denylist) == 0 && networkMode == "" && len(sources) == 0 { + return nil + } + return &WorkerProvisionEnvironmentConfig{ + NetworkMode: networkMode, + AllowlistDomains: allowlist, + DenylistDomains: denylist, + Env: envMap, + Sources: sources, + } +} + +// StartWorkRequest carries parameters dispatched from Tier2Poller to a Tier-2 WorkerRPCClient. +type StartWorkRequest struct { + InteractionID string + EnvironmentID string + LeaseToken string + ControlPlaneURL string + GRPCTarget string +} + +// StartWorkResponse carries the dispatch status returned by a Tier-2 WorkerRPCClient. +type StartWorkResponse struct { + Status string +} + +// IdleStatus carries worker activity state returned by WorkerRPCClient.GetIdleStatus. +type IdleStatus struct { + Running bool + IdleDuration time.Duration +} + +// WorkerRPCClient abstracts StartWork and GetIdleStatus RPC calls to a provisioned Tier-2 worker. +type WorkerRPCClient interface { + StartWork(ctx context.Context, req *StartWorkRequest, headers map[string]string) (*StartWorkResponse, error) + GetIdleStatus(ctx context.Context, headers map[string]string) (*IdleStatus, error) + Close() error +} + +// WorkerDialer opens a WorkerRPCClient connection to uri. +type WorkerDialer func(uri string) (WorkerRPCClient, error) + +type provisionedWorkerState struct { + EnvironmentID string + WorkerID string + URI string + Headers map[string]string + LastActiveAt time.Time +} + +// Tier2PollerConfig configures a stateful Tier2Poller controller. +type Tier2PollerConfig struct { + Client EnvironmentClient + Provisioner WorkerProvisioner + WorkerProvisionCmd string + WorkerSuspendCmd string + DialWorker WorkerDialer + IdleTimeout time.Duration + ResumeGracePeriod time.Duration + EnvironmentID string + EnvironmentKey string + WorkerID string + ControlPlaneURL string + GRPCTarget string + MaxConcurrentInteractions int + PollInterval time.Duration + LeaseTTL time.Duration + Labels map[string]string + MaxHandoffRetries int + HandoffRetryBackoff time.Duration + PoisonCooldown time.Duration + SweepInterval time.Duration + DrainTimeout time.Duration + Now func() time.Time + NewTimer func(d time.Duration) *time.Timer + Logger *slog.Logger + Observer Observer +} + +// Defaults populates default values for unset fields in Tier2PollerConfig. +func (c *Tier2PollerConfig) Defaults() { + if c.MaxConcurrentInteractions <= 0 { + c.MaxConcurrentInteractions = DefaultPollerMaxConcurrentInteractions + } + if c.PollInterval <= 0 { + c.PollInterval = DefaultPollBackoff + } + if c.LeaseTTL <= 0 { + c.LeaseTTL = DefaultLeaseTTL + } + if strings.TrimSpace(c.WorkerID) == "" { + host, _ := os.Hostname() + c.WorkerID = cmp.Or(host, "poller-controller") + } + if c.MaxHandoffRetries <= 0 { + c.MaxHandoffRetries = DefaultMaxHandoffRetries + } + if c.HandoffRetryBackoff <= 0 { + c.HandoffRetryBackoff = DefaultHandoffRetryBackoff + } + if c.PoisonCooldown <= 0 { + c.PoisonCooldown = DefaultPoisonCooldown + } + if c.SweepInterval <= 0 { + c.SweepInterval = DefaultSweepInterval + } + if c.DrainTimeout <= 0 { + c.DrainTimeout = DefaultDrainTimeout + } + if c.IdleTimeout <= 0 { + c.IdleTimeout = DefaultWorkerIdleTimeout + } + if c.ResumeGracePeriod <= 0 { + c.ResumeGracePeriod = DefaultWorkerResumeGracePeriod + } + if c.Now == nil { + c.Now = time.Now + } + if c.NewTimer == nil { + c.NewTimer = time.NewTimer + } + if c.Logger == nil { + c.Logger = slog.Default() + } + if c.Observer == nil { + c.Observer = NopObserver{} + } +} + +type provFlight struct { + done chan struct{} + cancel context.CancelFunc + waiters int + shared bool + resp *WorkerProvisionResponse + err error +} + +// Tier2Poller orchestrates stateful multi-environment polling, singleflight worker provisioning, +// poison-pill cooldowns, idle worker sweeps, and graceful drain heartbeats. +type Tier2Poller struct { + cfg Tier2PollerConfig + client EnvironmentClient + provisioner WorkerProvisioner + dialWorker WorkerDialer + logger *slog.Logger + observer Observer + + mu sync.Mutex + handoffLeases map[string]ActiveInteractionLease + failureCounts map[string]int + cooldownUntil map[string]time.Time + deadLettered map[string]bool + activeWorkers map[string]*provisionedWorkerState + provFlights map[string]*provFlight + dialGroup Singleflight + workerConns map[string]WorkerRPCClient + draining bool + handoffWg sync.WaitGroup + slotReleasedCh chan struct{} +} + +// NewTier2Poller constructs a Tier2Poller from cfg. +func NewTier2Poller(cfg Tier2PollerConfig) (*Tier2Poller, error) { + if cfg.Client == nil { + return nil, &PathError{Kind: PathErrInvalidArgument, Message: "EnvironmentClient is required for PollerController"} + } + provisioner := cfg.Provisioner + if provisioner == nil && strings.TrimSpace(cfg.WorkerProvisionCmd) != "" { + cmdProv, err := NewDirectCommandProvisioner(cfg.WorkerProvisionCmd, cfg.WorkerSuspendCmd) + if err != nil { + return nil, err + } + provisioner = cmdProv + } + if provisioner == nil { + return nil, &PathError{Kind: PathErrInvalidArgument, Message: "WorkerProvisioner or WorkerProvisionCmd is required for PollerController"} + } + if cfg.DialWorker == nil { + return nil, &PathError{Kind: PathErrInvalidArgument, Message: "DialWorker is required for PollerController"} + } + cfg.Defaults() + + return &Tier2Poller{ + cfg: cfg, + client: cfg.Client, + provisioner: provisioner, + dialWorker: cfg.DialWorker, + logger: cfg.Logger, + observer: cfg.Observer, + handoffLeases: make(map[string]ActiveInteractionLease), + failureCounts: make(map[string]int), + cooldownUntil: make(map[string]time.Time), + deadLettered: make(map[string]bool), + activeWorkers: make(map[string]*provisionedWorkerState), + provFlights: make(map[string]*provFlight), + workerConns: make(map[string]WorkerRPCClient), + slotReleasedCh: make(chan struct{}, 1), + }, nil +} + +// Close closes any cached WorkerRPCClient connections. +func (c *Tier2Poller) Close() error { + c.mu.Lock() + conns := make([]WorkerRPCClient, 0, len(c.workerConns)) + for uri, conn := range c.workerConns { + conns = append(conns, conn) + delete(c.workerConns, uri) + } + c.mu.Unlock() + + var firstErr error + for _, conn := range conns { + if err := conn.Close(); err != nil && firstErr == nil { + firstErr = err + } + } + return firstErr +} + +func (c *Tier2Poller) now() time.Time { + return c.cfg.Now() +} + +func (c *Tier2Poller) newTimer(d time.Duration) *time.Timer { + return c.cfg.NewTimer(d) +} + +func (c *Tier2Poller) notifySlotReleased() { + select { + case c.slotReleasedCh <- struct{}{}: + default: + } +} + +// HandoffLeases returns a snapshot copy of leases currently held in handoffLeases. +func (c *Tier2Poller) HandoffLeases() map[string]ActiveInteractionLease { + c.mu.Lock() + defer c.mu.Unlock() + out := make(map[string]ActiveInteractionLease, len(c.handoffLeases)) + for k, v := range c.handoffLeases { + out[k] = v + } + return out +} + +// InFlightCount returns the number of interactions currently in handoff. +func (c *Tier2Poller) InFlightCount() int { + c.mu.Lock() + defer c.mu.Unlock() + return len(c.handoffLeases) +} + +// FailureCount returns the consecutive handoff failure count for interactionID. +func (c *Tier2Poller) FailureCount(interactionID string) int { + c.mu.Lock() + defer c.mu.Unlock() + return c.failureCounts[interactionID] +} + +// IsDeadLettered reports whether interactionID exhausted MaxHandoffRetries. +func (c *Tier2Poller) IsDeadLettered(interactionID string) bool { + c.mu.Lock() + defer c.mu.Unlock() + return c.deadLettered[interactionID] +} + +// InCooldown reports whether interactionID is currently in its poison-pill cooldown window. +func (c *Tier2Poller) InCooldown(interactionID string) bool { + c.mu.Lock() + defer c.mu.Unlock() + until, ok := c.cooldownUntil[interactionID] + return ok && c.now().Before(until) +} + +// WaitHandoffs blocks until all active handoff goroutines complete. +func (c *Tier2Poller) WaitHandoffs() { + c.handoffWg.Wait() +} + +// Run executes the Tier2Poller main loop and idle worker sweep loop until ctx is canceled, then drains in-flight handoffs. +func (c *Tier2Poller) Run(ctx context.Context) error { + defer func() { + _ = c.Close() + }() + + stopHandoffsCh := make(chan struct{}) + var stopOnce sync.Once + cancelHandoffs := func() { + stopOnce.Do(func() { + close(stopHandoffsCh) + }) + } + defer cancelHandoffs() + + bgCtx, cancelBg := context.WithCancel(ctx) + var bgWg sync.WaitGroup + bgWg.Add(1) + go func() { + defer bgWg.Done() + c.runSweepLoop(bgCtx) + }() + defer func() { + cancelBg() + bgWg.Wait() + }() + + timer := c.newTimer(0) + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + defer timer.Stop() + + for { + if ctx.Err() != nil { + break + } + + wasAtCapacity := c.InFlightCount() >= c.cfg.MaxConcurrentInteractions + nextBackoff, err := c.pollOnce(ctx, stopHandoffsCh) + if err != nil { + if ctx.Err() != nil { + break + } + c.logger.Warn("PollerController LeaseInteractions cycle failed", slog.Any("error", err)) + nextBackoff = c.cfg.PollInterval + } + if nextBackoff <= 0 || c.cfg.PollInterval < nextBackoff { + nextBackoff = c.cfg.PollInterval + } + if c.InFlightCount() > 0 && c.cfg.LeaseTTL > 0 { + if hbInterval := c.cfg.LeaseTTL / 3; hbInterval > 0 && hbInterval < nextBackoff { + nextBackoff = hbInterval + } + } + + timer.Reset(nextBackoff) + if wasAtCapacity { + select { + case <-ctx.Done(): + stopAndDrainTimer(timer) + case <-c.slotReleasedCh: + stopAndDrainTimer(timer) + case <-timer.C: + } + } else { + select { + case <-ctx.Done(): + stopAndDrainTimer(timer) + case <-timer.C: + } + } + } + + c.drainInFlightHandoffs(stopHandoffsCh, cancelHandoffs) + return nil +} + +func stopAndDrainTimer(t *time.Timer) { + if !t.Stop() { + select { + case <-t.C: + default: + } + } +} + +func (c *Tier2Poller) runSweepLoop(ctx context.Context) { + ticker := time.NewTicker(c.cfg.SweepInterval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + if err := c.sweepProvisionedWorkers(ctx); err != nil && ctx.Err() == nil { + c.logger.Warn("PollerController sweepProvisionedWorkers failed", slog.Any("error", err)) + } + } + } +} + +func cloneStringMap(m map[string]string) map[string]string { + if len(m) == 0 { + return nil + } + out := make(map[string]string, len(m)) + for k, v := range m { + out[k] = v + } + return out +} + +func (c *Tier2Poller) waitContext(ctx context.Context, d time.Duration) error { + if d <= 0 { + d = 10 * time.Millisecond + } + t := c.newTimer(d) + defer t.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-t.C: + return nil + } +} + +func (c *Tier2Poller) getOrDialWorkerClient(uri string) (WorkerRPCClient, error) { + uri = strings.TrimSpace(uri) + if uri == "" { + return nil, &PathError{Kind: PathErrInvalidArgument, Message: "worker uri is required"} + } + + c.mu.Lock() + if conn := c.workerConns[uri]; conn != nil { + c.mu.Unlock() + return conn, nil + } + c.mu.Unlock() + + val, _, err := c.dialGroup.Do(uri, func() (any, error) { + c.mu.Lock() + if existing := c.workerConns[uri]; existing != nil { + c.mu.Unlock() + return existing, nil + } + c.mu.Unlock() + + conn, err := c.dialWorker(uri) + if err != nil { + return nil, fmt.Errorf("failed to dial worker uri %q: %w", uri, err) + } + if conn == nil { + return nil, &PathError{Kind: PathErrInternal, Message: fmt.Sprintf("worker dialer returned nil client for %q", uri)} + } + + c.mu.Lock() + if existing := c.workerConns[uri]; existing != nil { + c.mu.Unlock() + _ = conn.Close() + return existing, nil + } + c.workerConns[uri] = conn + c.mu.Unlock() + return conn, nil + }) + if err != nil { + return nil, err + } + client, ok := val.(WorkerRPCClient) + if !ok || client == nil { + return nil, &PathError{Kind: PathErrInternal, Message: fmt.Sprintf("invalid worker client for %q", uri)} + } + return client, nil +} + +func (c *Tier2Poller) provisionEnvironmentWorker(ctx context.Context, stopHandoffsCh <-chan struct{}, envID, interactionID string, cfg *EnvironmentConfig) (*WorkerProvisionResponse, error) { + c.mu.Lock() + call := c.provFlights[envID] + if call != nil { + call.waiters++ + call.shared = true + c.mu.Unlock() + } else { + provCtx, provCancel := context.WithTimeout(context.WithoutCancel(ctx), DefaultHandoffAttemptTimeout) + if stopHandoffsCh != nil { + go func() { + select { + case <-stopHandoffsCh: + provCancel() + case <-provCtx.Done(): + } + }() + } + call = &provFlight{ + done: make(chan struct{}), + cancel: provCancel, + waiters: 1, + } + c.provFlights[envID] = call + c.mu.Unlock() + + go func() { + normalReturn := false + defer func() { + if !normalReturn { + if r := recover(); r != nil { + call.err = fmt.Errorf("provision worker panic: %v", r) + } else if call.err == nil { + call.err = errors.New("provision worker goroutine exited unexpectedly") + } + } + call.cancel() + c.mu.Lock() + if c.provFlights[envID] == call { + delete(c.provFlights, envID) + } + c.mu.Unlock() + close(call.done) + }() + call.resp, call.err = c.provisioner.ProvisionWorker(provCtx, WorkerProvisionRequest{ + SchemaVersion: WorkerProvisionSchemaVersionV1, + EnvironmentID: envID, + InteractionID: interactionID, + EnvironmentConfig: ConvertToProvisionEnvironmentConfig(cfg), + }) + if call.err == nil && call.resp == nil { + call.err = &PathError{Kind: PathErrInternal, Message: "worker provisioner returned nil response"} + } + normalReturn = true + }() + } + + select { + case <-call.done: + c.mu.Lock() + shared := call.shared + c.mu.Unlock() + c.observer.WorkerProvisioned(shared, call.err) + if call.err != nil { + return nil, call.err + } + return call.resp, nil + case <-ctx.Done(): + c.mu.Lock() + call.waiters-- + if call.waiters <= 0 { + if c.provFlights[envID] == call { + delete(c.provFlights, envID) + } + call.cancel() + } + c.mu.Unlock() + return nil, ctx.Err() + } +} + +// SweepOnce executes a single idle worker sweep pass, suspending workers idle beyond IdleTimeout. +func (c *Tier2Poller) SweepOnce(ctx context.Context) error { + return c.sweepProvisionedWorkers(ctx) +} + +func (c *Tier2Poller) sweepProvisionedWorkers(ctx context.Context) error { + now := c.now() + + c.mu.Lock() + inFlightEnvs := make(map[string]bool, len(c.handoffLeases)) + for _, lease := range c.handoffLeases { + if envID := strings.TrimSpace(lease.EnvironmentID); envID != "" { + inFlightEnvs[envID] = true + } + } + workers := make([]provisionedWorkerState, 0, len(c.activeWorkers)) + for _, w := range c.activeWorkers { + if w == nil || inFlightEnvs[w.EnvironmentID] { + continue + } + if now.Sub(w.LastActiveAt) < c.cfg.ResumeGracePeriod { + continue + } + workers = append(workers, provisionedWorkerState{ + EnvironmentID: w.EnvironmentID, + WorkerID: w.WorkerID, + URI: w.URI, + Headers: cloneStringMap(w.Headers), + LastActiveAt: w.LastActiveAt, + }) + } + c.mu.Unlock() + + var firstErr error + for _, w := range workers { + client, err := c.getOrDialWorkerClient(w.URI) + if err != nil { + if firstErr == nil { + firstErr = err + } + continue + } + idleResp, err := client.GetIdleStatus(ctx, w.Headers) + if err != nil { + if firstErr == nil { + firstErr = err + } + continue + } + if idleResp.Running || idleResp.IdleDuration < c.cfg.IdleTimeout { + continue + } + + c.mu.Lock() + current, ok := c.activeWorkers[w.EnvironmentID] + if !ok || !current.LastActiveAt.Equal(w.LastActiveAt) { + c.mu.Unlock() + continue + } + delete(c.activeWorkers, w.EnvironmentID) + uriStillUsed := false + for _, other := range c.activeWorkers { + if other.URI == w.URI { + uriStillUsed = true + break + } + } + var staleConn WorkerRPCClient + if !uriStillUsed { + staleConn = c.workerConns[w.URI] + delete(c.workerConns, w.URI) + } + c.mu.Unlock() + if staleConn != nil { + _ = staleConn.Close() + } + + if err := c.provisioner.SuspendWorker(ctx, WorkerSuspendRequest{ + SchemaVersion: WorkerProvisionSchemaVersionV1, + EnvironmentID: w.EnvironmentID, + WorkerID: w.WorkerID, + IdleDuration: idleResp.IdleDuration.String(), + }); err != nil { + if firstErr == nil { + firstErr = err + } + continue + } + c.logger.Info("PollerController suspended idle worker", + slog.String("environment_id", w.EnvironmentID), + slog.String("worker_id", w.WorkerID), + slog.Duration("idle_duration", idleResp.IdleDuration), + ) + } + return firstErr +} + +func (c *Tier2Poller) drainInFlightHandoffs(stopHandoffsCh <-chan struct{}, cancelHandoffs func()) { + c.mu.Lock() + c.draining = true + c.mu.Unlock() + + drainCtx, cancelDrain := context.WithTimeout(context.Background(), c.cfg.DrainTimeout) + defer cancelDrain() + + allDoneCh := make(chan struct{}) + go func() { + c.handoffWg.Wait() + close(allDoneCh) + }() + + waitDur := min(c.cfg.PollInterval, max(c.cfg.LeaseTTL/3, time.Millisecond)) + for c.InFlightCount() > 0 && drainCtx.Err() == nil { + _, _ = c.pollOnce(drainCtx, stopHandoffsCh) + if c.InFlightCount() == 0 { + break + } + timer := c.newTimer(waitDur) + select { + case <-allDoneCh: + stopAndDrainTimer(timer) + case <-drainCtx.Done(): + stopAndDrainTimer(timer) + case <-timer.C: + } + } + + cancelHandoffs() + <-allDoneCh +} + +// PollOnce executes a single LeaseInteractions poll and launches non-blocking handoff goroutines for claimed leases. +func (c *Tier2Poller) PollOnce(ctx context.Context) (time.Duration, error) { + return c.pollOnce(ctx, nil) +} + +func (c *Tier2Poller) pollOnce(ctx context.Context, stopHandoffsCh <-chan struct{}) (time.Duration, error) { + now := c.now() + + c.mu.Lock() + for id, until := range c.cooldownUntil { + if !now.Before(until) { + delete(c.cooldownUntil, id) + } + } + + draining := c.draining + inFlightCount := len(c.handoffLeases) + maxNew := 0 + if !draining { + maxNew = max(c.cfg.MaxConcurrentInteractions-inFlightCount, 0) + } + + activeLeases := make([]LeaseEntry, 0, len(c.handoffLeases)) + for _, lease := range c.handoffLeases { + activeLeases = append(activeLeases, LeaseEntry{ + InteractionID: lease.InteractionID, + LeaseToken: lease.LeaseToken, + }) + } + c.mu.Unlock() + + slices.SortFunc(activeLeases, func(a, b LeaseEntry) int { + return strings.Compare(a.InteractionID, b.InteractionID) + }) + + if draining && len(activeLeases) == 0 { + return 0, nil + } + + leaseCtx, cancelLease := context.WithTimeout(ctx, DefaultLeaseCallTimeout) + defer cancelLease() + + res, err := c.client.LeaseInteractions(leaseCtx, &LeaseInteractionsRequest{ + EnvironmentID: c.cfg.EnvironmentID, + EnvironmentKey: c.cfg.EnvironmentKey, + WorkerID: c.cfg.WorkerID, + ActiveLeases: activeLeases, + MaxNewInteractions: int32(maxNew), + SendZeroMaxNew: maxNew == 0, + RequestedLeaseTTL: c.cfg.LeaseTTL, + Labels: c.cfg.Labels, + }) + if err != nil { + c.observer.LeasePolled(0, err) + return 0, err + } + c.observer.LeasePolled(len(res.Leases), nil) + + for _, leased := range res.Leases { + id := strings.TrimSpace(leased.InteractionID) + if id == "" { + continue + } + envID := cmp.Or(strings.TrimSpace(leased.EnvironmentID), c.cfg.EnvironmentID) + + c.mu.Lock() + if existing, ok := c.handoffLeases[id]; ok { + if leased.LeaseToken != "" { + existing.LeaseToken = leased.LeaseToken + } + if envID != "" { + existing.EnvironmentID = envID + } + if leased.EnvironmentConfig != nil { + existing.EnvironmentConfig = leased.EnvironmentConfig + } + c.handoffLeases[id] = existing + c.mu.Unlock() + continue + } + + if draining { + c.mu.Unlock() + continue + } + if until, inCooldown := c.cooldownUntil[id]; inCooldown && c.now().Before(until) { + c.mu.Unlock() + continue + } + if len(c.handoffLeases) >= c.cfg.MaxConcurrentInteractions { + c.mu.Unlock() + continue + } + + active := ActiveInteractionLease{ + InteractionID: id, + EnvironmentID: envID, + LeaseToken: leased.LeaseToken, + EnvironmentConfig: leased.EnvironmentConfig, + } + c.handoffLeases[id] = active + c.handoffWg.Add(1) + c.mu.Unlock() + + c.logger.Info("PollerController claimed interaction lease", + slog.String("interaction_id", id), + slog.String("environment_id", envID), + ) + go c.runHandoff(ctx, stopHandoffsCh, active) + } + + return res.PollBackoff, nil +} + +func (c *Tier2Poller) runHandoff(parentCtx context.Context, stopHandoffsCh <-chan struct{}, initial ActiveInteractionLease) { + defer c.handoffWg.Done() + id := initial.InteractionID + + var ( + ctx context.Context + cancel context.CancelFunc + ) + if stopHandoffsCh != nil { + ctx, cancel = context.WithTimeout(context.WithoutCancel(parentCtx), DefaultHandoffAttemptTimeout) + go func() { + select { + case <-stopHandoffsCh: + cancel() + case <-ctx.Done(): + } + }() + } else { + ctx, cancel = context.WithTimeout(parentCtx, DefaultHandoffAttemptTimeout) + } + defer cancel() + + defer func() { + c.mu.Lock() + delete(c.handoffLeases, id) + c.mu.Unlock() + c.notifySlotReleased() + }() + + for { + if ctx.Err() != nil { + return + } + + c.mu.Lock() + current := c.handoffLeases[id] + c.mu.Unlock() + + envID := cmp.Or(strings.TrimSpace(current.EnvironmentID), c.cfg.EnvironmentID) + + var ( + resp *StartWorkResponse + provResp *WorkerProvisionResponse + err error + ) + provResp, err = c.provisionEnvironmentWorker(ctx, stopHandoffsCh, envID, id, current.EnvironmentConfig) + if err == nil { + var workerClient WorkerRPCClient + workerClient, err = c.getOrDialWorkerClient(provResp.URI) + if err == nil { + c.mu.Lock() + current = c.handoffLeases[id] + c.mu.Unlock() + resp, err = c.callStartWork(ctx, workerClient, &StartWorkRequest{ + InteractionID: id, + EnvironmentID: envID, + LeaseToken: current.LeaseToken, + ControlPlaneURL: c.cfg.ControlPlaneURL, + GRPCTarget: c.cfg.GRPCTarget, + }, provResp.Headers) + if err == nil { + c.mu.Lock() + c.activeWorkers[envID] = &provisionedWorkerState{ + EnvironmentID: envID, + WorkerID: cmp.Or(strings.TrimSpace(provResp.WorkerID), envID), + URI: provResp.URI, + Headers: cloneStringMap(provResp.Headers), + LastActiveAt: c.now(), + } + c.mu.Unlock() + } + } + } + if err == nil && (resp == nil || resp.Status == "" || strings.EqualFold(resp.Status, WorkerStatusAccepted)) { + c.mu.Lock() + delete(c.failureCounts, id) + delete(c.cooldownUntil, id) + delete(c.deadLettered, id) + c.mu.Unlock() + c.observer.HandoffFinished("accepted") + c.logger.Info("PollerController dispatched interaction to worker", + slog.String("interaction_id", id), + slog.String("environment_id", envID), + slog.String("worker_id", cmp.Or(strings.TrimSpace(provResp.WorkerID), envID)), + ) + return + } + + if err == nil { + err = fmt.Errorf("unexpected StartWork status %q", resp.Status) + } + if ctx.Err() != nil { + return + } + + c.mu.Lock() + c.failureCounts[id]++ + failures := c.failureCounts[id] + if failures >= c.cfg.MaxHandoffRetries { + c.cooldownUntil[id] = c.now().Add(c.cfg.PoisonCooldown) + c.deadLettered[id] = true + c.mu.Unlock() + c.observer.HandoffFinished("dead_lettered") + c.logger.Warn("PollerController dropping interaction lease after max handoff retries", + slog.String("interaction_id", id), + slog.String("environment_id", envID), + slog.Int("failures", failures), + slog.Any("error", err), + ) + return + } + c.mu.Unlock() + c.observer.HandoffFinished("retry") + c.logger.Warn("PollerController interaction handoff attempt failed; retrying", + slog.String("interaction_id", id), + slog.String("environment_id", envID), + slog.Int("failures", failures), + slog.Any("error", err), + ) + + backoff := min(RetryBackoff(c.cfg.HandoffRetryBackoff, failures), 2*time.Second) + if waitErr := c.waitContext(ctx, backoff); waitErr != nil { + return + } + } +} + +func (c *Tier2Poller) callStartWork(ctx context.Context, workerClient WorkerRPCClient, req *StartWorkRequest, headers map[string]string) (*StartWorkResponse, error) { + rpcCtx, cancelRPC := context.WithTimeout(ctx, DefaultLeaseCallTimeout) + defer cancelRPC() + return workerClient.StartWork(rpcCtx, req, headers) +} diff --git a/internal/cli/custom/envruntime/poller_test.go b/internal/cli/custom/envruntime/poller_test.go new file mode 100644 index 0000000..3dd14e1 --- /dev/null +++ b/internal/cli/custom/envruntime/poller_test.go @@ -0,0 +1,1260 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package envruntime + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "log" + "log/slog" + "net/http" + "os" + "path/filepath" + "runtime" + "slices" + "strconv" + "strings" + "sync" + "sync/atomic" + "syscall" + "testing" + "time" +) + +func TestSingleflightDoAndDoChan(t *testing.T) { + var sf Singleflight + var runs atomic.Int32 + entered := make(chan struct{}) + gate := make(chan struct{}) + + var wg sync.WaitGroup + var sharedCount atomic.Int32 + for i := 0; i < 4; i++ { + wg.Add(1) + go func() { + defer wg.Done() + val, shared, err := sf.Do("k1", func() (any, error) { + runs.Add(1) + close(entered) + <-gate + return "v1", nil + }) + if err != nil || val != "v1" { + t.Errorf("sf.Do() = (%v, %v), want (\"v1\", nil)", val, err) + } + if shared { + sharedCount.Add(1) + } + }() + } + <-entered + for { + sf.mu.Lock() + allJoined := sf.m["k1"] != nil && sf.m["k1"].dups == 3 + sf.mu.Unlock() + if allJoined { + break + } + runtime.Gosched() + } + close(gate) + wg.Wait() + + if runs.Load() != 1 || sharedCount.Load() != 4 { + t.Errorf("sf.Do runs=%d sharedCount=%d, want runs=1 sharedCount=4", runs.Load(), sharedCount.Load()) + } + + // Sequential call after completion is unshared. + val, shared, err := sf.Do("k1", func() (any, error) { return "v2", nil }) + if err != nil || val != "v2" || shared { + t.Errorf("sequential sf.Do() = (%v, %v, %v), want (\"v2\", false, nil)", val, shared, err) + } + + // DoChan deduplication & error propagation. + gate2 := make(chan struct{}) + boom := errors.New("provision boom") + ch1 := sf.DoChan("k2", func() (any, error) { + <-gate2 + return nil, boom + }) + ch2 := sf.DoChan("k2", func() (any, error) { + return "unexpected", nil + }) + close(gate2) + r1, r2 := <-ch1, <-ch2 + if !errors.Is(r1.Err, boom) || !errors.Is(r2.Err, boom) || !r1.Shared || !r2.Shared { + t.Errorf("DoChan results = (%+v, %+v), want shared boom errors", r1, r2) + } + + // Panic recovery on Do and DoChan must return an error without deadlocking or poisoning the key. + if _, _, pErr := sf.Do("k_panic", func() (any, error) { panic("boom-do") }); pErr == nil || !strings.Contains(pErr.Error(), "singleflight panic") { + t.Errorf("sf.Do(panic) = %v, want singleflight panic error", pErr) + } + pRes := <-sf.DoChan("k_panic", func() (any, error) { panic("boom-chan") }) + if pRes.Err == nil || !strings.Contains(pRes.Err.Error(), "singleflight panic") { + t.Errorf("sf.DoChan(panic) = %+v, want singleflight panic error", pRes) + } + if okVal, _, okErr := sf.Do("k_panic", func() (any, error) { return "recovered", nil }); okErr != nil || okVal != "recovered" { + t.Errorf("sf.Do(after panic) = (%v, %v), want (\"recovered\", nil)", okVal, okErr) + } +} + +func TestShellCommandProvisionerAndStatelessPoller(t *testing.T) { + if bin, args := ResolveProvisionShell("windows", "echo ok"); !strings.Contains(bin, "sh") || len(args) != 2 || args[0] != "-c" { + t.Errorf("ResolveProvisionShell(windows with PATH sh) = (%q, %v), want sh -c", bin, args) + } + origPath := os.Getenv("PATH") + t.Setenv("PATH", "") + if bin, args := ResolveProvisionShell("windows", "echo ok"); bin != "cmd.exe" || len(args) != 2 || args[0] != "/C" { + t.Errorf("ResolveProvisionShell(windows without PATH sh) = (%q, %v), want cmd.exe /C", bin, args) + } + if bin, args := ResolveProvisionShell("linux", "echo ok"); bin != "sh" || len(args) != 2 || args[0] != "-c" { + t.Errorf("ResolveProvisionShell(linux without PATH sh) = (%q, %v), want sh -c", bin, args) + } + t.Setenv("PATH", origPath) + + emptyProv := NewShellCommandProvisioner("") + if err := emptyProv.ProvisionAndDispatch(t.Context(), ProvisionInput{InteractionID: "i1", EnvironmentID: "e1"}); err == nil { + t.Error("ProvisionAndDispatch(empty cmd) = nil, want error") + } + + tempDir := t.TempDir() + outFile := filepath.Join(tempDir, "stdin.json") + statFile := filepath.Join(tempDir, "stat.txt") + t.Setenv("GEMINI_API_KEY", "secret-must-scrub") + okProv := NewShellCommandProvisioner(`if [ -n "$GEMINI_API_KEY" ]; then exit 9; fi; cat /proc/self/stat > ` + statFile + `; cat > ` + outFile) + if err := okProv.ProvisionAndDispatch(t.Context(), ProvisionInput{ + InteractionID: "int_p1", + EnvironmentID: "env_p1", + LeaseToken: "tok_secret_123", + ServerURL: "https://generativelanguage.googleapis.com", + EnvironmentConfig: json.RawMessage(`{"network_mode":"OFF"}`), + }); err != nil { + t.Fatalf("ProvisionAndDispatch(ok) failed: %v", err) + } + written, err := os.ReadFile(outFile) + if err != nil { + t.Fatalf("ReadFile(outFile) failed: %v", err) + } + var gotIn ProvisionInput + if err := json.Unmarshal(written, &gotIn); err != nil || gotIn.SchemaVersion != WorkerProvisionSchemaVersionV1 || gotIn.InteractionID != "int_p1" { + t.Errorf("ProvisionInput stdin = (%+v, %v)", gotIn, err) + } + if statBytes, err := os.ReadFile(statFile); err == nil { + if idx := strings.LastIndexByte(string(statBytes), ')'); idx >= 0 { + fields := strings.Fields(string(statBytes[idx+1:])) + if len(fields) >= 3 { + if childPgid, err := strconv.Atoi(fields[2]); err == nil && childPgid == syscall.Getpgrp() { + t.Errorf("childPgid = %d, want != parent pgid %d (ConfigureSubprocessGroup must set new process group)", childPgid, syscall.Getpgrp()) + } + } + } + } + + // Verify LeaseToken is redacted from stderr on failure, and empty stderr omits parentheses. + failProv := NewShellCommandProvisioner(`echo "leaked tok_secret_999" >&2; exit 2`) + err = failProv.ProvisionAndDispatch(t.Context(), ProvisionInput{ + InteractionID: "int_f", + EnvironmentID: "env_f", + LeaseToken: "tok_secret_999", + }) + if err == nil || strings.Contains(err.Error(), "tok_secret_999") || !strings.Contains(err.Error(), "[REDACTED]") { + t.Errorf("ProvisionAndDispatch(stderr redaction) = %v, want [REDACTED] without tok_secret_999", err) + } + silentFailProv := NewShellCommandProvisioner(`exit 3`) + if err := silentFailProv.ProvisionAndDispatch(t.Context(), ProvisionInput{InteractionID: "int_sf", EnvironmentID: "env_sf"}); err == nil || strings.Contains(err.Error(), "():") { + t.Errorf("ProvisionAndDispatch(empty stderr) = %v, want error without '():'", err) + } + + // Test RunLeasePollLoop error backoff, reset after success, StartMsg logging, inFlight deduplication, and non-retryable termination. + var loopLog bytes.Buffer + var loopObs recordingObserver + var loopPoll atomic.Int32 + var dupHandled atomic.Int32 + dupGate := make(chan struct{}) + loopErr := RunLeasePollLoop(t.Context(), LeasePollLoopConfig{ + Client: &fakeEnvClient{ + leaseFn: func(_ context.Context, req *LeaseInteractionsRequest) (*LeaseInteractionsResult, error) { + switch loopPoll.Add(1) { + case 1: + if req.MaxNewInteractions != 2 { + t.Errorf("poll 1 MaxNewInteractions = %d, want 2", req.MaxNewInteractions) + } + return nil, errors.New("transient 1") + case 2: + return nil, errors.New("transient 2") + case 3: + return &LeaseInteractionsResult{ + Leases: []LeasedInteraction{ + {InteractionID: "int_dup", LeaseToken: "t1"}, + {InteractionID: "int_dup", LeaseToken: "t1"}, + }, + }, nil + case 4: + close(dupGate) + return nil, errors.New("transient after reset") + default: + return nil, &TransportError{StatusCode: http.StatusForbidden, Message: "fatal 403"} + } + }, + }, + MaxSlots: 2, + PollInterval: 10 * time.Millisecond, + Logger: log.New(&loopLog, "", 0), + Observer: &loopObs, + ErrLabel: "TestLoopErr", + StartMsg: "TestLoopStartedBanner", + }, func(_ context.Context, _ *LeasedInteraction) { + dupHandled.Add(1) + <-dupGate + }) + if loopErr == nil || !IsNonRetryableError(loopErr) { + t.Errorf("RunLeasePollLoop() = %v, want non-retryable 403 error", loopErr) + } + if dupHandled.Load() != 1 { + t.Errorf("dupHandled = %d, want 1 (inFlight must deduplicate concurrent identical lease IDs)", dupHandled.Load()) + } + logStr := loopLog.String() + // Poll 1 -> 10ms, Poll 2 -> 20ms, Poll 3 -> resets consecutiveErrors & logs StartMsg, Poll 4 -> 10ms (not 40ms!). + if !strings.Contains(logStr, "TestLoopStartedBanner") || !strings.Contains(logStr, "(retrying in 20ms): transient 2") || !strings.Contains(logStr, "(retrying in 10ms): transient after reset") { + t.Errorf("unexpected RunLeasePollLoop log:\n%s", logStr) + } + loopObs.mu.Lock() + if len(loopObs.leasePolledCounts) < 4 || loopObs.leasePolledErrs[0] == nil || loopObs.leasePolledCounts[2] != 2 { + t.Errorf("loopObs = counts:%v errs:%v", loopObs.leasePolledCounts, loopObs.leasePolledErrs) + } + loopObs.mu.Unlock() + + // StatelessPoller defaults, custom fields, ServerURL precedence, recordDispatchFailure shift, and dispatchLease recovery. + customSP := NewStatelessPoller(StatelessPollerConfig{ + PollerID: "custom-poller-id", + MaxInteractions: 4, + }) + if customSP.pollerID != "custom-poller-id" || customSP.maxInteractions != 4 { + t.Errorf("customSP = (pollerID=%q, maxInteractions=%d)", customSP.pollerID, customSP.maxInteractions) + } + + var hbCalls atomic.Int32 + provisionAborted := make(chan error, 1) + var polledOnce atomic.Bool + var firstSPMaxNew atomic.Int32 + pollerCtx, cancelPoller := context.WithCancel(t.Context()) + defer cancelPoller() + + spObs := &recordingObserver{} + sp := NewStatelessPoller(StatelessPollerConfig{ + Client: &fakeEnvClient{ + leaseFn: func(_ context.Context, req *LeaseInteractionsRequest) (*LeaseInteractionsResult, error) { + if !polledOnce.Swap(true) { + firstSPMaxNew.Store(req.MaxNewInteractions) + return &LeaseInteractionsResult{ + Leases: []LeasedInteraction{{ + InteractionID: "int_cancel_me", + EnvironmentID: "env_c", + LeaseToken: "tok_c", + }}, + }, nil + } + return &LeaseInteractionsResult{}, nil + }, + }, + Heartbeat: func(context.Context, string, string, string, string) (*HeartbeatResult, error) { + if hbCalls.Add(1) >= 2 { + return &HeartbeatResult{ShouldCancel: true}, nil + } + return &HeartbeatResult{Renewed: true}, nil + }, + ProvisionFunc: func(ctx context.Context, _ ProvisionInput) error { + <-ctx.Done() + provisionAborted <- ctx.Err() + return ctx.Err() + }, + Observer: spObs, + PollInterval: 15 * time.Millisecond, + HandoffHeartbeatInterval: 15 * time.Millisecond, + ProvisionTimeout: 5 * time.Second, + Logger: log.New(io.Discard, "", 0), + }) + if !strings.HasPrefix(sp.pollerID, "poller-") || sp.maxInteractions != 10 { + t.Errorf("sp defaults = (pollerID=%q, maxInteractions=%d)", sp.pollerID, sp.maxInteractions) + } + + // Seed an expired failure entry (>10m old) to verify recordDispatchFailure evicts it. + sp.failures["stale_int"] = &statelessFailureState{ + consecutive: 1, + cooldownEnd: time.Now().Add(-11 * time.Minute), + } + + runErrCh := make(chan error, 1) + go func() { + runErrCh <- sp.Run(pollerCtx) + }() + + select { + case abortErr := <-provisionAborted: + if !errors.Is(abortErr, context.Canceled) { + t.Errorf("ProvisionFunc abort err = %v, want context.Canceled", abortErr) + } + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for Heartbeat ShouldCancel to abort ProvisionFunc") + } + cancelPoller() + <-runErrCh + + if firstSPMaxNew.Load() != 10 { + t.Errorf("firstSPMaxNew = %d, want 10", firstSPMaxNew.Load()) + } + sp.mu.Lock() + _, staleStillPresent := sp.failures["stale_int"] + sp.mu.Unlock() + if staleStillPresent { + t.Error("recordDispatchFailure did not evict >10m stale failure entry") + } + + // Verify recordDispatchFailure backoff shift for consecutive=2 (~1s, not 2s) and cleanup on subsequent success. + sp.recordDispatchFailure("int_shift") + sp.recordDispatchFailure("int_shift") + sp.mu.Lock() + remShift := time.Until(sp.failures["int_shift"].cooldownEnd) + sp.failures["int_shift"].cooldownEnd = time.Now().Add(-time.Millisecond) + sp.mu.Unlock() + if remShift <= 500*time.Millisecond || remShift > 1300*time.Millisecond { + t.Errorf("recordDispatchFailure(consecutive=2) remaining cooldown = %v, want ~1s", remShift) + } + + // Test StatelessPoller with nil Heartbeat (falling back to Client), ServerURL preservation vs HTTPClient fallback, and successful dispatch clearing failures. + var gotServerURL string + spSuccess := NewStatelessPoller(StatelessPollerConfig{ + Client: NewHTTPClient(HTTPClientConfig{BaseURL: "https://fallback.example.com"}), + ServerURL: "https://explicit.example.com", + ProvisionFunc: func(_ context.Context, in ProvisionInput) error { + gotServerURL = in.ServerURL + return nil + }, + Observer: spObs, + Logger: log.New(io.Discard, "", 0), + }) + spSuccess.failures["int_recover"] = &statelessFailureState{consecutive: 1, cooldownEnd: time.Now().Add(-time.Second)} + spSuccess.dispatchLease(t.Context(), &LeasedInteraction{InteractionID: "int_recover", EnvironmentID: "env_r", LeaseToken: "tok_r"}) + if gotServerURL != "https://explicit.example.com" { + t.Errorf("dispatchLease(explicit ServerURL) = %q, want https://explicit.example.com", gotServerURL) + } + if _, stillFailed := spSuccess.failures["int_recover"]; stillFailed { + t.Error("expected successful dispatchLease to delete failures[int_recover]") + } + spSuccess.cfg.ServerURL = "" + spSuccess.dispatchLease(t.Context(), &LeasedInteraction{InteractionID: "int_recover_2", EnvironmentID: "env_r", LeaseToken: "tok_r"}) + if gotServerURL != "https://fallback.example.com" { + t.Errorf("dispatchLease(fallback ServerURL) = %q, want https://fallback.example.com", gotServerURL) + } + // Verify heartbeatOnce when p.cfg.Heartbeat == nil invokes Client rather than panicking. + spNilHB := NewStatelessPoller(StatelessPollerConfig{ + Client: &fakeEnvClient{ + leaseFn: func(_ context.Context, _ *LeaseInteractionsRequest) (*LeaseInteractionsResult, error) { + return &LeaseInteractionsResult{Leases: []LeasedInteraction{{InteractionID: "int_hb"}}}, nil + }, + }, + Logger: log.New(io.Discard, "", 0), + }) + if hbRes, err := spNilHB.heartbeatOnce(t.Context(), "int_hb", "env_1", "tok_1"); err != nil || !hbRes.Renewed { + t.Errorf("spNilHB.heartbeatOnce() = (%+v, %v), want Renewed=true", hbRes, err) + } + + spObs.mu.Lock() + statuses := append([]string(nil), spObs.handoffStatuses...) + spObs.mu.Unlock() + if !slices.Contains(statuses, "error") || !slices.Contains(statuses, "accepted") { + t.Errorf("spObs.handoffStatuses = %v, want both 'error' and 'accepted'", statuses) + } +} + +func TestDirectCommandProvisionerAndConfigConversion(t *testing.T) { + if _, err := NewDirectCommandProvisioner(" ", ""); err == nil { + t.Error("NewDirectCommandProvisioner(empty) = nil, want error") + } + + t.Setenv("GEMINI_API_KEY", "secret-must-scrub") + dir := t.TempDir() + provScript := filepath.Join(dir, "prov.sh") + suspScript := filepath.Join(dir, "susp.sh") + provStdin := filepath.Join(dir, "prov_stdin.json") + suspLog := filepath.Join(dir, "suspended.json") + + provBody := "#!/bin/sh\n" + + "if [ \"$1\" != \"--prov-arg\" ] || [ -n \"$GEMINI_API_KEY\" ]; then echo \"bad prov env or arg\" >&2; exit 8; fi\n" + + "cat > " + provStdin + "\n" + + "if grep -q '\"env_fail\"' " + provStdin + "; then echo \"prov_stderr_detail\" >&2; exit 2; fi\n" + + "echo '{\"uri\":\" http://127.0.0.1:9001 \",\"worker_id\":\"w-99\",\"headers\":{\"X-Auth\":\"ok\"}}'\n" + if err := os.WriteFile(provScript, []byte(provBody), 0o755); err != nil { + t.Fatalf("WriteFile(provScript) failed: %v", err) + } + + suspBody := "#!/bin/sh\n" + + "if [ \"$1\" != \"--susp-arg\" ] || [ -n \"$GEMINI_API_KEY\" ]; then echo \"bad susp env or arg\" >&2; exit 8; fi\n" + + "cat > " + suspLog + "\n" + + "if grep -q '\"env_susp_fail\"' " + suspLog + "; then echo \"susp_stderr_detail\" >&2; exit 2; fi\n" + if err := os.WriteFile(suspScript, []byte(suspBody), 0o755); err != nil { + t.Fatalf("WriteFile(suspScript) failed: %v", err) + } + + dcp, err := NewDirectCommandProvisioner(provScript+" --prov-arg", suspScript+" --susp-arg") + if err != nil { + t.Fatalf("NewDirectCommandProvisioner() failed: %v", err) + } + if _, err := dcp.ProvisionWorker(t.Context(), WorkerProvisionRequest{}); err == nil { + t.Error("ProvisionWorker(empty EnvironmentID) = nil, want error") + } + resp, err := dcp.ProvisionWorker(t.Context(), WorkerProvisionRequest{ + EnvironmentID: " env_d1 ", + InteractionID: " int_d1 ", + }) + if err != nil || resp.URI != "http://127.0.0.1:9001" || resp.WorkerID != "w-99" || resp.SchemaVersion != WorkerProvisionSchemaVersionV1 { + t.Fatalf("ProvisionWorker() = (%+v, %v)", resp, err) + } + rawProvIn, err := os.ReadFile(provStdin) + if err != nil { + t.Fatalf("ReadFile(provStdin) failed: %v", err) + } + var gotProvReq WorkerProvisionRequest + if err := json.Unmarshal(rawProvIn, &gotProvReq); err != nil || gotProvReq.SchemaVersion != WorkerProvisionSchemaVersionV1 || gotProvReq.EnvironmentID != "env_d1" || gotProvReq.InteractionID != "int_d1" { + t.Errorf("ProvisionWorker stdin = (%+v, %v)", gotProvReq, err) + } + if _, err := dcp.ProvisionWorker(t.Context(), WorkerProvisionRequest{EnvironmentID: "env_fail"}); err == nil || !strings.Contains(err.Error(), "prov_stderr_detail") { + t.Errorf("ProvisionWorker(env_fail) = %v, want prov_stderr_detail", err) + } + + if err := dcp.SuspendWorker(t.Context(), WorkerSuspendRequest{}); err == nil { + t.Error("SuspendWorker(empty ids) = nil, want error") + } + if err := dcp.SuspendWorker(t.Context(), WorkerSuspendRequest{ + EnvironmentID: " env_d1 ", + WorkerID: " w-99 ", + IdleDuration: " 65s ", + }); err != nil { + t.Fatalf("SuspendWorker() failed: %v", err) + } + rawSusp, err := os.ReadFile(suspLog) + if err != nil { + t.Fatalf("ReadFile(suspLog) failed: %v", err) + } + var gotSuspReq WorkerSuspendRequest + if err := json.Unmarshal(rawSusp, &gotSuspReq); err != nil || gotSuspReq.SchemaVersion != WorkerProvisionSchemaVersionV1 || gotSuspReq.EnvironmentID != "env_d1" || gotSuspReq.WorkerID != "w-99" || gotSuspReq.IdleDuration != "65s" { + t.Errorf("SuspendWorker stdin = (%+v, %v)", gotSuspReq, err) + } + // SuspendWorker succeeds when only EnvironmentID or only WorkerID is provided. + if err := dcp.SuspendWorker(t.Context(), WorkerSuspendRequest{EnvironmentID: "env_only"}); err != nil { + t.Errorf("SuspendWorker(EnvironmentID only) = %v, want nil", err) + } + if err := dcp.SuspendWorker(t.Context(), WorkerSuspendRequest{WorkerID: "w-only"}); err != nil { + t.Errorf("SuspendWorker(WorkerID only) = %v, want nil", err) + } + if err := dcp.SuspendWorker(t.Context(), WorkerSuspendRequest{EnvironmentID: "env_susp_fail", WorkerID: "w-99"}); err == nil || !strings.Contains(err.Error(), "susp_stderr_detail") { + t.Errorf("SuspendWorker(env_susp_fail) = %v, want susp_stderr_detail", err) + } + + if ConvertToProvisionEnvironmentConfig(nil) != nil { + t.Error("ConvertToProvisionEnvironmentConfig(nil) != nil") + } + if ConvertToProvisionEnvironmentConfig(&EnvironmentConfig{}) != nil { + t.Error("ConvertToProvisionEnvironmentConfig(empty) != nil") + } + if envOnly := ConvertToProvisionEnvironmentConfig(&EnvironmentConfig{Env: map[string]string{"ONLY_ENV": "V1"}}); envOnly == nil || envOnly.Env["ONLY_ENV"] != "V1" { + t.Errorf("ConvertToProvisionEnvironmentConfig(Env only) = %+v", envOnly) + } + if allowStructOnly := ConvertToProvisionEnvironmentConfig(&EnvironmentConfig{AllowlistDomains: []string{"struct.example.com"}}); allowStructOnly == nil || len(allowStructOnly.AllowlistDomains) != 1 || allowStructOnly.AllowlistDomains[0] != "struct.example.com" { + t.Errorf("ConvertToProvisionEnvironmentConfig(AllowlistDomains struct only) = %+v", allowStructOnly) + } + if allowRawOnly := ConvertToProvisionEnvironmentConfig(&EnvironmentConfig{RawJSON: json.RawMessage(`{"allowlist_domains":["raw.example.com"]}`)}); allowRawOnly == nil || len(allowRawOnly.AllowlistDomains) != 1 || allowRawOnly.AllowlistDomains[0] != "raw.example.com" { + t.Errorf("ConvertToProvisionEnvironmentConfig(RawJSON allowlist_domains only) = %+v", allowRawOnly) + } + if denyOnly := ConvertToProvisionEnvironmentConfig(&EnvironmentConfig{DenylistDomains: []string{"bad.example.com"}}); denyOnly == nil || len(denyOnly.DenylistDomains) != 1 { + t.Errorf("ConvertToProvisionEnvironmentConfig(DenylistDomains only) = %+v", denyOnly) + } + if modeOnly := ConvertToProvisionEnvironmentConfig(&EnvironmentConfig{NetworkMode: "OFF"}); modeOnly == nil || modeOnly.NetworkMode != "OFF" { + t.Errorf("ConvertToProvisionEnvironmentConfig(NetworkMode only) = %+v", modeOnly) + } + if modeRawOnly := ConvertToProvisionEnvironmentConfig(&EnvironmentConfig{RawJSON: json.RawMessage(`{"network_mode":"ALLOWLIST"}`)}); modeRawOnly == nil || modeRawOnly.NetworkMode != "ALLOWLIST" { + t.Errorf("ConvertToProvisionEnvironmentConfig(RawJSON network_mode only) = %+v", modeRawOnly) + } + if got := truncateCommandOutput(" short "); got != "short" { + t.Errorf("truncateCommandOutput(short) = %q, want short", got) + } + if got := truncateCommandOutput(strings.Repeat("x", maxCommandErrorBytes+10)); !strings.HasSuffix(got, "...(truncated)") { + t.Errorf("truncateCommandOutput(long) = %q, want ...(truncated)", got) + } + if srcOnly := ConvertToProvisionEnvironmentConfig(&EnvironmentConfig{Sources: []EnvironmentSource{{Type: "inline", Target: "main.py"}}}); srcOnly == nil || len(srcOnly.Sources) != 1 { + t.Errorf("ConvertToProvisionEnvironmentConfig(Sources only) = %+v", srcOnly) + } + parsed := ParseEnvironmentConfig(json.RawMessage(`{ + "network_mode": "ALLOWLIST", + "allowlist_domains": ["a.example.com"], + "network_allowlist": [{"domain": "b.example.com"}, {"domain": "bad.example.com", "mode": "deny"}], + "network_denylist": ["evil.example.com"], + "env": {"K1": "V1"}, + "env_vars": {"K2": {"value": "V2"}}, + "sources": [{"type": "inline", "target": "main.py", "content": "print(1)"}] + }`)) + conv := ConvertToProvisionEnvironmentConfig(parsed) + if conv == nil || len(conv.AllowlistDomains) != 2 || len(conv.DenylistDomains) != 2 || conv.Env["K1"] != "V1" || conv.Env["K2"] != "V2" || len(conv.Sources) != 1 { + t.Errorf("ConvertToProvisionEnvironmentConfig() = %+v", conv) + } +} + +type fakeWorkerProvisioner struct { + mu sync.Mutex + provisions int + suspensions []WorkerSuspendRequest + provFn func(req WorkerProvisionRequest) (*WorkerProvisionResponse, error) + provCtxFn func(ctx context.Context, req WorkerProvisionRequest) (*WorkerProvisionResponse, error) +} + +func (f *fakeWorkerProvisioner) ProvisionWorker(ctx context.Context, req WorkerProvisionRequest) (*WorkerProvisionResponse, error) { + f.mu.Lock() + f.provisions++ + ctxFn := f.provCtxFn + fn := f.provFn + f.mu.Unlock() + if ctxFn != nil { + return ctxFn(ctx, req) + } + if fn != nil { + return fn(req) + } + return &WorkerProvisionResponse{ + SchemaVersion: WorkerProvisionSchemaVersionV1, + URI: " worker://" + req.EnvironmentID + " ", + WorkerID: "wk-" + req.EnvironmentID, + Headers: map[string]string{"X-Worker": "1"}, + }, nil +} + +func (f *fakeWorkerProvisioner) SuspendWorker(_ context.Context, req WorkerSuspendRequest) error { + f.mu.Lock() + f.suspensions = append(f.suspensions, req) + f.mu.Unlock() + return nil +} + +type fakeWorkerRPCClient struct { + startFn func(ctx context.Context, req *StartWorkRequest, headers map[string]string) (*StartWorkResponse, error) + idleFn func(headers map[string]string) (*IdleStatus, error) + closed atomic.Bool +} + +func (f *fakeWorkerRPCClient) StartWork(ctx context.Context, req *StartWorkRequest, headers map[string]string) (*StartWorkResponse, error) { + return f.startFn(ctx, req, headers) +} + +func (f *fakeWorkerRPCClient) GetIdleStatus(_ context.Context, headers map[string]string) (*IdleStatus, error) { + if headers["X-Worker"] != "1" { + return nil, errors.New("missing cloned X-Worker header in GetIdleStatus") + } + if f.idleFn != nil { + return f.idleFn(headers) + } + return &IdleStatus{Running: false, IdleDuration: 2 * time.Minute}, nil +} + +func (f *fakeWorkerRPCClient) Close() error { + f.closed.Store(true) + return nil +} + +func TestTier2PollerSingleflightPoisonPillSweepAndDrain(t *testing.T) { + if _, err := NewTier2Poller(Tier2PollerConfig{}); err == nil { + t.Error("NewTier2Poller(missing Client) = nil, want error") + } + dummyClient := &fakeEnvClient{leaseFn: func(context.Context, *LeaseInteractionsRequest) (*LeaseInteractionsResult, error) { + return &LeaseInteractionsResult{}, nil + }} + if _, err := NewTier2Poller(Tier2PollerConfig{Client: dummyClient}); err == nil { + t.Error("NewTier2Poller(missing Provisioner) = nil, want error") + } + if _, err := NewTier2Poller(Tier2PollerConfig{Client: dummyClient, Provisioner: &fakeWorkerProvisioner{}}); err == nil { + t.Error("NewTier2Poller(missing DialWorker) = nil, want error") + } + + emptyCfg := Tier2PollerConfig{} + emptyCfg.Defaults() + if emptyCfg.MaxConcurrentInteractions != DefaultPollerMaxConcurrentInteractions || + emptyCfg.PollInterval != DefaultPollBackoff || + emptyCfg.LeaseTTL != DefaultLeaseTTL || + emptyCfg.WorkerID == "" || + emptyCfg.MaxHandoffRetries != DefaultMaxHandoffRetries || + emptyCfg.HandoffRetryBackoff != DefaultHandoffRetryBackoff || + emptyCfg.PoisonCooldown != DefaultPoisonCooldown || + emptyCfg.SweepInterval != DefaultSweepInterval || + emptyCfg.DrainTimeout != DefaultDrainTimeout || + emptyCfg.IdleTimeout != DefaultWorkerIdleTimeout || + emptyCfg.ResumeGracePeriod != DefaultWorkerResumeGracePeriod || + emptyCfg.Now == nil || + emptyCfg.NewTimer == nil || + emptyCfg.Logger == nil || + emptyCfg.Observer == nil { + t.Errorf("emptyCfg.Defaults() did not populate all defaults: %+v", emptyCfg) + } + + var customTimerCalled bool + customCfg := Tier2PollerConfig{ + MaxConcurrentInteractions: 7, + PollInterval: 3 * time.Second, + LeaseTTL: 12 * time.Second, + WorkerID: "explicit-worker-id", + MaxHandoffRetries: 9, + HandoffRetryBackoff: 4 * time.Second, + PoisonCooldown: 20 * time.Second, + SweepInterval: 11 * time.Second, + DrainTimeout: 13 * time.Second, + IdleTimeout: 14 * time.Second, + ResumeGracePeriod: 15 * time.Second, + NewTimer: func(d time.Duration) *time.Timer { + customTimerCalled = true + return time.NewTimer(d) + }, + } + customCfg.Defaults() + if customCfg.MaxConcurrentInteractions != 7 || + customCfg.PollInterval != 3*time.Second || + customCfg.LeaseTTL != 12*time.Second || + customCfg.WorkerID != "explicit-worker-id" || + customCfg.MaxHandoffRetries != 9 || + customCfg.HandoffRetryBackoff != 4*time.Second || + customCfg.PoisonCooldown != 20*time.Second || + customCfg.SweepInterval != 11*time.Second || + customCfg.DrainTimeout != 13*time.Second || + customCfg.IdleTimeout != 14*time.Second || + customCfg.ResumeGracePeriod != 15*time.Second || + customCfg.Observer == nil { + t.Errorf("customCfg.Defaults() overwrote custom values: %+v", customCfg) + } + tm := customCfg.NewTimer(time.Millisecond) + tm.Stop() + if !customTimerCalled { + t.Error("customCfg.Defaults() overwrote custom NewTimer") + } + + nowMu := sync.Mutex{} + fakeNow := time.Date(2026, 10, 26, 12, 0, 0, 0, time.UTC) + getNow := func() time.Time { + nowMu.Lock() + defer nowMu.Unlock() + return fakeNow + } + advanceNow := func(d time.Duration) { + nowMu.Lock() + fakeNow = fakeNow.Add(d) + nowMu.Unlock() + } + + prov := &fakeWorkerProvisioner{ + provFn: func(req WorkerProvisionRequest) (*WorkerProvisionResponse, error) { + uri := "worker://shared" + if req.EnvironmentID == "env_poison" { + uri = "worker://poison" + } + return &WorkerProvisionResponse{ + SchemaVersion: WorkerProvisionSchemaVersionV1, + URI: uri, + WorkerID: "wk-" + req.EnvironmentID, + Headers: map[string]string{"X-Worker": "1", "X-Env": req.EnvironmentID}, + }, nil + }, + } + var dials atomic.Int32 + var accepted atomic.Int32 + var allowPoison atomic.Bool + var env2Busy atomic.Bool + var env2ShortIdle atomic.Bool + env2Busy.Store(true) + inFlightStarted := make(chan struct{}) + inFlightGate := make(chan struct{}) + var closeInFlightOnce sync.Once + drainStarted := make(chan struct{}) + var closeDrainOnce sync.Once + + sharedConn := &fakeWorkerRPCClient{ + startFn: func(ctx context.Context, req *StartWorkRequest, headers map[string]string) (*StartWorkResponse, error) { + if req.InteractionID == "int_poison" && !allowPoison.Load() { + return nil, errors.New("worker rejected start_work") + } + if headers["X-Worker"] != "1" { + return nil, errors.New("missing worker header") + } + if req.InteractionID == "int_ok_1" { + closeInFlightOnce.Do(func() { close(inFlightStarted) }) + <-inFlightGate + } + if req.InteractionID == "int_drain" { + closeDrainOnce.Do(func() { close(drainStarted) }) + time.Sleep(25 * time.Millisecond) + if err := ctx.Err(); err != nil { + return nil, err + } + } + accepted.Add(1) + return &StartWorkResponse{Status: WorkerStatusAccepted}, nil + }, + idleFn: func(headers map[string]string) (*IdleStatus, error) { + if headers["X-Env"] == "env_shared_2" { + if env2Busy.Load() { + return &IdleStatus{Running: true, IdleDuration: 2 * time.Minute}, nil + } + if env2ShortIdle.Load() { + return &IdleStatus{Running: false, IdleDuration: 10 * time.Second}, nil + } + } + return &IdleStatus{Running: false, IdleDuration: 2 * time.Minute}, nil + }, + } + poisonConn := &fakeWorkerRPCClient{ + startFn: sharedConn.startFn, + idleFn: func(map[string]string) (*IdleStatus, error) { + return &IdleStatus{Running: true}, nil + }, + } + + var pollStep atomic.Int32 + var sawDrainZeroMaxNew atomic.Bool + client := &fakeEnvClient{ + leaseFn: func(_ context.Context, req *LeaseInteractionsRequest) (*LeaseInteractionsResult, error) { + switch pollStep.Add(1) { + case 1: + // Two concurrent interactions in env_shared + one in env_shared_2 (sharing worker://shared URI) + one poison interaction! + return &LeaseInteractionsResult{ + Leases: []LeasedInteraction{ + {InteractionID: "int_ok_1", EnvironmentID: "env_shared", LeaseToken: "tok_1"}, + {InteractionID: "int_ok_2", EnvironmentID: "env_shared", LeaseToken: "tok_2"}, + {InteractionID: "int_ok_3", EnvironmentID: "env_shared_2", LeaseToken: "tok_3"}, + {InteractionID: "int_poison", EnvironmentID: "env_poison", LeaseToken: "tok_bad"}, + }, + PollBackoff: 20 * time.Millisecond, + }, nil + case 2: + // Poll 2 occurs while int_ok_1 is still inFlight -> must update token without spawning a second handoff! + return &LeaseInteractionsResult{ + Leases: []LeasedInteraction{ + {InteractionID: "int_ok_1", EnvironmentID: "env_shared", LeaseToken: "tok_1_renewed"}, + }, + PollBackoff: 20 * time.Millisecond, + }, nil + case 3: + // Poll 3 occurs after cooldown expires and allowPoison=true -> int_poison succeeds and clears failure/cooldown/deadLetter state! + return &LeaseInteractionsResult{ + Leases: []LeasedInteraction{ + {InteractionID: "int_poison", EnvironmentID: "env_poison", LeaseToken: "tok_recovered"}, + }, + PollBackoff: 20 * time.Millisecond, + }, nil + case 4: + // Poll 4 occurs inside Run() -> claims int_drain so drainInFlightHandoffs is exercised with an active handoff! + return &LeaseInteractionsResult{ + Leases: []LeasedInteraction{ + {InteractionID: "int_drain", EnvironmentID: "env_shared", LeaseToken: "tok_drain"}, + }, + PollBackoff: 5 * time.Second, + }, nil + default: + if req.MaxNewInteractions == 0 && req.SendZeroMaxNew { + sawDrainZeroMaxNew.Store(true) + } + return &LeaseInteractionsResult{PollBackoff: 20 * time.Millisecond}, nil + } + }, + } + + tpObs := &recordingObserver{} + var poisonDials atomic.Int32 + tp, err := NewTier2Poller(Tier2PollerConfig{ + Client: client, + Provisioner: prov, + Observer: tpObs, + DialWorker: func(uri string) (WorkerRPCClient, error) { + if uri != strings.TrimSpace(uri) { + return nil, errors.New("untrimmed worker uri: " + uri) + } + dials.Add(1) + if uri == "worker://poison" { + if poisonDials.Add(1) == 1 { + return nil, errors.New("dial poison failed on attempt 1") + } + return poisonConn, nil + } + return sharedConn, nil + }, + ControlPlaneURL: "https://generativelanguage.googleapis.com", + GRPCTarget: "generativelanguage.googleapis.com:443", + MaxHandoffRetries: 2, + HandoffRetryBackoff: 5 * time.Millisecond, + PoisonCooldown: 30 * time.Second, + IdleTimeout: 60 * time.Second, + ResumeGracePeriod: 3 * time.Second, + PollInterval: 5 * time.Second, + SweepInterval: 1 * time.Hour, + DrainTimeout: 2 * time.Second, + Now: getNow, + Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + }) + if err != nil { + t.Fatalf("NewTier2Poller() failed: %v", err) + } + if tp.cfg.LeaseTTL != DefaultLeaseTTL || tp.cfg.WorkerID == "" { + t.Errorf("tp.cfg defaults: LeaseTTL=%v WorkerID=%q", tp.cfg.LeaseTTL, tp.cfg.WorkerID) + } + + if _, err := tp.PollOnce(t.Context()); err != nil { + t.Fatalf("PollOnce(1) failed: %v", err) + } + // Wait for int_ok_1 to enter StartWork, then run PollOnce(2) while int_ok_1 is still inFlight[id]==true! + <-inFlightStarted + if _, err := tp.PollOnce(t.Context()); err != nil { + t.Fatalf("PollOnce(2) failed: %v", err) + } + // Seed a future cooldown on int_ok_1 before releasing inFlightGate so we verify accepted handoff deletes cooldownUntil[id]. + tp.mu.Lock() + tp.cooldownUntil["int_ok_1"] = getNow().Add(time.Hour) + tp.mu.Unlock() + close(inFlightGate) + tp.WaitHandoffs() + + // If inFlight["int_ok_1"] was not tracked, PollOnce(2) would have spawned a duplicate handoff and accepted would be 4 instead of 3! + if accepted.Load() != 3 { + t.Errorf("accepted = %d, want 3", accepted.Load()) + } + // Verify getOrDialWorkerClient trims whitespace around URI and hits cached connection. + if cachedConn, err := tp.getOrDialWorkerClient(" worker://shared "); err != nil || cachedConn != sharedConn { + t.Errorf("getOrDialWorkerClient(whitespace URI) = (%v, %v), want cached sharedConn", cachedConn, err) + } + if tp.InCooldown("int_ok_1") { + t.Error("expected accepted handoff to clear cooldownUntil[int_ok_1]") + } + if tp.InFlightCount() != 0 || len(tp.HandoffLeases()) != 0 { + t.Errorf("InFlightCount=%d HandoffLeases=%v, want 0", tp.InFlightCount(), tp.HandoffLeases()) + } + if !tp.IsDeadLettered("int_poison") || !tp.InCooldown("int_poison") || tp.FailureCount("int_poison") != 2 { + t.Errorf("int_poison state: dead=%v cooldown=%v failures=%d", tp.IsDeadLettered("int_poison"), tp.InCooldown("int_poison"), tp.FailureCount("int_poison")) + } + + // Within ResumeGracePeriod (0s elapsed < 3s), SweepOnce must NOT suspend active workers. + if err := tp.SweepOnce(t.Context()); err != nil { + t.Fatalf("SweepOnce(grace period) failed: %v", err) + } + prov.mu.Lock() + if len(prov.suspensions) != 0 { + t.Errorf("suspensions during grace period = %v, want empty", prov.suspensions) + } + prov.mu.Unlock() + + // Advance clock past ResumeGracePeriod & PoisonCooldown. + advanceNow(40 * time.Second) + if tp.InCooldown("int_poison") { + t.Error("InCooldown(int_poison) = true after 40s, want false") + } + // Even after cooldown expires, deadLettered["int_poison"] must remain true until a handoff succeeds! + if !tp.IsDeadLettered("int_poison") { + t.Error("IsDeadLettered(int_poison) = false after cooldown expiry, want true") + } + + // Sweep 1: env_shared is idle, but env_shared_2 is still Running=true (sharing worker://shared URI), so sharedConn must NOT be closed yet! + if err := tp.SweepOnce(t.Context()); err != nil { + t.Fatalf("SweepOnce(1 after grace) failed: %v", err) + } + prov.mu.Lock() + suspAfter1 := len(prov.suspensions) + prov.mu.Unlock() + if suspAfter1 != 1 || sharedConn.closed.Load() { + t.Errorf("after Sweep 1: suspensions=%d sharedConn.closed=%v, want 1 and false (uriStillUsed=true)", suspAfter1, sharedConn.closed.Load()) + } + + // Sweep 1b: env_shared_2 has Running=false but IdleDuration (10s) < IdleTimeout (60s) -> must NOT be suspended! + env2Busy.Store(false) + env2ShortIdle.Store(true) + if err := tp.SweepOnce(t.Context()); err != nil { + t.Fatalf("SweepOnce(1b short idle) failed: %v", err) + } + prov.mu.Lock() + suspAfter1b := len(prov.suspensions) + prov.mu.Unlock() + if suspAfter1b != 1 || sharedConn.closed.Load() { + t.Errorf("after Sweep 1b: suspensions=%d sharedConn.closed=%v, want 1 and false (IdleDuration < IdleTimeout)", suspAfter1b, sharedConn.closed.Load()) + } + + // Poll 3: allow int_poison (worker://poison) to succeed before Sweep 2 so activeWorkers contains a worker with a different URI, + // and verify failureCounts and deadLettered are cleared! + allowPoison.Store(true) + if _, err := tp.PollOnce(t.Context()); err != nil { + t.Fatalf("PollOnce(3 recovery) failed: %v", err) + } + tp.WaitHandoffs() + if tp.FailureCount("int_poison") != 0 || tp.IsDeadLettered("int_poison") { + t.Errorf("after recovery: FailureCount=%d IsDeadLettered=%v, want 0 and false", tp.FailureCount("int_poison"), tp.IsDeadLettered("int_poison")) + } + + // Sweep 2: env_shared_2 becomes idle (>= IdleTimeout) while env_poison (worker://poison) is Running=true -> + // env_shared_2 is suspended and sharedConn is closed because no remaining active worker uses worker://shared! + env2ShortIdle.Store(false) + if err := tp.SweepOnce(t.Context()); err != nil { + t.Fatalf("SweepOnce(2) failed: %v", err) + } + prov.mu.Lock() + suspAfter2 := len(prov.suspensions) + prov.mu.Unlock() + if suspAfter2 != 2 || !sharedConn.closed.Load() { + t.Errorf("after Sweep 2: suspensions=%d sharedConn.closed=%v, want 2 and true", suspAfter2, sharedConn.closed.Load()) + } + + // Sweep 3: suspended workers were deleted from activeWorkers and env_poison is Running=true, so no further suspensions occur. + if err := tp.SweepOnce(t.Context()); err != nil { + t.Fatalf("SweepOnce(3) failed: %v", err) + } + prov.mu.Lock() + suspAfter3 := len(prov.suspensions) + prov.mu.Unlock() + if suspAfter3 != 2 { + t.Errorf("after Sweep 3: suspensions=%d, want 2 (deleted activeWorkers must not be re-suspended)", suspAfter3) + } + + // Run() + in-flight drain + Close(). + sharedConn.closed.Store(false) + runCtx, cancel := context.WithCancel(t.Context()) + done := make(chan error, 1) + go func() { + done <- tp.Run(runCtx) + }() + // Wait until int_drain is mid-handoff, then cancel runCtx to trigger drainInFlightHandoffs! + select { + case <-drainStarted: + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for int_drain handoff to start") + } + cancel() + select { + case err := <-done: + if err != nil { + t.Errorf("tp.Run() = %v, want nil", err) + } + case <-time.After(1 * time.Second): + t.Fatal("tp.Run() did not finish draining rapidly when stateChangedCh fired") + } + if accepted.Load() != 5 { + t.Errorf("accepted after drain = %d, want 5 (drainInFlightHandoffs must wait for in-flight handoff)", accepted.Load()) + } + if !sawDrainZeroMaxNew.Load() { + t.Error("expected drainInFlightHandoffs pollOnce to set MaxNewInteractions=0 (c.draining=true)") + } + if !sharedConn.closed.Load() || !poisonConn.closed.Load() { + t.Error("expected Close() to close cached WorkerRPCClient connections") + } + // Verify Close() unlocked c.mu by calling a lock-acquiring method after Run() returns. + if tp.InFlightCount() != 0 { + t.Errorf("tp.InFlightCount() after Run() = %d, want 0", tp.InFlightCount()) + } + + tpObs.mu.Lock() + if len(tpObs.workerProvisioned) == 0 || len(tpObs.leasePolledCounts) == 0 || !slices.Contains(tpObs.handoffStatuses, "accepted") { + t.Errorf("tpObs = prov:%v polled:%v handoff:%v", tpObs.workerProvisioned, tpObs.leasePolledCounts, tpObs.handoffStatuses) + } + tpObs.mu.Unlock() +} + +func TestTier2PollerHandoffCancellationAndTokenRefresh(t *testing.T) { + // 1. Public PollOnce(pollCtx) without Run(): canceling pollCtx must cancel in-flight handoff immediately + // without dead-lettering or waiting for DefaultHandoffAttemptTimeout. + pollOnceStarted := make(chan struct{}) + tpDirect, err := NewTier2Poller(Tier2PollerConfig{ + Client: &fakeEnvClient{ + leaseFn: func(context.Context, *LeaseInteractionsRequest) (*LeaseInteractionsResult, error) { + return &LeaseInteractionsResult{ + Leases: []LeasedInteraction{{InteractionID: "int_direct_cancel", EnvironmentID: "env_d", LeaseToken: "tok_d"}}, + }, nil + }, + }, + Provisioner: &fakeWorkerProvisioner{ + provFn: func(WorkerProvisionRequest) (*WorkerProvisionResponse, error) { + return &WorkerProvisionResponse{SchemaVersion: WorkerProvisionSchemaVersionV1, URI: "worker://d"}, nil + }, + }, + DialWorker: func(string) (WorkerRPCClient, error) { + return &fakeWorkerRPCClient{ + startFn: func(ctx context.Context, _ *StartWorkRequest, _ map[string]string) (*StartWorkResponse, error) { + close(pollOnceStarted) + <-ctx.Done() + return nil, ctx.Err() + }, + }, nil + }, + Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + }) + if err != nil { + t.Fatalf("NewTier2Poller(tpDirect) failed: %v", err) + } + pollCtx, cancelPoll := context.WithCancel(t.Context()) + if _, err := tpDirect.PollOnce(pollCtx); err != nil { + t.Fatalf("tpDirect.PollOnce() failed: %v", err) + } + <-pollOnceStarted + cancelPoll() + waitDone := make(chan struct{}) + go func() { + tpDirect.WaitHandoffs() + close(waitDone) + }() + select { + case <-waitDone: + case <-time.After(500 * time.Millisecond): + t.Fatal("PollOnce handoff did not respect caller context cancellation promptly") + } + tpDirect.mu.Lock() + directDead := tpDirect.deadLettered["int_direct_cancel"] + tpDirect.mu.Unlock() + if directDead { + t.Error("expected canceled PollOnce handoff not to be dead-lettered") + } + + // 2. Run(runCtx) with short DrainTimeout: when StartWork blocks past DrainTimeout, closing stopHandoffsCh + // cancels the detached handoff context so Run() and WaitHandoffs() return promptly. + runBlockedStarted := make(chan struct{}) + var closeRunBlockedOnce sync.Once + tpRunTimeout, err := NewTier2Poller(Tier2PollerConfig{ + Client: &fakeEnvClient{ + leaseFn: func(_ context.Context, req *LeaseInteractionsRequest) (*LeaseInteractionsResult, error) { + if req.MaxNewInteractions == 0 { + return &LeaseInteractionsResult{}, nil + } + return &LeaseInteractionsResult{ + Leases: []LeasedInteraction{{InteractionID: "int_drain_timeout", EnvironmentID: "env_dt", LeaseToken: "tok_dt"}}, + PollBackoff: 5 * time.Second, + }, nil + }, + }, + Provisioner: &fakeWorkerProvisioner{ + provFn: func(WorkerProvisionRequest) (*WorkerProvisionResponse, error) { + return &WorkerProvisionResponse{SchemaVersion: WorkerProvisionSchemaVersionV1, URI: "worker://dt"}, nil + }, + }, + DialWorker: func(string) (WorkerRPCClient, error) { + return &fakeWorkerRPCClient{ + startFn: func(ctx context.Context, _ *StartWorkRequest, _ map[string]string) (*StartWorkResponse, error) { + closeRunBlockedOnce.Do(func() { close(runBlockedStarted) }) + <-ctx.Done() + return nil, ctx.Err() + }, + }, nil + }, + DrainTimeout: 25 * time.Millisecond, + SweepInterval: time.Hour, + Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + }) + if err != nil { + t.Fatalf("NewTier2Poller(tpRunTimeout) failed: %v", err) + } + runCtx, cancelRun := context.WithCancel(t.Context()) + runDone := make(chan error, 1) + go func() { + runDone <- tpRunTimeout.Run(runCtx) + }() + <-runBlockedStarted + cancelRun() + select { + case runErr := <-runDone: + if runErr != nil { + t.Errorf("tpRunTimeout.Run() = %v, want nil", runErr) + } + case <-time.After(500 * time.Millisecond): + t.Fatal("tpRunTimeout.Run() did not cancel blocked handoff after DrainTimeout") + } + waitRunDone := make(chan struct{}) + go func() { + tpRunTimeout.WaitHandoffs() + close(waitRunDone) + }() + select { + case <-waitRunDone: + case <-time.After(500 * time.Millisecond): + t.Fatal("tpRunTimeout.WaitHandoffs() remained blocked after stopHandoffsCh closed") + } + + // 3. Cold-start provisioning token rotation: if a concurrent PollOnce refreshes handoffLeases[id] + // while ProvisionWorker is in flight, StartWork must receive the refreshed LeaseToken. + provEntered := make(chan struct{}) + provRelease := make(chan struct{}) + var coldPollCount atomic.Int32 + var gotStartToken atomic.Value + tpCold, err := NewTier2Poller(Tier2PollerConfig{ + Client: &fakeEnvClient{ + leaseFn: func(context.Context, *LeaseInteractionsRequest) (*LeaseInteractionsResult, error) { + if coldPollCount.Add(1) == 1 { + return &LeaseInteractionsResult{ + Leases: []LeasedInteraction{{InteractionID: "int_cold", EnvironmentID: "env_cold", LeaseToken: "tok_v1"}}, + }, nil + } + return &LeaseInteractionsResult{ + Leases: []LeasedInteraction{{InteractionID: "int_cold", EnvironmentID: "env_cold", LeaseToken: "tok_v2_refreshed"}}, + }, nil + }, + }, + Provisioner: &fakeWorkerProvisioner{ + provFn: func(WorkerProvisionRequest) (*WorkerProvisionResponse, error) { + close(provEntered) + <-provRelease + return &WorkerProvisionResponse{SchemaVersion: WorkerProvisionSchemaVersionV1, URI: "worker://cold"}, nil + }, + }, + DialWorker: func(string) (WorkerRPCClient, error) { + return &fakeWorkerRPCClient{ + startFn: func(_ context.Context, req *StartWorkRequest, _ map[string]string) (*StartWorkResponse, error) { + gotStartToken.Store(req.LeaseToken) + return &StartWorkResponse{Status: WorkerStatusAccepted}, nil + }, + }, nil + }, + Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + }) + if err != nil { + t.Fatalf("NewTier2Poller(tpCold) failed: %v", err) + } + if _, err := tpCold.PollOnce(t.Context()); err != nil { + t.Fatalf("tpCold.PollOnce(1) failed: %v", err) + } + <-provEntered + if _, err := tpCold.PollOnce(t.Context()); err != nil { + t.Fatalf("tpCold.PollOnce(2) failed: %v", err) + } + close(provRelease) + tpCold.WaitHandoffs() + if got, _ := gotStartToken.Load().(string); got != "tok_v2_refreshed" { + t.Errorf("StartWork LeaseToken after cold-start refresh = %q, want %q", got, "tok_v2_refreshed") + } + + // 4. Shared singleflight provisioning caller-cancellation isolation: + // if waiter 1 cancels while waiter 2 is still waiting on the same envID, ProvisionWorker continues and waiter 2 succeeds; + // if the sole waiter cancels, ProvisionWorker's context cancels promptly. + sfProvStarted := make(chan struct{}) + sfProvRelease := make(chan struct{}) + var sfPollNum atomic.Int32 + var sfProvCtxErrOnRelease atomic.Value + var sfAcceptedIDs []string + var sfMu sync.Mutex + tpSharedProv, err := NewTier2Poller(Tier2PollerConfig{ + Client: &fakeEnvClient{ + leaseFn: func(context.Context, *LeaseInteractionsRequest) (*LeaseInteractionsResult, error) { + switch sfPollNum.Add(1) { + case 1: + return &LeaseInteractionsResult{ + Leases: []LeasedInteraction{{InteractionID: "int_sf_1", EnvironmentID: "env_sf", LeaseToken: "t1"}}, + }, nil + case 2: + return &LeaseInteractionsResult{ + Leases: []LeasedInteraction{{InteractionID: "int_sf_2", EnvironmentID: "env_sf", LeaseToken: "t2"}}, + }, nil + default: + return &LeaseInteractionsResult{ + Leases: []LeasedInteraction{{InteractionID: "int_sf_sole", EnvironmentID: "env_sole", LeaseToken: "t3"}}, + }, nil + } + }, + }, + Provisioner: &fakeWorkerProvisioner{ + provCtxFn: func(ctx context.Context, req WorkerProvisionRequest) (*WorkerProvisionResponse, error) { + if req.EnvironmentID == "env_sole" { + close(sfProvStarted) + <-ctx.Done() + return nil, ctx.Err() + } + close(sfProvStarted) + <-sfProvRelease + if ctx.Err() != nil { + sfProvCtxErrOnRelease.Store(ctx.Err()) + return nil, ctx.Err() + } + return &WorkerProvisionResponse{SchemaVersion: WorkerProvisionSchemaVersionV1, URI: "worker://sf"}, nil + }, + }, + DialWorker: func(string) (WorkerRPCClient, error) { + return &fakeWorkerRPCClient{ + startFn: func(_ context.Context, req *StartWorkRequest, _ map[string]string) (*StartWorkResponse, error) { + sfMu.Lock() + sfAcceptedIDs = append(sfAcceptedIDs, req.InteractionID) + sfMu.Unlock() + return &StartWorkResponse{Status: WorkerStatusAccepted}, nil + }, + }, nil + }, + Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + }) + if err != nil { + t.Fatalf("NewTier2Poller(tpSharedProv) failed: %v", err) + } + ctxW1, cancelW1 := context.WithCancel(t.Context()) + if _, err := tpSharedProv.PollOnce(ctxW1); err != nil { + t.Fatalf("tpSharedProv.PollOnce(W1) failed: %v", err) + } + <-sfProvStarted + ctxW2, cancelW2 := context.WithCancel(t.Context()) + defer cancelW2() + if _, err := tpSharedProv.PollOnce(ctxW2); err != nil { + t.Fatalf("tpSharedProv.PollOnce(W2) failed: %v", err) + } + for { + tpSharedProv.mu.Lock() + bothWaiting := tpSharedProv.provFlights["env_sf"] != nil && tpSharedProv.provFlights["env_sf"].waiters == 2 + tpSharedProv.mu.Unlock() + if bothWaiting { + break + } + runtime.Gosched() + } + cancelW1() + for { + tpSharedProv.mu.Lock() + oneRemaining := tpSharedProv.provFlights["env_sf"] != nil && tpSharedProv.provFlights["env_sf"].waiters == 1 + tpSharedProv.mu.Unlock() + if oneRemaining { + break + } + runtime.Gosched() + } + close(sfProvRelease) + tpSharedProv.WaitHandoffs() + if gotErr := sfProvCtxErrOnRelease.Load(); gotErr != nil { + t.Errorf("shared ProvisionWorker ctx canceled when waiter 1 canceled while waiter 2 remained: %v", gotErr) + } + sfMu.Lock() + gotAccepted := slices.Clone(sfAcceptedIDs) + sfMu.Unlock() + if !slices.Equal(gotAccepted, []string{"int_sf_2"}) { + t.Errorf("sfAcceptedIDs = %v, want [\"int_sf_2\"]", gotAccepted) + } + + sfProvStarted = make(chan struct{}) + ctxSole, cancelSole := context.WithCancel(t.Context()) + if _, err := tpSharedProv.PollOnce(ctxSole); err != nil { + t.Fatalf("tpSharedProv.PollOnce(sole) failed: %v", err) + } + <-sfProvStarted + cancelSole() + soleDone := make(chan struct{}) + go func() { + tpSharedProv.WaitHandoffs() + close(soleDone) + }() + select { + case <-soleDone: + case <-time.After(500 * time.Millisecond): + t.Fatal("sole waiter cancellation did not abort ProvisionWorker promptly") + } +} diff --git a/internal/cli/custom/runtime_proc_unix.go b/internal/cli/custom/envruntime/proc_unix.go similarity index 62% rename from internal/cli/custom/runtime_proc_unix.go rename to internal/cli/custom/envruntime/proc_unix.go index 4f9987c..37f363c 100644 --- a/internal/cli/custom/runtime_proc_unix.go +++ b/internal/cli/custom/envruntime/proc_unix.go @@ -14,23 +14,32 @@ //go:build unix -package custom +package envruntime import ( + "errors" "os" "os/exec" "syscall" ) -func runtimeShutdownSignals() []os.Signal { +// ShutdownSignals returns the OS signals that trigger graceful runtime shutdown. +func ShutdownSignals() []os.Signal { return []os.Signal{os.Interrupt, syscall.SIGTERM} } -func configureSubprocessGroup(cmd *exec.Cmd) { +// ConfigureSubprocessGroup places cmd in its own process group and configures +// cmd.Cancel to terminate the entire process group on context cancellation. +func ConfigureSubprocessGroup(cmd *exec.Cmd) { cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} cmd.Cancel = func() error { if cmd.Process != nil && cmd.Process.Pid > 0 { - return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL) + if err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL); err != nil { + if errors.Is(err, syscall.ESRCH) { + return os.ErrProcessDone + } + return err + } } return nil } diff --git a/internal/cli/custom/envruntime/proc_windows.go b/internal/cli/custom/envruntime/proc_windows.go new file mode 100644 index 0000000..c7867bc --- /dev/null +++ b/internal/cli/custom/envruntime/proc_windows.go @@ -0,0 +1,66 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build windows + +package envruntime + +import ( + "context" + "os" + "os/exec" + "path/filepath" + "strconv" + "strings" + "syscall" + "time" +) + +// ShutdownSignals returns the OS signals that trigger graceful runtime shutdown. +func ShutdownSignals() []os.Signal { + return []os.Signal{os.Interrupt, syscall.SIGTERM} +} + +// ConfigureSubprocessGroup places cmd in a new Windows process group, preserves +// cmd.exe /C quoting, and terminates child process trees on cancellation. +func ConfigureSubprocessGroup(cmd *exec.Cmd) { + sysAttr := &syscall.SysProcAttr{ + CreationFlags: syscall.CREATE_NEW_PROCESS_GROUP, + } + if len(cmd.Args) == 3 && strings.EqualFold(cmd.Args[1], "/C") { + base := strings.ToLower(filepath.Base(cmd.Args[0])) + if base == "cmd.exe" || base == "cmd" { + sysAttr.CmdLine = `cmd.exe /d /s /c "` + cmd.Args[2] + `"` + } + } + cmd.SysProcAttr = sysAttr + cmd.Cancel = func() error { + if cmd.Process == nil { + return nil + } + if cmd.Process.Pid > 0 { + taskkillPath := "taskkill.exe" + if sysRoot := strings.TrimSpace(os.Getenv("SystemRoot")); sysRoot != "" { + taskkillPath = filepath.Join(sysRoot, "System32", "taskkill.exe") + } + killCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + err := exec.CommandContext(killCtx, taskkillPath, "/T", "/F", "/PID", strconv.Itoa(cmd.Process.Pid)).Run() + cancel() + if err == nil { + return nil + } + } + return cmd.Process.Kill() + } +} diff --git a/internal/cli/custom/envruntime/sandbox_fs.go b/internal/cli/custom/envruntime/sandbox_fs.go new file mode 100644 index 0000000..0451b37 --- /dev/null +++ b/internal/cli/custom/envruntime/sandbox_fs.go @@ -0,0 +1,1242 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package envruntime + +import ( + "bufio" + "errors" + "fmt" + "io" + "io/fs" + "os" + "path/filepath" + "regexp" + "slices" + "strings" + "time" +) + +const ( + // MaxReadFileBytes caps full-file and single-call range reads at 4 MiB unless UnlimitedRead is set. + MaxReadFileBytes = 4 * 1024 * 1024 + + // MaxGrepMatches bounds grep_search match count per call when bounded reads are enabled. + MaxGrepMatches = 1000 + + // MaxGrepOutputBytes bounds grep_search output bytes per call when bounded reads are enabled. + MaxGrepOutputBytes = 4 * 1024 * 1024 + + // MaxFindFiles bounds find_by_name results per call when bounded reads are enabled. + MaxFindFiles = 5000 +) + +var errSearchTruncated = errors.New("search results truncated") + +// PathErrorKind classifies filesystem confinement and I/O errors for callers that map to status codes. +type PathErrorKind int + +const ( + // PathErrInternal indicates an unexpected filesystem I/O failure. + PathErrInternal PathErrorKind = iota + // PathErrInvalidArgument indicates an invalid path or file type mismatch (such as reading a directory as a file). + PathErrInvalidArgument + // PathErrPermissionDenied indicates a confinement escape, symlink rejection, or OS permission error. + PathErrPermissionDenied + // PathErrNotFound indicates that the target path or parent directory does not exist. + PathErrNotFound + // PathErrAlreadyExists indicates that an exclusive create encountered an existing entry. + PathErrAlreadyExists +) + +// PathError is a structured filesystem error carrying a PathErrorKind classification. +type PathError struct { + Kind PathErrorKind + Path string + Message string + Err error +} + +// Error implements error. +func (e *PathError) Error() string { + if e.Message != "" { + if e.Err != nil { + return fmt.Sprintf("%s: %v", e.Message, e.Err) + } + return e.Message + } + if e.Err != nil { + return e.Err.Error() + } + return fmt.Sprintf("filesystem error on %q", e.Path) +} + +// Unwrap returns the underlying error. +func (e *PathError) Unwrap() error { + return e.Err +} + +// ClassifyOSError maps an error returned by os filesystem calls to a *PathError, or nil when err is nil. +func ClassifyOSError(err error, path string) error { + if err == nil { + return nil + } + var pe *PathError + if errors.As(err, &pe) { + return pe + } + switch { + case os.IsNotExist(err): + return &PathError{ + Kind: PathErrNotFound, + Path: path, + Message: fmt.Sprintf("path not found: %s", path), + Err: err, + } + case os.IsPermission(err): + return &PathError{ + Kind: PathErrPermissionDenied, + Path: path, + Message: fmt.Sprintf("permission denied: %s", path), + Err: err, + } + case os.IsExist(err): + return &PathError{ + Kind: PathErrAlreadyExists, + Path: path, + Message: fmt.Sprintf("path already exists: %s", path), + Err: err, + } + default: + return &PathError{ + Kind: PathErrInternal, + Path: path, + Message: fmt.Sprintf("operation failed on %s", path), + Err: err, + } + } +} + +// ConfinementMode selects how SandboxFS confines file operations under RootDir. +type ConfinementMode int + +const ( + // ConfinementOSRoot canonicalizes symlinks, rejects escaping or dangling symlinks, and opens files through *os.Root. + ConfinementOSRoot ConfinementMode = iota + // ConfinementLexical confines paths lexically via filepath.Clean under RootDir (for isolated container filesystems). + ConfinementLexical +) + +// FileKind classifies POSIX file types returned by SandboxFS.Stat and SandboxFS.ReadDir. +type FileKind int + +const ( + // FileKindRegular denotes a regular file. + FileKindRegular FileKind = iota + // FileKindDirectory denotes a directory. + FileKindDirectory + // FileKindSymlink denotes a symbolic link. + FileKindSymlink + // FileKindFIFOPipe denotes a named pipe (FIFO). + FileKindFIFOPipe + // FileKindSocket denotes a Unix domain socket. + FileKindSocket + // FileKindCharDevice denotes a character device. + FileKindCharDevice + // FileKindBlockDevice denotes a block device. + FileKindBlockDevice +) + +// FileMetadata holds POSIX metadata for a filesystem entry. +type FileMetadata struct { + Name string + Kind FileKind + IsDir bool + Size int64 + ModTime time.Time + Permissions uint32 + SymlinkTarget string +} + +// SandboxFSConfig configures a SandboxFS instance. +type SandboxFSConfig struct { + RootDir string + WorkDir string + Mode ConfinementMode + UnlimitedRead bool +} + +// SandboxFS provides path-confined filesystem operations for agent tools and gRPC filesystem services. +type SandboxFS struct { + rootDir string + workDir string + mode ConfinementMode + unlimitedRead bool +} + +// NewSandboxFS creates a SandboxFS using ConfinementOSRoot and bounded 4 MiB reads. +func NewSandboxFS(rootDir, workDir string) *SandboxFS { + return NewSandboxFSWithConfig(SandboxFSConfig{ + RootDir: rootDir, + WorkDir: workDir, + Mode: ConfinementOSRoot, + }) +} + +// NewSandboxFSWithConfig creates a SandboxFS with the given configuration. +func NewSandboxFSWithConfig(cfg SandboxFSConfig) *SandboxFS { + if cfg.Mode == ConfinementLexical { + root := filepath.Clean(cfg.RootDir) + if cfg.RootDir == "" { + root = string(filepath.Separator) + } else if !isVirtualRootPath(cfg.RootDir) { + if abs, err := filepath.Abs(cfg.RootDir); err == nil { + root = filepath.Clean(abs) + } + } + work := root + if strings.TrimSpace(cfg.WorkDir) != "" { + work = filepath.Clean(cfg.WorkDir) + if !isVirtualRootPath(cfg.WorkDir) { + if absWork, err := filepath.Abs(cfg.WorkDir); err == nil && isWithinRoot(root, absWork) { + work = filepath.Clean(absWork) + } else if !isWithinRoot(root, work) { + work = filepath.Join(root, work) + } + } + if !isWithinRoot(root, work) { + work = root + } + } + return &SandboxFS{ + rootDir: root, + workDir: work, + mode: ConfinementLexical, + unlimitedRead: cfg.UnlimitedRead, + } + } + + rootDir := cfg.RootDir + if abs, err := filepath.Abs(rootDir); err == nil { + rootDir = canonicalizeLongestPrefix(abs) + } + cleanRoot := filepath.Clean(rootDir) + cleanWork := cleanRoot + if strings.TrimSpace(cfg.WorkDir) != "" { + if absWork, err := filepath.Abs(cfg.WorkDir); err == nil { + cleanWork = canonicalizeLongestPrefix(absWork) + } else { + cleanWork = filepath.Clean(cfg.WorkDir) + } + if !isWithinRoot(cleanRoot, cleanWork) { + cleanWork = cleanRoot + } + } + return &SandboxFS{ + rootDir: cleanRoot, + workDir: cleanWork, + mode: ConfinementOSRoot, + unlimitedRead: cfg.UnlimitedRead, + } +} + +// RootDir returns the normalized confinement root directory. +func (s *SandboxFS) RootDir() string { + return s.rootDir +} + +// WorkDir returns the normalized default working directory for relative paths. +func (s *SandboxFS) WorkDir() string { + return s.workDir +} + +// Mode returns the configured ConfinementMode. +func (s *SandboxFS) Mode() ConfinementMode { + return s.mode +} + +// ResolveSandboxDirs validates rootDir and workDir confinement (creating both directories when create is true). +func ResolveSandboxDirs(rootDir, workDir string, create bool) (string, string, error) { + if strings.TrimSpace(rootDir) == "" { + rootDir = "." + } + absRoot, err := filepath.Abs(rootDir) + if err != nil { + return "", "", fmt.Errorf("resolving root-dir %q: %w", rootDir, err) + } + origAbsRoot := filepath.Clean(absRoot) + absRoot = canonicalizeLongestPrefix(origAbsRoot) + + if strings.TrimSpace(workDir) == "" || workDir == "." { + workDir = absRoot + } + absWork, err := filepath.Abs(workDir) + if err != nil { + return "", "", fmt.Errorf("resolving workdir %q: %w", workDir, err) + } + if isWithinRoot(origAbsRoot, absWork) { + if rel, relErr := filepath.Rel(origAbsRoot, absWork); relErr == nil { + absWork = filepath.Join(absRoot, rel) + } + } + if !isWithinRoot(absRoot, absWork) { + return "", "", fmt.Errorf("workdir %q escapes root-dir %q", absWork, absRoot) + } + resolvedWork, err := NewSandboxFS(absRoot, absRoot).ResolvePath(absWork) + if err != nil { + return "", "", fmt.Errorf("workdir %q escapes root-dir %q: %w", absWork, absRoot, err) + } + absWork = resolvedWork + + if !create { + return absRoot, absWork, nil + } + if err := os.MkdirAll(absRoot, 0o755); err != nil { + return "", "", fmt.Errorf("creating root-dir %q: %w", absRoot, err) + } + if err := os.MkdirAll(absWork, 0o755); err != nil { + return "", "", fmt.Errorf("creating workdir %q: %w", absWork, err) + } + return absRoot, absWork, nil +} + +func isWithinRoot(root, target string) bool { + if root == string(filepath.Separator) && filepath.IsAbs(target) { + return true + } + rel, err := filepath.Rel(root, target) + if err != nil { + return false + } + return rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator)) +} + +func canonicalizeLongestPrefix(absPath string) string { + cur := filepath.Clean(absPath) + var suffix []string + for { + if eval, err := filepath.EvalSymlinks(cur); err == nil { + for i := len(suffix) - 1; i >= 0; i-- { + eval = filepath.Join(eval, suffix[i]) + } + return filepath.Clean(eval) + } + parent := filepath.Dir(cur) + if parent == cur { + break + } + suffix = append(suffix, filepath.Base(cur)) + cur = parent + } + return filepath.Clean(absPath) +} + +func rejectNULByte(path string) error { + if strings.Contains(path, "\x00") { + return &PathError{ + Kind: PathErrPermissionDenied, + Path: path, + Message: "access denied: path contains NUL byte", + } + } + return nil +} + +func isDriveRelativePath(p string) bool { + return filepath.VolumeName(p) != "" && !filepath.IsAbs(p) +} + +func isVirtualRootPath(p string) bool { + if filepath.IsAbs(p) { + return true + } + return filepath.VolumeName(p) == "" && (strings.HasPrefix(p, "/") || strings.HasPrefix(p, `\`)) +} + +// ResolvePath resolves relPath within RootDir according to Mode. +func (s *SandboxFS) ResolvePath(relPath string) (string, error) { + if s.mode == ConfinementLexical { + return s.resolveLexicalPath(relPath) + } + return s.resolveOSRootPath(relPath) +} + +func (s *SandboxFS) resolveLexicalPath(clientPath string) (string, error) { + if err := rejectNULByte(clientPath); err != nil { + return "", err + } + if isDriveRelativePath(clientPath) { + return "", &PathError{ + Kind: PathErrPermissionDenied, + Path: clientPath, + Message: fmt.Sprintf("access denied: drive-relative path %q is not permitted", clientPath), + } + } + root := filepath.Clean(s.rootDir) + if root == "" { + root = string(filepath.Separator) + } + work := s.workDir + if work == "" { + work = root + } + cleaned := filepath.Clean(clientPath) + + var candidate string + if isVirtualRootPath(cleaned) { + switch { + case root == string(filepath.Separator): + candidate = cleaned + case isWithinRoot(root, cleaned): + candidate = cleaned + default: + trimmed := strings.TrimPrefix(cleaned, filepath.VolumeName(cleaned)) + trimmed = strings.TrimLeft(trimmed, `/\`) + candidate = filepath.Join(root, trimmed) + } + } else { + candidate = filepath.Join(work, cleaned) + } + + if !isWithinRoot(root, candidate) { + return "", &PathError{ + Kind: PathErrPermissionDenied, + Path: clientPath, + Message: fmt.Sprintf("access denied: path %q escapes root prefix %q", clientPath, root), + } + } + return candidate, nil +} + +func (s *SandboxFS) lexicalOSRootCandidate(relPath string) (string, error) { + if err := rejectNULByte(relPath); err != nil { + return "", err + } + if isDriveRelativePath(relPath) { + return "", &PathError{ + Kind: PathErrPermissionDenied, + Path: relPath, + Message: fmt.Sprintf("access denied: drive-relative path %q is not permitted", relPath), + } + } + var candidate string + if isVirtualRootPath(relPath) { + cleaned := filepath.Clean(relPath) + if isWithinRoot(s.rootDir, cleaned) { + candidate = cleaned + } else { + trimmed := strings.TrimPrefix(cleaned, filepath.VolumeName(cleaned)) + trimmed = strings.TrimLeft(trimmed, `/\`) + candidate = filepath.Join(s.rootDir, trimmed) + } + } else { + candidate = filepath.Join(s.workDir, relPath) + } + if !isWithinRoot(s.rootDir, candidate) { + return "", &PathError{ + Kind: PathErrPermissionDenied, + Path: relPath, + Message: fmt.Sprintf("access denied: path %q escapes root directory", relPath), + } + } + return candidate, nil +} + +func (s *SandboxFS) resolveOSRootPath(relPath string) (string, error) { + candidate, err := s.lexicalOSRootCandidate(relPath) + if err != nil { + return "", err + } + + cur := candidate + var suffix []string + for { + if eval, err := filepath.EvalSymlinks(cur); err == nil { + resolved := eval + for i := len(suffix) - 1; i >= 0; i-- { + resolved = filepath.Join(resolved, suffix[i]) + if fi, lerr := os.Lstat(resolved); lerr == nil && fi.Mode()&os.ModeSymlink != 0 { + return "", &PathError{ + Kind: PathErrPermissionDenied, + Path: relPath, + Message: fmt.Sprintf("access denied: dangling symlink %q is not permitted", relPath), + } + } + } + if !isWithinRoot(s.rootDir, resolved) { + return "", &PathError{ + Kind: PathErrPermissionDenied, + Path: relPath, + Message: fmt.Sprintf("access denied: symlink target for %q escapes root directory", relPath), + } + } + return resolved, nil + } + if fi, lerr := os.Lstat(cur); lerr == nil && fi.Mode()&os.ModeSymlink != 0 { + return "", &PathError{ + Kind: PathErrPermissionDenied, + Path: relPath, + Message: fmt.Sprintf("access denied: dangling symlink %q is not permitted", relPath), + } + } + parent := filepath.Dir(cur) + if parent == cur { + break + } + suffix = append(suffix, filepath.Base(cur)) + cur = parent + } + return candidate, nil +} + +func reanchorPathError(err error, origPath, internalPath string) error { + var pe *PathError + if errors.As(err, &pe) { + return &PathError{ + Kind: pe.Kind, + Path: origPath, + Message: strings.ReplaceAll(pe.Message, internalPath, origPath), + Err: pe.Err, + } + } + return err +} + +func (s *SandboxFS) resolveOSRootNoFollowLeaf(relPath string) (string, error) { + candidate, err := s.lexicalOSRootCandidate(relPath) + if err != nil { + return "", err + } + if filepath.Clean(candidate) == s.rootDir { + return s.rootDir, nil + } + parentDir := filepath.Dir(candidate) + parentResolved, err := s.resolveOSRootPath(parentDir) + if err != nil { + return "", reanchorPathError(err, relPath, parentDir) + } + leafFull := filepath.Join(parentResolved, filepath.Base(candidate)) + if fi, lerr := os.Lstat(leafFull); lerr == nil && fi.Mode()&os.ModeSymlink != 0 { + if _, err := s.resolveOSRootPath(leafFull); err != nil { + return "", reanchorPathError(err, relPath, leafFull) + } + } + return leafFull, nil +} + +func (s *SandboxFS) openInRoot(relPath string) (*os.Root, string, error) { + resolved, err := s.resolveOSRootPath(relPath) + if err != nil { + return nil, "", err + } + rel, _ := filepath.Rel(s.rootDir, resolved) + root, err := os.OpenRoot(s.rootDir) + if err != nil { + return nil, "", ClassifyOSError(err, relPath) + } + return root, rel, nil +} + +func (s *SandboxFS) openInRootNoFollowLeaf(relPath string) (*os.Root, string, error) { + resolved, err := s.resolveOSRootNoFollowLeaf(relPath) + if err != nil { + return nil, "", err + } + rel, _ := filepath.Rel(s.rootDir, resolved) + root, err := os.OpenRoot(s.rootDir) + if err != nil { + return nil, "", ClassifyOSError(err, relPath) + } + return root, rel, nil +} + +// OpenRead opens path for reading after validating confinement and checking that it is a regular file. +func (s *SandboxFS) OpenRead(path string) (*os.File, error) { + if path == "" { + return nil, &PathError{Kind: PathErrInvalidArgument, Path: path, Message: "path is required"} + } + var ( + f *os.File + err error + ) + if s.mode == ConfinementLexical { + targetPath, resErr := s.resolveLexicalPath(path) + if resErr != nil { + return nil, resErr + } + if fi, statErr := os.Stat(targetPath); statErr == nil { + if fi.IsDir() { + return nil, &PathError{Kind: PathErrInvalidArgument, Path: path, Message: fmt.Sprintf("cannot read directory %q as file", path)} + } + if fi.Mode()&(os.ModeNamedPipe|os.ModeSocket) != 0 { + return nil, &PathError{Kind: PathErrInvalidArgument, Path: path, Message: fmt.Sprintf("cannot read non-regular file %q", path)} + } + } + f, err = os.Open(targetPath) + } else { + root, rel, rootErr := s.openInRoot(path) + if rootErr != nil { + return nil, rootErr + } + defer root.Close() + if fi, statErr := root.Lstat(rel); statErr == nil && !fi.Mode().IsRegular() { + return nil, &PathError{Kind: PathErrInvalidArgument, Path: path, Message: fmt.Sprintf("cannot read non-regular file %q", path)} + } + f, err = root.Open(rel) + } + if err != nil { + return nil, ClassifyOSError(err, path) + } + info, err := f.Stat() + if err != nil { + _ = f.Close() + return nil, ClassifyOSError(err, path) + } + if s.mode == ConfinementLexical { + if info.IsDir() { + _ = f.Close() + return nil, &PathError{Kind: PathErrInvalidArgument, Path: path, Message: fmt.Sprintf("cannot read directory %q as file", path)} + } + if info.Mode()&(os.ModeNamedPipe|os.ModeSocket) != 0 { + _ = f.Close() + return nil, &PathError{Kind: PathErrInvalidArgument, Path: path, Message: fmt.Sprintf("cannot read non-regular file %q", path)} + } + } else if !info.Mode().IsRegular() { + _ = f.Close() + return nil, &PathError{Kind: PathErrInvalidArgument, Path: path, Message: fmt.Sprintf("cannot read non-regular file %q", path)} + } + return f, nil +} + +// OpenWrite opens path for writing, creating parent directories with 0755 permissions as needed. +// When perm is non-zero, the file's permissions are explicitly set to perm (including when overwriting an existing file). +func (s *SandboxFS) OpenWrite(path string, perm os.FileMode, overwrite bool) (*os.File, error) { + return s.openWriteFile(path, perm, overwrite, false) +} + +func (s *SandboxFS) openWriteFile(path string, perm os.FileMode, overwrite, appendMode bool) (*os.File, error) { + if path == "" { + return nil, &PathError{Kind: PathErrInvalidArgument, Path: path, Message: "path is required"} + } + effectivePerm := perm + if effectivePerm == 0 { + effectivePerm = 0o644 + } + flags := os.O_CREATE | os.O_WRONLY + switch { + case appendMode: + flags |= os.O_APPEND + case overwrite: + flags |= os.O_TRUNC + default: + flags |= os.O_EXCL + } + + var ( + f *os.File + err error + ) + if s.mode == ConfinementLexical { + safePath, resErr := s.resolveLexicalPath(path) + if resErr != nil { + return nil, resErr + } + if mkErr := os.MkdirAll(filepath.Dir(safePath), 0o755); mkErr != nil { + return nil, ClassifyOSError(mkErr, path) + } + f, err = os.OpenFile(safePath, flags, effectivePerm) + } else { + root, rel, rootErr := s.openInRoot(path) + if rootErr != nil { + return nil, rootErr + } + defer root.Close() + if fi, statErr := root.Lstat(rel); statErr == nil && !fi.Mode().IsRegular() { + return nil, &PathError{Kind: PathErrInvalidArgument, Path: path, Message: fmt.Sprintf("cannot write non-regular file %q", path)} + } + if mkErr := root.MkdirAll(filepath.Dir(rel), 0o755); mkErr != nil { + return nil, ClassifyOSError(mkErr, path) + } + f, err = root.OpenFile(rel, flags, effectivePerm) + } + if err != nil { + return nil, ClassifyOSError(err, path) + } + if perm != 0 { + if chmodErr := f.Chmod(effectivePerm); chmodErr != nil { + _ = f.Close() + return nil, ClassifyOSError(chmodErr, path) + } + } + return f, nil +} + +// ReadRegularFile reads a file under confinement, enforcing MaxReadFileBytes unless UnlimitedRead is set. +func (s *SandboxFS) ReadRegularFile(path string) ([]byte, error) { + f, err := s.OpenRead(path) + if err != nil { + return nil, err + } + defer f.Close() + + if s.unlimitedRead { + data, err := io.ReadAll(f) + if err != nil { + return nil, ClassifyOSError(err, path) + } + return data, nil + } + + data, exceeded, err := readLimitedAll(f, MaxReadFileBytes) + if err != nil { + return nil, ClassifyOSError(err, path) + } + if exceeded { + return nil, fmt.Errorf("file %q exceeds maximum full-read size (%d bytes); specify start_line and end_line", path, MaxReadFileBytes) + } + return data, nil +} + +func readLimitedAll(r io.Reader, maxBytes int) ([]byte, bool, error) { + data, err := io.ReadAll(io.LimitReader(r, int64(maxBytes)+1)) + if err != nil { + return nil, false, err + } + if len(data) > maxBytes { + return nil, true, nil + } + return data, false, nil +} + +// ReadFileRange reads 1-indexed inclusive lines [startLine, endLine] from path. +func (s *SandboxFS) ReadFileRange(path string, startLine, endLine int) (string, error) { + if s.unlimitedRead { + raw, err := s.ReadRegularFile(path) + if err != nil { + return "", err + } + return sliceLines(string(raw), startLine, endLine), nil + } + + f, err := s.OpenRead(path) + if err != nil { + return "", err + } + defer f.Close() + + reader := bufio.NewReader(f) + var selected []string + totalBytes := 0 + for lineNum := 1; endLine <= 0 || lineNum <= endLine; lineNum++ { + part, readErr := reader.ReadString('\n') + if readErr != nil && !errors.Is(readErr, io.EOF) { + return "", ClassifyOSError(readErr, path) + } + if part == "" { + break + } + line := strings.TrimSuffix(strings.TrimSuffix(part, "\n"), "\r") + if lineNum >= startLine { + totalBytes += len(line) + 1 + if totalBytes > MaxReadFileBytes { + return "", fmt.Errorf("line range [%d, %d] of %q exceeds maximum read size (%d bytes); narrow start_line and end_line", max(startLine, 1), endLine, path, MaxReadFileBytes) + } + selected = append(selected, line) + } + } + return strings.Join(selected, "\n"), nil +} + +func sliceLines(text string, startLine, endLine int) string { + lines := strings.Split(text, "\n") + if startLine <= 1 { + startLine = 1 + } + if startLine > len(lines) { + return "" + } + if endLine <= 0 || endLine > len(lines) { + endLine = len(lines) + } + if endLine < startLine { + return "" + } + return strings.Join(lines[startLine-1:endLine], "\n") +} + +// WriteRegularFile writes content to path with optional overwrite or append semantics. +func (s *SandboxFS) WriteRegularFile(path, content string, overwrite, appendMode bool) (int, error) { + f, err := s.openWriteFile(path, 0, overwrite, appendMode) + if err != nil { + return 0, err + } + n, writeErr := f.WriteString(content) + closeErr := f.Close() + if writeErr != nil { + return n, ClassifyOSError(writeErr, path) + } + if closeErr != nil { + return n, ClassifyOSError(closeErr, path) + } + return n, nil +} + +// Stat returns FileMetadata for path without following a leaf symlink. +func (s *SandboxFS) Stat(path string) (FileMetadata, error) { + if path == "" { + return FileMetadata{}, &PathError{Kind: PathErrInvalidArgument, Path: path, Message: "path is required"} + } + if s.mode == ConfinementLexical { + targetPath, err := s.resolveLexicalPath(path) + if err != nil { + return FileMetadata{}, err + } + info, err := os.Lstat(targetPath) + if err != nil { + return FileMetadata{}, ClassifyOSError(err, path) + } + return buildFileMetadata(info, func() (string, error) { return os.Readlink(targetPath) }), nil + } + + root, rel, err := s.openInRootNoFollowLeaf(path) + if err != nil { + return FileMetadata{}, err + } + defer root.Close() + + info, err := root.Lstat(rel) + if err != nil { + return FileMetadata{}, ClassifyOSError(err, path) + } + return buildFileMetadata(info, func() (string, error) { return root.Readlink(rel) }), nil +} + +// Mkdir creates a directory (or directory tree when recursive is true) under confinement. +func (s *SandboxFS) Mkdir(path string, perm os.FileMode, recursive bool) error { + if path == "" { + return &PathError{Kind: PathErrInvalidArgument, Path: path, Message: "path is required"} + } + if perm == 0 { + perm = 0o755 + } + if s.mode == ConfinementLexical { + targetPath, err := s.resolveLexicalPath(path) + if err != nil { + return err + } + if recursive { + return ClassifyOSError(os.MkdirAll(targetPath, perm), path) + } + return ClassifyOSError(os.Mkdir(targetPath, perm), path) + } + + root, rel, err := s.openInRoot(path) + if err != nil { + return err + } + defer root.Close() + + if recursive { + return ClassifyOSError(root.MkdirAll(rel, perm), path) + } + return ClassifyOSError(root.Mkdir(rel, perm), path) +} + +// Remove deletes a file or directory under confinement, rejecting removal of RootDir or WorkDir. +func (s *SandboxFS) Remove(path string, recursive bool) error { + if strings.TrimSpace(path) == "" { + return &PathError{ + Kind: PathErrPermissionDenied, + Path: path, + Message: "cannot remove empty path", + } + } + if s.mode == ConfinementLexical { + targetPath, err := s.resolveLexicalPath(path) + if err != nil { + return err + } + if targetPath == s.rootDir || targetPath == s.workDir { + return &PathError{ + Kind: PathErrPermissionDenied, + Path: path, + Message: fmt.Sprintf("cannot remove workspace root or working directory: %s", path), + } + } + fi, statErr := os.Lstat(targetPath) + if statErr != nil { + return ClassifyOSError(statErr, path) + } + if recursive && fi.Mode()&os.ModeSymlink == 0 { + return ClassifyOSError(os.RemoveAll(targetPath), path) + } + return ClassifyOSError(os.Remove(targetPath), path) + } + + root, rel, err := s.openInRootNoFollowLeaf(path) + if err != nil { + return err + } + defer root.Close() + if rel == "." || filepath.Join(s.rootDir, rel) == s.workDir { + return &PathError{ + Kind: PathErrPermissionDenied, + Path: path, + Message: fmt.Sprintf("cannot remove workspace root or working directory: %s", path), + } + } + fi, statErr := root.Lstat(rel) + if statErr != nil { + return ClassifyOSError(statErr, path) + } + if recursive && fi.Mode()&os.ModeSymlink == 0 { + return ClassifyOSError(root.RemoveAll(rel), path) + } + return ClassifyOSError(root.Remove(rel), path) +} + +// ReadDir lists FileMetadata entries in directory path in lexicographical order. +func (s *SandboxFS) ReadDir(path string) ([]FileMetadata, error) { + if s.mode == ConfinementLexical { + targetPath, err := s.resolveLexicalPath(path) + if err != nil { + return nil, err + } + dirEntries, err := os.ReadDir(targetPath) + if err != nil { + return nil, ClassifyOSError(err, path) + } + out := make([]FileMetadata, 0, len(dirEntries)) + for _, entry := range dirEntries { + info, err := entry.Info() + if err != nil { + continue + } + fullEntryPath := filepath.Join(targetPath, entry.Name()) + out = append(out, buildFileMetadata(info, func() (string, error) { return os.Readlink(fullEntryPath) })) + } + return out, nil + } + + root, rel, err := s.openInRoot(path) + if err != nil { + return nil, err + } + defer root.Close() + + dirFile, err := root.Open(rel) + if err != nil { + return nil, ClassifyOSError(err, path) + } + defer dirFile.Close() + + dirEntries, err := dirFile.ReadDir(-1) + if err != nil { + return nil, ClassifyOSError(err, path) + } + slices.SortFunc(dirEntries, func(a, b os.DirEntry) int { + return strings.Compare(a.Name(), b.Name()) + }) + out := make([]FileMetadata, 0, len(dirEntries)) + for _, entry := range dirEntries { + info, err := entry.Info() + if err != nil { + continue + } + entryRel := filepath.Join(rel, entry.Name()) + out = append(out, buildFileMetadata(info, func() (string, error) { return root.Readlink(entryRel) })) + } + return out, nil +} + +// Rename moves or renames oldPath to newPath under confinement, creating destination parent directories as needed. +func (s *SandboxFS) Rename(oldPath, newPath string) error { + if oldPath == "" || newPath == "" { + return &PathError{Kind: PathErrInvalidArgument, Path: oldPath, Message: "old_path and new_path are required"} + } + if s.mode == ConfinementLexical { + resolvedOld, err := s.resolveLexicalPath(oldPath) + if err != nil { + return err + } + resolvedNew, err := s.resolveLexicalPath(newPath) + if err != nil { + return err + } + if resolvedOld == s.rootDir || resolvedNew == s.rootDir || resolvedOld == s.workDir || resolvedNew == s.workDir { + return &PathError{ + Kind: PathErrPermissionDenied, + Path: oldPath, + Message: "cannot rename workspace root or working directory", + } + } + if _, statErr := os.Lstat(resolvedOld); statErr != nil { + return ClassifyOSError(statErr, oldPath) + } + if err := os.MkdirAll(filepath.Dir(resolvedNew), 0o755); err != nil { + return ClassifyOSError(err, newPath) + } + return ClassifyOSError(os.Rename(resolvedOld, resolvedNew), oldPath) + } + + resolvedOld, err := s.resolveOSRootNoFollowLeaf(oldPath) + if err != nil { + return err + } + resolvedNew, err := s.resolveOSRootNoFollowLeaf(newPath) + if err != nil { + return err + } + relOld, _ := filepath.Rel(s.rootDir, resolvedOld) + relNew, _ := filepath.Rel(s.rootDir, resolvedNew) + if relOld == "." || relNew == "." || resolvedOld == s.workDir || resolvedNew == s.workDir { + return &PathError{ + Kind: PathErrPermissionDenied, + Path: oldPath, + Message: "cannot rename workspace root or working directory", + } + } + root, err := os.OpenRoot(s.rootDir) + if err != nil { + return ClassifyOSError(err, oldPath) + } + defer root.Close() + if _, statErr := root.Lstat(relOld); statErr != nil { + return ClassifyOSError(statErr, oldPath) + } + if err := root.MkdirAll(filepath.Dir(relNew), 0o755); err != nil { + return ClassifyOSError(err, newPath) + } + return ClassifyOSError(root.Rename(relOld, relNew), oldPath) +} + +func buildFileMetadata(fi os.FileInfo, readlink func() (string, error)) FileMetadata { + var kind FileKind + var symlinkTarget string + mode := fi.Mode() + + switch { + case mode&os.ModeSymlink != 0: + kind = FileKindSymlink + if readlink != nil { + if target, err := readlink(); err == nil { + symlinkTarget = target + } + } + case fi.IsDir(): + kind = FileKindDirectory + case mode&os.ModeNamedPipe != 0: + kind = FileKindFIFOPipe + case mode&os.ModeSocket != 0: + kind = FileKindSocket + case mode&os.ModeDevice != 0: + if mode&os.ModeCharDevice != 0 { + kind = FileKindCharDevice + } else { + kind = FileKindBlockDevice + } + default: + kind = FileKindRegular + } + + return FileMetadata{ + Name: fi.Name(), + Kind: kind, + IsDir: fi.IsDir(), + Size: fi.Size(), + ModTime: fi.ModTime(), + Permissions: uint32(mode.Perm()), + SymlinkTarget: symlinkTarget, + } +} + +// ListEntries returns tool-formatted directory entries (name, is_dir, size) for path. +func (s *SandboxFS) ListEntries(path string) ([]any, error) { + entries, err := s.ReadDir(path) + if err != nil { + return nil, err + } + out := make([]any, 0, len(entries)) + for _, entry := range entries { + out = append(out, map[string]any{ + "name": entry.Name, + "is_dir": entry.IsDir, + "size": entry.Size, + }) + } + return out, nil +} + +func (s *SandboxFS) readGrepCandidate(osRoot *os.Root, p string) ([]byte, bool) { + if osRoot == nil { + data, err := s.ReadRegularFile(p) + return data, err == nil + } + resolved, err := s.resolveOSRootPath(p) + if err != nil { + return nil, false + } + rel, err := filepath.Rel(s.rootDir, resolved) + if err != nil { + return nil, false + } + fi, err := osRoot.Lstat(rel) + if err != nil || !fi.Mode().IsRegular() { + return nil, false + } + f, err := osRoot.Open(rel) + if err != nil { + return nil, false + } + defer f.Close() + if s.unlimitedRead { + data, readErr := io.ReadAll(f) + return data, readErr == nil + } + data, exceeded, readErr := readLimitedAll(f, MaxReadFileBytes) + return data, readErr == nil && !exceeded +} + +// Grep searches regular files under searchPath for query, returning matches and whether output was truncated. +func (s *SandboxFS) Grep(searchPath, query string, caseInsensitive, isRegex bool) ([]any, bool, error) { + fullPath, err := s.ResolvePath(searchPath) + if err != nil { + return nil, false, err + } + if query == "" { + return nil, false, errors.New("query is required") + } + pattern := query + if !isRegex { + pattern = regexp.QuoteMeta(query) + } + if caseInsensitive { + pattern = "(?i)" + pattern + } + re, err := regexp.Compile(pattern) + if err != nil { + return nil, false, fmt.Errorf("invalid search pattern: %w", err) + } + + var osRoot *os.Root + if s.mode != ConfinementLexical { + osRoot, err = os.OpenRoot(s.rootDir) + if err != nil { + return nil, false, ClassifyOSError(err, searchPath) + } + defer osRoot.Close() + } + + var matches []any + totalBytes := 0 + err = filepath.WalkDir(fullPath, func(p string, d fs.DirEntry, walkErr error) error { + if walkErr != nil || d.IsDir() { + return nil + } + if d.Type()&(fs.ModeNamedPipe|fs.ModeSocket|fs.ModeDevice|fs.ModeCharDevice) != 0 { + return nil + } + data, ok := s.readGrepCandidate(osRoot, p) + if !ok { + return nil + } + + rel, relErr := filepath.Rel(fullPath, p) + if relErr != nil || rel == "." { + rel = filepath.Base(p) + } + slashRel := filepath.ToSlash(rel) + for idx, rawLine := range strings.Split(string(data), "\n") { + line := strings.TrimSuffix(rawLine, "\r") + if re.MatchString(line) { + entryBytes := len(slashRel) + len(line) + 32 + if !s.unlimitedRead && (len(matches) >= MaxGrepMatches || (len(matches) > 0 && totalBytes+entryBytes > MaxGrepOutputBytes)) { + return errSearchTruncated + } + totalBytes += entryBytes + matches = append(matches, map[string]any{ + "file": slashRel, + "line_number": idx + 1, + "line": line, + }) + } + } + return nil + }) + if errors.Is(err, errSearchTruncated) { + return matches, true, nil + } + return matches, false, err +} + +// FindFiles walks searchDir and returns relative paths matching pattern, optional maxDepth, and optional extensions. +func (s *SandboxFS) FindFiles(searchDir, pattern string, maxDepth int, extensions ...string) ([]any, bool, error) { + fullPath, err := s.ResolvePath(searchDir) + if err != nil { + return nil, false, err + } + var extAllow map[string]bool + for _, rawExt := range extensions { + clean := strings.TrimPrefix(strings.ToLower(strings.TrimSpace(rawExt)), ".") + if clean != "" { + if extAllow == nil { + extAllow = make(map[string]bool, len(extensions)) + } + extAllow[clean] = true + } + } + patternLower := strings.ToLower(pattern) + var files []any + err = filepath.WalkDir(fullPath, func(p string, d fs.DirEntry, walkErr error) error { + if walkErr != nil || p == fullPath { + return nil + } + if _, resErr := s.ResolvePath(p); resErr != nil { + return nil + } + rel, relErr := filepath.Rel(fullPath, p) + if relErr != nil { + return nil + } + slashRel := filepath.ToSlash(rel) + if maxDepth > 0 && strings.Count(slashRel, "/")+1 > maxDepth { + if d.IsDir() { + return filepath.SkipDir + } + return nil + } + name := d.Name() + if len(extAllow) > 0 { + if d.IsDir() { + return nil + } + fileExt := strings.TrimPrefix(strings.ToLower(filepath.Ext(name)), ".") + if !extAllow[fileExt] { + return nil + } + } + matched, _ := filepath.Match(pattern, name) + if matched || patternLower == "" || strings.Contains(strings.ToLower(name), patternLower) { + if !s.unlimitedRead && len(files) >= MaxFindFiles { + return errSearchTruncated + } + files = append(files, slashRel) + } + return nil + }) + if errors.Is(err, errSearchTruncated) { + return files, true, nil + } + return files, false, err +} diff --git a/internal/cli/custom/envruntime/sandbox_test.go b/internal/cli/custom/envruntime/sandbox_test.go new file mode 100644 index 0000000..080e100 --- /dev/null +++ b/internal/cli/custom/envruntime/sandbox_test.go @@ -0,0 +1,991 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package envruntime + +import ( + "context" + "errors" + "io" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "testing" + "time" +) + +func mustSymlink(t *testing.T, target, link string) { + t.Helper() + if err := os.Symlink(target, link); err != nil { + if runtime.GOOS == "windows" { + t.Skipf("os.Symlink unsupported in this Windows environment: %v", err) + } + t.Fatalf("os.Symlink(%q, %q) failed: %v", target, link, err) + } +} + +func TestResolveSandboxDirs(t *testing.T) { + base := t.TempDir() + rootDir := filepath.Join(base, "root") + workDir := filepath.Join(rootDir, "proj") + // create=false validates confinement without mutating disk (used before --dry-run). + if _, _, err := ResolveSandboxDirs(rootDir, workDir, false); err != nil { + t.Fatalf("ResolveSandboxDirs(non-existent, create=false) = %v, want nil", err) + } + if _, statErr := os.Stat(rootDir); !os.IsNotExist(statErr) { + t.Errorf("ResolveSandboxDirs(create=false) unexpectedly created %q", rootDir) + } + gotRoot, gotWork, err := ResolveSandboxDirs(rootDir, workDir, true) + if err != nil { + t.Fatalf("ResolveSandboxDirs(create=true) failed: %v", err) + } + if !strings.HasPrefix(gotWork, gotRoot) { + t.Errorf("gotWork %q does not reside inside gotRoot %q", gotWork, gotRoot) + } + + outside := t.TempDir() + if _, _, err := ResolveSandboxDirs(gotRoot, outside, false); err == nil { + t.Error("ResolveSandboxDirs(outside workdir) = nil, want escape error") + } + + // Sibling prefix directory (/tmp/.../root2 vs /tmp/.../root) must be rejected. + siblingPrefix := gotRoot + "2" + if err := os.MkdirAll(siblingPrefix, 0o755); err != nil { + t.Fatalf("MkdirAll(siblingPrefix) failed: %v", err) + } + if _, _, err := ResolveSandboxDirs(gotRoot, siblingPrefix, false); err == nil { + t.Errorf("ResolveSandboxDirs(%q, %q) = nil, want sibling-prefix escape error", gotRoot, siblingPrefix) + } + + // Symlinked rootDir with workDir under symlinked rootDir canonicalizes both paths. + symlinkRoot := filepath.Join(base, "sym-root") + mustSymlink(t, gotRoot, symlinkRoot) + canonRoot, canonWork, err := ResolveSandboxDirs(symlinkRoot, filepath.Join(symlinkRoot, "proj"), false) + if err != nil || canonRoot != gotRoot || canonWork != gotWork { + t.Errorf("ResolveSandboxDirs(symlink) = (%q, %q, %v), want (%q, %q, nil)", canonRoot, canonWork, err, gotRoot, gotWork) + } + symFS := NewSandboxFS(symlinkRoot, filepath.Join(symlinkRoot, "proj")) + if symFS.RootDir() != gotRoot || symFS.WorkDir() != gotWork { + t.Errorf("NewSandboxFS(symlink) = (%q, %q), want (%q, %q)", symFS.RootDir(), symFS.WorkDir(), gotRoot, gotWork) + } + + // Symlink inside rootDir pointing to a subdirectory inside rootDir canonicalizes workDir. + innerWorkLink := filepath.Join(gotRoot, "proj_link") + mustSymlink(t, gotWork, innerWorkLink) + _, canonInnerWork, err := ResolveSandboxDirs(gotRoot, innerWorkLink, false) + if err != nil || canonInnerWork != gotWork { + t.Errorf("ResolveSandboxDirs(innerWorkLink) = (%q, %v), want (%q, nil)", canonInnerWork, err, gotWork) + } + + // Symlink inside rootDir pointing outside rootDir cannot be used as workDir. + escapeWorkLink := filepath.Join(gotRoot, "sym_escape_work") + mustSymlink(t, outside, escapeWorkLink) + if _, _, err := ResolveSandboxDirs(gotRoot, escapeWorkLink, false); err == nil { + t.Error("ResolveSandboxDirs(workDir symlink escaping rootDir) = nil, want escape error") + } + + // Dangling symlink as rootDir (both direct and nested under dangling prefix) must be rejected. + danglingRoot := filepath.Join(base, "dangling_root") + mustSymlink(t, filepath.Join(base, "does_not_exist"), danglingRoot) + for _, tc := range []struct { + name string + root string + create bool + }{ + {name: "direct_no_create", root: danglingRoot, create: false}, + {name: "direct_create", root: danglingRoot, create: true}, + {name: "nested_no_create", root: filepath.Join(danglingRoot, "nested"), create: false}, + } { + t.Run(tc.name, func(t *testing.T) { + if _, _, err := ResolveSandboxDirs(tc.root, "", tc.create); err == nil || !strings.Contains(err.Error(), "dangling symlink") { + t.Errorf("ResolveSandboxDirs(%q, create=%v) = %v, want dangling symlink error", tc.root, tc.create, err) + } + }) + } +} + +func TestSandboxFSConfinementModesAndOperations(t *testing.T) { + modes := []struct { + name string + mode ConfinementMode + }{ + {name: "os_root", mode: ConfinementOSRoot}, + {name: "lexical", mode: ConfinementLexical}, + } + + for _, tc := range modes { + t.Run(tc.name, func(t *testing.T) { + rootDir := t.TempDir() + workDir := filepath.Join(rootDir, "sub") + if err := os.MkdirAll(workDir, 0o755); err != nil { + t.Fatalf("MkdirAll failed: %v", err) + } + outsideDir := t.TempDir() + outsideSecret := filepath.Join(outsideDir, "secret.txt") + if err := os.WriteFile(outsideSecret, []byte("top-secret"), 0o600); err != nil { + t.Fatalf("WriteFile(outsideSecret) failed: %v", err) + } + + fs := NewSandboxFSWithConfig(SandboxFSConfig{ + RootDir: rootDir, + WorkDir: workDir, + Mode: tc.mode, + }) + if fs.Mode() != tc.mode { + t.Errorf("fs.Mode() = %v, want %v", fs.Mode(), tc.mode) + } + if fs.RootDir() == "" || fs.WorkDir() == "" { + t.Errorf("RootDir=%q WorkDir=%q, want non-empty", fs.RootDir(), fs.WorkDir()) + } + + // Valid parent traversal within rootDir (`../root_note.txt` from `rootDir/sub`). + if _, err := fs.WriteRegularFile("../root_note.txt", "in-root\n", false, false); err != nil { + t.Fatalf("WriteRegularFile(../root_note.txt) failed: %v", err) + } + if _, err := fs.WriteRegularFile("../root_note.txt", "dup", false, false); err == nil { + t.Error("WriteRegularFile(overwrite=false existing) = nil, want PathErrAlreadyExists") + } else { + var pe *PathError + if !errors.As(err, &pe) || pe.Kind != PathErrAlreadyExists { + t.Errorf("existing file error = %v, want PathErrAlreadyExists", err) + } + } + if _, err := fs.WriteRegularFile("../root_note.txt", "more\n", false, true); err != nil { + t.Fatalf("WriteRegularFile(appendMode=true) failed: %v", err) + } + data, err := fs.ReadRegularFile("/root_note.txt") + if err != nil || string(data) != "in-root\nmore\n" { + t.Errorf("ReadRegularFile(/root_note.txt) = (%q, %v), want \"in-root\\nmore\\n\"", data, err) + } + + // Relative traversal outside rootDir must fail with PathErrPermissionDenied. + if _, err := fs.ReadRegularFile("../../escape.txt"); err == nil { + t.Error("ReadRegularFile(../../escape.txt) succeeded, want confinement error") + } + if _, err := fs.WriteRegularFile("../../escape.txt", "x", true, false); err == nil { + t.Error("WriteRegularFile(../../escape.txt) succeeded, want confinement error") + } + + // Absolute paths outside rootDir (including sibling-prefix paths like rootDir+"2/...") are mapped into the virtual root inside rootDir, never touching host files. + if _, err := fs.ReadRegularFile(outsideSecret); err == nil { + t.Errorf("ReadRegularFile(%q) succeeded, want not-found inside virtual root", outsideSecret) + } else { + var pe *PathError + if !errors.As(err, &pe) || pe.Kind != PathErrNotFound { + t.Errorf("ReadRegularFile(missing) error = %v, want PathErrNotFound", err) + } + } + if _, err := fs.WriteRegularFile("/virtual_etc/config", "virtual", true, false); err != nil { + t.Fatalf("WriteRegularFile(/virtual_etc/config) failed: %v", err) + } + if _, statErr := os.Stat(filepath.Join(fs.RootDir(), "virtual_etc", "config")); statErr != nil { + t.Errorf("virtual absolute write did not land inside rootDir: %v", statErr) + } + siblingTarget := fs.RootDir() + "2" + string(filepath.Separator) + "sib.txt" + if _, err := fs.WriteRegularFile(siblingTarget, "virtual-sib", true, false); err != nil { + t.Fatalf("WriteRegularFile(siblingTarget) failed: %v", err) + } + if _, statErr := os.Stat(fs.RootDir() + "2"); !os.IsNotExist(statErr) { + t.Errorf("sibling-prefix path escaped rootDir onto host disk: statErr=%v", statErr) + } + + // NUL byte and empty-path guards. + if _, err := fs.ResolvePath("bad\x00name"); err == nil { + t.Error("ResolvePath(NUL) = nil, want PathErrPermissionDenied") + } + if _, err := fs.OpenRead(""); err == nil { + t.Error("OpenRead(\"\") = nil, want PathErrInvalidArgument") + } + if _, err := fs.OpenWrite("", 0o644, true); err == nil { + t.Error("OpenWrite(\"\") = nil, want PathErrInvalidArgument") + } + if _, err := fs.Stat(""); err == nil { + t.Error("Stat(\"\") = nil, want PathErrInvalidArgument") + } + if err := fs.Mkdir("", 0o755, true); err == nil { + t.Error("Mkdir(\"\") = nil, want PathErrInvalidArgument") + } + if err := fs.Rename("", "dst"); err == nil { + t.Error("Rename(\"\", \"dst\") = nil, want PathErrInvalidArgument") + } + if err := fs.Rename("src", ""); err == nil { + t.Error("Rename(\"src\", \"\") = nil, want PathErrInvalidArgument") + } + + // In ConfinementOSRoot mode, symlinks escaping rootDir and dangling symlinks are rejected across all operations. + if tc.mode == ConfinementOSRoot { + escapeLink := filepath.Join(workDir, "escape_link") + mustSymlink(t, outsideDir, escapeLink) + danglingLink := filepath.Join(workDir, "dangling_link") + mustSymlink(t, filepath.Join(outsideDir, "missing_target"), danglingLink) + + realTargetDir := filepath.Join(rootDir, "real_target") + if err := os.MkdirAll(realTargetDir, 0o755); err != nil { + t.Fatalf("MkdirAll(realTargetDir) failed: %v", err) + } + mustSymlink(t, filepath.Join(rootDir, "missing_leaf"), filepath.Join(realTargetDir, "dangling_leaf")) + + loopA := filepath.Join(workDir, "loop_a") + loopB := filepath.Join(workDir, "loop_b") + mustSymlink(t, loopB, loopA) + mustSymlink(t, loopA, loopB) + + for _, esc := range []struct { + name string + path string + }{ + {name: "symlink_existing_target", path: "escape_link/secret.txt"}, + {name: "symlink_nonexistent_leaf", path: "escape_link/new_subdir/pwn.txt"}, + {name: "dangling_symlink", path: "dangling_link"}, + {name: "dangling_suffix_under_real_dir", path: filepath.Join("..", "real_target", "dangling_leaf", "sub")}, + {name: "symlink_loop", path: "loop_a/file.txt"}, + } { + if got, err := fs.ResolvePath(esc.path); err == nil { + t.Errorf("ResolvePath(%s=%q) = %q, want confinement/symlink error", esc.name, esc.path, got) + } + } + if _, err := fs.ReadRegularFile("escape_link/secret.txt"); err == nil { + t.Error("ReadRegularFile(escape symlink) succeeded, want error") + } + if _, err := fs.WriteRegularFile("escape_link/new.txt", "bad", true, false); err == nil { + t.Error("WriteRegularFile(escape symlink) succeeded, want error") + } + if _, err := fs.WriteRegularFile("dangling_link", "bad", true, false); err == nil { + t.Error("WriteRegularFile(dangling symlink) succeeded, want error") + } + if _, err := fs.OpenWrite("/real_target", 0o644, true); err == nil || !strings.Contains(err.Error(), "cannot write non-regular file") { + t.Errorf("OpenWrite(directory in ConfinementOSRoot) = %v, want non-regular file error", err) + } + if _, err := fs.OpenRead("/real_target"); err == nil || !strings.Contains(err.Error(), "cannot read non-regular file") { + t.Errorf("OpenRead(directory in ConfinementOSRoot) = %v, want non-regular file error", err) + } + if _, err := fs.Stat("escape_link/secret.txt"); err == nil { + t.Error("Stat(escape symlink) succeeded in ConfinementOSRoot, want error") + } + if err := fs.Mkdir("escape_link/escaped_dir", 0o755, true); err == nil { + t.Error("Mkdir(escape symlink) succeeded in ConfinementOSRoot, want error") + } + if err := fs.Remove("escape_link/secret.txt", false); err == nil { + t.Error("Remove(escape symlink) succeeded in ConfinementOSRoot, want error") + } + if _, statErr := os.Stat(outsideSecret); statErr != nil { + t.Errorf("outsideSecret was unexpectedly removed: %v", statErr) + } + if err := fs.Rename("escape_link/secret.txt", "stolen.txt"); err == nil { + t.Error("Rename(from escape symlink) succeeded in ConfinementOSRoot, want error") + } + if err := fs.Rename("../root_note.txt", "escape_link/stolen.txt"); err == nil { + t.Error("Rename(to escape symlink) succeeded in ConfinementOSRoot, want error") + } + if _, err := fs.ReadDir("escape_link"); err == nil { + t.Error("ReadDir(escape symlink) succeeded in ConfinementOSRoot, want error") + } + + // Intra-root symlinks in ConfinementOSRoot: Stat reports FileKindSymlink + SymlinkTarget, + // Rename moves the symlink itself, and Remove(recursive=true) unlinks a directory symlink without deleting its target directory. + intraTargetFile := filepath.Join(workDir, "intra_target.txt") + if err := os.WriteFile(intraTargetFile, []byte("keep-me"), 0o644); err != nil { + t.Fatalf("WriteFile(intraTargetFile) failed: %v", err) + } + intraFileLink := filepath.Join(workDir, "intra_file_link.txt") + mustSymlink(t, "intra_target.txt", intraFileLink) + symMeta, err := fs.Stat("intra_file_link.txt") + if err != nil || symMeta.Kind != FileKindSymlink || symMeta.SymlinkTarget != "intra_target.txt" { + t.Errorf("Stat(intra_file_link.txt) = (%+v, %v), want FileKindSymlink target=intra_target.txt", symMeta, err) + } + if err := fs.Rename("intra_file_link.txt", "renamed_file_link.txt"); err != nil { + t.Fatalf("Rename(intra_file_link.txt) failed: %v", err) + } + if _, lerr := os.Lstat(filepath.Join(workDir, "intra_file_link.txt")); !os.IsNotExist(lerr) { + t.Errorf("expected intra_file_link.txt moved, lstat err=%v", lerr) + } + if err := fs.Remove("renamed_file_link.txt", false); err != nil { + t.Fatalf("Remove(renamed_file_link.txt) failed: %v", err) + } + if gotKeep, rerr := os.ReadFile(intraTargetFile); rerr != nil || string(gotKeep) != "keep-me" { + t.Errorf("intraTargetFile after link rename/remove = (%q, %v)", gotKeep, rerr) + } + + keptDir := filepath.Join(workDir, "kept_dir") + if err := os.MkdirAll(keptDir, 0o755); err != nil { + t.Fatalf("MkdirAll(keptDir) failed: %v", err) + } + keptFile := filepath.Join(keptDir, "inside.txt") + if err := os.WriteFile(keptFile, []byte("survives"), 0o644); err != nil { + t.Fatalf("WriteFile(keptFile) failed: %v", err) + } + mustSymlink(t, "kept_dir", filepath.Join(workDir, "dir_link")) + if err := fs.Remove("dir_link", true); err != nil { + t.Fatalf("Remove(dir_link, recursive=true) failed: %v", err) + } + if _, statKept := os.Stat(keptFile); statKept != nil { + t.Errorf("Remove(dir_link, recursive=true) deleted symlink target contents: %v", statKept) + } + _ = os.RemoveAll(keptDir) + _ = os.Remove(intraTargetFile) + } else { + // ConfinementLexical intentionally permits symlinks inside already-isolated container sandboxes + // (such as /bin -> /usr/bin) while still rejecting lexical '..' traversal outside rootDir. + containerTarget := filepath.Join(outsideDir, "secret.txt") + lexicalLink := filepath.Join(workDir, "container_link.txt") + mustSymlink(t, containerTarget, lexicalLink) + if got, err := fs.ReadRegularFile("container_link.txt"); err != nil || string(got) != "top-secret" { + t.Errorf("ConfinementLexical ReadRegularFile(symlink) = (%q, %v), want (\"top-secret\", nil)", got, err) + } + if _, err := fs.WriteRegularFile("container_link.txt", "top-secret", true, false); err != nil { + t.Errorf("ConfinementLexical WriteRegularFile(symlink) = %v, want nil", err) + } + extDirLink := filepath.Join(workDir, "ext_dir_link") + mustSymlink(t, outsideDir, extDirLink) + if err := fs.Mkdir("ext_dir_link/lex_sub", 0o755, true); err != nil { + t.Errorf("ConfinementLexical Mkdir(symlink) = %v, want nil", err) + } + extWrite, err := fs.OpenWrite("ext_dir_link/lex_sub/out.txt", 0o644, true) + if err != nil { + t.Fatalf("ConfinementLexical OpenWrite(symlink) = %v, want nil", err) + } + _, _ = extWrite.Write([]byte("lex-data")) + _ = extWrite.Close() + if meta, err := fs.Stat("ext_dir_link/lex_sub/out.txt"); err != nil || meta.Size != 8 { + t.Errorf("ConfinementLexical Stat(symlink) = (%+v, %v), want Size=8", meta, err) + } + if entries, err := fs.ReadDir("ext_dir_link/lex_sub"); err != nil || len(entries) != 1 { + t.Errorf("ConfinementLexical ReadDir(symlink) = (%+v, %v), want 1 entry", entries, err) + } + if err := fs.Rename("ext_dir_link/lex_sub/out.txt", "ext_dir_link/lex_sub/renamed.txt"); err != nil { + t.Errorf("ConfinementLexical Rename(symlink) = %v, want nil", err) + } + if err := fs.Remove("ext_dir_link/lex_sub/renamed.txt", false); err != nil { + t.Errorf("ConfinementLexical Remove(symlink) = %v, want nil", err) + } + var dirWritePE *PathError + if _, err := fs.OpenWrite("ext_dir_link/lex_sub", 0o644, true); !errors.As(err, &dirWritePE) || dirWritePE.Kind != PathErrInternal { + t.Errorf("ConfinementLexical OpenWrite(directory) = %v, want PathErrInternal", err) + } + var dirReadPE *PathError + if _, err := fs.OpenRead("ext_dir_link/lex_sub"); !errors.As(err, &dirReadPE) || dirReadPE.Kind != PathErrInvalidArgument { + t.Errorf("ConfinementLexical OpenRead(directory) = %v, want PathErrInvalidArgument", err) + } + } + + // OpenWrite & OpenRead round-trip + explicit permission updates on overwrite. + wFile, err := fs.OpenWrite("streamed.txt", 0o644, false) + if err != nil { + t.Fatalf("OpenWrite(streamed.txt) failed: %v", err) + } + _, _ = wFile.Write([]byte("stream-payload")) + _ = wFile.Close() + if _, err := fs.OpenWrite("streamed.txt", 0o644, false); err == nil { + t.Error("OpenWrite(overwrite=false existing) = nil, want error") + } + rFile, err := fs.OpenRead("streamed.txt") + if err != nil { + t.Fatalf("OpenRead(streamed.txt) failed: %v", err) + } + streamRead, _ := io.ReadAll(rFile) + _ = rFile.Close() + if string(streamRead) != "stream-payload" { + t.Errorf("OpenRead(streamed.txt) = %q, want \"stream-payload\"", streamRead) + } + if runtime.GOOS != "windows" { + cFile, err := fs.OpenWrite("chmod_target.sh", 0o600, false) + if err != nil { + t.Fatalf("OpenWrite(chmod_target.sh, 0600) failed: %v", err) + } + _ = cFile.Close() + if m600, err := fs.Stat("chmod_target.sh"); err != nil || m600.Permissions != 0o600 { + t.Errorf("Stat(chmod_target.sh initial) = (%#o, %v), want 0o600", m600.Permissions, err) + } + cFile755, err := fs.OpenWrite("chmod_target.sh", 0o755, true) + if err != nil { + t.Fatalf("OpenWrite(chmod_target.sh, 0755, overwrite=true) failed: %v", err) + } + _ = cFile755.Close() + if m755, err := fs.Stat("chmod_target.sh"); err != nil || m755.Permissions != 0o755 { + t.Errorf("Stat(chmod_target.sh after 0755 overwrite) = (%#o, %v), want 0o755", m755.Permissions, err) + } + cFileZero, err := fs.OpenWrite("chmod_target.sh", 0, true) + if err != nil { + t.Fatalf("OpenWrite(chmod_target.sh, 0, overwrite=true) failed: %v", err) + } + _ = cFileZero.Close() + if _, err := fs.WriteRegularFile("chmod_target.sh", "#!/bin/sh\necho ok\n", true, false); err != nil { + t.Fatalf("WriteRegularFile(chmod_target.sh, overwrite=true) failed: %v", err) + } + if mPreserved, err := fs.Stat("chmod_target.sh"); err != nil || mPreserved.Permissions != 0o755 { + t.Errorf("Stat(chmod_target.sh after perm=0/WriteRegularFile overwrite) = (%#o, %v), want preserved 0o755", mPreserved.Permissions, err) + } + } + + // CRLF ReadFileRange, Grep, FindFiles, Stat, Mkdir, Rename, ListEntries, Remove. + crlfText := "alpha\r\nbeta = 42\r\ngamma\r\n" + if _, err := fs.WriteRegularFile("cfg/app.ini", crlfText, true, false); err != nil { + t.Fatalf("WriteRegularFile(cfg/app.ini) failed: %v", err) + } + ranged, err := fs.ReadFileRange("cfg/app.ini", 2, 3) + if err != nil || ranged != "beta = 42\ngamma" { + t.Errorf("ReadFileRange() = (%q, %v), want (\"beta = 42\\ngamma\", nil)", ranged, err) + } + if emptyRange, err := fs.ReadFileRange("cfg/app.ini", 3, 2); err != nil || emptyRange != "" { + t.Errorf("ReadFileRange(start > end) = (%q, %v), want (\"\", nil)", emptyRange, err) + } + + matches, truncated, err := fs.Grep(".", "BETA = 42$", true, true) + if err != nil || truncated || len(matches) != 1 { + t.Errorf("Grep() = (%v, %v, %v), want 1 non-truncated match", matches, truncated, err) + } + if litMatches, _, err := fs.Grep(".", "beta = 42$", false, false); err != nil || len(litMatches) != 0 { + t.Errorf("Grep(isRegex=false, \"beta = 42$\") = (%v, %v), want 0 literal matches", litMatches, err) + } + if _, _, err := fs.Grep(".", "[", false, true); err == nil { + t.Error("Grep(invalid regex) = nil, want error") + } + if _, _, err := fs.Grep(".", "", false, false); err == nil { + t.Error("Grep(empty query) = nil, want error") + } + + found, _, err := fs.FindFiles(".", "*", 3, "INI") + if err != nil || len(found) != 1 || found[0] != "cfg/app.ini" { + t.Errorf("FindFiles(INI) = (%v, %v), want [cfg/app.ini]", found, err) + } + foundEmptyExt, _, err := fs.FindFiles(".", "app.ini", 3, "", " . ") + if err != nil || len(foundEmptyExt) != 1 { + t.Errorf("FindFiles(empty extensions) = (%v, %v), want 1 match", foundEmptyExt, err) + } + + meta, err := fs.Stat("cfg/app.ini") + if err != nil || meta.IsDir || meta.Kind != FileKindRegular { + t.Errorf("Stat(cfg/app.ini) = (%+v, %v)", meta, err) + } + if err := fs.Mkdir("missing_parent_nonrec/child", 0o755, false); err == nil { + t.Error("Mkdir(missing parent, recursive=false) = nil, want error") + } + if err := fs.Mkdir("default_perm_dir", 0, false); err != nil { + t.Fatalf("Mkdir(perm=0) failed: %v", err) + } + if _, err := fs.WriteRegularFile("default_perm_dir/probe.txt", "ok", true, false); err != nil { + t.Fatalf("WriteRegularFile(default_perm_dir/probe.txt) failed: %v", err) + } + if err := fs.Mkdir("nested/a/b", 0o755, true); err != nil { + t.Fatalf("Mkdir(recursive) failed: %v", err) + } + if err := fs.Rename("no_such_src.txt", "ghost_parent/dst.txt"); err == nil { + t.Error("Rename(non-existent source) = nil, want error") + } + if _, statErr := os.Stat(filepath.Join(workDir, "ghost_parent")); !os.IsNotExist(statErr) { + t.Errorf("Rename(non-existent source) created ghost_parent: statErr=%v", statErr) + } + if err := fs.Rename("cfg/app.ini", "moved/app.ini"); err != nil { + t.Fatalf("Rename() failed: %v", err) + } + if _, err := fs.WriteRegularFile("moved/aaa.ini", "first", true, false); err != nil { + t.Fatalf("WriteRegularFile(moved/aaa.ini) failed: %v", err) + } + entries, err := fs.ListEntries("moved") + if err != nil || len(entries) != 2 || entries[0].(map[string]any)["name"] != "aaa.ini" || entries[1].(map[string]any)["name"] != "app.ini" { + t.Errorf("ListEntries(moved) = (%v, %v), want sorted [aaa.ini, app.ini]", entries, err) + } + if _, ok := entries[0].(map[string]any)["size"].(int64); !ok { + t.Errorf("ListEntries size type = %T, want int64", entries[0].(map[string]any)["size"]) + } + dirMetas, err := fs.ReadDir("moved") + if err != nil || len(dirMetas) != 2 || dirMetas[0].Name != "aaa.ini" || dirMetas[1].Name != "app.ini" { + t.Errorf("ReadDir(moved) = (%+v, %v), want sorted [aaa.ini, app.ini]", dirMetas, err) + } + if err := fs.Remove("moved", false); err == nil { + t.Error("Remove(non-empty dir, recursive=false) = nil, want error") + } + if err := fs.Remove("moved", true); err != nil { + t.Errorf("Remove(moved, recursive=true) failed: %v", err) + } + if err := fs.Remove("no_such_entry", true); err == nil { + t.Error("Remove(non-existent, recursive=true) = nil, want error") + } + if err := fs.Remove("", true); err == nil { + t.Error("Remove(\"\", recursive=true) = nil, want PathErrInvalidArgument") + } + if err := fs.Remove(".", true); err == nil { + t.Error("Remove(\".\") = nil, want workDir removal protection error") + } + if err := fs.Rename(".", "moved_work"); err == nil { + t.Error("Rename(\".\", ...) = nil, want workDir rename protection error") + } + if err := fs.Remove("..", true); err == nil { + t.Error("Remove(..) = nil, want root removal protection error") + } + if runtime.GOOS != "windows" { + fifoPath := filepath.Join(workDir, "test_pipe.fifo") + if mkErr := exec.Command("mkfifo", fifoPath).Run(); mkErr == nil { + fifoDone := make(chan error, 1) + go func() { + _, rerr := fs.OpenRead("test_pipe.fifo") + fifoDone <- rerr + }() + select { + case fifoErr := <-fifoDone: + if fifoErr == nil { + t.Error("OpenRead(FIFO) = nil, want non-regular file error without blocking") + } + case <-time.After(time.Second): + t.Fatal("OpenRead(FIFO) blocked on named pipe without writer") + } + grepDone := make(chan error, 1) + go func() { + _, _, gerr := fs.Grep(".", "anything", false, false) + grepDone <- gerr + }() + select { + case gerr := <-grepDone: + if gerr != nil { + t.Errorf("Grep(dir containing FIFO) = %v, want nil", gerr) + } + case <-time.After(time.Second): + t.Fatal("Grep() blocked on named pipe") + } + _ = os.Remove(fifoPath) + } + } + }) + } + + // Verify ConfinementOSRoot returns structured errors when rootDir does not exist on disk. + uncreatedFS := NewSandboxFS(filepath.Join(t.TempDir(), "uncreated_root"), "") + if _, err := uncreatedFS.OpenRead("x.txt"); err == nil { + t.Error("uncreatedFS.OpenRead() = nil, want error") + } + if err := uncreatedFS.Rename("a.txt", "b.txt"); err == nil { + t.Error("uncreatedFS.Rename() = nil, want error") + } + + // Verify bounded vs UnlimitedRead behavior above MaxReadFileBytes. + bigDir := t.TempDir() + bigPayload := strings.Repeat("x", MaxReadFileBytes+128) + if err := os.WriteFile(filepath.Join(bigDir, "big.bin"), []byte(bigPayload), 0o644); err != nil { + t.Fatalf("WriteFile(big.bin) failed: %v", err) + } + boundedFS := NewSandboxFS(bigDir, bigDir) + if _, err := boundedFS.ReadRegularFile("big.bin"); err == nil { + t.Error("boundedFS.ReadRegularFile(>MaxReadFileBytes) = nil, want size limit error") + } + limitSrc := strings.NewReader("0123456789ABCDEF") + if _, exceeded, err := readLimitedAll(limitSrc, 10); err != nil || !exceeded || limitSrc.Len() != 5 { + t.Errorf("readLimitedAll(10) = (exceeded=%v, err=%v, remaining=%d), want (true, nil, 5)", exceeded, err, limitSrc.Len()) + } + + halfLine := strings.Repeat("a", MaxReadFileBytes/2) + twoHalfLines := halfLine + "\n" + halfLine + "\n" + if err := os.WriteFile(filepath.Join(bigDir, "two_lines.txt"), []byte(twoHalfLines), 0o644); err != nil { + t.Fatalf("WriteFile(two_lines.txt) failed: %v", err) + } + if _, err := boundedFS.ReadFileRange("two_lines.txt", 1, 2); err == nil { + t.Error("boundedFS.ReadFileRange(two half-max lines + newlines > MaxReadFileBytes) = nil, want size limit error") + } + exactHalfLine := strings.Repeat("b", MaxReadFileBytes/2-1) + exactTwoLines := exactHalfLine + "\n" + exactHalfLine + "\n" + if err := os.WriteFile(filepath.Join(bigDir, "exact_two_lines.txt"), []byte(exactTwoLines), 0o644); err != nil { + t.Fatalf("WriteFile(exact_two_lines.txt) failed: %v", err) + } + if gotExact, err := boundedFS.ReadFileRange("exact_two_lines.txt", 1, 2); err != nil || len(gotExact) != MaxReadFileBytes-1 { + t.Errorf("boundedFS.ReadFileRange(exact MaxReadFileBytes) len=%d err=%v, want len=%d", len(gotExact), err, MaxReadFileBytes-1) + } + unlimitedFS := NewSandboxFSWithConfig(SandboxFSConfig{ + RootDir: bigDir, + WorkDir: bigDir, + Mode: ConfinementLexical, + UnlimitedRead: true, + }) + gotBig, err := unlimitedFS.ReadRegularFile("big.bin") + if err != nil || len(gotBig) != len(bigPayload) { + t.Errorf("unlimitedFS.ReadRegularFile(big.bin) len=%d err=%v, want len=%d", len(gotBig), err, len(bigPayload)) + } + + // Verify ClassifyOSError mappings (including true untyped nil error) and PathError.Error() formatting. + if err := ClassifyOSError(nil, "p"); err != nil { + t.Errorf("ClassifyOSError(nil) = %#v, want nil error interface", err) + } + var pe *PathError + if err := ClassifyOSError(os.ErrNotExist, "missing"); !errors.As(err, &pe) || pe.Kind != PathErrNotFound || !errors.Is(err, os.ErrNotExist) || !strings.Contains(err.Error(), "missing") { + t.Errorf("ClassifyOSError(ErrNotExist) = %+v", err) + } + if err := ClassifyOSError(os.ErrExist, "dup"); !errors.As(err, &pe) || pe.Kind != PathErrAlreadyExists { + t.Errorf("ClassifyOSError(ErrExist) = %+v", err) + } + if err := ClassifyOSError(os.ErrPermission, "denied"); !errors.As(err, &pe) || pe.Kind != PathErrPermissionDenied { + t.Errorf("ClassifyOSError(ErrPermission) = %+v", err) + } + if got := (&PathError{Message: "msg-with-err", Err: errors.New("inner")}).Error(); got != "msg-with-err: inner" { + t.Errorf("PathError(Message+Err).Error() = %q, want \"msg-with-err: inner\"", got) + } + if got := (&PathError{Message: "msg-only"}).Error(); got != "msg-only" { + t.Errorf("PathError(Message only).Error() = %q, want \"msg-only\"", got) + } + if got := (&PathError{Err: errors.New("err-only")}).Error(); got != "err-only" { + t.Errorf("PathError(Err only).Error() = %q, want \"err-only\"", got) + } + if got := (&PathError{Path: "p.txt"}).Error(); got != `filesystem error on "p.txt"` { + t.Errorf("PathError(Path only).Error() = %q, want %q", got, `filesystem error on "p.txt"`) + } +} + +func TestSandboxShellEnvBuffersAndSubprocessKill(t *testing.T) { + var defaultMaxBuf CappedBuffer + if n, err := defaultMaxBuf.Write([]byte("abc")); err != nil || n != 3 || defaultMaxBuf.String() != "abc" { + t.Errorf("CappedBuffer(default MaxBytes=0) = (%d, %v, %q), want (3, nil, \"abc\")", n, err, defaultMaxBuf.String()) + } + + var exactBuf CappedBuffer + exactBuf.MaxBytes = 5 + _, _ = exactBuf.Write([]byte("12345")) + if got := exactBuf.String(); got != "12345" { + t.Errorf("CappedBuffer(exact MaxBytes) = %q, want \"12345\"", got) + } + _, _ = exactBuf.Write([]byte("6")) + if got := exactBuf.String(); got != "12345\n...[output truncated at 4 MiB]..." { + t.Errorf("CappedBuffer(overflow after exact MaxBytes) = %q, want \"12345\\n...[output truncated at 4 MiB]...\"", got) + } + + var singleOverflowBuf CappedBuffer + singleOverflowBuf.MaxBytes = 5 + _, _ = singleOverflowBuf.Write([]byte("123456")) + if got := singleOverflowBuf.String(); got != "12345\n...[output truncated at 4 MiB]..." { + t.Errorf("CappedBuffer(single write > MaxBytes) = %q, want \"12345\\n...[output truncated at 4 MiB]...\"", got) + } + + allowEnv := BuildAllowlistEnv([]string{ + "PATH=/usr/bin", + "LC_ALL=C.UTF-8", + "COMPUTERNAME=WIN-NODE", + "USERNAME=runner", + "OS=Windows_NT", + "SYSTEMDRIVE=C:", + "PSMODULEPATH=C:\\Modules", + "GEMINI_API_KEY=secret", + "UNLISTED=drop", + }) + if len(allowEnv) != 7 { + t.Errorf("BuildAllowlistEnv() = %v, want 7 entries", allowEnv) + } + + if emptyScrub := FilterScrubbedEnv(nil); emptyScrub == nil || len(emptyScrub) != 0 { + t.Errorf("FilterScrubbedEnv(nil) = %#v, want non-nil empty slice", emptyScrub) + } + if allFiltered := FilterScrubbedEnv([]string{"GEMINI_API_KEY=secret", "GEMINI_LEASE_TOKEN=tok"}); allFiltered == nil || len(allFiltered) != 0 { + t.Errorf("FilterScrubbedEnv(all credentials) = %#v, want non-nil empty slice", allFiltered) + } + + scrubbed := FilterScrubbedEnv([]string{ + "CUSTOM_SAFE=keep", + "GEMINI_API_KEY=drop1", + "GEMINI_ENVIRONMENT_KEY=drop2", + "GOOGLE_API_KEY=drop3", + "GEMINI_LEASE_TOKEN=drop4", + }, "EXTRA_VAR=ok") + joined := strings.Join(scrubbed, "\n") + if !strings.Contains(joined, "CUSTOM_SAFE=keep") || !strings.Contains(joined, "EXTRA_VAR=ok") || strings.Contains(joined, "drop") { + t.Errorf("FilterScrubbedEnv() = %v", scrubbed) + } + + origPath := os.Getenv("PATH") + t.Setenv("PATH", t.TempDir()) + if bin, args := ResolveShellCommand("windows", "dir"); bin != "cmd.exe" || len(args) != 2 || args[0] != "/C" { + t.Errorf("ResolveShellCommand(windows, empty PATH) = (%q, %v), want (\"cmd.exe\", [\"/C\", \"dir\"])", bin, args) + } + if !isWindowsWSLStubBash("windows", `C:\Windows\System32\bash.exe`) || !isWindowsWSLStubBash("windows", `c:/windows/sysnative/bash.exe`) || isWindowsWSLStubBash("windows", `C:\Program Files\Git\bin\bash.exe`) || isWindowsWSLStubBash("linux", "/usr/bin/bash") { + t.Error("isWindowsWSLStubBash classification mismatch") + } + if bin, args := resolveExecShell("windows", true, "dir"); bin != "cmd.exe" || len(args) != 2 || args[0] != "/C" { + t.Errorf("resolveExecShell(windows, useBinSh=true, empty PATH) = (%q, %v), want (\"cmd.exe\", [\"/C\", \"dir\"])", bin, args) + } + if bin, _ := ResolveShellCommand("linux", "ls"); bin != "sh" { + t.Errorf("ResolveShellCommand(linux, empty PATH) = %q, want sh", bin) + } + t.Setenv("PATH", origPath) + if len(ShutdownSignals()) == 0 { + t.Error("ShutdownSignals() returned empty slice") + } + + workDir := filepath.Join(t.TempDir(), "auto-work") + sh := NewSandboxShellWithConfig(SandboxShellConfig{ + WorkDir: workDir, + EnvPolicy: EnvPolicyScrubCredentials, + ExtraEnv: []string{"INJECTED_FLAG=present"}, + AutoCreateWorkDir: true, + UseBinSh: true, + }) + if sh.WorkDir() != workDir { + t.Errorf("sh.WorkDir() = %q, want %q", sh.WorkDir(), workDir) + } + + t.Setenv("GEMINI_API_KEY", "should-not-leak") + stdout, stderr, exitCode, execErr := sh.ExecuteBash(t.Context(), `printf "%s|%s|%s" "$INJECTED_FLAG" "$GEMINI_API_KEY" "$(pwd)"`, "", 5*time.Second) + if execErr != nil || stdout != "present||"+workDir || stderr != "" || exitCode != 0 { + t.Errorf("ExecuteBash(env+cwd check) = (%q, %q, %d, %v), want stdout=%q exitCode=0", stdout, stderr, exitCode, execErr, "present||"+workDir) + } + scrubNoExtraSh := NewSandboxShellWithConfig(SandboxShellConfig{ + WorkDir: workDir, + EnvPolicy: EnvPolicyScrubCredentials, + UseBinSh: true, + }) + noExtraOut, _, noExtraCode, noExtraErr := scrubNoExtraSh.ExecuteBash(t.Context(), `printf "%s" "$GEMINI_API_KEY"`, "", 5*time.Second) + if noExtraErr != nil || noExtraCode != 0 || noExtraOut != "" { + t.Errorf("scrubNoExtraSh.ExecuteBash() = (%q, %d, %v), want (\"\", 0, nil)", noExtraOut, noExtraCode, noExtraErr) + } + + // Verify AutoCreateWorkDir=false fails when workDir does not exist, EnvPolicyAllowlist drops unlisted env vars, + // and UseBinSh=false uses ResolveShellCommand (preferring bash on PATH). + noAutoSh := NewSandboxShell(filepath.Join(t.TempDir(), "does-not-exist")) + if _, _, _, err := noAutoSh.ExecuteBash(t.Context(), "echo hi", "", 5*time.Second); err == nil { + t.Error("ExecuteBash(AutoCreateWorkDir=false, missing workDir) = nil, want error") + } + allowSh := NewSandboxShellWithConfig(SandboxShellConfig{ + WorkDir: workDir, + EnvPolicy: EnvPolicyAllowlist, + ExtraEnv: []string{"EXTRA_ALLOW=present"}, + }) + t.Setenv("CUSTOM_UNLISTED_ENV", "secret-val") + if runtime.GOOS != "windows" { + fakeBinDir := t.TempDir() + fakeBash := filepath.Join(fakeBinDir, "bash") + if err := os.WriteFile(fakeBash, []byte("#!/bin/sh\nprintf \"via-bash:%s|%s\" \"$CUSTOM_UNLISTED_ENV\" \"$EXTRA_ALLOW\"\n"), 0o755); err != nil { + t.Fatalf("WriteFile(fakeBash) failed: %v", err) + } + t.Setenv("PATH", fakeBinDir+string(os.PathListSeparator)+origPath) + if bin, _ := resolveExecShell("linux", true, "ls"); bin != "/bin/sh" { + t.Errorf("resolveExecShell(linux, useBinSh=true) = %q, want /bin/sh", bin) + } + if bin, _ := resolveExecShell("linux", false, "ls"); bin != fakeBash { + t.Errorf("resolveExecShell(linux, useBinSh=false) = %q, want %q", bin, fakeBash) + } + allowOut, _, allowCode, allowErr := allowSh.ExecuteBash(t.Context(), "ignored", "", 5*time.Second) + if allowErr != nil || allowCode != 0 || allowOut != "via-bash:|present" { + t.Errorf("allowSh.ExecuteBash() = (%q, %d, %v), want (\"via-bash:|present\", 0, nil)", allowOut, allowCode, allowErr) + } + binShOut, _, binShCode, binShErr := sh.ExecuteBash(t.Context(), `printf "direct-bin-sh"`, "", 5*time.Second) + if binShErr != nil || binShCode != 0 || binShOut != "direct-bin-sh" { + t.Errorf("sh.ExecuteBash(UseBinSh=true with fakeBash on PATH) = (%q, %d, %v), want (\"direct-bin-sh\", 0, nil)", binShOut, binShCode, binShErr) + } + t.Setenv("PATH", origPath) + } + + // Verify already-canceled or already-expired parent ctx returns ctx.Err() without formatting as a command timeout. + canceledCtx, cancelNow := context.WithCancel(t.Context()) + cancelNow() + _, _, cancelCode, cancelErr := sh.ExecuteBash(canceledCtx, "echo hi", "", 5*time.Second) + if !errors.Is(cancelErr, context.Canceled) || cancelCode != -1 || strings.Contains(cancelErr.Error(), "timed out") { + t.Errorf("ExecuteBash(canceledCtx) = (code=%d, err=%v), want context.Canceled without 'timed out'", cancelCode, cancelErr) + } + expiredCtx, cancelExpired := context.WithDeadline(t.Context(), time.Now().Add(-time.Second)) + defer cancelExpired() + _, _, expCode, expErr := sh.ExecuteBash(expiredCtx, "echo hi", "", 5*time.Second) + if !errors.Is(expErr, context.DeadlineExceeded) || expCode != -1 || strings.Contains(expErr.Error(), "timed out") { + t.Errorf("ExecuteBash(expiredCtx) = (code=%d, err=%v), want context.DeadlineExceeded without 'timed out'", expCode, expErr) + } + + _, failStderr, failCode, failErr := sh.ExecuteBash(t.Context(), "echo err >&2; exit 7", "", 5*time.Second) + if failErr != nil || failCode != 7 || !strings.Contains(failStderr, "err") { + t.Errorf("ExecuteBash(exit 7) = (stderr=%q, code=%d, err=%v)", failStderr, failCode, failErr) + } + + // Verify timeout kills background grandchild holding stdout pipe without hanging. + start := time.Now() + _, _, toCode, toErr := sh.ExecuteBash(t.Context(), "(sleep 30) & sleep 30", "", 80*time.Millisecond) + if toErr == nil || toCode != -1 || !strings.Contains(toErr.Error(), "timed out") { + t.Errorf("ExecuteBash(timeout) = (code=%d, err=%v), want timed out error", toCode, toErr) + } + if elapsed := time.Since(start); elapsed > 5*time.Second { + t.Errorf("ExecuteBash(timeout) took %v, subprocess group did not terminate promptly", elapsed) + } +} + +func TestToolDispatcherToolsAndExtendedFS(t *testing.T) { + rootDir := t.TempDir() + fs := NewSandboxFS(rootDir, rootDir) + sh := NewSandboxShell(rootDir) + + disabled := NewToolDispatcher(ToolDispatcherConfig{ + FS: fs, + Shell: sh, + DisableBash: true, + }) + if res, isErr := disabled.Dispatch(t.Context(), "execute_bash", map[string]any{"command": "echo hi"}); !isErr || !strings.Contains(res["error"].(string), "disabled") { + t.Errorf("Dispatch(execute_bash disabled) = (%v, %v), want disabled error", res, isErr) + } + if _, isErr := disabled.Dispatch(t.Context(), "stat_file", map[string]any{"path": "."}); !isErr { + t.Error("Dispatch(stat_file with EnableExtendedFSTools=false) = false, want error") + } + + var gotTimeout time.Duration + td := NewToolDispatcher(ToolDispatcherConfig{ + FS: fs, + Shell: sh, + ShellExec: func(_ context.Context, cmd, cwd string, timeout time.Duration) (string, string, int, error) { + gotTimeout = timeout + return cmd + "@" + filepath.Base(cwd), "", 0, nil + }, + CommandTimeout: 2 * time.Second, + EnableExtendedFSTools: true, + DefaultWriteToFileOverwrite: true, + RespectWaitMsBeforeAsync: true, + }) + + bashRes, bashErr := td.Dispatch(t.Context(), "run_shell_command", map[string]any{ + "CommandLine": "ls", + "WaitMsBeforeAsync": float64(500), + }) + wantBashOut := "ls@" + filepath.Base(rootDir) + if bashErr || gotTimeout != 500*time.Millisecond || bashRes["stdout"] != wantBashOut { + t.Errorf("Dispatch(run_shell_command) = (%v, %v), timeout=%v, want stdout=%q", bashRes, bashErr, gotTimeout, wantBashOut) + } + if _, escErr := td.Dispatch(t.Context(), "bash", map[string]any{"command": "pwd", "cwd": "../../outside"}); !escErr { + t.Error("Dispatch(bash escaping cwd) = false, want error") + } + + // File tools: write_to_file, view_file, replace_file_content, list_dir, grep_search, find_by_name, extended tools. + // Note: "dup" appears on BOTH line 1 and line 2 with AllowMultiple=false (default), so string-typed + // StartLine="2" MUST be parsed to 2 rather than falling back to 0 (which would fail on 2 occurrences). + if res, isErr := td.Dispatch(t.Context(), "write_to_file", map[string]any{ + "TargetFile": "dir/code.py", + "CodeContent": "dup\r\ndup\r\nline3\r\n", + }); isErr { + t.Fatalf("write_to_file failed: %v", res) + } + // Verify write_file defaults to overwrite=true even when DefaultWriteToFileOverwrite is false, + // whereas write_to_file defaults to overwrite=false when DefaultWriteToFileOverwrite is false. + if res, isErr := disabled.Dispatch(t.Context(), "write_to_file", map[string]any{ + "TargetFile": "dir/code.py", + "CodeContent": "should_fail", + }); !isErr { + t.Errorf("disabled.Dispatch(write_to_file existing with DefaultWriteToFileOverwrite=false) = (%v, false), want error", res) + } + if res, isErr := disabled.Dispatch(t.Context(), "write_file", map[string]any{ + "TargetFile": "dir/code.py", + "CodeContent": "old", + }); isErr { + t.Fatalf("disabled.Dispatch(write_file existing with DefaultWriteToFileOverwrite=false) failed: %v", res) + } + if res, isErr := td.Dispatch(t.Context(), "write_to_file", map[string]any{ + "TargetFile": "dir/code.py", + "CodeContent": "dup\r\ndup\r\nline3\r\n", + }); isErr { + t.Fatalf("td.Dispatch(write_to_file existing with DefaultWriteToFileOverwrite=true) failed: %v", res) + } + if res, isErr := td.Dispatch(t.Context(), "replace_file_content", map[string]any{ + "TargetFile": "dir/code.py", + "TargetContent": "dup", + "ReplacementContent": "line2_edited", + "StartLine": "2", + "EndLine": float64(2), + }); isErr { + t.Fatalf("replace_file_content(string StartLine=2) failed: %v", res) + } + viewRes, viewErr := td.Dispatch(t.Context(), "view_file", map[string]any{"AbsolutePath": "dir/code.py"}) + if viewErr || viewRes["content"] != "dup\r\nline2_edited\r\nline3\r\n" { + t.Errorf("view_file(full) = (%v, %v)", viewRes, viewErr) + } + rangeRes, rangeErr := td.Dispatch(t.Context(), "view_file", map[string]any{ + "AbsolutePath": "dir/code.py", + "StartLine": 1, + "EndLine": 2, + }) + if rangeErr || rangeRes["content"] != "dup\nline2_edited" { + t.Errorf("view_file(range) = (%v, %v)", rangeRes, rangeErr) + } + startOnlyRes, startOnlyErr := td.Dispatch(t.Context(), "view_file", map[string]any{ + "AbsolutePath": "dir/code.py", + "StartLine": 2, + }) + if startOnlyErr || startOnlyRes["content"] != "line2_edited\nline3" { + t.Errorf("view_file(StartLine only) = (%v, %v), want \"line2_edited\\nline3\"", startOnlyRes, startOnlyErr) + } + endOnlyRes, endOnlyErr := td.Dispatch(t.Context(), "view_file", map[string]any{ + "AbsolutePath": "dir/code.py", + "EndLine": 1, + }) + if endOnlyErr || endOnlyRes["content"] != "dup" { + t.Errorf("view_file(EndLine only) = (%v, %v), want \"dup\"", endOnlyRes, endOnlyErr) + } + + // Verify ReplaceInText preserves LF-only line endings, normalizes multi-line LF targets against CRLF content, and handles explicit CRLF targets. + lfOut, lfCount, lfErr := ReplaceInText("alpha\nbeta\n", "alpha", "GAMMA", 0, 1, false) + if lfErr != nil || lfCount != 1 || lfOut != "GAMMA\nbeta\n" { + t.Errorf("ReplaceInText(LF-only, startLine=0, endLine=1) = (%q, %d, %v), want (\"GAMMA\\nbeta\\n\", 1, nil)", lfOut, lfCount, lfErr) + } + crlfOut, crlfCount, crlfErr := ReplaceInText("line1\r\nline2\r\nline3\r\n", "line1\nline2", "new1\r\nnew2", 1, 2, false) + if crlfErr != nil || crlfCount != 1 || crlfOut != "new1\r\nnew2\r\nline3\r\n" { + t.Errorf("ReplaceInText(multi-line LF target in CRLF content) = (%q, %d, %v), want (\"new1\\r\\nnew2\\r\\nline3\\r\\n\", 1, nil)", crlfOut, crlfCount, crlfErr) + } + rawCRLFOut, rawCRLFCount, rawCRLFErr := ReplaceInText("line1\r\nline2\r\n", "line1\r\nline2", "new1\r\nnew2", 0, 0, false) + if rawCRLFErr != nil || rawCRLFCount != 1 || rawCRLFOut != "new1\r\nnew2\r\n" { + t.Errorf("ReplaceInText(explicit CRLF target in CRLF content) = (%q, %d, %v), want (\"new1\\r\\nnew2\\r\\n\", 1, nil)", rawCRLFOut, rawCRLFCount, rawCRLFErr) + } + + if res, isErr := td.Dispatch(t.Context(), "stat_file", map[string]any{"path": "dir/code.py"}); isErr || res["is_dir"] != false { + t.Errorf("stat_file = (%v, %v)", res, isErr) + } else if _, ok := res["size"].(int64); !ok { + t.Errorf("stat_file size type = %T, want int64", res["size"]) + } + nilFSTools := NewToolDispatcher(ToolDispatcherConfig{EnableExtendedFSTools: true}) + if res, isErr := nilFSTools.Dispatch(t.Context(), "view_file", map[string]any{"path": "a.txt"}); !isErr || !strings.Contains(res["error"].(string), "not configured") { + t.Errorf("nilFSTools.Dispatch(view_file) = (%v, %v), want not configured error", res, isErr) + } + // toolMkdir defaults recursive=true, so passing string "false" on a missing parent MUST fail if string coercion works. + if res, isErr := td.Dispatch(t.Context(), "mkdir", map[string]any{"path": "missing_parent/child", "recursive": "false"}); !isErr { + t.Errorf("mkdir(recursive=\"false\" with missing parent) = (%v, false), want error", res) + } + if res, isErr := td.Dispatch(t.Context(), "mkdir", map[string]any{"path": "newdir/sub", "recursive": "true"}); isErr { + t.Errorf("mkdir = (%v, %v)", res, isErr) + } + if res, isErr := td.Dispatch(t.Context(), "rename_file", map[string]any{"old_path": "dir/code.py", "new_path": "newdir/code.py"}); isErr { + t.Errorf("rename_file = (%v, %v)", res, isErr) + } + if res, isErr := td.Dispatch(t.Context(), "list_dir", map[string]any{"DirectoryPath": "newdir"}); isErr { + t.Errorf("list_dir = (%v, %v)", res, isErr) + } + if res, isErr := td.Dispatch(t.Context(), "grep_search", map[string]any{"SearchPath": ".", "Query": "line2_edited"}); isErr { + t.Errorf("grep_search = (%v, %v)", res, isErr) + } + if res, isErr := td.Dispatch(t.Context(), "find_by_name", map[string]any{"SearchDirectory": ".", "Extensions": "rs,py"}); isErr { + t.Errorf("find_by_name(comma Extensions=\"rs,py\") = (%v, %v)", res, isErr) + } else if files, ok := res["files"].([]any); !ok || len(files) != 1 || files[0] != "newdir/code.py" { + t.Errorf("find_by_name(comma Extensions=\"rs,py\") files = %v, want [newdir/code.py]", res["files"]) + } + if res, isErr := td.Dispatch(t.Context(), "remove_file", map[string]any{}); !isErr { + t.Errorf("remove_file(empty path) = (%v, false), want error", res) + } + if res, isErr := td.Dispatch(t.Context(), "remove_file", map[string]any{"TargetFile": "newdir", "recursive": "true"}); isErr { + t.Errorf("remove_file(TargetFile) = (%v, %v)", res, isErr) + } + if _, isErr := td.Dispatch(t.Context(), "unknown_tool", nil); !isErr { + t.Error("Dispatch(unknown_tool) = false, want error") + } + + // Direct unit verification of intArg and boolArg string coercion and fallback behavior. + if got := intArg(map[string]any{"n": " 17 "}, 3, "n"); got != 17 { + t.Errorf("intArg(\" 17 \", fallback=3) = %d, want 17", got) + } + if got := intArg(map[string]any{"n": "not-an-int"}, 3, "n"); got != 3 { + t.Errorf("intArg(\"not-an-int\", fallback=3) = %d, want 3", got) + } + if got := boolArg(map[string]any{"b": "false"}, true, "b"); got { + t.Error("boolArg(\"false\", fallback=true) = true, want false") + } + if got := boolArg(map[string]any{"b": " TRUE "}, false, "b"); !got { + t.Error("boolArg(\" TRUE \", fallback=false) = false, want true") + } + if got := boolArg(map[string]any{"b": "not-a-bool"}, true, "b"); !got { + t.Error("boolArg(\"not-a-bool\", fallback=true) = false, want true") + } +} diff --git a/internal/cli/custom/envruntime/shell.go b/internal/cli/custom/envruntime/shell.go new file mode 100644 index 0000000..b667a0b --- /dev/null +++ b/internal/cli/custom/envruntime/shell.go @@ -0,0 +1,278 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package envruntime + +import ( + "bytes" + "context" + "errors" + "fmt" + "os" + "os/exec" + "runtime" + "strings" + "time" +) + +// MaxBashOutputBytes caps stdout and stderr capture per command execution at 4 MiB. +const MaxBashOutputBytes = 4 * 1024 * 1024 + +// EnvPolicy selects how SandboxShell filters environment variables passed to child processes. +type EnvPolicy int + +const ( + // EnvPolicyAllowlist passes only baseline OS/locale variables (recommended for host CLI workers). + EnvPolicyAllowlist EnvPolicy = iota + // EnvPolicyScrubCredentials passes the process environment with Gemini/Google credential variables removed plus ExtraEnv. + EnvPolicyScrubCredentials +) + +// CappedBuffer is a bytes.Buffer wrapper that bounds retained output to MaxBytes. +type CappedBuffer struct { + buf bytes.Buffer + MaxBytes int + truncated bool +} + +// Write appends up to MaxBytes into the buffer and discards overflow while reporting full write length. +func (c *CappedBuffer) Write(p []byte) (int, error) { + origLen := len(p) + maxBytes := c.MaxBytes + if maxBytes <= 0 { + maxBytes = MaxBashOutputBytes + } + remaining := maxBytes - c.buf.Len() + if remaining <= 0 { + c.truncated = true + return origLen, nil + } + if len(p) > remaining { + c.buf.Write(p[:remaining]) + c.truncated = true + return origLen, nil + } + c.buf.Write(p) + return origLen, nil +} + +// String returns the captured text, appending a truncation marker when MaxBytes was exceeded. +func (c *CappedBuffer) String() string { + if c.truncated { + return c.buf.String() + "\n...[output truncated at 4 MiB]..." + } + return c.buf.String() +} + +var scrubbedCredentialEnvKeys = map[string]bool{ + "GEMINI_ENVIRONMENT_KEY": true, + "ENVIRONMENT_KEY": true, + "GEMINI_API_KEY": true, + "GOOGLE_API_KEY": true, + "GEMINI_ACCESS_TOKEN": true, + "ACCESS_TOKEN": true, + "GOOGLE_OAUTH_ACCESS_TOKEN": true, + "CLOUDSDK_AUTH_ACCESS_TOKEN": true, + "GEMINI_LEASE_TOKEN": true, + "LEASE_TOKEN": true, +} + +// FilterScrubbedEnv removes Gemini and Google credential environment variables from baseEnv and appends extraEnv. +func FilterScrubbedEnv(baseEnv []string, extraEnv ...string) []string { + // Must be non-nil: a nil exec.Cmd.Env inherits the parent environment. + out := make([]string, 0, len(baseEnv)+len(extraEnv)) + for _, kv := range baseEnv { + key, _, _ := strings.Cut(kv, "=") + if scrubbedCredentialEnvKeys[strings.ToUpper(strings.TrimSpace(key))] { + continue + } + out = append(out, kv) + } + return append(out, extraEnv...) +} + +var allowedWorkerBashEnvKeys = map[string]bool{ + "PATH": true, + "HOME": true, + "USER": true, + "LOGNAME": true, + "SHELL": true, + "LANG": true, + "TERM": true, + "TMPDIR": true, + "SYSTEMROOT": true, + "WINDIR": true, + "COMSPEC": true, + "PATHEXT": true, + "TEMP": true, + "TMP": true, + "USERPROFILE": true, + "APPDATA": true, + "LOCALAPPDATA": true, + "PROGRAMFILES": true, + "PROGRAMFILES(X86)": true, + "PROGRAMDATA": true, + "HOMEDRIVE": true, + "HOMEPATH": true, + "COMPUTERNAME": true, + "USERNAME": true, + "USERDOMAIN": true, + "OS": true, + "SYSTEMDRIVE": true, + "PROCESSOR_ARCHITECTURE": true, + "NUMBER_OF_PROCESSORS": true, + "PSMODULEPATH": true, +} + +// BuildAllowlistEnv retains only baseline OS and locale variables from baseEnv. +func BuildAllowlistEnv(baseEnv []string) []string { + out := make([]string, 0, 12) + for _, kv := range baseEnv { + key, _, ok := strings.Cut(kv, "=") + if !ok { + continue + } + upper := strings.ToUpper(strings.TrimSpace(key)) + if allowedWorkerBashEnvKeys[upper] || strings.HasPrefix(upper, "LC_") { + out = append(out, kv) + } + } + return out +} + +func isWindowsWSLStubBash(goos, p string) bool { + if goos != "windows" { + return false + } + normalized := strings.ToLower(strings.TrimRight(strings.ReplaceAll(p, "/", `\`), `\`)) + return strings.HasSuffix(normalized, `\system32\bash.exe`) || strings.HasSuffix(normalized, `\sysnative\bash.exe`) +} + +// ResolveShellCommand selects bash, sh, or cmd.exe to execute command on goos. +func ResolveShellCommand(goos, command string) (string, []string) { + if p, err := exec.LookPath("bash"); err == nil && !isWindowsWSLStubBash(goos, p) { + return p, []string{"-c", command} + } + if p, err := exec.LookPath("sh"); err == nil { + return p, []string{"-c", command} + } + if goos == "windows" { + return "cmd.exe", []string{"/C", command} + } + return "sh", []string{"-c", command} +} + +func resolveExecShell(goos string, useBinSh bool, command string) (string, []string) { + if useBinSh && goos != "windows" { + return "/bin/sh", []string{"-c", command} + } + return ResolveShellCommand(goos, command) +} + +// SandboxShellConfig configures a SandboxShell instance. +type SandboxShellConfig struct { + WorkDir string + EnvPolicy EnvPolicy + ExtraEnv []string + AutoCreateWorkDir bool + UseBinSh bool +} + +// SandboxShell executes shell commands with process-group termination, output caps, and environment filtering. +type SandboxShell struct { + workDir string + envPolicy EnvPolicy + extraEnv []string + autoCreateWorkDir bool + useBinSh bool +} + +// NewSandboxShell creates a SandboxShell using EnvPolicyAllowlist. +func NewSandboxShell(workDir string) *SandboxShell { + return NewSandboxShellWithConfig(SandboxShellConfig{ + WorkDir: workDir, + EnvPolicy: EnvPolicyAllowlist, + }) +} + +// NewSandboxShellWithConfig creates a SandboxShell with the provided configuration. +func NewSandboxShellWithConfig(cfg SandboxShellConfig) *SandboxShell { + return &SandboxShell{ + workDir: cfg.WorkDir, + envPolicy: cfg.EnvPolicy, + extraEnv: append([]string(nil), cfg.ExtraEnv...), + autoCreateWorkDir: cfg.AutoCreateWorkDir, + useBinSh: cfg.UseBinSh, + } +} + +// WorkDir returns the default working directory configured on the shell. +func (s *SandboxShell) WorkDir() string { + return s.workDir +} + +// ExecuteBash executes command in workDir (or s.WorkDir() when empty) with an optional timeout. +// When timeout <= 0, no additional deadline is imposed beyond ctx. +func (s *SandboxShell) ExecuteBash(ctx context.Context, command, workDir string, timeout time.Duration) (string, string, int, error) { + runCtx := ctx + var cancel context.CancelFunc + if timeout > 0 { + runCtx, cancel = context.WithTimeout(ctx, timeout) + defer cancel() + } + + dir := s.workDir + if workDir != "" { + dir = workDir + } + if s.autoCreateWorkDir && dir != "" { + _ = os.MkdirAll(dir, 0o755) + } + + shellBin, shellArgs := resolveExecShell(runtime.GOOS, s.useBinSh, command) + + cmd := exec.CommandContext(runCtx, shellBin, shellArgs...) + cmd.Dir = dir + ConfigureSubprocessGroup(cmd) + cmd.WaitDelay = 5 * time.Second + if s.envPolicy == EnvPolicyScrubCredentials { + cmd.Env = FilterScrubbedEnv(os.Environ(), s.extraEnv...) + } else { + cmd.Env = append(BuildAllowlistEnv(os.Environ()), s.extraEnv...) + } + + stdoutBuf := &CappedBuffer{MaxBytes: MaxBashOutputBytes} + stderrBuf := &CappedBuffer{MaxBytes: MaxBashOutputBytes} + cmd.Stdout = stdoutBuf + cmd.Stderr = stderrBuf + + err := cmd.Run() + stdout := stdoutBuf.String() + stderr := stderrBuf.String() + + if runCtx.Err() != nil { + if ctx.Err() == nil { + return stdout, stderr, -1, fmt.Errorf("command timed out after %v: %w", timeout, runCtx.Err()) + } + return stdout, stderr, -1, runCtx.Err() + } + if err != nil { + var exitErr *exec.ExitError + if errors.As(err, &exitErr) { + return stdout, stderr, exitErr.ExitCode(), nil + } + return stdout, stderr, -1, err + } + return stdout, stderr, 0, nil +} diff --git a/internal/cli/custom/envruntime/tools.go b/internal/cli/custom/envruntime/tools.go new file mode 100644 index 0000000..4af4655 --- /dev/null +++ b/internal/cli/custom/envruntime/tools.go @@ -0,0 +1,452 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package envruntime + +import ( + "cmp" + "context" + "fmt" + "strconv" + "strings" + "time" +) + +// ShellExecFunc executes a shell command and returns stdout, stderr, exitCode, and any execution error. +type ShellExecFunc func(ctx context.Context, command, workDir string, timeout time.Duration) (string, string, int, error) + +// ToolDispatcherConfig configures built-in sandbox tool dispatch. +type ToolDispatcherConfig struct { + FS *SandboxFS + Shell *SandboxShell + ShellExec ShellExecFunc + DisableBash bool + CommandTimeout time.Duration + EnableExtendedFSTools bool + DefaultWriteToFileOverwrite bool + RespectWaitMsBeforeAsync bool +} + +// ToolDispatcher dispatches client-side filesystem, search, and shell tool invocations. +type ToolDispatcher struct { + cfg ToolDispatcherConfig +} + +// NewToolDispatcher constructs a ToolDispatcher from cfg. +func NewToolDispatcher(cfg ToolDispatcherConfig) *ToolDispatcher { + return &ToolDispatcher{cfg: cfg} +} + +// Dispatch executes toolName with args and returns the structured JSON-like map result and isError flag. +func (d *ToolDispatcher) Dispatch(ctx context.Context, toolName string, args map[string]any) (map[string]any, bool) { + normalized := strings.ToLower(strings.TrimSpace(toolName)) + switch normalized { + case "execute_bash", "run_command", "run_shell_command", "bash", "shell": + return d.toolExecuteBash(ctx, args) + case "view_file", "read_file": + return d.dispatchFS(func() (map[string]any, bool) { return d.toolViewFile(args) }) + case "write_to_file", "write_file": + return d.dispatchFS(func() (map[string]any, bool) { return d.toolWriteFile(normalized, args) }) + case "replace_file_content", "edit_file": + return d.dispatchFS(func() (map[string]any, bool) { return d.toolReplaceFileContent(args) }) + case "list_dir", "read_dir", "list_directory": + return d.dispatchFS(func() (map[string]any, bool) { return d.toolListDir(args) }) + case "grep_search": + return d.dispatchFS(func() (map[string]any, bool) { return d.toolGrepSearch(args) }) + case "find_by_name", "find_file": + return d.dispatchFS(func() (map[string]any, bool) { return d.toolFindByName(args) }) + case "stat_file", "stat": + if d.cfg.EnableExtendedFSTools { + return d.dispatchFS(func() (map[string]any, bool) { return d.toolStatFile(args) }) + } + case "mkdir": + if d.cfg.EnableExtendedFSTools { + return d.dispatchFS(func() (map[string]any, bool) { return d.toolMkdir(args) }) + } + case "remove_file", "remove": + if d.cfg.EnableExtendedFSTools { + return d.dispatchFS(func() (map[string]any, bool) { return d.toolRemoveFile(args) }) + } + case "rename_file", "rename": + if d.cfg.EnableExtendedFSTools { + return d.dispatchFS(func() (map[string]any, bool) { return d.toolRenameFile(args) }) + } + } + return map[string]any{"error": fmt.Sprintf("unsupported sandbox tool %q", toolName)}, true +} + +func (d *ToolDispatcher) dispatchFS(fn func() (map[string]any, bool)) (map[string]any, bool) { + if d.cfg.FS == nil { + return map[string]any{"error": "filesystem is not configured"}, true + } + return fn() +} + +func (d *ToolDispatcher) toolExecuteBash(ctx context.Context, args map[string]any) (map[string]any, bool) { + if d.cfg.DisableBash { + return map[string]any{"error": "execute_bash is disabled (--disable-bash)"}, true + } + cmdLine := stringArg(args, "command", "CommandLine", "cmd") + rawWorkDir := stringArg(args, "Cwd", "cwd", "working_dir") + resolvedWorkDir := rawWorkDir + if d.cfg.FS != nil && (rawWorkDir != "" || d.cfg.FS.Mode() == ConfinementOSRoot) { + resolved, err := d.cfg.FS.ResolvePath(rawWorkDir) + if err != nil { + return map[string]any{"error": err.Error()}, true + } + resolvedWorkDir = resolved + } + + timeout := d.cfg.CommandTimeout + timeoutKeys := []string{"timeout_ms"} + if d.cfg.RespectWaitMsBeforeAsync { + timeoutKeys = append(timeoutKeys, "WaitMsBeforeAsync") + } + if timeoutMS := intArg(args, 0, timeoutKeys...); timeoutMS > 0 { + timeout = time.Duration(timeoutMS) * time.Millisecond + } else if timeoutSec := intArg(args, 0, "timeout_seconds", "timeout"); timeoutSec > 0 { + timeout = time.Duration(timeoutSec) * time.Second + } + if d.cfg.CommandTimeout > 0 && timeout > d.cfg.CommandTimeout { + timeout = d.cfg.CommandTimeout + } + + var ( + stdout string + stderr string + exitCode int + err error + ) + switch { + case d.cfg.ShellExec != nil: + stdout, stderr, exitCode, err = d.cfg.ShellExec(ctx, cmdLine, resolvedWorkDir, timeout) + case d.cfg.Shell != nil: + stdout, stderr, exitCode, err = d.cfg.Shell.ExecuteBash(ctx, cmdLine, resolvedWorkDir, timeout) + default: + return map[string]any{"error": "shell executor is not configured"}, true + } + if err != nil { + return map[string]any{ + "error": err.Error(), + "stdout": stdout, + "stderr": stderr, + "exit_code": exitCode, + }, true + } + return map[string]any{ + "stdout": stdout, + "stderr": stderr, + "exit_code": exitCode, + }, false +} + +func (d *ToolDispatcher) toolViewFile(args map[string]any) (map[string]any, bool) { + path := stringArg(args, "AbsolutePath", "path", "file_path") + if path == "" { + return map[string]any{"error": "path is required"}, true + } + startLine := intArg(args, 0, "StartLine", "start_line") + endLine := intArg(args, 0, "EndLine", "end_line") + var content string + if startLine > 0 || endLine > 0 { + sliced, err := d.cfg.FS.ReadFileRange(path, startLine, endLine) + if err != nil { + return map[string]any{"error": err.Error()}, true + } + content = sliced + } else { + data, err := d.cfg.FS.ReadRegularFile(path) + if err != nil { + return map[string]any{"error": err.Error()}, true + } + content = string(data) + } + return map[string]any{ + "path": path, + "content": content, + }, false +} + +func (d *ToolDispatcher) toolWriteFile(normalizedToolName string, args map[string]any) (map[string]any, bool) { + path := stringArg(args, "TargetFile", "path", "file_path") + if path == "" { + return map[string]any{"error": "path is required"}, true + } + content := stringArg(args, "CodeContent", "content") + defaultOverwrite := normalizedToolName == "write_file" || d.cfg.DefaultWriteToFileOverwrite + overwrite := boolArg(args, defaultOverwrite, "Overwrite", "overwrite") + appendMode := boolArg(args, false, "Append", "append") + n, err := d.cfg.FS.WriteRegularFile(path, content, overwrite, appendMode) + if err != nil { + return map[string]any{"error": err.Error()}, true + } + return map[string]any{ + "status": "ok", + "path": path, + "bytes_written": n, + }, false +} + +func (d *ToolDispatcher) toolReplaceFileContent(args map[string]any) (map[string]any, bool) { + path := stringArg(args, "TargetFile", "path", "file_path") + if path == "" { + return map[string]any{"error": "TargetFile is required"}, true + } + target := stringArg(args, "TargetContent", "old_string") + if target == "" { + return map[string]any{"error": "TargetContent is required"}, true + } + replacement := stringArg(args, "ReplacementContent", "new_string") + allowMultiple := boolArg(args, false, "AllowMultiple", "allow_multiple") + startLine := intArg(args, 0, "StartLine", "start_line") + endLine := intArg(args, 0, "EndLine", "end_line") + + raw, err := d.cfg.FS.ReadRegularFile(path) + if err != nil { + return map[string]any{"error": err.Error()}, true + } + updated, count, err := ReplaceInText(string(raw), target, replacement, startLine, endLine, allowMultiple) + if err != nil { + return map[string]any{"error": err.Error()}, true + } + if _, err := d.cfg.FS.WriteRegularFile(path, updated, true, false); err != nil { + return map[string]any{"error": err.Error()}, true + } + return map[string]any{ + "status": "ok", + "path": path, + "replacements": count, + }, false +} + +func (d *ToolDispatcher) toolListDir(args map[string]any) (map[string]any, bool) { + path := stringArg(args, "DirectoryPath", "path", "dir_path") + if path == "" { + path = "." + } + entries, err := d.cfg.FS.ListEntries(path) + if err != nil { + return map[string]any{"error": err.Error()}, true + } + return map[string]any{ + "path": path, + "entries": entries, + }, false +} + +func (d *ToolDispatcher) toolStatFile(args map[string]any) (map[string]any, bool) { + path := stringArg(args, "AbsolutePath", "path", "file_path") + info, err := d.cfg.FS.Stat(path) + if err != nil { + return map[string]any{"error": err.Error()}, true + } + return map[string]any{ + "name": info.Name, + "is_dir": info.IsDir, + "size": info.Size, + "permissions": int(info.Permissions), + }, false +} + +func (d *ToolDispatcher) toolMkdir(args map[string]any) (map[string]any, bool) { + path := stringArg(args, "DirectoryPath", "path", "dir_path") + recursive := boolArg(args, true, "recursive", "Recursive") + if err := d.cfg.FS.Mkdir(path, 0o755, recursive); err != nil { + return map[string]any{"error": err.Error()}, true + } + return map[string]any{"status": "ok", "path": path}, false +} + +func (d *ToolDispatcher) toolRemoveFile(args map[string]any) (map[string]any, bool) { + path := stringArg(args, "AbsolutePath", "TargetFile", "path", "file_path") + if strings.TrimSpace(path) == "" { + return map[string]any{"error": "path is required"}, true + } + recursive := boolArg(args, false, "recursive", "Recursive") + if err := d.cfg.FS.Remove(path, recursive); err != nil { + return map[string]any{"error": err.Error()}, true + } + return map[string]any{"status": "ok", "path": path}, false +} + +func (d *ToolDispatcher) toolRenameFile(args map[string]any) (map[string]any, bool) { + oldPath := stringArg(args, "old_path", "source_path", "OldPath") + newPath := stringArg(args, "new_path", "destination_path", "NewPath") + if err := d.cfg.FS.Rename(oldPath, newPath); err != nil { + return map[string]any{"error": err.Error()}, true + } + return map[string]any{"status": "ok", "old_path": oldPath, "new_path": newPath}, false +} + +func (d *ToolDispatcher) toolGrepSearch(args map[string]any) (map[string]any, bool) { + searchPath := cmp.Or(stringArg(args, "SearchPath", "path"), ".") + query := stringArg(args, "Query", "query") + caseInsensitive := boolArg(args, false, "CaseInsensitive", "case_insensitive") + isRegex := boolArg(args, false, "IsRegex", "is_regex") + matches, truncated, err := d.cfg.FS.Grep(searchPath, query, caseInsensitive, isRegex) + if err != nil { + return map[string]any{"error": err.Error()}, true + } + res := map[string]any{"matches": matches} + if truncated { + res["truncated"] = true + } + return res, false +} + +func (d *ToolDispatcher) toolFindByName(args map[string]any) (map[string]any, bool) { + searchDir := cmp.Or(stringArg(args, "SearchDirectory", "path"), ".") + pattern := stringArg(args, "Pattern", "pattern") + maxDepth := intArg(args, 0, "MaxDepth", "max_depth") + extensions := stringSliceArg(args, "Extensions", "extensions") + files, truncated, err := d.cfg.FS.FindFiles(searchDir, pattern, maxDepth, extensions...) + if err != nil { + return map[string]any{"error": err.Error()}, true + } + res := map[string]any{"files": files} + if truncated { + res["truncated"] = true + } + return res, false +} + +// ReplaceInText replaces target with replacement within 1-indexed inclusive [startLine, endLine] while preserving CRLF line endings. +func ReplaceInText(content, target, replacement string, startLine, endLine int, allowMultiple bool) (string, int, error) { + preserveCRLF := strings.Contains(content, "\r\n") && !strings.Contains(target, "\r\n") + if preserveCRLF { + content = strings.ReplaceAll(content, "\r\n", "\n") + replacement = strings.ReplaceAll(replacement, "\r\n", "\n") + } + restore := func(s string) string { + if preserveCRLF { + return strings.ReplaceAll(s, "\n", "\r\n") + } + return s + } + + if startLine <= 0 && endLine <= 0 { + count := strings.Count(content, target) + if count == 0 { + return "", 0, fmt.Errorf("target content not found") + } + if count > 1 && !allowMultiple { + return "", 0, fmt.Errorf("found %d occurrences of target content when AllowMultiple is false", count) + } + return restore(strings.ReplaceAll(content, target, replacement)), count, nil + } + + lines := strings.Split(content, "\n") + if startLine <= 0 { + startLine = 1 + } + if endLine <= 0 || endLine > len(lines) { + endLine = len(lines) + } + if startLine > endLine || startLine > len(lines) { + return "", 0, fmt.Errorf("invalid line range [%d, %d]", startLine, endLine) + } + + region := strings.Join(lines[startLine-1:endLine], "\n") + count := strings.Count(region, target) + if count == 0 { + return "", 0, fmt.Errorf("target content not found in lines [%d, %d]", startLine, endLine) + } + if count > 1 && !allowMultiple { + return "", 0, fmt.Errorf("found %d occurrences in lines [%d, %d] when AllowMultiple is false", count, startLine, endLine) + } + + var parts []string + if startLine > 1 { + parts = append(parts, strings.Join(lines[:startLine-1], "\n")) + } + parts = append(parts, strings.ReplaceAll(region, target, replacement)) + if endLine < len(lines) { + parts = append(parts, strings.Join(lines[endLine:], "\n")) + } + return restore(strings.Join(parts, "\n")), count, nil +} + +func stringArg(args map[string]any, keys ...string) string { + for _, k := range keys { + if v := args[k]; v != nil { + if s, ok := v.(string); ok { + return s + } + return fmt.Sprint(v) + } + } + return "" +} + +func intArg(args map[string]any, fallback int, keys ...string) int { + for _, k := range keys { + switch val := args[k].(type) { + case int: + return val + case int32: + return int(val) + case int64: + return int(val) + case float64: + return int(val) + case float32: + return int(val) + case string: + if n, err := strconv.Atoi(strings.TrimSpace(val)); err == nil { + return n + } + } + } + return fallback +} + +func boolArg(args map[string]any, fallback bool, keys ...string) bool { + for _, k := range keys { + switch val := args[k].(type) { + case bool: + return val + case string: + if b, err := strconv.ParseBool(strings.TrimSpace(val)); err == nil { + return b + } + } + } + return fallback +} + +func stringSliceArg(args map[string]any, keys ...string) []string { + for _, k := range keys { + switch val := args[k].(type) { + case []string: + return val + case []any: + out := make([]string, 0, len(val)) + for _, item := range val { + if item != nil { + out = append(out, fmt.Sprint(item)) + } + } + return out + case string: + var out []string + for _, part := range strings.Split(val, ",") { + if trimmed := strings.TrimSpace(part); trimmed != "" { + out = append(out, trimmed) + } + } + return out + } + } + return nil +} diff --git a/internal/cli/custom/envruntime/wire.go b/internal/cli/custom/envruntime/wire.go new file mode 100644 index 0000000..0aef281 --- /dev/null +++ b/internal/cli/custom/envruntime/wire.go @@ -0,0 +1,1051 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package envruntime + +import ( + "encoding/binary" + "fmt" + "io" + "math" + "net/url" + "sort" + "strings" +) + +const ( + wireVarint = 0 + wireFixed64 = 1 + wireBytes = 2 + wireFixed32 = 5 + + // MaxGRPCFrameBytes bounds uncompressed gRPC message frames at 32 MiB. + MaxGRPCFrameBytes = 32 << 20 + + // MaxProtoNestingDepth bounds recursive Struct/Value/ListValue decoding. + MaxProtoNestingDepth = 64 +) + +// GRPCWriteFrame writes msg as an uncompressed 5-byte length-prefixed gRPC frame to w. +func GRPCWriteFrame(w io.Writer, msg []byte) error { + if len(msg) > MaxGRPCFrameBytes { + return fmt.Errorf("outbound gRPC frame size %d exceeds limit %d", len(msg), MaxGRPCFrameBytes) + } + var hdr [5]byte + binary.BigEndian.PutUint32(hdr[1:5], uint32(len(msg))) + if _, err := w.Write(hdr[:]); err != nil { + return err + } + if len(msg) > 0 { + if _, err := w.Write(msg); err != nil { + return err + } + } + return nil +} + +// GRPCReadFrame reads an uncompressed 5-byte length-prefixed gRPC frame from r. +func GRPCReadFrame(r io.Reader) ([]byte, error) { + var hdr [5]byte + if _, err := io.ReadFull(r, hdr[:]); err != nil { + return nil, err + } + if hdr[0] != 0 { + return nil, fmt.Errorf("unsupported compressed gRPC frame flag %d", hdr[0]) + } + n := binary.BigEndian.Uint32(hdr[1:5]) + if n > MaxGRPCFrameBytes { + return nil, fmt.Errorf("inbound gRPC frame size %d exceeds limit %d", n, MaxGRPCFrameBytes) + } + if n == 0 { + return nil, nil + } + buf := make([]byte, n) + if _, err := io.ReadFull(r, buf); err != nil { + if err == io.EOF { + return nil, io.ErrUnexpectedEOF + } + return nil, err + } + return buf, nil +} + +// DecodeGRPCMessage percent-decodes a grpc-message header or trailer value. +func DecodeGRPCMessage(raw string) string { + trimmed := strings.TrimSpace(raw) + if unescaped, err := url.PathUnescape(trimmed); err == nil { + return unescaped + } + return trimmed +} + +func appendTag(dst []byte, fieldNum int, wireType int) []byte { + return binary.AppendUvarint(dst, uint64(fieldNum<<3|wireType)) +} + +func appendVarintField(dst []byte, fieldNum int, v uint64) []byte { + dst = appendTag(dst, fieldNum, wireVarint) + return binary.AppendUvarint(dst, v) +} + +func appendBoolField(dst []byte, fieldNum int, v bool) []byte { + if !v { + return dst + } + return appendVarintField(dst, fieldNum, 1) +} + +func appendStringField(dst []byte, fieldNum int, s string) []byte { + if s == "" { + return dst + } + dst = appendTag(dst, fieldNum, wireBytes) + dst = binary.AppendUvarint(dst, uint64(len(s))) + return append(dst, s...) +} + +func appendBytesField(dst []byte, fieldNum int, b []byte) []byte { + dst = appendTag(dst, fieldNum, wireBytes) + dst = binary.AppendUvarint(dst, uint64(len(b))) + return append(dst, b...) +} + +// MarshalClientEventBatch encodes b into protobuf wire format matching genai.ClientEventBatch. +// All wire field numbers in this file must stay synchronized with the canonical +// genai.ClientEventBatch and genai.ServerEventBatch protobuf definitions (verified by +// differential protobuf round-trip tests). +func MarshalClientEventBatch(b *ClientEventBatch) []byte { + if b == nil { + return nil + } + var out []byte + for _, ev := range b.Events { + if ev != nil { + out = appendBytesField(out, 1, marshalClientEvent(ev)) + } + } + if b.AckOutputSeqNum != 0 { + out = appendVarintField(out, 2, uint64(b.AckOutputSeqNum)) + } + if r := b.ReconnectInteraction; r != nil { + var sub []byte + sub = appendStringField(sub, 1, r.InteractionID) + if r.NextOutputSeqNum != 0 { + sub = appendVarintField(sub, 2, uint64(r.NextOutputSeqNum)) + } + sub = appendBoolField(sub, 3, r.Background) + sub = appendStringField(sub, 5, r.ClientID) + out = appendBytesField(out, 8, sub) + } + if b.Disconnect != nil { + out = appendBytesField(out, 9, nil) + } + return out +} + +// MarshalClientEvent encodes a single ClientEvent into protobuf wire format. +func MarshalClientEvent(e *ClientEvent) []byte { + if e == nil { + return nil + } + return marshalClientEvent(e) +} + +func marshalClientEvent(e *ClientEvent) []byte { + var out []byte + if p := e.Payload; p != nil { + out = appendBytesField(out, 3, marshalPayload(p)) + } + if e.SeqNum != 0 { + out = appendVarintField(out, 10, uint64(e.SeqNum)) + } + if e.TimestampMicros != 0 { + out = appendVarintField(out, 11, uint64(e.TimestampMicros)) + } + return out +} + +func marshalPayload(p *EventPayload) []byte { + var sub []byte + sub = appendStringField(sub, 6, p.SubInteractionID) + if st := p.Step; st != nil { + var stepBytes []byte + if tc := st.ToolCall; tc != nil { + var tcb []byte + tcb = appendStringField(tcb, 1, tc.ID) + if fc := tc.FunctionCall; fc != nil { + var fcb []byte + fcb = appendStringField(fcb, 1, fc.Name) + fcb = appendBytesField(fcb, 2, EncodeProtoStruct(fc.Arguments)) + tcb = appendBytesField(tcb, 3, fcb) + } + stepBytes = appendBytesField(stepBytes, 4, tcb) + } + if tr := st.ToolResult; tr != nil { + var tb []byte + tb = appendStringField(tb, 1, tr.CallID) + if fnRes := tr.FunctionResult; fnRes != nil { + var frb []byte + frb = appendStringField(frb, 2, fnRes.Name) + frb = appendBoolField(frb, 3, fnRes.IsError) + frb = appendBytesField(frb, 7, EncodeProtoValue(fnRes.Result)) + tb = appendBytesField(tb, 3, frb) + } + stepBytes = appendBytesField(stepBytes, 5, tb) + } + if len(stepBytes) > 0 { + sub = appendBytesField(sub, 11, stepBytes) + } + } + if fi := p.FunctionInvocation; fi != nil { + var fib []byte + fib = appendStringField(fib, 1, fi.CallID) + fib = appendStringField(fib, 2, fi.Name) + fib = appendBytesField(fib, 3, EncodeProtoStruct(fi.Arguments)) + sub = appendBytesField(sub, 14, fib) + } + if req := p.CallHookRequest; req != nil { + if len(req.RawWire) > 0 { + sub = appendBytesField(sub, 15, req.RawWire) + } else { + var hrb []byte + hrb = appendStringField(hrb, 1, req.RequestID) + hrb = appendStringField(hrb, 2, req.Name) + if req.HookType != 0 { + hrb = appendVarintField(hrb, 7, uint64(req.HookType)) + } + sub = appendBytesField(sub, 15, hrb) + } + } + if resp := p.CallHookResponse; resp != nil { + if len(resp.RawWire) > 0 { + sub = appendBytesField(sub, 16, resp.RawWire) + } else { + var hresp []byte + hresp = appendStringField(hresp, 1, resp.RequestID) + if resp.EmptyResult { + hresp = appendBytesField(hresp, 4, nil) + } + hresp = appendStringField(hresp, 5, resp.ErrorMessage) + sub = appendBytesField(sub, 16, hresp) + } + } + if fr := p.FunctionResult; fr != nil { + var frb []byte + frb = appendStringField(frb, 1, fr.CallID) + frb = appendStringField(frb, 2, fr.Name) + frb = appendBoolField(frb, 3, fr.IsError) + frb = appendBytesField(frb, 4, EncodeProtoValue(fr.Result)) + sub = appendBytesField(sub, 18, frb) + } + return sub +} + +// MarshalServerEventBatch encodes b into protobuf wire format matching genai.ServerEventBatch. +func MarshalServerEventBatch(b *ServerEventBatch) []byte { + if b == nil { + return nil + } + var out []byte + for _, ev := range b.Events { + if ev != nil { + var evBytes []byte + if ev.Payload != nil { + evBytes = appendBytesField(evBytes, 3, marshalPayload(ev.Payload)) + } + if ev.SeqNum != 0 { + evBytes = appendVarintField(evBytes, 10, uint64(ev.SeqNum)) + } + if ev.TimestampMicros != 0 { + evBytes = appendVarintField(evBytes, 11, uint64(ev.TimestampMicros)) + } + out = appendBytesField(out, 1, evBytes) + } + } + if b.AckInputSeqNum != 0 { + out = appendVarintField(out, 2, uint64(b.AckInputSeqNum)) + } + if b.InteractionCompleted != nil { + var comp []byte + if st := b.InteractionCompleted.Status; st != nil { + var stb []byte + if st.Code != 0 { + stb = appendVarintField(stb, 1, uint64(st.Code)) + } + stb = appendStringField(stb, 2, st.Message) + comp = appendBytesField(comp, 1, stb) + } + out = appendBytesField(out, 6, comp) + } + if b.InteractionReconnected != nil { + var rec []byte + if b.InteractionReconnected.NextInputSeqNum != 0 { + rec = appendVarintField(rec, 1, uint64(b.InteractionReconnected.NextInputSeqNum)) + } + out = appendBytesField(out, 8, rec) + } + if b.ReconnectionRequired != nil { + out = appendBytesField(out, 9, nil) + } + if b.Yield != nil { + var yb []byte + if b.Yield.NextInputSeqNum != 0 { + yb = appendVarintField(yb, 1, uint64(b.Yield.NextInputSeqNum)) + } + out = appendBytesField(out, 10, yb) + } + return out +} + +// EncodeProtoValue encodes a Go value into protobuf wire format matching genai.Value. +func EncodeProtoValue(v any) []byte { + var out []byte + switch val := v.(type) { + case nil: + return appendVarintField(out, 1, 0) + case bool: + out = appendTag(out, 4, wireVarint) + if val { + return binary.AppendUvarint(out, 1) + } + return binary.AppendUvarint(out, 0) + case int: + return encodeNumberValue(float64(val)) + case int32: + return encodeNumberValue(float64(val)) + case int64: + return encodeNumberValue(float64(val)) + case uint: + return encodeNumberValue(float64(val)) + case uint32: + return encodeNumberValue(float64(val)) + case uint64: + return encodeNumberValue(float64(val)) + case float32: + return encodeNumberValue(float64(val)) + case float64: + return encodeNumberValue(val) + case string: + out = appendTag(out, 3, wireBytes) + out = binary.AppendUvarint(out, uint64(len(val))) + return append(out, val...) + case []any: + var listBuf []byte + for _, elem := range val { + listBuf = appendBytesField(listBuf, 1, EncodeProtoValue(elem)) + } + return appendBytesField(out, 6, listBuf) + case []string: + var listBuf []byte + for _, elem := range val { + listBuf = appendBytesField(listBuf, 1, EncodeProtoValue(elem)) + } + return appendBytesField(out, 6, listBuf) + case map[string]any: + return appendBytesField(out, 5, EncodeProtoStruct(val)) + default: + s := fmt.Sprint(val) + out = appendTag(out, 3, wireBytes) + out = binary.AppendUvarint(out, uint64(len(s))) + return append(out, s...) + } +} + +func encodeNumberValue(f float64) []byte { + out := appendTag(nil, 2, wireFixed64) + return binary.LittleEndian.AppendUint64(out, math.Float64bits(f)) +} + +// EncodeProtoStruct encodes a Go map into deterministic protobuf wire format matching genai.Struct. +func EncodeProtoStruct(m map[string]any) []byte { + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + sort.Strings(keys) + var out []byte + for _, k := range keys { + var entry []byte + entry = appendTag(entry, 1, wireBytes) + entry = binary.AppendUvarint(entry, uint64(len(k))) + entry = append(entry, k...) + entry = appendBytesField(entry, 2, EncodeProtoValue(m[k])) + out = appendBytesField(out, 1, entry) + } + return out +} + +func consumeVarint(b []byte) (uint64, int, error) { + v, n := binary.Uvarint(b) + if n == 0 { + return 0, 0, io.ErrUnexpectedEOF + } + if n < 0 { + return 0, 0, fmt.Errorf("protobuf varint overflows uint64") + } + return v, n, nil +} + +// ConsumeField decodes a single protobuf wire field from b, returning its field number, wire type, value, and byte length. +func ConsumeField(b []byte) (fieldNum int, wireType int, valVarint uint64, valBytes []byte, consumed int, err error) { + tag, n, err := consumeVarint(b) + if err != nil { + return 0, 0, 0, nil, 0, err + } + rawFieldNum := tag >> 3 + if rawFieldNum == 0 || rawFieldNum > (1<<29)-1 { + return 0, 0, 0, nil, 0, fmt.Errorf("invalid protobuf field number %d", rawFieldNum) + } + fieldNum = int(rawFieldNum) + wireType = int(tag & 0x7) + rest := b[n:] + switch wireType { + case wireVarint: + v, m, err := consumeVarint(rest) + if err != nil { + return 0, 0, 0, nil, 0, err + } + return fieldNum, wireType, v, nil, n + m, nil + case wireFixed64: + if len(rest) < 8 { + return 0, 0, 0, nil, 0, io.ErrUnexpectedEOF + } + return fieldNum, wireType, binary.LittleEndian.Uint64(rest[:8]), nil, n + 8, nil + case wireBytes: + l, m, err := consumeVarint(rest) + if err != nil { + return 0, 0, 0, nil, 0, err + } + after := rest[m:] + if l > uint64(len(after)) { + return 0, 0, 0, nil, 0, io.ErrUnexpectedEOF + } + return fieldNum, wireType, 0, after[:int(l)], n + m + int(l), nil + case wireFixed32: + if len(rest) < 4 { + return 0, 0, 0, nil, 0, io.ErrUnexpectedEOF + } + return fieldNum, wireType, uint64(binary.LittleEndian.Uint32(rest[:4])), nil, n + 4, nil + default: + return 0, 0, 0, nil, 0, fmt.Errorf("unsupported protobuf wire type %d", wireType) + } +} + +// UnmarshalServerEventBatch decodes protobuf wire bytes into a ServerEventBatch. +// All wire field numbers decoded below must stay synchronized with the canonical +// genai.ServerEventBatch protobuf definition. +func UnmarshalServerEventBatch(data []byte) (*ServerEventBatch, error) { + out := &ServerEventBatch{} + for len(data) > 0 { + num, wt, vv, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 1 && wt == wireBytes: + ev, err := unmarshalServerEvent(vb) + if err != nil { + return nil, err + } + out.Events = append(out.Events, ev) + case num == 2 && wt == wireVarint: + out.AckInputSeqNum = int64(vv) + case num == 6 && wt == wireBytes: + comp, err := unmarshalInteractionCompleted(vb) + if err != nil { + return nil, err + } + out.InteractionCompleted = comp + case num == 8 && wt == wireBytes: + rec, err := unmarshalInteractionReconnected(vb) + if err != nil { + return nil, err + } + out.InteractionReconnected = rec + case num == 9 && wt == wireBytes: + out.ReconnectionRequired = &ReconnectionRequiredEvent{} + case num == 10 && wt == wireBytes: + y, err := unmarshalYield(vb) + if err != nil { + return nil, err + } + out.Yield = y + } + } + return out, nil +} + +// UnmarshalClientEventBatch decodes protobuf wire bytes into a ClientEventBatch. +func UnmarshalClientEventBatch(data []byte) (*ClientEventBatch, error) { + out := &ClientEventBatch{} + for len(data) > 0 { + num, wt, vv, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 1 && wt == wireBytes: + ev, err := unmarshalClientEvent(vb) + if err != nil { + return nil, err + } + out.Events = append(out.Events, ev) + case num == 2 && wt == wireVarint: + out.AckOutputSeqNum = int64(vv) + case num == 8 && wt == wireBytes: + rec, err := unmarshalReconnectInteraction(vb) + if err != nil { + return nil, err + } + out.ReconnectInteraction = rec + case num == 9 && wt == wireBytes: + out.Disconnect = &DisconnectEvent{} + } + } + return out, nil +} + +func unmarshalReconnectInteraction(data []byte) (*ReconnectInteractionEvent, error) { + out := &ReconnectInteractionEvent{} + for len(data) > 0 { + num, wt, vv, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 1 && wt == wireBytes: + out.InteractionID = string(vb) + case num == 2 && wt == wireVarint: + out.NextOutputSeqNum = int64(vv) + case num == 3 && wt == wireVarint: + out.Background = vv != 0 + case (num == 5 || (num == 4 && out.ClientID == "")) && wt == wireBytes: + out.ClientID = string(vb) + } + } + return out, nil +} + +func unmarshalClientEvent(data []byte) (*ClientEvent, error) { + out := &ClientEvent{} + for len(data) > 0 { + num, wt, vv, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 3 && wt == wireBytes: + p, err := unmarshalPayload(vb) + if err != nil { + return nil, err + } + out.Payload = p + case num == 10 && wt == wireVarint: + out.SeqNum = int64(vv) + case num == 11 && wt == wireVarint: + out.TimestampMicros = int64(vv) + } + } + return out, nil +} + +func unmarshalInteractionCompleted(data []byte) (*InteractionCompletedEvent, error) { + out := &InteractionCompletedEvent{} + for len(data) > 0 { + num, wt, _, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + if num == 1 && wt == wireBytes { + st, err := unmarshalRPCStatus(vb) + if err != nil { + return nil, err + } + out.Status = st + } + } + return out, nil +} + +func unmarshalRPCStatus(data []byte) (*RPCStatus, error) { + out := &RPCStatus{} + for len(data) > 0 { + num, wt, vv, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 1 && wt == wireVarint: + out.Code = int32(vv) + case num == 2 && wt == wireBytes: + out.Message = string(vb) + } + } + return out, nil +} + +func unmarshalInteractionReconnected(data []byte) (*InteractionReconnectedEvent, error) { + out := &InteractionReconnectedEvent{} + var legacySeq int64 + for len(data) > 0 { + num, wt, vv, _, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 1 && wt == wireVarint: + out.NextInputSeqNum = int64(vv) + case num == 2 && wt == wireVarint: + legacySeq = int64(vv) + } + } + if out.NextInputSeqNum == 0 && legacySeq != 0 { + out.NextInputSeqNum = legacySeq + } + return out, nil +} + +func unmarshalYield(data []byte) (*YieldEvent, error) { + out := &YieldEvent{} + for len(data) > 0 { + num, wt, vv, _, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + if num == 1 && wt == wireVarint { + out.NextInputSeqNum = int64(vv) + } + } + return out, nil +} + +func unmarshalServerEvent(data []byte) (*ServerEvent, error) { + out := &ServerEvent{} + for len(data) > 0 { + num, wt, vv, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 3 && wt == wireBytes: + p, err := unmarshalPayload(vb) + if err != nil { + return nil, err + } + out.Payload = p + case num == 10 && wt == wireVarint: + out.SeqNum = int64(vv) + case num == 11 && wt == wireVarint: + out.TimestampMicros = int64(vv) + } + } + return out, nil +} + +func unmarshalPayload(data []byte) (*EventPayload, error) { + out := &EventPayload{} + for len(data) > 0 { + num, wt, _, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 6 && wt == wireBytes: + out.SubInteractionID = string(vb) + case num == 11 && wt == wireBytes: + st, err := unmarshalStep(vb) + if err != nil { + return nil, err + } + out.Step = st + case num == 14 && wt == wireBytes: + fi, err := unmarshalFunctionInvocation(vb) + if err != nil { + return nil, err + } + out.FunctionInvocation = fi + case num == 15 && wt == wireBytes: + hr, err := unmarshalCallHookRequest(vb) + if err != nil { + return nil, err + } + out.CallHookRequest = hr + case num == 16 && wt == wireBytes: + resp, err := unmarshalCallHookResponse(vb) + if err != nil { + return nil, err + } + out.CallHookResponse = resp + case num == 18 && wt == wireBytes: + fr, err := unmarshalFunctionResultEvent(vb) + if err != nil { + return nil, err + } + out.FunctionResult = fr + } + } + return out, nil +} + +func unmarshalStep(data []byte) (*Step, error) { + out := &Step{} + for len(data) > 0 { + num, wt, _, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 4 && wt == wireBytes: + tc, err := unmarshalToolCallStep(vb) + if err != nil { + return nil, err + } + out.ToolCall = tc + case num == 5 && wt == wireBytes: + tr, err := unmarshalToolResultStep(vb) + if err != nil { + return nil, err + } + out.ToolResult = tr + } + } + return out, nil +} + +func unmarshalToolCallStep(data []byte) (*ToolCallStep, error) { + out := &ToolCallStep{} + for len(data) > 0 { + num, wt, _, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 1 && wt == wireBytes: + out.ID = string(vb) + case num == 3 && wt == wireBytes: + fc, err := unmarshalFunctionCallStep(vb) + if err != nil { + return nil, err + } + out.FunctionCall = fc + } + } + return out, nil +} + +func unmarshalFunctionCallStep(data []byte) (*FunctionCallStep, error) { + out := &FunctionCallStep{} + for len(data) > 0 { + num, wt, _, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 1 && wt == wireBytes: + out.Name = string(vb) + case num == 2 && wt == wireBytes: + m, err := DecodeProtoStruct(vb, 0) + if err != nil { + return nil, err + } + out.Arguments = m + } + } + return out, nil +} + +func unmarshalToolResultStep(data []byte) (*ToolResultStep, error) { + out := &ToolResultStep{} + for len(data) > 0 { + num, wt, _, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 1 && wt == wireBytes: + out.CallID = string(vb) + case num == 3 && wt == wireBytes: + fr, err := unmarshalFunctionResultStep(vb) + if err != nil { + return nil, err + } + out.FunctionResult = fr + } + } + return out, nil +} + +func unmarshalFunctionResultStep(data []byte) (*FunctionResultStep, error) { + out := &FunctionResultStep{} + for len(data) > 0 { + num, wt, vv, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 2 && wt == wireBytes: + out.Name = string(vb) + case num == 3 && wt == wireVarint: + out.IsError = vv != 0 + case num == 7 && wt == wireBytes: + v, err := DecodeProtoValue(vb, 0) + if err != nil { + return nil, err + } + out.Result = v + } + } + return out, nil +} + +func unmarshalFunctionInvocation(data []byte) (*FunctionInvocation, error) { + out := &FunctionInvocation{} + for len(data) > 0 { + num, wt, _, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 1 && wt == wireBytes: + out.CallID = string(vb) + case num == 2 && wt == wireBytes: + out.Name = string(vb) + case num == 3 && wt == wireBytes: + m, err := DecodeProtoStruct(vb, 0) + if err != nil { + return nil, err + } + out.Arguments = m + } + } + return out, nil +} + +func unmarshalFunctionResultEvent(data []byte) (*FunctionResultEvent, error) { + out := &FunctionResultEvent{} + for len(data) > 0 { + num, wt, vv, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 1 && wt == wireBytes: + out.CallID = string(vb) + case num == 2 && wt == wireBytes: + out.Name = string(vb) + case num == 3 && wt == wireVarint: + out.IsError = vv != 0 + case num == 4 && wt == wireBytes: + v, err := DecodeProtoValue(vb, 0) + if err != nil { + return nil, err + } + out.Result = v + } + } + return out, nil +} + +func unmarshalCallHookRequest(data []byte) (*CallHookRequest, error) { + out := &CallHookRequest{ + RawWire: append([]byte(nil), data...), + } + for len(data) > 0 { + num, wt, vv, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 1 && wt == wireBytes: + out.RequestID = string(vb) + case num == 2 && wt == wireBytes: + out.Name = string(vb) + case num == 7 && wt == wireVarint: + out.HookType = int32(vv) + } + } + return out, nil +} + +func unmarshalCallHookResponse(data []byte) (*CallHookResponse, error) { + out := &CallHookResponse{ + RawWire: append([]byte(nil), data...), + } + for len(data) > 0 { + num, wt, _, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 1 && wt == wireBytes: + out.RequestID = string(vb) + case num == 4 && wt == wireBytes: + out.EmptyResult = true + case num == 5 && wt == wireBytes: + out.ErrorMessage = string(vb) + } + } + return out, nil +} + +// DecodeProtoStruct decodes protobuf wire bytes matching genai.Struct into a Go map. +func DecodeProtoStruct(data []byte, depth int) (map[string]any, error) { + if depth > MaxProtoNestingDepth { + return nil, fmt.Errorf("protobuf struct nesting exceeds maximum depth %d", MaxProtoNestingDepth) + } + out := make(map[string]any) + for len(data) > 0 { + num, wt, _, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + if num != 1 || wt != wireBytes { + continue + } + var entryKey string + var entryVal any + entryData := vb + for len(entryData) > 0 { + enum, ewt, _, evb, en, err := ConsumeField(entryData) + if err != nil { + return nil, err + } + entryData = entryData[en:] + switch { + case enum == 1 && ewt == wireBytes: + entryKey = string(evb) + case enum == 2 && ewt == wireBytes: + v, err := DecodeProtoValue(evb, depth+1) + if err != nil { + return nil, err + } + entryVal = v + } + } + out[entryKey] = entryVal + } + return out, nil +} + +// DecodeProtoValue decodes protobuf wire bytes matching genai.Value into a native Go value. +func DecodeProtoValue(data []byte, depth int) (any, error) { + if depth > MaxProtoNestingDepth { + return nil, fmt.Errorf("protobuf value nesting exceeds maximum depth %d", MaxProtoNestingDepth) + } + var res any + for len(data) > 0 { + num, wt, vv, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + switch { + case num == 1 && wt == wireVarint: + res = nil + case num == 2 && wt == wireFixed64: + res = math.Float64frombits(vv) + case num == 3 && wt == wireBytes: + res = string(vb) + case num == 4 && wt == wireVarint: + res = vv != 0 + case num == 5 && wt == wireBytes: + m, err := DecodeProtoStruct(vb, depth+1) + if err != nil { + return nil, err + } + res = m + case num == 6 && wt == wireBytes: + items, err := decodeProtoList(vb, depth+1) + if err != nil { + return nil, err + } + res = items + case num == 7 && wt == wireBytes: + txt, err := decodeContentText(vb) + if err != nil { + return nil, err + } + res = txt + } + } + return res, nil +} + +func decodeContentText(data []byte) (any, error) { + for len(data) > 0 { + num, wt, _, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + if num == 10 && wt == wireBytes { + inner := vb + var textVal string + for len(inner) > 0 { + inum, iwt, _, ivb, in, ierr := ConsumeField(inner) + if ierr != nil { + return nil, ierr + } + inner = inner[in:] + if inum == 3 && iwt == wireBytes { + textVal = string(ivb) + } + } + return textVal, nil + } + } + return nil, nil +} + +func decodeProtoList(data []byte, depth int) ([]any, error) { + if depth > MaxProtoNestingDepth { + return nil, fmt.Errorf("protobuf list nesting exceeds maximum depth %d", MaxProtoNestingDepth) + } + var out []any + for len(data) > 0 { + num, wt, _, vb, n, err := ConsumeField(data) + if err != nil { + return nil, err + } + data = data[n:] + if num == 1 && wt == wireBytes { + v, err := DecodeProtoValue(vb, depth+1) + if err != nil { + return nil, err + } + out = append(out, v) + } + } + return out, nil +} diff --git a/internal/cli/custom/envruntime/wire_test.go b/internal/cli/custom/envruntime/wire_test.go new file mode 100644 index 0000000..cff1b94 --- /dev/null +++ b/internal/cli/custom/envruntime/wire_test.go @@ -0,0 +1,722 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package envruntime + +import ( + "bytes" + "encoding/binary" + "errors" + "io" + "math" + "reflect" + "strings" + "testing" +) + +type failAfterNWriter struct { + limit int + wrote int + err error +} + +func (w *failAfterNWriter) Write(p []byte) (int, error) { + if w.wrote+len(p) > w.limit { + return 0, w.err + } + w.wrote += len(p) + return len(p), nil +} + +type errReader struct { + prefix []byte + err error + used bool +} + +func (r *errReader) Read(p []byte) (int, error) { + if !r.used && len(r.prefix) > 0 { + r.used = true + n := copy(p, r.prefix) + return n, nil + } + return 0, r.err +} + +type customStringer struct { + label string +} + +func (s customStringer) String() string { + return s.label +} + +func TestGRPCFramingAndDecodeMessage(t *testing.T) { + var buf bytes.Buffer + wantPayload := []byte("hello-grpc-frame") + if err := GRPCWriteFrame(&buf, wantPayload); err != nil { + t.Fatalf("GRPCWriteFrame(non-empty) failed: %v", err) + } + if err := GRPCWriteFrame(&buf, nil); err != nil { + t.Fatalf("GRPCWriteFrame(empty) failed: %v", err) + } + + got1, err := GRPCReadFrame(&buf) + if err != nil || !bytes.Equal(got1, wantPayload) { + t.Fatalf("GRPCReadFrame(1) = (%q, %v), want (%q, nil)", got1, err, wantPayload) + } + got2, err := GRPCReadFrame(&buf) + if err != nil || len(got2) != 0 { + t.Fatalf("GRPCReadFrame(2) = (%x, %v), want (empty, nil)", got2, err) + } + + oversized := make([]byte, MaxGRPCFrameBytes+1) + if err := GRPCWriteFrame(&buf, oversized); err == nil { + t.Error("GRPCWriteFrame(MaxGRPCFrameBytes+1) = nil, want error") + } + + writeBoom := errors.New("write boom") + if err := GRPCWriteFrame(&failAfterNWriter{limit: 2, err: writeBoom}, []byte("abc")); !errors.Is(err, writeBoom) { + t.Errorf("GRPCWriteFrame(header fail) = %v, want %v", err, writeBoom) + } + if err := GRPCWriteFrame(&failAfterNWriter{limit: 5, err: writeBoom}, []byte("abc")); !errors.Is(err, writeBoom) { + t.Errorf("GRPCWriteFrame(payload fail) = %v, want %v", err, writeBoom) + } + + var compressedHdr [5]byte + compressedHdr[0] = 1 + if _, err := GRPCReadFrame(bytes.NewReader(compressedHdr[:])); err == nil { + t.Error("GRPCReadFrame(compressed) = nil, want error") + } + + var oversizedHdr [5]byte + binary.BigEndian.PutUint32(oversizedHdr[1:5], MaxGRPCFrameBytes+1) + if _, err := GRPCReadFrame(bytes.NewReader(oversizedHdr[:])); err == nil { + t.Error("GRPCReadFrame(oversized) = nil, want error") + } + + if _, err := GRPCReadFrame(bytes.NewReader([]byte{0, 0, 1})); err == nil { + t.Error("GRPCReadFrame(short header) = nil, want error") + } + + var truncatedHdr [5]byte + binary.BigEndian.PutUint32(truncatedHdr[1:5], 10) + if _, err := GRPCReadFrame(bytes.NewReader(append(truncatedHdr[:], 1, 2, 3))); !errors.Is(err, io.ErrUnexpectedEOF) { + t.Errorf("GRPCReadFrame(truncated payload) = %v, want ErrUnexpectedEOF", err) + } + + readBoom := errors.New("read boom") + var validHdr [5]byte + binary.BigEndian.PutUint32(validHdr[1:5], 4) + if _, err := GRPCReadFrame(&errReader{prefix: validHdr[:], err: readBoom}); !errors.Is(err, readBoom) { + t.Errorf("GRPCReadFrame(payload read error) = %v, want %v", err, readBoom) + } + + if got := DecodeGRPCMessage(" invalid%20lease%20token "); got != "invalid lease token" { + t.Errorf("DecodeGRPCMessage(escaped) = %q, want %q", got, "invalid lease token") + } + if got := DecodeGRPCMessage(" bad%zzescape "); got != "bad%zzescape" { + t.Errorf("DecodeGRPCMessage(malformed escape) = %q, want %q", got, "bad%zzescape") + } +} + +func TestMarshalAndUnmarshalClientEventBatch(t *testing.T) { + if got := MarshalClientEventBatch(nil); got != nil { + t.Errorf("MarshalClientEventBatch(nil) = %x, want nil", got) + } + if got := MarshalClientEventBatch(&ClientEventBatch{}); len(got) != 0 { + t.Errorf("MarshalClientEventBatch(empty) = %x, want empty", got) + } + if got := MarshalClientEvent(nil); got != nil { + t.Errorf("MarshalClientEvent(nil) = %x, want nil", got) + } + if got := MarshalClientEvent(&ClientEvent{}); len(got) != 0 { + t.Errorf("MarshalClientEvent(empty) = %x, want empty", got) + } + + emptyStep := MarshalClientEvent(&ClientEvent{Payload: &EventPayload{Step: &Step{}}}) + if !bytes.Equal(emptyStep, []byte{0x1a, 0x00}) { + t.Errorf("MarshalClientEvent(empty Step) = %x, want 1a00", emptyStep) + } + noFnStep := MarshalClientEvent(&ClientEvent{Payload: &EventPayload{Step: &Step{ToolCall: &ToolCallStep{ID: "tc_0"}}}}) + if len(noFnStep) == 0 { + t.Error("MarshalClientEvent(ToolCall with nil FunctionCall) returned empty") + } + + batch := &ClientEventBatch{ + AckOutputSeqNum: 7, + ReconnectInteraction: &ReconnectInteractionEvent{ + InteractionID: "int_123", + NextOutputSeqNum: 8, + Background: true, + ClientID: DefaultEnvironmentClientID, + }, + Disconnect: &DisconnectEvent{}, + Events: []*ClientEvent{ + nil, + { + SeqNum: 1, + TimestampMicros: 1700000001, + Payload: &EventPayload{ + SubInteractionID: "sub_1", + Step: &Step{ + ToolCall: &ToolCallStep{ + ID: "tc_1", + FunctionCall: &FunctionCallStep{ + Name: "view_file", + Arguments: map[string]any{"path": "README.md"}, + }, + }, + ToolResult: &ToolResultStep{ + CallID: "tc_1", + FunctionResult: &FunctionResultStep{ + Name: "view_file", + IsError: true, + Result: map[string]any{"error": "missing"}, + }, + }, + }, + }, + }, + { + SeqNum: 2, + Payload: &EventPayload{ + FunctionInvocation: &FunctionInvocation{ + CallID: "fi_1", + Name: "list_dir", + Arguments: map[string]any{"dir": "."}, + }, + FunctionResult: &FunctionResultEvent{ + CallID: "fi_1", + Name: "list_dir", + IsError: true, + Result: map[string]any{"entries": []any{"a.txt"}}, + }, + }, + }, + { + SeqNum: 3, + Payload: &EventPayload{ + CallHookRequest: &CallHookRequest{ + RequestID: "hr_1", + Name: "pre_tool", + HookType: 2, + }, + CallHookResponse: &CallHookResponse{ + RequestID: "hr_1", + EmptyResult: true, + ErrorMessage: "hook err", + }, + }, + }, + { + SeqNum: 4, + Payload: &EventPayload{ + CallHookRequest: &CallHookRequest{RawWire: appendStringField(appendStringField(nil, 1, "hr_raw"), 2, "raw_hook")}, + CallHookResponse: &CallHookResponse{RawWire: appendBoolField(appendStringField(nil, 1, "hr_raw"), 2, true)}, + }, + }, + }, + } + + raw := MarshalClientEventBatch(batch) + + // Verify ReconnectInteraction encodes ClientID on field 5 on the wire. + rest := raw + sawReconnectClientIDField5 := false + for len(rest) > 0 { + num, _, _, sub, n, err := ConsumeField(rest) + if err != nil { + t.Fatalf("ConsumeField(ClientEventBatch) failed: %v", err) + } + if num == 8 { + inner := sub + for len(inner) > 0 { + inum, _, _, ib, in, ierr := ConsumeField(inner) + if ierr != nil { + t.Fatalf("ConsumeField(ReconnectInteraction) failed: %v", ierr) + } + if inum == 5 && string(ib) == DefaultEnvironmentClientID { + sawReconnectClientIDField5 = true + } + if inum == 4 { + t.Errorf("ReconnectInteraction unexpectedly wrote field 4 on wire") + } + inner = inner[in:] + } + } + rest = rest[n:] + } + if !sawReconnectClientIDField5 { + t.Error("MarshalClientEventBatch did not encode ReconnectInteraction.ClientID on field 5") + } + + decoded, err := UnmarshalClientEventBatch(raw) + if err != nil { + t.Fatalf("UnmarshalClientEventBatch() failed: %v", err) + } + if decoded.AckOutputSeqNum != 7 || decoded.Disconnect == nil || decoded.ReconnectInteraction == nil { + t.Fatalf("decoded control fields mismatch: %+v", decoded) + } + if got := decoded.ReconnectInteraction; got.InteractionID != "int_123" || got.NextOutputSeqNum != 8 || !got.Background || got.ClientID != DefaultEnvironmentClientID { + t.Errorf("decoded ReconnectInteraction = %+v", got) + } + if len(decoded.Events) != 4 { + t.Fatalf("len(decoded.Events) = %d, want 4", len(decoded.Events)) + } + if decoded.Events[0].SeqNum != 1 || decoded.Events[0].TimestampMicros != 1700000001 { + t.Errorf("decoded.Events[0] seq/ts = (%d, %d), want (1, 1700000001)", decoded.Events[0].SeqNum, decoded.Events[0].TimestampMicros) + } + if got := decoded.Events[0].Payload; got.SubInteractionID != "sub_1" || + got.Step.ToolCall.ID != "tc_1" || + got.Step.ToolCall.FunctionCall.Name != "view_file" || + got.Step.ToolCall.FunctionCall.Arguments["path"] != "README.md" || + got.Step.ToolResult.CallID != "tc_1" || + got.Step.ToolResult.FunctionResult.Name != "view_file" || + !got.Step.ToolResult.FunctionResult.IsError || + !reflect.DeepEqual(got.Step.ToolResult.FunctionResult.Result, map[string]any{"error": "missing"}) { + t.Errorf("decoded Events[0] = %+v", got) + } + if got := decoded.Events[1].Payload; got.FunctionInvocation.CallID != "fi_1" || + got.FunctionInvocation.Name != "list_dir" || + got.FunctionInvocation.Arguments["dir"] != "." || + got.FunctionResult.CallID != "fi_1" || + got.FunctionResult.Name != "list_dir" || + !got.FunctionResult.IsError || + !reflect.DeepEqual(got.FunctionResult.Result, map[string]any{"entries": []any{"a.txt"}}) { + t.Errorf("decoded Events[1] = %+v", got) + } + if got := decoded.Events[2].Payload; got.CallHookRequest.RequestID != "hr_1" || + got.CallHookRequest.Name != "pre_tool" || + got.CallHookRequest.HookType != 2 || + got.CallHookResponse.RequestID != "hr_1" || + !got.CallHookResponse.EmptyResult || + got.CallHookResponse.ErrorMessage != "hook err" { + t.Errorf("decoded Events[2] = %+v", got) + } + if got := decoded.Events[3].Payload; got.CallHookRequest.RequestID != "hr_raw" || got.CallHookResponse.RequestID != "hr_raw" { + t.Errorf("decoded Events[3] RawWire = %+v", got) + } + + // Verify legacy field 4 ClientID fallback and field 5 precedence. + legacyReconnect := appendStringField(nil, 1, "int_legacy") + legacyReconnect = appendStringField(legacyReconnect, 4, "legacy-client") + legacyBatch, err := UnmarshalClientEventBatch(appendBytesField(nil, 8, legacyReconnect)) + if err != nil || legacyBatch.ReconnectInteraction.ClientID != "legacy-client" { + t.Errorf("legacy field 4 ReconnectInteraction = (%+v, %v), want ClientID=legacy-client", legacyBatch, err) + } + bothReconnect := appendStringField(legacyReconnect, 5, "canonical-client") + bothBatch, err := UnmarshalClientEventBatch(appendBytesField(nil, 8, bothReconnect)) + if err != nil || bothBatch.ReconnectInteraction.ClientID != "canonical-client" { + t.Errorf("field 5 precedence ReconnectInteraction = (%+v, %v), want ClientID=canonical-client", bothBatch, err) + } + + if _, err := UnmarshalClientEventBatch([]byte{0x0a, 0x05, 0x01}); err == nil { + t.Error("UnmarshalClientEventBatch(truncated) = nil, want error") + } +} + +func TestMarshalAndUnmarshalServerEventBatchAndHandshakeCheck(t *testing.T) { + if got := MarshalServerEventBatch(nil); got != nil { + t.Errorf("MarshalServerEventBatch(nil) = %x, want nil", got) + } + + batch := &ServerEventBatch{ + AckInputSeqNum: 11, + InteractionCompleted: &InteractionCompletedEvent{ + Status: &RPCStatus{Code: 13, Message: "backend failure"}, + }, + InteractionReconnected: &InteractionReconnectedEvent{NextInputSeqNum: 12}, + ReconnectionRequired: &ReconnectionRequiredEvent{}, + Yield: &YieldEvent{NextInputSeqNum: 15}, + Events: []*ServerEvent{ + nil, + { + SeqNum: 9, + TimestampMicros: 1700000099, + Payload: &EventPayload{ + Step: &Step{ + ToolCall: &ToolCallStep{ + ID: "call_9", + FunctionCall: &FunctionCallStep{ + Name: "execute_bash", + Arguments: map[string]any{"command": "pwd"}, + }, + }, + }, + }, + }, + }, + } + + raw := MarshalServerEventBatch(batch) + decoded, err := UnmarshalServerEventBatch(raw) + if err != nil { + t.Fatalf("UnmarshalServerEventBatch() failed: %v", err) + } + if decoded.AckInputSeqNum != 11 || decoded.ReconnectionRequired == nil || decoded.Yield == nil || decoded.Yield.NextInputSeqNum != 15 { + t.Fatalf("decoded ServerEventBatch mismatch: %+v", decoded) + } + if decoded.InteractionCompleted == nil || decoded.InteractionCompleted.Status == nil || decoded.InteractionCompleted.Status.Code != 13 || decoded.InteractionCompleted.Status.Message != "backend failure" { + t.Errorf("decoded InteractionCompleted = %+v", decoded.InteractionCompleted) + } + if decoded.InteractionReconnected == nil || decoded.InteractionReconnected.NextInputSeqNum != 12 { + t.Errorf("decoded InteractionReconnected = %+v", decoded.InteractionReconnected) + } + if len(decoded.Events) != 1 || + decoded.Events[0].SeqNum != 9 || + decoded.Events[0].TimestampMicros != 1700000099 || + decoded.Events[0].Payload.Step.ToolCall.FunctionCall.Name != "execute_bash" || + decoded.Events[0].Payload.Step.ToolCall.FunctionCall.Arguments["command"] != "pwd" { + t.Errorf("decoded Events = %+v", decoded.Events) + } + + // Verify ServerEvent with zero SeqNum and TimestampMicros does not emit non-zero values. + zeroEvBatch, err := UnmarshalServerEventBatch(MarshalServerEventBatch(&ServerEventBatch{ + Events: []*ServerEvent{{Payload: &EventPayload{SubInteractionID: "s0"}}}, + })) + if err != nil || len(zeroEvBatch.Events) != 1 || zeroEvBatch.Events[0].SeqNum != 0 || zeroEvBatch.Events[0].TimestampMicros != 0 { + t.Errorf("zero SeqNum/TimestampMicros ServerEvent = (%+v, %v)", zeroEvBatch, err) + } + zeroCodeWire := MarshalServerEventBatch(&ServerEventBatch{ + InteractionCompleted: &InteractionCompletedEvent{Status: &RPCStatus{Code: 0, Message: "ok"}}, + }) + if bytes.Contains(zeroCodeWire, []byte{0x08, 0x00}) { + t.Errorf("MarshalServerEventBatch(Status.Code=0) unexpectedly wrote varint field 1: %x", zeroCodeWire) + } + zeroRecWire := MarshalServerEventBatch(&ServerEventBatch{ + InteractionReconnected: &InteractionReconnectedEvent{NextInputSeqNum: 0}, + }) + if !bytes.Equal(zeroRecWire, []byte{0x42, 0x00}) { + t.Errorf("MarshalServerEventBatch(InteractionReconnected NextInputSeqNum=0) = %x, want 4200", zeroRecWire) + } + + // Verify legacy field 2 fallback and field 1 precedence on InteractionReconnected, plus unknown field tolerance on Completed/Yield. + legacyReconnected := appendBytesField(nil, 8, appendVarintField(nil, 2, 42)) + decLegacy, err := UnmarshalServerEventBatch(legacyReconnected) + if err != nil || decLegacy.InteractionReconnected == nil || decLegacy.InteractionReconnected.NextInputSeqNum != 42 { + t.Errorf("legacy field 2 InteractionReconnected = (%+v, %v), want NextInputSeqNum=42", decLegacy, err) + } + bothReconnected := appendBytesField(nil, 8, appendVarintField(appendVarintField(nil, 1, 12), 2, 42)) + decBothRec, err := UnmarshalServerEventBatch(bothReconnected) + if err != nil || decBothRec.InteractionReconnected == nil || decBothRec.InteractionReconnected.NextInputSeqNum != 12 { + t.Errorf("field 1+2 InteractionReconnected = (%+v, %v), want NextInputSeqNum=12", decBothRec, err) + } + compWithExtra := appendBytesField(nil, 1, appendStringField(appendVarintField(nil, 1, 7), 2, "msg")) + compWithExtra = appendStringField(compWithExtra, 99, "ignored-extra-bytes") + yieldWithExtra := appendVarintField(appendVarintField(nil, 1, 15), 99, 999) + extraBatchWire := append(appendBytesField(nil, 6, compWithExtra), appendBytesField(nil, 10, yieldWithExtra)...) + decExtra, err := UnmarshalServerEventBatch(extraBatchWire) + if err != nil || decExtra.InteractionCompleted.Status.Code != 7 || decExtra.Yield.NextInputSeqNum != 15 { + t.Errorf("UnmarshalServerEventBatch(extra fields) = (%+v, %v)", decExtra, err) + } + + negStatusWire := MarshalServerEventBatch(&ServerEventBatch{ + InteractionCompleted: &InteractionCompletedEvent{Status: &RPCStatus{Code: -1, Message: "neg-code"}}, + }) + decNeg, err := UnmarshalServerEventBatch(negStatusWire) + if err != nil || decNeg.InteractionCompleted == nil || decNeg.InteractionCompleted.Status == nil || decNeg.InteractionCompleted.Status.Code != -1 { + t.Errorf("negative RPCStatus.Code round-trip = (%+v, %v), want Code=-1", decNeg, err) + } + // Standard protobuf negative int32 varint uses 10 bytes with sign extension, not a 5-byte uint32 varint. + if len(negStatusWire) < 12 { + t.Errorf("MarshalServerEventBatch(Status.Code=-1) wire len = %d (%x), want >= 12 bytes with 10-byte int32 varint", len(negStatusWire), negStatusWire) + } + + if _, err := UnmarshalServerEventBatch([]byte{0x0a, 0x09, 0x01}); err == nil { + t.Error("UnmarshalServerEventBatch(truncated) = nil, want error") + } +} + +func TestProtoStructAndValueRoundTripAndDepthLimits(t *testing.T) { + input := map[string]any{ + "nil_val": nil, + "bool_true": true, + "bool_false": false, + "int_val": int(-7), + "int32_val": int32(12), + "int64_val": int64(99), + "uint_val": uint(5), + "f32_val": float32(1.5), + "f64_val": float64(2.25), + "str_val": "hello", + "str_slice": []string{"x", "y"}, + "any_slice": []any{"a", float64(3)}, + "nested": map[string]any{"inner": "ok"}, + "stringer": customStringer{label: "custom-val"}, + } + + encoded1 := EncodeProtoStruct(input) + encoded2 := EncodeProtoStruct(input) + if !bytes.Equal(encoded1, encoded2) { + t.Error("EncodeProtoStruct is not deterministic across invocations") + } + orderedKeys := []string{"any_slice", "bool_false", "bool_true", "f32_val", "f64_val", "int32_val", "int64_val", "int_val", "nested", "nil_val", "str_slice", "str_val", "stringer", "uint_val"} + lastPos := -1 + for _, k := range orderedKeys { + pos := bytes.Index(encoded1, []byte(k)) + if pos <= lastPos { + t.Fatalf("EncodeProtoStruct key %q at byte %d is not in sorted order (lastPos=%d)", k, pos, lastPos) + } + lastPos = pos + } + + decoded, err := DecodeProtoStruct(encoded1, 0) + if err != nil { + t.Fatalf("DecodeProtoStruct() failed: %v", err) + } + if decoded["bool_true"] != true || decoded["bool_false"] != false || decoded["nil_val"] != nil { + t.Errorf("decoded booleans/nil mismatch: %v", decoded) + } + if decoded["int_val"] != float64(-7) || decoded["uint_val"] != float64(5) || decoded["f32_val"] != float64(1.5) { + t.Errorf("decoded numbers mismatch: %v", decoded) + } + if decoded["stringer"] != "custom-val" { + t.Errorf("decoded stringer = %v, want \"custom-val\"", decoded["stringer"]) + } + if !reflect.DeepEqual(decoded["str_slice"], []any{"x", "y"}) { + t.Errorf("decoded str_slice = %v, want [x y]", decoded["str_slice"]) + } + if !reflect.DeepEqual(decoded["any_slice"], []any{"a", float64(3)}) { + t.Errorf("decoded any_slice = %v, want [a 3]", decoded["any_slice"]) + } + if !reflect.DeepEqual(decoded["nested"], map[string]any{"inner": "ok"}) { + t.Errorf("decoded nested = %v, want map[inner:ok]", decoded["nested"]) + } + // Verify oneof override where null_value (tag 1) follows string_value (tag 3). + overriddenNull, err := DecodeProtoValue(append(EncodeProtoValue("old"), EncodeProtoValue(nil)...), 0) + if err != nil || overriddenNull != nil { + t.Errorf("DecodeProtoValue(string then null) = (%v, %v), want (nil, nil)", overriddenNull, err) + } + + // Verify tag 7 Content -> tag 10 TextContent -> tag 3 text string decoding (ignoring extra fields). + textContent := appendStringField(nil, 1, "ignored-inner-1") + textContent = appendStringField(textContent, 3, "extracted-content-text") + textContent = appendStringField(textContent, 9, "ignored-inner-9") + contentMsg := appendStringField(nil, 1, "ignored-outer-1") + contentMsg = appendBytesField(contentMsg, 10, textContent) + contentVal := appendBytesField(nil, 7, contentMsg) + gotText, err := DecodeProtoValue(contentVal, 0) + if err != nil || gotText != "extracted-content-text" { + t.Errorf("DecodeProtoValue(content_value) = (%v, %v), want (\"extracted-content-text\", nil)", gotText, err) + } + + // Verify exact MaxProtoNestingDepth boundary and depth+1 propagation across Struct, Value, and ListValue. + if _, err := DecodeProtoStruct(nil, MaxProtoNestingDepth); err != nil { + t.Errorf("DecodeProtoStruct(depth=MaxProtoNestingDepth) = %v, want nil", err) + } + if _, err := DecodeProtoStruct(nil, MaxProtoNestingDepth+1); err == nil { + t.Error("DecodeProtoStruct(depth=MaxProtoNestingDepth+1) = nil, want error") + } + structWithScalar := EncodeProtoStruct(map[string]any{"k": "v"}) + if _, err := DecodeProtoStruct(structWithScalar, MaxProtoNestingDepth-1); err != nil { + t.Errorf("DecodeProtoStruct(structWithScalar, MaxProtoNestingDepth-1) = %v, want nil", err) + } + if _, err := DecodeProtoStruct(structWithScalar, MaxProtoNestingDepth); err == nil { + t.Error("DecodeProtoStruct(structWithScalar, MaxProtoNestingDepth) = nil, want depth error") + } + valWithStruct := EncodeProtoValue(map[string]any{"k": "v"}) + if _, err := DecodeProtoValue(valWithStruct, MaxProtoNestingDepth-2); err != nil { + t.Errorf("DecodeProtoValue(valWithStruct, MaxProtoNestingDepth-2) = %v, want nil", err) + } + if _, err := DecodeProtoValue(valWithStruct, MaxProtoNestingDepth-1); err == nil { + t.Error("DecodeProtoValue(valWithStruct, MaxProtoNestingDepth-1) = nil, want depth error") + } + valWithEmptyList := EncodeProtoValue([]any{}) + if _, err := DecodeProtoValue(valWithEmptyList, MaxProtoNestingDepth-1); err != nil { + t.Errorf("DecodeProtoValue(valWithEmptyList, MaxProtoNestingDepth-1) = %v, want nil", err) + } + if _, err := DecodeProtoValue(valWithEmptyList, MaxProtoNestingDepth); err == nil { + t.Error("DecodeProtoValue(valWithEmptyList, MaxProtoNestingDepth) = nil, want depth error") + } + valWithOneElemList := EncodeProtoValue([]any{"leaf"}) + if _, err := DecodeProtoValue(valWithOneElemList, MaxProtoNestingDepth-2); err != nil { + t.Errorf("DecodeProtoValue(valWithOneElemList, MaxProtoNestingDepth-2) = %v, want nil", err) + } + if _, err := DecodeProtoValue(valWithOneElemList, MaxProtoNestingDepth-1); err == nil { + t.Error("DecodeProtoValue(valWithOneElemList, MaxProtoNestingDepth-1) = nil, want depth error") + } + scalarVal := EncodeProtoValue("leaf") + if got, err := DecodeProtoValue(scalarVal, MaxProtoNestingDepth); err != nil || got != "leaf" { + t.Errorf("DecodeProtoValue(depth=MaxProtoNestingDepth) = (%v, %v), want (\"leaf\", nil)", got, err) + } + if _, err := DecodeProtoValue(scalarVal, MaxProtoNestingDepth+1); err == nil { + t.Error("DecodeProtoValue(depth=MaxProtoNestingDepth+1) = nil, want error") + } + + // Verify CallHookRequest with HookType==0 omits field 7 on the wire, and decodeProtoList ignores + // field 1 when wireType != wireBytes and wireBytes when field != 1. + hookZeroBytes := MarshalClientEvent(&ClientEvent{ + Payload: &EventPayload{ + CallHookRequest: &CallHookRequest{RequestID: "r", Name: "n", HookType: 0}, + }, + }) + if bytes.Contains(hookZeroBytes, []byte{0x38, 0x00}) { + t.Errorf("MarshalClientEvent(CallHookRequest HookType=0) encoded zero field 7: %x", hookZeroBytes) + } + + var listWithIgnoredFields []byte + listWithIgnoredFields = appendBytesField(listWithIgnoredFields, 1, EncodeProtoValue("only-item")) + listWithIgnoredFields = appendVarintField(listWithIgnoredFields, 1, 99) + listWithIgnoredFields = appendBytesField(listWithIgnoredFields, 2, []byte{0xFF}) + if gotList, err := DecodeProtoValue(appendBytesField(nil, 6, listWithIgnoredFields), 0); err != nil || !reflect.DeepEqual(gotList, []any{"only-item"}) { + t.Errorf("DecodeProtoValue(listWithIgnoredFields) = (%v, %v), want [only-item]", gotList, err) + } + + var deepStruct []byte + for i := 0; i < MaxProtoNestingDepth+2; i++ { + val := appendBytesField(nil, 5, deepStruct) + entry := appendStringField(nil, 1, "k") + entry = appendBytesField(entry, 2, val) + deepStruct = appendBytesField(nil, 1, entry) + } + if _, err := DecodeProtoStruct(deepStruct, 0); err == nil || !strings.Contains(err.Error(), "maximum depth") { + t.Errorf("DecodeProtoStruct(deepStruct) = %v, want maximum depth error", err) + } +} + +func TestConsumeFieldWireTypesAndMalformedInputs(t *testing.T) { + // wireFixed64 (type 1) + var f64 []byte + f64 = appendTag(f64, 2, wireFixed64) + f64 = binary.LittleEndian.AppendUint64(f64, math.Float64bits(3.5)) + num, wt, val, _, n, err := ConsumeField(f64) + if err != nil || num != 2 || wt != wireFixed64 || math.Float64frombits(val) != 3.5 || n != len(f64) { + t.Errorf("ConsumeField(wireFixed64) = (%d, %d, %d, %d, %v)", num, wt, val, n, err) + } + if _, _, _, _, _, err := ConsumeField(f64[:len(f64)-1]); !errors.Is(err, io.ErrUnexpectedEOF) { + t.Errorf("ConsumeField(truncated fixed64) = %v, want ErrUnexpectedEOF", err) + } + + // wireFixed32 (type 5) + var f32 []byte + f32 = appendTag(f32, 3, wireFixed32) + f32 = binary.LittleEndian.AppendUint32(f32, 0x12345678) + num, wt, val, _, n, err = ConsumeField(f32) + if err != nil || num != 3 || wt != wireFixed32 || val != 0x12345678 || n != len(f32) { + t.Errorf("ConsumeField(wireFixed32) = (%d, %d, %x, %d, %v)", num, wt, val, n, err) + } + if _, _, _, _, _, err := ConsumeField(f32[:len(f32)-1]); !errors.Is(err, io.ErrUnexpectedEOF) { + t.Errorf("ConsumeField(truncated fixed32) = %v, want ErrUnexpectedEOF", err) + } + + // Field 0, field > (1<<29)-1, unsupported wire type 3, truncated varint, overflow varint. + maxValidFieldWire := appendVarintField(nil, (1<<29)-1, 7) + if num, _, val, _, _, err := ConsumeField(maxValidFieldWire); err != nil || num != (1<<29)-1 || val != 7 { + t.Errorf("ConsumeField(max valid field (1<<29)-1) = (num=%d, val=%d, err=%v)", num, val, err) + } + var tooLargeFieldWire []byte + tooLargeFieldWire = binary.AppendUvarint(tooLargeFieldWire, uint64(1<<29)<<3) + tooLargeFieldWire = append(tooLargeFieldWire, 0x01) + badInputs := [][]byte{ + {0x00, 0x01}, + tooLargeFieldWire, + {0x0b}, + {0x08, 0x80}, + {0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x02}, + } + for i, in := range badInputs { + if _, _, _, _, _, err := ConsumeField(in); err == nil { + t.Errorf("ConsumeField(badInputs[%d]=%x) = nil, want error", i, in) + } + } +} + +func TestEventPayloadCloneAndNopObserver(t *testing.T) { + var nilPayload *EventPayload + if nilPayload.Clone() != nil { + t.Error("(*EventPayload)(nil).Clone() != nil") + } + + // Verify Clone handles nil FunctionCall and nil FunctionResult inside Step without panicking. + nilStepChildren := (&EventPayload{ + Step: &Step{ + ToolCall: &ToolCallStep{ID: "c0"}, + ToolResult: &ToolResultStep{CallID: "c0"}, + }, + }).Clone() + if nilStepChildren.Step.ToolCall.FunctionCall != nil || nilStepChildren.Step.ToolResult.FunctionResult != nil { + t.Errorf("Clone(nil FunctionCall/FunctionResult) = %+v, want nil children", nilStepChildren.Step) + } + + orig := &EventPayload{ + SubInteractionID: "sub", + Step: &Step{ + ToolCall: &ToolCallStep{ + ID: "c1", + FunctionCall: &FunctionCallStep{ + Name: "fn", + Arguments: map[string]any{"nested": map[string]any{"k": "v"}, "arr": []any{map[string]any{"inner": "x"}}}, + }, + }, + ToolResult: &ToolResultStep{ + CallID: "c1", + FunctionResult: &FunctionResultStep{ + Name: "fn", + Result: map[string]any{"out": "ok"}, + }, + }, + }, + FunctionInvocation: &FunctionInvocation{ + CallID: "f1", + Name: "fn2", + Arguments: map[string]any{"a": "b"}, + }, + FunctionResult: &FunctionResultEvent{ + CallID: "f1", + Name: "fn2", + Result: []string{"s1"}, + }, + CallHookRequest: &CallHookRequest{RequestID: "h1", RawWire: []byte{1, 2}}, + CallHookResponse: &CallHookResponse{RequestID: "h1", RawWire: []byte{3, 4}}, + } + + cloned := orig.Clone() + cloned.Step.ToolCall.FunctionCall.Arguments["nested"].(map[string]any)["k"] = "mutated" + cloned.Step.ToolCall.FunctionCall.Arguments["arr"].([]any)[0].(map[string]any)["inner"] = "mutated" + cloned.Step.ToolResult.FunctionResult.Name = "mutated_fn" + cloned.Step.ToolResult.FunctionResult.Result.(map[string]any)["out"] = "mutated" + cloned.FunctionInvocation.Name = "mutated_fn2" + cloned.FunctionInvocation.Arguments["a"] = "mutated" + cloned.FunctionResult.Name = "mutated_fn2" + cloned.FunctionResult.Result.([]string)[0] = "mutated" + cloned.CallHookRequest.RawWire[0] = 99 + cloned.CallHookResponse.RawWire[0] = 88 + + if orig.Step.ToolCall.FunctionCall.Arguments["nested"].(map[string]any)["k"] != "v" { + t.Error("Clone() shared nested map reference with original") + } + if orig.Step.ToolCall.FunctionCall.Arguments["arr"].([]any)[0].(map[string]any)["inner"] != "x" { + t.Error("Clone() shared nested slice-element map reference with original") + } + if orig.Step.ToolResult.FunctionResult.Name != "fn" || orig.Step.ToolResult.FunctionResult.Result.(map[string]any)["out"] != "ok" { + t.Errorf("Clone() shared Step.ToolResult.FunctionResult with original: %+v", orig.Step.ToolResult.FunctionResult) + } + if orig.FunctionInvocation.Name != "fn2" || orig.FunctionInvocation.Arguments["a"] != "b" { + t.Errorf("Clone() shared FunctionInvocation with original: %+v", orig.FunctionInvocation) + } + if orig.FunctionResult.Name != "fn2" || orig.FunctionResult.Result.([]string)[0] != "s1" { + t.Errorf("Clone() shared FunctionResult with original: %+v", orig.FunctionResult) + } + if orig.CallHookRequest.RawWire[0] != 1 || orig.CallHookResponse.RawWire[0] != 3 { + t.Error("Clone() shared RawWire byte slice with original") + } + + var obs NopObserver + obs.LeasePolled(1, nil) + obs.WorkerProvisioned(true, nil) + obs.HandoffFinished("accepted") + obs.InteractionStarted("int_1") + obs.InteractionFinished("int_1", nil) +} diff --git a/internal/cli/custom/envruntime/worker.go b/internal/cli/custom/envruntime/worker.go new file mode 100644 index 0000000..ffecd53 --- /dev/null +++ b/internal/cli/custom/envruntime/worker.go @@ -0,0 +1,786 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package envruntime + +import ( + "cmp" + "context" + "errors" + "fmt" + "io" + "log" + "os" + "strings" + "sync" + "time" +) + +const ( + // DefaultMaxStreamReconnects is the default maximum stream reconnect attempts per leased interaction. + DefaultMaxStreamReconnects = 3 + + // DefaultCommandTimeout is the default timeout for execute_bash tool calls when PullWorker builds a fallback ToolDispatcher. + DefaultCommandTimeout = 2 * time.Minute + + // MaxToolResultWireBytes bounds a single serialized ClientEvent below MaxGRPCFrameBytes. + MaxToolResultWireBytes = MaxGRPCFrameBytes - 4096 +) + +// PullWorkerConfig configures a PullWorker for single-node polling, active-lease renewal, or one-shot serve execution. +type PullWorkerConfig struct { + Client EnvironmentClient + Tools *ToolDispatcher + FS *SandboxFS + Shell *SandboxShell + Hooks HookHandler + Observer Observer + EnvironmentID string + EnvironmentKey string + WorkerID string + RegisterEnvironment func(ctx context.Context) (envID, envKey string, err error) + DirectInteractionID string + DirectLeaseToken string + MaxConcurrent int + PollInterval time.Duration + LeaseTTL time.Duration + HeartbeatInterval time.Duration + CommandTimeout time.Duration + DisableBash bool + Labels map[string]string + // TrackActiveLeases renews active_leases on each LeaseInteractions poll cycle and cancels + // in-flight interaction sessions if the control plane stops returning their lease. + TrackActiveLeases bool + InBandStreamHeartbeat bool + RetryConnectErrors bool + NoRetries bool + MaxReconnectAttempts int + NonRetryableCheck func(error) bool + Logger *log.Logger +} + +// CLIPullWorkerDefaults returns the baseline PullWorkerConfig for CLI sandbox worker commands +// (4 concurrent sessions, 2s poll interval, 2m command timeout, in-band stream heartbeats, and connect retries). +func CLIPullWorkerDefaults() PullWorkerConfig { + return PullWorkerConfig{ + MaxConcurrent: 4, + PollInterval: 2 * time.Second, + HeartbeatInterval: DefaultHeartbeatInterval, + CommandTimeout: DefaultCommandTimeout, + InBandStreamHeartbeat: true, + RetryConnectErrors: true, + } +} + +// TrackedPullWorkerDefaults returns the baseline PullWorkerConfig for standalone container workers +// that renew active_leases on each poll cycle (1 concurrent session, 5s poll backoff, and 2m command timeout). +func TrackedPullWorkerDefaults() PullWorkerConfig { + return PullWorkerConfig{ + MaxConcurrent: 1, + PollInterval: DefaultPollBackoff, + LeaseTTL: DefaultLeaseTTL, + CommandTimeout: DefaultCommandTimeout, + TrackActiveLeases: true, + } +} + +type activeWorkerSession struct { + interactionID string + leaseToken string + cancel context.CancelFunc +} + +// PullWorker leases pending interactions from the control plane and executes tools and hooks over bidirectional streams. +type PullWorker struct { + cfg PullWorkerConfig + client EnvironmentClient + tools *ToolDispatcher + hooks HookHandler + observer Observer + workerID string + maxConcurrent int + pollInterval time.Duration + leaseTTL time.Duration + heartbeatInterval time.Duration + maxReconnects int + logger *log.Logger + + mu sync.Mutex + environmentID string + environmentKey string + activeSessions map[string]*activeWorkerSession + slotReleasedCh chan struct{} +} + +// NewPullWorker constructs a PullWorker from cfg. +func NewPullWorker(cfg PullWorkerConfig) *PullWorker { + workerID := strings.TrimSpace(cfg.WorkerID) + if workerID == "" { + workerID = DefaultRuntimeID("worker") + } + maxConcurrent := cfg.MaxConcurrent + if maxConcurrent <= 0 { + maxConcurrent = 4 + } + pollInterval := cfg.PollInterval + if pollInterval <= 0 { + pollInterval = 2 * time.Second + } + leaseTTL := cfg.LeaseTTL + if leaseTTL <= 0 { + leaseTTL = DefaultLeaseTTL + } + hbInterval := cfg.HeartbeatInterval + if hbInterval <= 0 && cfg.InBandStreamHeartbeat { + hbInterval = DefaultHeartbeatInterval + } + tools := cfg.Tools + if tools == nil { + tools = NewToolDispatcher(ToolDispatcherConfig{ + FS: cfg.FS, + Shell: cfg.Shell, + DisableBash: cfg.DisableBash, + CommandTimeout: cmp.Or(cfg.CommandTimeout, DefaultCommandTimeout), + }) + } + maxReconnects := cfg.MaxReconnectAttempts + if maxReconnects <= 0 { + maxReconnects = DefaultMaxStreamReconnects + } + if cfg.NoRetries { + maxReconnects = 0 + } + obs := cfg.Observer + if obs == nil { + obs = NopObserver{} + } + logger := cfg.Logger + if logger == nil { + logger = log.New(os.Stderr, "", log.LstdFlags) + } + + return &PullWorker{ + cfg: cfg, + client: cfg.Client, + tools: tools, + hooks: cfg.Hooks, + observer: obs, + workerID: workerID, + maxConcurrent: maxConcurrent, + pollInterval: pollInterval, + leaseTTL: leaseTTL, + heartbeatInterval: hbInterval, + maxReconnects: maxReconnects, + logger: logger, + environmentID: strings.TrimSpace(cfg.EnvironmentID), + environmentKey: strings.TrimSpace(cfg.EnvironmentKey), + activeSessions: make(map[string]*activeWorkerSession), + slotReleasedCh: make(chan struct{}, 1), + } +} + +// EnvironmentID returns the worker's current environment identifier. +func (w *PullWorker) EnvironmentID() string { + w.mu.Lock() + defer w.mu.Unlock() + return w.environmentID +} + +// EnvironmentKey returns the worker's current environment key. +func (w *PullWorker) EnvironmentKey() string { + w.mu.Lock() + defer w.mu.Unlock() + return w.environmentKey +} + +func (w *PullWorker) isNonRetryable(err error) bool { + if IsNonRetryableError(err) { + return true + } + return w.cfg.NonRetryableCheck != nil && w.cfg.NonRetryableCheck(err) +} + +// ServeInteraction validates lease identifiers and executes a single pre-leased interaction session. +func (w *PullWorker) ServeInteraction(ctx context.Context, lease LeasedInteraction) error { + interactionID := strings.TrimSpace(lease.InteractionID) + if interactionID == "" { + return &PathError{Kind: PathErrInvalidArgument, Message: "interaction_id is required"} + } + if !IsValidIdentifier(interactionID) { + return &PathError{Kind: PathErrInvalidArgument, Message: "interaction_id contains invalid characters"} + } + if envID := strings.TrimSpace(lease.EnvironmentID); envID != "" && !IsValidIdentifier(envID) { + return &PathError{Kind: PathErrInvalidArgument, Message: "environment_id contains invalid characters"} + } + lease.InteractionID = interactionID + return w.ExecuteLease(ctx, &lease) +} + +// ServeLeases executes each pre-leased interaction sequentially, logging attachment and completion. +func (w *PullWorker) ServeLeases(ctx context.Context, leases []*LeasedInteraction) error { + for _, lease := range leases { + if lease == nil { + continue + } + w.logger.Printf("[%s] Attached interaction session", lease.InteractionID) + if err := w.ExecuteLease(ctx, lease); err != nil { + return err + } + w.logger.Printf("[%s] Completed interaction session", lease.InteractionID) + } + return nil +} + +// Run executes the pull-mode worker loop until ctx is canceled or a fatal non-retryable error occurs. +func (w *PullWorker) Run(ctx context.Context) error { + if w.cfg.TrackActiveLeases { + return w.runTrackedLoop(ctx) + } + return RunLeasePollLoop(ctx, LeasePollLoopConfig{ + Client: w.client, + CallerID: w.workerID, + EnvironmentID: w.EnvironmentID(), + EnvironmentKey: w.EnvironmentKey(), + MaxSlots: w.maxConcurrent, + PollInterval: w.pollInterval, + NoRetries: w.cfg.NoRetries, + NonRetryableCheck: w.isNonRetryable, + Logger: w.logger, + Observer: w.observer, + ErrLabel: "LeaseInteractions error", + StartMsg: fmt.Sprintf("Worker %s polling (maxConcurrent=%d)", w.workerID, w.maxConcurrent), + }, func(ctx context.Context, item *LeasedInteraction) { + w.logger.Printf("[%s] Claimed interaction lease (env=%s)", item.InteractionID, cmp.Or(item.EnvironmentID, "default")) + if err := w.ExecuteLease(ctx, item); err != nil && ctx.Err() == nil { + w.logger.Printf("[%s] session ended with error: %v", item.InteractionID, err) + return + } + if ctx.Err() == nil { + w.logger.Printf("[%s] Completed interaction session", item.InteractionID) + } + }) +} + +func (w *PullWorker) runTrackedLoop(ctx context.Context) error { + if w.cfg.RegisterEnvironment != nil && w.EnvironmentKey() == "" { + envID, envKey, err := w.cfg.RegisterEnvironment(ctx) + if err != nil { + return fmt.Errorf("failed to register environment on startup: %w", err) + } + w.mu.Lock() + w.environmentID = envID + w.environmentKey = envKey + w.mu.Unlock() + w.logger.Printf("PullWorker registered environment %q", envID) + } + + var wg sync.WaitGroup + defer wg.Wait() + + loopCtx, cancelLoop := context.WithCancel(ctx) + defer cancelLoop() + + if w.cfg.DirectInteractionID != "" { + w.startTrackedSession(loopCtx, &wg, LeasedInteraction{ + InteractionID: w.cfg.DirectInteractionID, + EnvironmentID: w.EnvironmentID(), + LeaseToken: w.cfg.DirectLeaseToken, + EnvironmentKey: w.EnvironmentKey(), + }) + if w.EnvironmentKey() == "" { + wg.Wait() + return nil + } + } + + backoff := w.pollInterval + consecutiveErrors := 0 + for { + if ctx.Err() != nil { + return nil + } + nextBackoff, err := w.leaseTrackedOnce(loopCtx, &wg) + if err != nil { + if ctx.Err() != nil { + return nil + } + if w.isNonRetryable(err) || w.cfg.NoRetries { + return err + } + consecutiveErrors++ + nextBackoff = RetryBackoff(w.pollInterval, consecutiveErrors) + w.logger.Printf("PullWorker lease cycle failed (retrying in %v): %v", nextBackoff, err) + } else { + consecutiveErrors = 0 + } + if nextBackoff <= 0 { + nextBackoff = backoff + } + + select { + case <-ctx.Done(): + return nil + case <-w.slotReleasedCh: + case <-time.After(nextBackoff): + } + } +} + +func (w *PullWorker) leaseTrackedOnce(ctx context.Context, wg *sync.WaitGroup) (time.Duration, error) { + w.mu.Lock() + envID := w.environmentID + envKey := w.environmentKey + activeLeases := make([]LeaseEntry, 0, len(w.activeSessions)) + for id, sess := range w.activeSessions { + activeLeases = append(activeLeases, LeaseEntry{ + InteractionID: id, + LeaseToken: sess.leaseToken, + }) + } + availSlots := max(w.maxConcurrent-len(w.activeSessions), 0) + w.mu.Unlock() + + if availSlots == 0 && len(activeLeases) == 0 { + return w.pollInterval, nil + } + + leaseRes, err := w.client.LeaseInteractions(ctx, &LeaseInteractionsRequest{ + EnvironmentID: envID, + EnvironmentKey: envKey, + WorkerID: w.workerID, + ActiveLeases: activeLeases, + MaxNewInteractions: int32(availSlots), + SendZeroMaxNew: availSlots == 0, + RequestedLeaseTTL: w.leaseTTL, + Labels: w.cfg.Labels, + }) + if err != nil { + w.observer.LeasePolled(0, err) + return 0, err + } + w.observer.LeasePolled(len(leaseRes.Leases), nil) + + granted := make(map[string]LeasedInteraction, len(leaseRes.Leases)) + for _, lease := range leaseRes.Leases { + granted[lease.InteractionID] = lease + } + + w.mu.Lock() + for id, sess := range w.activeSessions { + if updated, ok := granted[id]; ok { + if updated.LeaseToken != "" { + sess.leaseToken = updated.LeaseToken + } + continue + } + if sess.leaseToken != "" { + w.logger.Printf("PullWorker lease lost for interaction %q; canceling local execution", id) + sess.cancel() + } + } + w.mu.Unlock() + + if availSlots > 0 { + for _, lease := range leaseRes.Leases { + if lease.InteractionID != "" { + w.startTrackedSession(ctx, wg, lease) + } + } + } + return leaseRes.PollBackoff, nil +} + +func (w *PullWorker) startTrackedSession(parentCtx context.Context, wg *sync.WaitGroup, lease LeasedInteraction) { + w.mu.Lock() + if _, exists := w.activeSessions[lease.InteractionID]; exists || len(w.activeSessions) >= w.maxConcurrent { + w.mu.Unlock() + return + } + sessCtx, cancel := context.WithCancel(parentCtx) + w.activeSessions[lease.InteractionID] = &activeWorkerSession{ + interactionID: lease.InteractionID, + leaseToken: lease.LeaseToken, + cancel: cancel, + } + w.mu.Unlock() + + wg.Add(1) + go func() { + defer wg.Done() + defer func() { + cancel() + w.mu.Lock() + delete(w.activeSessions, lease.InteractionID) + w.mu.Unlock() + select { + case w.slotReleasedCh <- struct{}{}: + default: + } + }() + cp := lease + if err := w.ExecuteLease(sessCtx, &cp); err != nil && sessCtx.Err() == nil { + w.logger.Printf("PullWorker session for interaction %q ended with error: %v", lease.InteractionID, err) + } + }() +} + +type sessionState struct { + interactionID string + mu sync.Mutex + lastAckSeq int64 + nextClientSeq int64 + executedCalls map[string]*EventPayload +} + +func (s *sessionState) ackSeq() int64 { + s.mu.Lock() + defer s.mu.Unlock() + return s.lastAckSeq +} + +// ExecuteLease attaches to lease's bidirectional interaction stream, handling reconnects, tools, hooks, and Yield. +func (w *PullWorker) ExecuteLease(ctx context.Context, lease *LeasedInteraction) (retErr error) { + if lease == nil { + return &PathError{Kind: PathErrInvalidArgument, Message: "lease is required"} + } + w.observer.InteractionStarted(lease.InteractionID) + defer func() { + w.observer.InteractionFinished(lease.InteractionID, retErr) + }() + + state := &sessionState{ + interactionID: lease.InteractionID, + nextClientSeq: 1, + executedCalls: make(map[string]*EventPayload), + } + + envKey := cmp.Or(lease.EnvironmentKey, w.EnvironmentKey()) + envID := cmp.Or(lease.EnvironmentID, w.EnvironmentID()) + var reqHB time.Duration + if !w.cfg.InBandStreamHeartbeat { + reqHB = cmp.Or(lease.HeartbeatInterval, w.heartbeatInterval) + if reqHB <= 0 && envKey == "" && lease.LeaseToken != "" { + reqHB = DefaultHeartbeatInterval + } + } + + for attempt := 0; attempt <= w.maxReconnects; attempt++ { + if err := ctx.Err(); err != nil { + return err + } + + leaseToken := lease.LeaseToken + w.mu.Lock() + if sess, ok := w.activeSessions[lease.InteractionID]; ok && sess.leaseToken != "" { + leaseToken = sess.leaseToken + } + w.mu.Unlock() + + ack := state.ackSeq() + stream, err := w.client.ConnectInteraction(ctx, &ConnectInteractionRequest{ + InteractionID: lease.InteractionID, + EnvironmentID: envID, + EnvironmentKey: envKey, + LeaseToken: leaseToken, + WorkerID: w.workerID, + ClientID: DefaultEnvironmentClientID, + NextOutputSeqNum: ack + 1, + AckOutputSeqNum: ack, + HeartbeatInterval: reqHB, + }) + if err != nil { + if !w.cfg.RetryConnectErrors || w.isNonRetryable(err) || attempt >= w.maxReconnects { + if !w.cfg.RetryConnectErrors { + return fmt.Errorf("ConnectInteraction(%q) failed: %w", lease.InteractionID, err) + } + return err + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(RetryBackoff(200*time.Millisecond, attempt+1)): + continue + } + } + + reconnect, runErr := w.consumeStream(ctx, stream, state) + _ = stream.Close() + if runErr != nil { + if !w.cfg.RetryConnectErrors || w.isNonRetryable(runErr) || attempt >= w.maxReconnects { + return runErr + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(RetryBackoff(200*time.Millisecond, attempt+1)): + continue + } + } + if !reconnect { + return nil + } + } + return fmt.Errorf("exceeded max stream reconnect attempts (%d) for interaction %q", w.maxReconnects, lease.InteractionID) +} + +func (w *PullWorker) consumeStream(ctx context.Context, stream InteractionStream, state *sessionState) (bool, error) { + var sendMu sync.Mutex + sendBatch := func(sendCtx context.Context, batch *ClientEventBatch) error { + sendMu.Lock() + defer sendMu.Unlock() + return stream.Send(sendCtx, batch) + } + closeSend := func() error { + sendMu.Lock() + defer sendMu.Unlock() + return stream.CloseSend() + } + + stopHeartbeat := func() {} + if w.cfg.InBandStreamHeartbeat { + hbCtx, cancelHB := context.WithCancel(ctx) + hbDone := make(chan struct{}) + go func() { + defer close(hbDone) + ticker := time.NewTicker(w.heartbeatInterval) + defer ticker.Stop() + for { + select { + case <-hbCtx.Done(): + return + case <-ticker.C: + _ = sendBatch(hbCtx, &ClientEventBatch{AckOutputSeqNum: state.ackSeq()}) + } + } + }() + stopHeartbeat = func() { + cancelHB() + <-hbDone + } + defer stopHeartbeat() + } + + for { + batch, err := stream.Recv(ctx) + if errors.Is(err, io.EOF) { + return false, nil + } + if err != nil { + return false, err + } + + if batch.InteractionReconnected != nil { + state.mu.Lock() + if batch.InteractionReconnected.NextInputSeqNum > state.nextClientSeq { + state.nextClientSeq = batch.InteractionReconnected.NextInputSeqNum + } + state.mu.Unlock() + } + + state.mu.Lock() + pendingAck := state.lastAckSeq + startClientSeq := state.nextClientSeq + state.mu.Unlock() + + var outEvents []*ClientEvent + ackUpdated := false + for _, ev := range batch.Events { + if ev == nil { + continue + } + if ev.SeqNum > pendingAck { + pendingAck = ev.SeqNum + ackUpdated = true + } + if respPayload := w.handleServerEvent(ctx, ev, state); respPayload != nil { + state.mu.Lock() + seq := state.nextClientSeq + state.nextClientSeq++ + state.mu.Unlock() + outEvents = append(outEvents, &ClientEvent{ + SeqNum: seq, + TimestampMicros: time.Now().UnixMicro(), + Payload: respPayload, + }) + } + } + + if len(outEvents) > 0 || ackUpdated { + if err := sendBatch(ctx, &ClientEventBatch{ + AckOutputSeqNum: pendingAck, + Events: outEvents, + }); err != nil { + state.mu.Lock() + state.nextClientSeq = startClientSeq + state.mu.Unlock() + return false, err + } + state.mu.Lock() + state.lastAckSeq = pendingAck + state.mu.Unlock() + } + + if batch.InteractionCompleted != nil { + stopHeartbeat() + _ = closeSend() + return false, nil + } + if batch.ReconnectionRequired != nil { + return true, nil + } + if batch.Yield != nil { + state.mu.Lock() + nextSeq := state.nextClientSeq + lastAck := state.lastAckSeq + state.mu.Unlock() + if len(outEvents) == 0 || batch.Yield.NextInputSeqNum >= nextSeq { + stopHeartbeat() + _ = sendBatch(ctx, &ClientEventBatch{ + AckOutputSeqNum: lastAck, + Disconnect: &DisconnectEvent{}, + }) + _ = closeSend() + drainCtx, cancelDrain := context.WithTimeout(ctx, 2*time.Second) + defer cancelDrain() + _, _ = stream.Recv(drainCtx) + return false, nil + } + } + } +} + +func (w *PullWorker) handleServerEvent(ctx context.Context, ev *ServerEvent, state *sessionState) *EventPayload { + if ev.Payload == nil { + return nil + } + + if fi := ev.Payload.FunctionInvocation; fi != nil { + callKey := "fn:" + cmp.Or(fi.CallID, fmt.Sprintf("seq-%d", ev.SeqNum)) + if cached := w.loadCachedCall(state, callKey); cached != nil { + return cached + } + w.logger.Printf("[%s] Executing function %s", state.interactionID, fi.Name) + resMap, isErr := w.tools.Dispatch(ctx, fi.Name, fi.Arguments) + out := &EventPayload{ + FunctionResult: &FunctionResultEvent{ + CallID: fi.CallID, + Name: fi.Name, + IsError: isErr, + Result: resMap, + }, + } + if exceedsMaxToolResultWire(out) { + out = &EventPayload{ + FunctionResult: &FunctionResultEvent{ + CallID: fi.CallID, + Name: fi.Name, + IsError: true, + Result: map[string]any{ + "error": fmt.Sprintf("tool %q result exceeds maximum frame size (%d bytes); narrow the query or specify start_line and end_line", fi.Name, MaxGRPCFrameBytes), + }, + }, + } + } + w.storeCachedCall(state, callKey, out) + return out + } + + if ev.Payload.Step != nil && ev.Payload.Step.ToolCall != nil && ev.Payload.Step.ToolCall.FunctionCall != nil { + tc := ev.Payload.Step.ToolCall + fnCall := tc.FunctionCall + callKey := "step:" + cmp.Or(tc.ID, fmt.Sprintf("seq-%d", ev.SeqNum)) + if cached := w.loadCachedCall(state, callKey); cached != nil { + return cached + } + w.logger.Printf("[%s] Executing tool %s", state.interactionID, fnCall.Name) + resMap, isErr := w.tools.Dispatch(ctx, fnCall.Name, fnCall.Arguments) + out := &EventPayload{ + Step: &Step{ + ToolResult: &ToolResultStep{ + CallID: tc.ID, + FunctionResult: &FunctionResultStep{ + Name: fnCall.Name, + IsError: isErr, + Result: resMap, + }, + }, + }, + } + if exceedsMaxToolResultWire(out) { + out = &EventPayload{ + Step: &Step{ + ToolResult: &ToolResultStep{ + CallID: tc.ID, + FunctionResult: &FunctionResultStep{ + Name: fnCall.Name, + IsError: true, + Result: map[string]any{ + "error": fmt.Sprintf("tool %q result exceeds maximum frame size (%d bytes); narrow the query or specify start_line and end_line", fnCall.Name, MaxGRPCFrameBytes), + }, + }, + }, + }, + } + } + w.storeCachedCall(state, callKey, out) + return out + } + + if hr := ev.Payload.CallHookRequest; hr != nil && w.hooks != nil { + callKey := "hook:" + cmp.Or(hr.RequestID, fmt.Sprintf("seq-%d", ev.SeqNum)) + if cached := w.loadCachedCall(state, callKey); cached != nil { + return cached + } + resp, err := w.hooks.CallHook(ctx, hr) + if err != nil { + resp = &CallHookResponse{ + RequestID: hr.RequestID, + ErrorMessage: err.Error(), + } + } else if resp == nil { + resp = &CallHookResponse{ + RequestID: hr.RequestID, + EmptyResult: true, + } + } + out := &EventPayload{CallHookResponse: resp} + w.storeCachedCall(state, callKey, out) + return out + } + + return nil +} + +func (w *PullWorker) loadCachedCall(state *sessionState, key string) *EventPayload { + state.mu.Lock() + defer state.mu.Unlock() + if cached, ok := state.executedCalls[key]; ok { + return cached.Clone() + } + return nil +} + +func (w *PullWorker) storeCachedCall(state *sessionState, key string, payload *EventPayload) { + state.mu.Lock() + defer state.mu.Unlock() + state.executedCalls[key] = payload.Clone() +} + +func exceedsMaxToolResultWire(payload *EventPayload) bool { + return len(MarshalClientEvent(&ClientEvent{ + SeqNum: 1, + TimestampMicros: 1, + Payload: payload, + })) > MaxToolResultWireBytes +} diff --git a/internal/cli/custom/envruntime/worker_test.go b/internal/cli/custom/envruntime/worker_test.go new file mode 100644 index 0000000..a891bac --- /dev/null +++ b/internal/cli/custom/envruntime/worker_test.go @@ -0,0 +1,850 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package envruntime + +import ( + "context" + "errors" + "io" + "log" + "net/http" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +type scriptedStream struct { + mu sync.Mutex + sent []*ClientEventBatch + recvFn func(call int) (*ServerEventBatch, error) + sendFn func(call int, batch *ClientEventBatch) error + recvCall int + sendCall int + closed bool +} + +func (s *scriptedStream) Send(_ context.Context, batch *ClientEventBatch) error { + s.mu.Lock() + s.sendCall++ + call := s.sendCall + s.sent = append(s.sent, batch) + fn := s.sendFn + s.mu.Unlock() + if fn != nil { + return fn(call, batch) + } + return nil +} + +func (s *scriptedStream) Recv(ctx context.Context) (*ServerEventBatch, error) { + s.mu.Lock() + s.recvCall++ + call := s.recvCall + fn := s.recvFn + s.mu.Unlock() + if fn != nil { + return fn(call) + } + <-ctx.Done() + return nil, ctx.Err() +} + +func (s *scriptedStream) CloseSend() error { return nil } + +func (s *scriptedStream) Close() error { + s.mu.Lock() + s.closed = true + s.mu.Unlock() + return nil +} + +type fakeEnvClient struct { + leaseFn func(ctx context.Context, req *LeaseInteractionsRequest) (*LeaseInteractionsResult, error) + connectFn func(ctx context.Context, req *ConnectInteractionRequest) (InteractionStream, error) +} + +func (f *fakeEnvClient) LeaseInteractions(ctx context.Context, req *LeaseInteractionsRequest) (*LeaseInteractionsResult, error) { + return f.leaseFn(ctx, req) +} + +func (f *fakeEnvClient) ConnectInteraction(ctx context.Context, req *ConnectInteractionRequest) (InteractionStream, error) { + return f.connectFn(ctx, req) +} + +type hookHandlerFunc func(ctx context.Context, req *CallHookRequest) (*CallHookResponse, error) + +func (f hookHandlerFunc) CallHook(ctx context.Context, req *CallHookRequest) (*CallHookResponse, error) { + return f(ctx, req) +} + +type recordingObserver struct { + mu sync.Mutex + leasePolledCounts []int + leasePolledErrs []error + interactionStarted []string + interactionFinished []string + workerProvisioned []bool + handoffStatuses []string +} + +func (o *recordingObserver) LeasePolled(leasedCount int, err error) { + o.mu.Lock() + o.leasePolledCounts = append(o.leasePolledCounts, leasedCount) + o.leasePolledErrs = append(o.leasePolledErrs, err) + o.mu.Unlock() +} + +func (o *recordingObserver) InteractionStarted(interactionID string) { + o.mu.Lock() + o.interactionStarted = append(o.interactionStarted, interactionID) + o.mu.Unlock() +} + +func (o *recordingObserver) InteractionFinished(interactionID string, _ error) { + o.mu.Lock() + o.interactionFinished = append(o.interactionFinished, interactionID) + o.mu.Unlock() +} + +func (o *recordingObserver) WorkerProvisioned(shared bool, _ error) { + o.mu.Lock() + o.workerProvisioned = append(o.workerProvisioned, shared) + o.mu.Unlock() +} + +func (o *recordingObserver) HandoffFinished(status string) { + o.mu.Lock() + o.handoffStatuses = append(o.handoffStatuses, status) + o.mu.Unlock() +} + +func TestPullWorkerReconnectReplayHooksYieldAndOversizedResult(t *testing.T) { + var toolCalls atomic.Int32 + var fnCalls atomic.Int32 + var hookCalls atomic.Int32 + var connectReqs []*ConnectInteractionRequest + var attempt3Stream, attempt4Stream *scriptedStream + var mu sync.Mutex + + tools := NewToolDispatcher(ToolDispatcherConfig{ + ShellExec: func(_ context.Context, cmd, _ string, _ time.Duration) (string, string, int, error) { + if cmd == "huge" { + return strings.Repeat("Z", MaxToolResultWireBytes+8192), "", 0, nil + } + if cmd == "fn_once" { + fnCalls.Add(1) + return "fn:" + cmd, "", 0, nil + } + toolCalls.Add(1) + return "ran:" + cmd, "", 0, nil + }, + }) + + var attempt atomic.Int32 + client := &fakeEnvClient{ + connectFn: func(_ context.Context, req *ConnectInteractionRequest) (InteractionStream, error) { + mu.Lock() + reqCopy := *req + connectReqs = append(connectReqs, &reqCopy) + mu.Unlock() + + switch attempt.Add(1) { + case 1: + // Attempt 1: transient connect error -> retried because RetryConnectErrors=true. + return nil, errors.New("transient dial error") + case 2: + // Attempt 2: delivers empty payload + FunctionInvocation("fi_once") + Step.ToolCall("call_1"), Send fails -> rolls back seq while caching both! + return &scriptedStream{ + sendFn: func(call int, _ *ClientEventBatch) error { + if call == 1 { + return errors.New("broken pipe during tool reply") + } + return nil + }, + recvFn: func(call int) (*ServerEventBatch, error) { + if call == 1 { + return &ServerEventBatch{ + Events: []*ServerEvent{ + nil, + {Payload: nil}, + {SeqNum: 1, Payload: &EventPayload{}}, + {SeqNum: 1, Payload: &EventPayload{Step: &Step{}}}, + {SeqNum: 1, Payload: &EventPayload{Step: &Step{ToolCall: &ToolCallStep{ID: "no_fn"}}}}, + { + SeqNum: 2, + Payload: &EventPayload{ + FunctionInvocation: &FunctionInvocation{ + CallID: "fi_once", + Name: "execute_bash", + Arguments: map[string]any{"command": "fn_once"}, + }, + }, + }, + { + SeqNum: 3, + Payload: &EventPayload{ + Step: &Step{ + ToolCall: &ToolCallStep{ + ID: "call_1", + FunctionCall: &FunctionCallStep{Name: "execute_bash", Arguments: map[string]any{"command": "once"}}, + }, + }, + }, + }, + }, + }, nil + } + return nil, io.EOF + }, + }, nil + case 3: + // Attempt 3: + // - Call 1: ack-only event (SeqNum=1, Payload=nil). + // - Call 2: replays "fi_once" (SeqNum=2) and "call_1" (SeqNum=3) verifying rolled-back client seq 1 & 2. + // - Call 3: InteractionReconnected{NextInputSeqNum: 10} + CallHookRequest("hk_1") + FunctionInvocation("fi_big", "huge") + non-terminal Yield{NextInputSeqNum: 1}. + // - Call 4: ReconnectionRequired. + s := &scriptedStream{ + recvFn: func(call int) (*ServerEventBatch, error) { + switch call { + case 1: + return &ServerEventBatch{ + Events: []*ServerEvent{ + {SeqNum: 1, Payload: nil}, + }, + }, nil + case 2: + return &ServerEventBatch{ + Events: []*ServerEvent{ + { + SeqNum: 2, + Payload: &EventPayload{ + FunctionInvocation: &FunctionInvocation{ + CallID: "fi_once", + Name: "execute_bash", + Arguments: map[string]any{"command": "fn_once"}, + }, + }, + }, + { + SeqNum: 3, + Payload: &EventPayload{ + Step: &Step{ + ToolCall: &ToolCallStep{ + ID: "call_1", + FunctionCall: &FunctionCallStep{Name: "execute_bash", Arguments: map[string]any{"command": "once"}}, + }, + }, + }, + }, + }, + }, nil + case 3: + return &ServerEventBatch{ + InteractionReconnected: &InteractionReconnectedEvent{NextInputSeqNum: 10}, + Events: []*ServerEvent{ + { + SeqNum: 4, + Payload: &EventPayload{ + CallHookRequest: &CallHookRequest{RequestID: "hk_1", Name: "audit"}, + }, + }, + { + SeqNum: 5, + Payload: &EventPayload{ + FunctionInvocation: &FunctionInvocation{ + CallID: "fi_big", + Name: "execute_bash", + Arguments: map[string]any{"command": "huge"}, + }, + }, + }, + }, + Yield: &YieldEvent{NextInputSeqNum: 1}, + }, nil + case 4: + return &ServerEventBatch{ + ReconnectionRequired: &ReconnectionRequiredEvent{}, + }, nil + default: + return nil, io.EOF + } + }, + } + mu.Lock() + attempt3Stream = s + mu.Unlock() + return s, nil + default: + // Attempt 4: + // - Call 1: waits for at least one in-band heartbeat tick, replays hk_1 with SeqNum=0 (verifying len(outEvents) > 0 sends even when ackUpdated==false), and sends Yield{NextInputSeqNum: 15} (verifying batch.Yield.NextInputSeqNum >= nextSeq disconnects when len(outEvents) > 0). + s := &scriptedStream{ + recvFn: func(call int) (*ServerEventBatch, error) { + if call == 1 { + time.Sleep(25 * time.Millisecond) + return &ServerEventBatch{ + Events: []*ServerEvent{{ + SeqNum: 0, + Payload: &EventPayload{CallHookRequest: &CallHookRequest{RequestID: "hk_1", Name: "audit"}}, + }}, + Yield: &YieldEvent{NextInputSeqNum: 15}, + }, nil + } + return nil, io.EOF + }, + } + mu.Lock() + attempt4Stream = s + mu.Unlock() + return s, nil + } + }, + } + + obs := &recordingObserver{} + worker := NewPullWorker(PullWorkerConfig{ + Client: client, + Tools: tools, + Hooks: hookHandlerFunc(func(_ context.Context, req *CallHookRequest) (*CallHookResponse, error) { + hookCalls.Add(1) + return &CallHookResponse{RequestID: req.RequestID, EmptyResult: true}, nil + }), + Observer: obs, + EnvironmentKey: "ek_w", + WorkerID: "w-1", + RetryConnectErrors: true, + InBandStreamHeartbeat: true, + HeartbeatInterval: 10 * time.Millisecond, + Logger: log.New(io.Discard, "", 0), + }) + + if err := worker.ServeInteraction(t.Context(), LeasedInteraction{ + InteractionID: " int_replay ", + LeaseToken: "tok_replay", + }); err != nil { + t.Fatalf("ServeInteraction() failed: %v", err) + } + + if got := toolCalls.Load(); got != 1 { + t.Errorf("side-effecting tool execution count = %d, want 1 (replay cache must prevent duplicate execution)", got) + } + if got := fnCalls.Load(); got != 1 { + t.Errorf("FunctionInvocation execution count = %d, want 1 (replay cache must prevent duplicate execution)", got) + } + if got := hookCalls.Load(); got != 1 { + t.Errorf("hook execution count = %d, want 1 (replay cache must prevent duplicate hook run)", got) + } + obs.mu.Lock() + startedCopy := append([]string(nil), obs.interactionStarted...) + obs.mu.Unlock() + if len(startedCopy) != 1 || startedCopy[0] != "int_replay" { + t.Errorf("observer.interactionStarted = %v, want [\"int_replay\"]", startedCopy) + } + + mu.Lock() + defer mu.Unlock() + if len(connectReqs) != 4 { + t.Fatalf("len(connectReqs) = %d, want 4", len(connectReqs)) + } + if connectReqs[0].InteractionID != "int_replay" || connectReqs[0].HeartbeatInterval != 0 { + t.Errorf("attempt 1 req = %+v, want trimmed InteractionID and HeartbeatInterval=0 when InBandStreamHeartbeat=true", connectReqs[0]) + } + // Attempt 3 followed failed Send on attempt 2, so AckOutputSeqNum must have rolled back to 0! + if connectReqs[2].AckOutputSeqNum != 0 || connectReqs[2].NextOutputSeqNum != 1 { + t.Errorf("attempt 3 seq after rolled-back Send = (ack=%d, next=%d), want (0, 1)", connectReqs[2].AckOutputSeqNum, connectReqs[2].NextOutputSeqNum) + } + // Attempt 4 followed acked seq 5, so AckOutputSeqNum=5, NextOutputSeqNum=6! + if connectReqs[3].AckOutputSeqNum != 5 || connectReqs[3].NextOutputSeqNum != 6 { + t.Errorf("attempt 4 seq = (ack=%d, next=%d), want (5, 6)", connectReqs[3].AckOutputSeqNum, connectReqs[3].NextOutputSeqNum) + } + + // Inspect attempt3Stream sent batches: + // - Batch 0 (or after any in-band heartbeat): ack-only batch with AckOutputSeqNum=1 and len(Events)==0. + // - Replay batch: Events[0].SeqNum==1 (FunctionResult IsError=false), Events[1].SeqNum==2. + // - Reconnected batch: Events[0].SeqNum==10 (CallHookResponse), Events[1].SeqNum==11 (FunctionResult IsError=true for oversized output). + var sawAckOnly, sawRolledBackSeq, sawReconnectedSeqAndOversized bool + for _, b := range attempt3Stream.sent { + if b.AckOutputSeqNum == 1 && len(b.Events) == 0 && b.Disconnect == nil { + sawAckOnly = true + } + if len(b.Events) == 2 && b.Events[0].SeqNum == 1 && b.Events[1].SeqNum == 2 { + if b.Events[0].Payload != nil && b.Events[0].Payload.FunctionResult != nil && !b.Events[0].Payload.FunctionResult.IsError { + sawRolledBackSeq = true + } + } + if len(b.Events) == 2 && b.Events[0].SeqNum == 10 && b.Events[1].SeqNum == 11 { + if b.Events[1].Payload != nil && b.Events[1].Payload.FunctionResult != nil && b.Events[1].Payload.FunctionResult.IsError { + sawReconnectedSeqAndOversized = true + } + } + } + if !sawAckOnly || !sawRolledBackSeq || !sawReconnectedSeqAndOversized { + t.Errorf("attempt3Stream.sent checks: ackOnly=%v rolledBackSeq=%v reconnectedSeqAndOversized=%v", sawAckOnly, sawRolledBackSeq, sawReconnectedSeqAndOversized) + } + + var sawInBandHeartbeat, sawAttempt4HookReply, sawAttempt4Disconnect bool + for _, b := range attempt4Stream.sent { + if len(b.Events) == 0 && b.Disconnect == nil { + sawInBandHeartbeat = true + } + if len(b.Events) == 1 && b.Events[0].Payload != nil && b.Events[0].Payload.CallHookResponse != nil { + sawAttempt4HookReply = true + } + if b.Disconnect != nil { + sawAttempt4Disconnect = true + } + } + if !sawInBandHeartbeat || !sawAttempt4HookReply || !sawAttempt4Disconnect { + t.Errorf("attempt4Stream.sent: sawInBandHeartbeat=%v sawAttempt4HookReply=%v sawAttempt4Disconnect=%v, want all true", sawInBandHeartbeat, sawAttempt4HookReply, sawAttempt4Disconnect) + } +} + +func TestPullWorkerValidationServeLeasesAndActiveLeaseRenewal(t *testing.T) { + // 1. ServeInteraction identifier validation and NewPullWorker defaults. + wVal := NewPullWorker(PullWorkerConfig{Logger: log.New(io.Discard, "", 0)}) + if !strings.HasPrefix(wVal.workerID, "worker-") || wVal.leaseTTL != DefaultLeaseTTL || wVal.maxConcurrent != 4 || wVal.pollInterval != 2*time.Second { + t.Errorf("wVal defaults: workerID=%q leaseTTL=%v maxConcurrent=%d pollInterval=%v", wVal.workerID, wVal.leaseTTL, wVal.maxConcurrent, wVal.pollInterval) + } + wCustom := NewPullWorker(PullWorkerConfig{ + WorkerID: "custom-w", + LeaseTTL: 9 * time.Second, + Logger: log.New(io.Discard, "", 0), + }) + if wCustom.workerID != "custom-w" || wCustom.leaseTTL != 9*time.Second { + t.Errorf("wCustom = (workerID=%q, leaseTTL=%v)", wCustom.workerID, wCustom.leaseTTL) + } + for _, bad := range []struct { + lease LeasedInteraction + wantErr string + }{ + {lease: LeasedInteraction{}, wantErr: "interaction_id is required"}, + {lease: LeasedInteraction{InteractionID: "a..b"}, wantErr: "interaction_id contains invalid characters"}, + {lease: LeasedInteraction{InteractionID: "int_ok", EnvironmentID: "bad/env"}, wantErr: "environment_id contains invalid characters"}, + } { + if err := wVal.ServeInteraction(t.Context(), bad.lease); err == nil || !strings.Contains(err.Error(), bad.wantErr) { + t.Errorf("ServeInteraction(%+v) = %v, want %q", bad.lease, err, bad.wantErr) + } + } + + // 2. ServeLeases, NoRetries with RetryConnectErrors=true, bounded recv retries, and HeartbeatInterval resolution. + var connectCalls atomic.Int32 + var lastConnectHB time.Duration + var nonInBandStream *scriptedStream + errClient := &fakeEnvClient{ + connectFn: func(_ context.Context, req *ConnectInteractionRequest) (InteractionStream, error) { + connectCalls.Add(1) + lastConnectHB = req.HeartbeatInterval + if req.InteractionID == "int_broken" { + return nil, errors.New("stream broken") + } + if req.InteractionID == "int_recv_err" { + callN := connectCalls.Load() + return &scriptedStream{ + recvFn: func(int) (*ServerEventBatch, error) { + if callN == 1 { + return nil, errors.New("transient recv error 1") + } + return nil, errors.New("fatal custom recv error 2") + }, + }, nil + } + var s *scriptedStream + s = &scriptedStream{ + recvFn: func(call int) (*ServerEventBatch, error) { + if call == 1 { + if req.InteractionID == "int_hb_wait" { + time.Sleep(25 * time.Millisecond) + } + if req.InteractionID == "int_yield_empty" { + // Standalone Yield with NextInputSeqNum=0 (< nextSeq=5) and no outEvents -> must disconnect via len(outEvents)==0! + return &ServerEventBatch{ + InteractionReconnected: &InteractionReconnectedEvent{NextInputSeqNum: 5}, + Yield: &YieldEvent{NextInputSeqNum: 0}, + }, nil + } + return &ServerEventBatch{ + InteractionCompleted: &InteractionCompletedEvent{ + Status: &RPCStatus{Code: 13, Message: "model failed"}, + }, + }, nil + } + if req.InteractionID == "int_yield_empty" { + s.mu.Lock() + disconnected := len(s.sent) > 0 && s.sent[len(s.sent)-1].Disconnect != nil + s.mu.Unlock() + if !disconnected { + return nil, errors.New("standalone Yield did not send Disconnect") + } + } + return nil, io.EOF + }, + } + nonInBandStream = s + return s, nil + }, + } + wErr := NewPullWorker(PullWorkerConfig{ + Client: errClient, + NoRetries: true, + RetryConnectErrors: true, + MaxReconnectAttempts: 3, + Logger: log.New(io.Discard, "", 0), + }) + if err := wErr.ServeLeases(t.Context(), []*LeasedInteraction{ + nil, + {InteractionID: "int_done", LeaseToken: "tok"}, + {InteractionID: "int_yield_empty", LeaseToken: "tok_y"}, + {InteractionID: "int_broken", LeaseToken: "tok2"}, + }); err == nil || !strings.Contains(err.Error(), "stream broken") { + t.Errorf("ServeLeases(int_done, int_yield_empty, int_broken) = %v, want stream broken error", err) + } + if got := connectCalls.Load(); got != 3 { + t.Errorf("connectCalls = %d, want 3 (NoRetries=true must set maxReconnects=0)", got) + } + if lastConnectHB != DefaultHeartbeatInterval { + t.Errorf("lastConnectHB(no envKey, with LeaseToken) = %v, want %v", lastConnectHB, DefaultHeartbeatInterval) + } + // With EnvironmentKey set and HeartbeatInterval=0, reqHB must be 0; with LeaseToken="" and envKey="", reqHB must be 0. + _ = wErr.ExecuteLease(t.Context(), &LeasedInteraction{InteractionID: "int_done", EnvironmentKey: "ek_set", LeaseToken: "tok"}) + if lastConnectHB != 0 { + t.Errorf("lastConnectHB(with envKey) = %v, want 0", lastConnectHB) + } + _ = wErr.ExecuteLease(t.Context(), &LeasedInteraction{InteractionID: "int_done"}) + if lastConnectHB != 0 { + t.Errorf("lastConnectHB(without LeaseToken) = %v, want 0", lastConnectHB) + } + + // Verify InBandStreamHeartbeat=false with HeartbeatInterval>0 passes HeartbeatInterval in ConnectInteractionRequest and does NOT run consumeStream ticker. + wStreamHB := NewPullWorker(PullWorkerConfig{ + Client: errClient, + InBandStreamHeartbeat: false, + HeartbeatInterval: 10 * time.Millisecond, + NoRetries: true, + Logger: log.New(io.Discard, "", 0), + }) + if err := wStreamHB.ExecuteLease(t.Context(), &LeasedInteraction{InteractionID: "int_hb_wait", EnvironmentKey: "ek"}); err != nil { + t.Fatalf("wStreamHB.ExecuteLease() failed: %v", err) + } + if lastConnectHB != 10*time.Millisecond || len(nonInBandStream.sent) != 0 { + t.Errorf("wStreamHB: lastConnectHB=%v sent=%d, want 10ms and 0 consumeStream sends", lastConnectHB, len(nonInBandStream.sent)) + } + + // Verify RetryConnectErrors=true with MaxReconnectAttempts=1 stops after 2 attempts on recv/connect error and returns the underlying error, + // while RetryConnectErrors=false or custom NonRetryableCheck stops after 1 attempt. + connectCalls.Store(0) + wRecvRetry := NewPullWorker(PullWorkerConfig{ + Client: errClient, + RetryConnectErrors: true, + MaxReconnectAttempts: 1, + Logger: log.New(io.Discard, "", 0), + }) + if err := wRecvRetry.ExecuteLease(t.Context(), &LeasedInteraction{InteractionID: "int_recv_err", EnvironmentKey: "ek"}); err == nil || !strings.Contains(err.Error(), "fatal custom recv error 2") || connectCalls.Load() != 2 { + t.Errorf("wRecvRetry.ExecuteLease(int_recv_err) = %v, connectCalls=%d, want fatal custom recv error 2 after 2 attempts", err, connectCalls.Load()) + } + connectCalls.Store(0) + if err := wRecvRetry.ExecuteLease(t.Context(), &LeasedInteraction{InteractionID: "int_broken", EnvironmentKey: "ek"}); err == nil || !strings.Contains(err.Error(), "stream broken") || connectCalls.Load() != 2 { + t.Errorf("wRecvRetry.ExecuteLease(int_broken) = %v, connectCalls=%d, want stream broken after 2 attempts", err, connectCalls.Load()) + } + connectCalls.Store(0) + wRecvNoRetry := NewPullWorker(PullWorkerConfig{ + Client: errClient, + RetryConnectErrors: false, + MaxReconnectAttempts: 3, + Logger: log.New(io.Discard, "", 0), + }) + if err := wRecvNoRetry.ExecuteLease(t.Context(), &LeasedInteraction{InteractionID: "int_recv_err", EnvironmentKey: "ek"}); err == nil || !strings.Contains(err.Error(), "transient recv error 1") || connectCalls.Load() != 1 { + t.Errorf("wRecvNoRetry.ExecuteLease(int_recv_err) = %v, connectCalls=%d, want transient recv error 1 after 1 attempt", err, connectCalls.Load()) + } + connectCalls.Store(0) + if err := wRecvNoRetry.ExecuteLease(t.Context(), &LeasedInteraction{InteractionID: "int_broken", EnvironmentKey: "ek"}); err == nil || !strings.Contains(err.Error(), `ConnectInteraction("int_broken") failed: stream broken`) || connectCalls.Load() != 1 { + t.Errorf("wRecvNoRetry.ExecuteLease(int_broken) = %v, connectCalls=%d, want wrapped ConnectInteraction error after 1 attempt", err, connectCalls.Load()) + } + connectCalls.Store(0) + wRecvCustomNonRetry := NewPullWorker(PullWorkerConfig{ + Client: errClient, + RetryConnectErrors: true, + MaxReconnectAttempts: 3, + NonRetryableCheck: func(err error) bool { + return strings.Contains(err.Error(), "fatal custom recv") + }, + Logger: log.New(io.Discard, "", 0), + }) + if err := wRecvCustomNonRetry.ExecuteLease(t.Context(), &LeasedInteraction{InteractionID: "int_recv_err", EnvironmentKey: "ek"}); err == nil || connectCalls.Load() != 2 { + t.Errorf("wRecvCustomNonRetry.ExecuteLease() = %v, connectCalls=%d, want error after 2 attempts (attempt 1 retryable, attempt 2 non-retryable)", err, connectCalls.Load()) + } + + // 3. TrackActiveLeases=true mode: RegisterEnvironment + active lease renewal + rotated token on reconnect + lost-lease cancellation! + streamStarted := make(chan struct{}) + streamCanceled := make(chan struct{}) + tokenRotatedCh := make(chan struct{}) + var leasePollCount atomic.Int32 + var sawActiveRenewal atomic.Bool + var firstPollMaxNew atomic.Int32 + var activeConnectTokens []string + var activeMu sync.Mutex + + activeClient := &fakeEnvClient{ + leaseFn: func(_ context.Context, req *LeaseInteractionsRequest) (*LeaseInteractionsResult, error) { + n := leasePollCount.Add(1) + switch n { + case 1: + firstPollMaxNew.Store(req.MaxNewInteractions) + return &LeaseInteractionsResult{ + Leases: []LeasedInteraction{{ + InteractionID: "int_active_1", + EnvironmentID: req.EnvironmentID, + LeaseToken: "tok_v1", + }}, + PollBackoff: 15 * time.Millisecond, + }, nil + case 2: + hasDirect := false + hasActiveToken := false + for _, al := range req.ActiveLeases { + if al.InteractionID == "int_direct_notoken" && al.LeaseToken == "" { + hasDirect = true + } + if al.InteractionID == "int_active_1" && al.LeaseToken == "tok_v1" { + hasActiveToken = true + } + } + if len(req.ActiveLeases) == 2 && hasDirect && hasActiveToken && req.MaxNewInteractions == 0 && req.SendZeroMaxNew { + sawActiveRenewal.Store(true) + } + defer close(tokenRotatedCh) + return &LeaseInteractionsResult{ + Leases: []LeasedInteraction{{ + InteractionID: "int_active_1", + EnvironmentID: req.EnvironmentID, + LeaseToken: "tok_v2_rotated", + }}, + PollBackoff: 15 * time.Millisecond, + }, nil + default: + // Omit int_active_1 -> worker must cancel lost session! + return &LeaseInteractionsResult{PollBackoff: 15 * time.Millisecond}, nil + } + }, + connectFn: func(_ context.Context, req *ConnectInteractionRequest) (InteractionStream, error) { + if req.InteractionID == "int_direct_notoken" { + return &ctxBlockingStream{started: make(chan struct{}), canceled: make(chan struct{})}, nil + } + activeMu.Lock() + activeConnectTokens = append(activeConnectTokens, req.LeaseToken) + callNum := len(activeConnectTokens) + activeMu.Unlock() + if callNum == 1 { + // First stream waits until poll 2 rotates the token, then triggers ReconnectionRequired so attempt 2 uses rotated token! + return &scriptedStream{ + recvFn: func(call int) (*ServerEventBatch, error) { + if call == 1 { + close(streamStarted) + <-tokenRotatedCh + return &ServerEventBatch{ReconnectionRequired: &ReconnectionRequiredEvent{}}, nil + } + return nil, io.EOF + }, + }, nil + } + return &ctxBlockingStream{started: make(chan struct{}), canceled: streamCanceled}, nil + }, + } + + runCtx, cancelRun := context.WithCancel(t.Context()) + defer cancelRun() + + activeObs := &recordingObserver{} + wActive := NewPullWorker(PullWorkerConfig{ + Client: activeClient, + Observer: activeObs, + TrackActiveLeases: true, + DirectInteractionID: "int_direct_notoken", + MaxConcurrent: 2, + PollInterval: 15 * time.Millisecond, + LeaseTTL: 45 * time.Millisecond, + RegisterEnvironment: func(context.Context) (string, string, error) { + return "env_registered", "ek_registered", nil + }, + Logger: log.New(io.Discard, "", 0), + }) + + runDone := make(chan error, 1) + go func() { + runDone <- wActive.Run(runCtx) + }() + + select { + case <-streamStarted: + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for active-lease session stream to start") + } + select { + case <-streamCanceled: + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for lost lease to cancel session stream") + } + cancelRun() + if err := <-runDone; err != nil && !errors.Is(err, context.Canceled) { + t.Errorf("wActive.Run() = %v, want nil or context.Canceled", err) + } + if wActive.EnvironmentID() != "env_registered" || wActive.EnvironmentKey() != "ek_registered" { + t.Errorf("registered env = (%q, %q)", wActive.EnvironmentID(), wActive.EnvironmentKey()) + } + if firstPollMaxNew.Load() != 1 { + t.Errorf("firstPollMaxNew = %d, want 1", firstPollMaxNew.Load()) + } + if !sawActiveRenewal.Load() { + t.Error("sawActiveRenewal = false, want ActiveLeases sent on poll 2") + } + activeMu.Lock() + tokensCopy := append([]string(nil), activeConnectTokens...) + activeMu.Unlock() + if len(tokensCopy) < 2 || tokensCopy[0] != "tok_v1" || tokensCopy[1] != "tok_v2_rotated" { + t.Errorf("activeConnectTokens = %v, want [\"tok_v1\", \"tok_v2_rotated\"]", tokensCopy) + } + activeObs.mu.Lock() + polledCounts := append([]int(nil), activeObs.leasePolledCounts...) + activeObs.mu.Unlock() + if len(polledCounts) == 0 || polledCounts[0] != 1 { + t.Errorf("activeObs.leasePolledCounts = %v, want first poll count 1", polledCounts) + } + + // Verify RegisterEnvironment is skipped when EnvironmentKey is already configured, and nil RegisterEnvironment with empty EnvironmentKey does not panic. + var regCalled bool + wPresetKey := NewPullWorker(PullWorkerConfig{ + Client: errClient, + TrackActiveLeases: true, + EnvironmentKey: "ek_already_set", + DirectInteractionID: "int_done", + RegisterEnvironment: func(context.Context) (string, string, error) { + regCalled = true + return "env_no", "ek_no", nil + }, + Logger: log.New(io.Discard, "", 0), + }) + presetCtx, cancelPreset := context.WithCancel(t.Context()) + cancelPreset() + _ = wPresetKey.Run(presetCtx) + if regCalled { + t.Error("RegisterEnvironment was called when EnvironmentKey was already non-empty") + } + + var directSessionCompleted atomic.Bool + wNilReg := NewPullWorker(PullWorkerConfig{ + Client: &fakeEnvClient{ + leaseFn: func(context.Context, *LeaseInteractionsRequest) (*LeaseInteractionsResult, error) { + return nil, &TransportError{StatusCode: http.StatusForbidden, Message: "unexpected lease poll without environment key"} + }, + connectFn: func(ctx context.Context, req *ConnectInteractionRequest) (InteractionStream, error) { + if ctx.Err() != nil { + return nil, ctx.Err() + } + if req.InteractionID != "int_done" || req.LeaseToken != "lt_direct_1" { + return nil, errors.New("unexpected direct connect req") + } + return &scriptedStream{ + recvFn: func(call int) (*ServerEventBatch, error) { + if call == 1 { + directSessionCompleted.Store(true) + return &ServerEventBatch{InteractionCompleted: &InteractionCompletedEvent{}}, nil + } + return nil, io.EOF + }, + }, nil + }, + }, + TrackActiveLeases: true, + DirectInteractionID: "int_done", + DirectLeaseToken: "lt_direct_1", + Logger: log.New(io.Discard, "", 0), + }) + if err := wNilReg.Run(t.Context()); err != nil { + t.Errorf("wNilReg.Run(DirectInteractionID without EnvironmentKey) = %v, want nil", err) + } + if !directSessionCompleted.Load() { + t.Error("directSessionCompleted = false, want direct session to complete before Run returns") + } +} + +func TestPullWorkerPresetsAndDecoupledTrackActiveLeases(t *testing.T) { + cliCfg := CLIPullWorkerDefaults() + if cliCfg.MaxConcurrent != 4 || cliCfg.PollInterval != 2*time.Second || cliCfg.CommandTimeout != 2*time.Minute || !cliCfg.InBandStreamHeartbeat || !cliCfg.RetryConnectErrors || cliCfg.TrackActiveLeases { + t.Errorf("CLIPullWorkerDefaults() = %+v", cliCfg) + } + trackedCfg := TrackedPullWorkerDefaults() + if trackedCfg.MaxConcurrent != 1 || trackedCfg.PollInterval != DefaultPollBackoff || trackedCfg.LeaseTTL != DefaultLeaseTTL || trackedCfg.CommandTimeout != DefaultCommandTimeout || !trackedCfg.TrackActiveLeases || trackedCfg.RetryConnectErrors { + t.Errorf("TrackedPullWorkerDefaults() = %+v", trackedCfg) + } + + // Verify runTrackedLoop returns immediately on non-retryable LeaseInteractions errors instead of spinning forever. + wFatalTracked := NewPullWorker(PullWorkerConfig{ + Client: &fakeEnvClient{ + leaseFn: func(context.Context, *LeaseInteractionsRequest) (*LeaseInteractionsResult, error) { + return nil, &TransportError{StatusCode: http.StatusForbidden, Message: "forbidden key"} + }, + }, + TrackActiveLeases: true, + EnvironmentKey: "ek_bad", + PollInterval: 10 * time.Millisecond, + Logger: log.New(io.Discard, "", 0), + }) + if err := wFatalTracked.Run(t.Context()); err == nil || !IsNonRetryableError(err) { + t.Errorf("wFatalTracked.Run(403) = %v, want non-retryable 403 error", err) + } + + // Enabling TrackActiveLeases alone must NOT implicitly enable extended FS tools or disable write_to_file overwrite, + // and omitted CommandTimeout must default to DefaultCommandTimeout rather than leaving bash unbounded. + rootDir := t.TempDir() + fs := NewSandboxFS(rootDir, rootDir) + wDecoupled := NewPullWorker(PullWorkerConfig{ + FS: fs, + TrackActiveLeases: true, + Logger: log.New(io.Discard, "", 0), + }) + if wDecoupled.tools.cfg.CommandTimeout != DefaultCommandTimeout { + t.Errorf("wDecoupled.tools.cfg.CommandTimeout = %v, want %v", wDecoupled.tools.cfg.CommandTimeout, DefaultCommandTimeout) + } + if _, isErr := wDecoupled.tools.Dispatch(t.Context(), "mkdir", map[string]any{"path": "sub"}); !isErr { + t.Error("Dispatch(mkdir) succeeded with TrackActiveLeases=true and default Tools, want unsupported tool error") + } + wCustomTimeout := NewPullWorker(PullWorkerConfig{ + FS: fs, + CommandTimeout: 7 * time.Second, + Logger: log.New(io.Discard, "", 0), + }) + if wCustomTimeout.tools.cfg.CommandTimeout != 7*time.Second { + t.Errorf("wCustomTimeout.tools.cfg.CommandTimeout = %v, want 7s", wCustomTimeout.tools.cfg.CommandTimeout) + } +} + +type ctxBlockingStream struct { + started chan struct{} + canceled chan struct{} + once sync.Once + cancOnce sync.Once +} + +func (s *ctxBlockingStream) Send(context.Context, *ClientEventBatch) error { return nil } + +func (s *ctxBlockingStream) Recv(ctx context.Context) (*ServerEventBatch, error) { + s.once.Do(func() { close(s.started) }) + <-ctx.Done() + s.cancOnce.Do(func() { close(s.canceled) }) + return nil, ctx.Err() +} + +func (s *ctxBlockingStream) CloseSend() error { return nil } +func (s *ctxBlockingStream) Close() error { return nil } diff --git a/internal/cli/custom/runtime_client.go b/internal/cli/custom/runtime_client.go index 2c708d4..3357ae0 100644 --- a/internal/cli/custom/runtime_client.go +++ b/internal/cli/custom/runtime_client.go @@ -15,23 +15,19 @@ package custom import ( - "bytes" + "cmp" "context" "crypto/tls" - "encoding/json" "errors" "fmt" "io" "net" "net/http" "net/url" - "os" - "regexp" - "strconv" "strings" - "sync" "time" + "github.com/google-gemini/gemini-api-cli/internal/cli/custom/envruntime" "github.com/google-gemini/gemini-api-cli/internal/client" "github.com/google-gemini/gemini-api-cli/internal/config" "github.com/google-gemini/gemini-api-cli/internal/flagutil" @@ -41,81 +37,11 @@ import ( "github.com/spf13/cobra" ) -const defaultHeartbeatInterval = 5 * time.Second +// LeasedInteraction aliases envruntime.LeasedInteraction for CLI runtime commands. +type LeasedInteraction = envruntime.LeasedInteraction -// LeasedInteraction represents an active interaction lease returned by LeaseInteractions. -type LeasedInteraction struct { - InteractionID string `json:"interaction_id"` - EnvironmentID string `json:"environment_id,omitempty"` - LeaseToken string `json:"lease_token"` - LeaseGeneration int64 `json:"lease_generation,omitempty"` - LeaseExpirationTime string `json:"lease_expiration_time,omitempty"` - EnvironmentConfig json.RawMessage `json:"environment_config,omitempty"` -} - -// UnmarshalJSON decodes LeasedInteraction, accepting ProtoJSON string or numeric lease_generation. -func (l *LeasedInteraction) UnmarshalJSON(data []byte) error { - type alias LeasedInteraction - var aux struct { - alias - LeaseGeneration json.RawMessage `json:"lease_generation"` - } - if err := json.Unmarshal(data, &aux); err != nil { - return err - } - *l = LeasedInteraction(aux.alias) - gen, err := parseProtoInt64(aux.LeaseGeneration) - if err != nil { - return fmt.Errorf("invalid lease_generation: %w", err) - } - l.LeaseGeneration = gen - return nil -} - -func parseProtoInt64(raw json.RawMessage) (int64, error) { - trimmed := bytes.TrimSpace(raw) - if len(trimmed) == 0 { - return 0, nil - } - if trimmed[0] == '"' { - var s string - if err := json.Unmarshal(trimmed, &s); err != nil { - return 0, err - } - if s == "" { - return 0, nil - } - return strconv.ParseInt(s, 10, 64) - } - var n int64 - if err := json.Unmarshal(trimmed, &n); err != nil { - return 0, err - } - return n, nil -} - -type activeInteractionLease struct { - InteractionID string `json:"interaction_id"` - LeaseToken string `json:"lease_token,omitempty"` -} - -type leaseInteractionsRequest struct { - WorkerID string `json:"worker_id,omitempty"` - ActiveLeases []activeInteractionLease `json:"active_leases,omitempty"` - MaxNewInteractions *int32 `json:"max_new_interactions,omitempty"` - EnvironmentID string `json:"environment_id,omitempty"` -} - -type leaseInteractionsResponse struct { - Leases []*LeasedInteraction `json:"leases"` - PollBackoff string `json:"poll_backoff,omitempty"` -} - -// HeartbeatResult summarizes the result of a lease renewal request during worker handoff. -type HeartbeatResult struct { - Renewed bool - ShouldCancel bool -} +// HeartbeatResult aliases envruntime.HeartbeatResult for CLI runtime commands. +type HeartbeatResult = envruntime.HeartbeatResult type runtimeTransport struct { baseURL string @@ -134,23 +60,6 @@ type runtimeTransport struct { errOut io.Writer } -var leaseTokenJSONPattern = regexp.MustCompile(`("lease_token"\s*:\s*)"(?:\\.|[^"\\])*"`) - -type leaseTokenRedactor struct { - mu sync.Mutex - inner io.Writer -} - -func (r *leaseTokenRedactor) Write(p []byte) (int, error) { - redacted := leaseTokenJSONPattern.ReplaceAll(p, []byte(`${1}"[REDACTED]"`)) - r.mu.Lock() - defer r.mu.Unlock() - if _, err := r.inner.Write(redacted); err != nil { - return 0, err - } - return len(p), nil -} - func newRuntimeTransport(cmd *cobra.Command, explicitRuntimeCred bool) (*runtimeTransport, error) { rt := &runtimeTransport{headers: make(map[string]string)} @@ -205,7 +114,7 @@ func newRuntimeTransport(cmd *cobra.Command, explicitRuntimeCred bool) (*runtime if key == "" || strings.ContainsAny(key, " \t\r\n\x00") { return nil, usageError(fmt.Sprintf("invalid header name %q", k)) } - if err := validateHeaderSecret("--header", val); err != nil { + if err := envruntime.ValidateHeaderSecret("--header", val); err != nil { return nil, usageError(err.Error()) } rt.headers[http.CanonicalHeaderKey(key)] = val @@ -225,19 +134,28 @@ func newRuntimeTransport(cmd *cobra.Command, explicitRuntimeCred bool) (*runtime } if client.IsDebug(cmd) { - redactor := &leaseTokenRedactor{inner: cmd.ErrOrStderr()} + redactor := envruntime.NewLeaseTokenRedactor(cmd.ErrOrStderr()) cmd.SetErr(redactor) rt.debug = true rt.errOut = redactor } - baseHTTP := &http.Client{ - CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }, - Transport: &http.Transport{ + var baseTransport *http.Transport + if dt, ok := http.DefaultTransport.(*http.Transport); ok { + baseTransport = dt.Clone() + baseTransport.Proxy = http.ProxyFromEnvironment + baseTransport.ForceAttemptHTTP2 = true + baseTransport.TLSClientConfig = &tls.Config{MinVersion: tls.VersionTLS12} + } else { + baseTransport = &http.Transport{ Proxy: http.ProxyFromEnvironment, ForceAttemptHTTP2: true, TLSClientConfig: &tls.Config{MinVersion: tls.VersionTLS12}, - }, + } + } + baseHTTP := &http.Client{ + CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }, + Transport: baseTransport, } var wrapped client.HTTPClient = baseHTTP @@ -250,8 +168,10 @@ func newRuntimeTransport(cmd *cobra.Command, explicitRuntimeCred bool) (*runtime // Bidirectional gRPC streams stream frames over an open io.Pipe body and must // NOT be wrapped in DebugClient (which calls io.ReadAll on req.Body and blocks). switch { - case rt.hasTestClient || isDryRun(cmd): + case isDryRun(cmd): rt.streamClient = rt.httpClient + case rt.hasTestClient: + rt.streamClient = wrapped case rt.scheme == "http": var h2cProtos http.Protocols h2cProtos.SetUnencryptedHTTP2(true) @@ -287,13 +207,25 @@ func validateRuntimeServerURL(raw string) error { return nil } -func validateHeaderSecret(flagName, val string) error { - for _, r := range val { - if r < 0x20 || r == 0x7f { - return fmt.Errorf("%s contains control characters", flagName) - } - } - return nil +func (rt *runtimeTransport) newEnvHTTPClient(envID, envKey string) *envruntime.HTTPClient { + return envruntime.NewHTTPClient(envruntime.HTTPClientConfig{ + BaseURL: rt.baseURL, + APIVersion: rt.apiVersion, + APIKey: rt.apiKey, + AccessToken: rt.accessToken, + EnvironmentID: envID, + EnvironmentKey: envKey, + UserProject: rt.userProject, + APIRevision: rt.apiRevision, + UserAgent: rawUserAgent, + Headers: rt.headers, + Timeout: rt.timeout, + HTTPDoer: rt.httpClient, + StreamDoer: rt.streamClient, + Debug: rt.debug, + DebugOut: rt.errOut, + ProbePlainHTTPOnStreamError: true, + }) } func (rt *runtimeTransport) newRequest(ctx context.Context, method, path string, body io.Reader) (*http.Request, error) { @@ -302,410 +234,115 @@ func (rt *runtimeTransport) newRequest(ctx context.Context, method, path string, if err != nil { return nil, err } - req.Header.Set("User-Agent", rawUserAgent) - for k, v := range rt.headers { - req.Header.Set(k, v) - } - switch { - case rt.apiKey != "": - req.Header.Set("x-goog-api-key", rt.apiKey) - case rt.accessToken != "": - req.Header.Set("Authorization", "Bearer "+rt.accessToken) - } - if rt.userProject != "" { - req.Header.Set("x-goog-user-project", rt.userProject) - } - if rt.apiRevision != "" { - req.Header.Set("Api-Revision", rt.apiRevision) + for k, vals := range rt.newEnvHTTPClient("", "").BuildHeaders("", "", "") { + for _, v := range vals { + req.Header.Add(k, v) + } } return req, nil } -// HTTPClientOptions configures an environment runtime HTTP/gRPC client. +// HTTPClientOptions configures a CLI environment runtime HTTP/gRPC client. type HTTPClientOptions struct { Transport *runtimeTransport EnvironmentID string EnvironmentKey string } -// HTTPClient communicates with the Gemini API LeaseInteractions and bidi stream endpoints. +// HTTPClient adapts envruntime.HTTPClient to surface CLI sdkerrors.SDKDefaultError diagnostics. type HTTPClient struct { - rt *runtimeTransport - environmentID string - environmentKey string + inner *envruntime.HTTPClient } // NewHTTPClient creates an HTTPClient bound to opts. func NewHTTPClient(opts HTTPClientOptions) *HTTPClient { - return &HTTPClient{ - rt: opts.Transport, - environmentID: strings.TrimSpace(opts.EnvironmentID), - environmentKey: strings.TrimSpace(opts.EnvironmentKey), + var inner *envruntime.HTTPClient + if opts.Transport != nil { + inner = opts.Transport.newEnvHTTPClient(opts.EnvironmentID, opts.EnvironmentKey) + } else { + inner = envruntime.NewHTTPClient(envruntime.HTTPClientConfig{ + EnvironmentID: opts.EnvironmentID, + EnvironmentKey: opts.EnvironmentKey, + ProbePlainHTTPOnStreamError: true, + }) } + return &HTTPClient{inner: inner} } // ServerURL returns the normalized base URL configured on the client. func (c *HTTPClient) ServerURL() string { - return c.rt.baseURL + return c.inner.ServerURL() } -func (c *HTTPClient) applyEnvironmentHeaders(req *http.Request, envID, envKey, leaseToken string) { - if envKey != "" { - req.Header.Set("x-goog-environment-key", envKey) - } - if envID != "" { - req.Header.Set("x-goog-environment-id", envID) - } - if leaseToken != "" { - req.Header.Set("x-goog-lease-token", leaseToken) +// LeaseInteractions implements envruntime.EnvironmentClient, mapping transport errors to SDKDefaultError. +func (c *HTTPClient) LeaseInteractions(ctx context.Context, req *envruntime.LeaseInteractionsRequest) (*envruntime.LeaseInteractionsResult, error) { + res, err := c.inner.LeaseInteractions(ctx, req) + if err != nil { + return nil, wrapTransportError(err) } + return res, nil } -// LeaseInteractions polls POST //environments:leaseInteractions for queued work. -func (c *HTTPClient) LeaseInteractions(ctx context.Context, workerID string, maxNew int32) ([]*LeasedInteraction, error) { - parsed, err := c.doLeaseRPC(ctx, c.environmentID, leaseInteractionsRequest{ - WorkerID: workerID, - MaxNewInteractions: &maxNew, - EnvironmentID: c.environmentID, - }, "LeaseInteractions failed") +// LeaseSimple polls POST //environments:leaseInteractions for up to maxNew queued interactions. +func (c *HTTPClient) LeaseSimple(ctx context.Context, workerID string, maxNew int32) ([]*LeasedInteraction, error) { + leases, err := c.inner.LeaseSimple(ctx, workerID, maxNew) if err != nil { - return nil, err + return nil, wrapTransportError(err) } - return parsed.Leases, nil + return leases, nil } // HeartbeatInteraction renews a lease over HTTP while a poller hands work off to a worker. func (c *HTTPClient) HeartbeatInteraction(ctx context.Context, pollerID, interactionID, envID, leaseToken string) (*HeartbeatResult, error) { - zero := int32(0) - resolvedEnvID := firstNonEmpty(envID, c.environmentID) - parsed, err := c.doLeaseRPC(ctx, resolvedEnvID, leaseInteractionsRequest{ - WorkerID: pollerID, - ActiveLeases: []activeInteractionLease{{ - InteractionID: interactionID, - LeaseToken: leaseToken, - }}, - MaxNewInteractions: &zero, - EnvironmentID: resolvedEnvID, - }, "LeaseInteractions heartbeat failed") - if err != nil { - var sdkErr *sdkerrors.SDKDefaultError - if errors.As(err, &sdkErr) { - switch sdkErr.StatusCode { - case http.StatusNotFound, http.StatusConflict, http.StatusPreconditionFailed: - return &HeartbeatResult{ShouldCancel: true}, nil - } - } - return nil, err - } - for _, l := range parsed.Leases { - if l != nil && l.InteractionID == interactionID { - return &HeartbeatResult{Renewed: true}, nil - } - } - return &HeartbeatResult{ShouldCancel: true}, nil -} - -func (c *HTTPClient) doLeaseRPC(ctx context.Context, envID string, reqBody leaseInteractionsRequest, errPrefix string) (*leaseInteractionsResponse, error) { - payload, err := json.Marshal(reqBody) - if err != nil { - return nil, err - } - reqCtx := ctx - var cancel context.CancelFunc - if c.rt.timeout > 0 { - reqCtx, cancel = context.WithTimeout(ctx, c.rt.timeout) - defer cancel() - } - req, err := c.rt.newRequest(reqCtx, http.MethodPost, c.rt.apiVersion+"/environments:leaseInteractions", bytes.NewReader(payload)) + res, err := c.inner.HeartbeatInteraction(ctx, pollerID, interactionID, envID, leaseToken) if err != nil { - return nil, err + return nil, wrapTransportError(err) } - req.Header.Set("Content-Type", "application/json") - c.applyEnvironmentHeaders(req, envID, c.environmentKey, "") - - resp, err := c.rt.httpClient.Do(req) - if err != nil { - return nil, err - } - defer resp.Body.Close() - - body, err := io.ReadAll(io.LimitReader(resp.Body, 8<<20)) - if err != nil { - return nil, err - } - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - return nil, sdkerrors.NewSDKDefaultError(errPrefix, resp.StatusCode, string(body), resp) - } - if len(bytes.TrimSpace(body)) == 0 { - return &leaseInteractionsResponse{}, nil - } - var parsed leaseInteractionsResponse - if err := json.Unmarshal(body, &parsed); err != nil { - return nil, fmt.Errorf("decoding LeaseInteractions response: %w", err) - } - return &parsed, nil -} - -type connectStreamRequest struct { - InteractionID string - EnvironmentID string - LeaseToken string - NextOutputSeqNum int64 - AckOutputSeqNum int64 -} - -func (c *HTTPClient) bidiStreamMethodPath() string { - return fmt.Sprintf("google.learning.gemini.api.interactions.%s.InteractionsService/CreateInteractionBidiStreamByEnvironmentKey", c.rt.apiVersion) + return res, nil } // PreviewConnect emits the dry-run preview for opening an interaction bidi stream without touching the network. func (c *HTTPClient) PreviewConnect(ctx context.Context, leaseToken string) error { - req, err := c.rt.newRequest(ctx, http.MethodPost, c.bidiStreamMethodPath(), http.NoBody) - if err != nil { - return err - } - req.Header.Set("Content-Type", "application/grpc") - req.Header.Set("TE", "trailers") - c.applyEnvironmentHeaders(req, c.environmentID, c.environmentKey, leaseToken) - resp, err := c.rt.httpClient.Do(req) - if err != nil { - return err - } - if resp != nil && resp.Body != nil { - _ = resp.Body.Close() - } - return nil -} - -type grpcInteractionStream struct { - resp *http.Response - pw *io.PipeWriter - sendMu sync.Mutex - closed bool + return wrapTransportError(c.inner.PreviewConnect(ctx, leaseToken)) } -func (s *grpcInteractionStream) Send(batch *clientEventBatch) error { - payload := batch.marshalWire() - s.sendMu.Lock() - defer s.sendMu.Unlock() - if s.closed { - return io.ErrClosedPipe - } - return grpcWriteFrame(s.pw, payload) -} - -func (s *grpcInteractionStream) Recv() (*serverEventBatch, error) { - frame, err := grpcReadFrame(s.resp.Body) +// ConnectInteraction implements envruntime.EnvironmentClient, wrapping HTTP and gRPC trailer errors into SDKDefaultError. +func (c *HTTPClient) ConnectInteraction(ctx context.Context, req *envruntime.ConnectInteractionRequest) (envruntime.InteractionStream, error) { + stream, err := c.inner.ConnectInteraction(ctx, req) if err != nil { - if errors.Is(err, io.EOF) { - if stErr := checkGRPCTrailers(s.resp); stErr != nil { - return nil, stErr - } - } - return nil, err + return nil, wrapTransportError(err) } - return unmarshalServerEventBatch(frame) + return &sdkErrorStream{InteractionStream: stream}, nil } -func (s *grpcInteractionStream) CloseSend() error { - s.sendMu.Lock() - defer s.sendMu.Unlock() - if s.closed { - return nil - } - s.closed = true - return s.pw.Close() -} - -func (s *grpcInteractionStream) Close() error { - _ = s.CloseSend() - if s.resp != nil && s.resp.Body != nil { - return s.resp.Body.Close() - } - return nil -} - -func (c *HTTPClient) connectStream(ctx context.Context, req connectStreamRequest) (*grpcInteractionStream, error) { - pr, pw := io.Pipe() - methodPath := c.bidiStreamMethodPath() - httpReq, err := c.rt.newRequest(ctx, http.MethodPost, methodPath, pr) - if err != nil { - _ = pw.Close() - _ = pr.Close() - return nil, err - } - httpReq.Header.Set("Content-Type", "application/grpc") - httpReq.Header.Set("TE", "trailers") - c.applyEnvironmentHeaders(httpReq, firstNonEmpty(req.EnvironmentID, c.environmentID), c.environmentKey, req.LeaseToken) - - if c.rt.debug && c.rt.errOut != nil { - _, _ = fmt.Fprintf(c.rt.errOut, "=== DEBUG STREAM ===\n%s %s\n", httpReq.Method, httpReq.URL.String()) - } - - handshake := (&clientEventBatch{ - AckOutputSeqNum: req.AckOutputSeqNum, - ReconnectInteraction: &reconnectInteractionEvent{ - InteractionID: req.InteractionID, - NextOutputSeqNum: req.NextOutputSeqNum, - ClientID: "environment", - }, - }).marshalWire() - - // Write the initial ReconnectInteraction frame immediately once the transport starts reading pr. - writeErrCh := make(chan error, 1) - go func() { - writeErrCh <- grpcWriteFrame(pw, handshake) - }() - - resp, err := c.rt.streamClient.Do(httpReq) - if err != nil { - _ = pw.CloseWithError(err) - _ = pr.Close() - if c.rt.scheme == "http" && ctx.Err() == nil { - if httpErr := c.probePlainHTTPError(ctx, methodPath, req); httpErr != nil { - return nil, httpErr - } - } - return nil, err - } - - if c.rt.debug && c.rt.errOut != nil { - _, _ = fmt.Fprintf(c.rt.errOut, "=== DEBUG STREAM RESPONSE ===\n%s\n", resp.Status) - } - - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - _ = pw.Close() - body, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<20)) - _ = resp.Body.Close() - return nil, sdkerrors.NewSDKDefaultError("ConnectInteraction failed", resp.StatusCode, string(body), resp) - } - if stErr := checkGRPCHeaders(resp); stErr != nil { - _ = pw.Close() - _ = resp.Body.Close() - return nil, stErr - } - if werr := <-writeErrCh; werr != nil { - _ = pw.Close() - _ = resp.Body.Close() - return nil, werr - } - - return &grpcInteractionStream{ - resp: resp, - pw: pw, - }, nil +type sdkErrorStream struct { + envruntime.InteractionStream } -func (c *HTTPClient) probePlainHTTPError(ctx context.Context, methodPath string, req connectStreamRequest) error { - probeReq, err := c.rt.newRequest(ctx, http.MethodPost, methodPath, http.NoBody) +func (s *sdkErrorStream) Recv(ctx context.Context) (*envruntime.ServerEventBatch, error) { + batch, err := s.InteractionStream.Recv(ctx) if err != nil { - return nil + return nil, wrapTransportError(err) } - probeReq.Header.Set("Content-Type", "application/grpc") - c.applyEnvironmentHeaders(probeReq, firstNonEmpty(req.EnvironmentID, c.environmentID), c.environmentKey, req.LeaseToken) - resp, err := c.rt.httpClient.Do(probeReq) - if err != nil { - return nil - } - defer resp.Body.Close() - if resp.StatusCode >= 200 && resp.StatusCode < 300 { - return nil - } - body, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<20)) - return sdkerrors.NewSDKDefaultError("ConnectInteraction failed", resp.StatusCode, string(body), resp) -} - -func checkGRPCHeaders(resp *http.Response) error { - return parseGRPCStatus(resp.Header.Get("Grpc-Status"), resp.Header.Get("Grpc-Message"), resp) + return batch, nil } -func checkGRPCTrailers(resp *http.Response) error { - return parseGRPCStatus(resp.Trailer.Get("Grpc-Status"), resp.Trailer.Get("Grpc-Message"), resp) -} - -func parseGRPCStatus(rawCode, rawMsg string, resp *http.Response) error { - trimmed := strings.TrimSpace(rawCode) - if trimmed == "" || trimmed == "0" { +func wrapTransportError(err error) error { + if err == nil { return nil } - code, err := strconv.Atoi(trimmed) - if err != nil { - code = 2 - } - msg := decodeGRPCMessage(rawMsg) - if msg == "" { - msg = fmt.Sprintf("gRPC status %d", code) - } - httpStatus := grpcCodeToHTTPStatus(code) - errBody, _ := json.Marshal(map[string]any{ - "error": map[string]any{ - "code": httpStatus, - "message": msg, - }, - }) - return sdkerrors.NewSDKDefaultError(msg, httpStatus, string(errBody), resp) -} - -func grpcCodeToHTTPStatus(code int) int { - switch code { - case 3, 9, 11: // INVALID_ARGUMENT, FAILED_PRECONDITION, OUT_OF_RANGE - return http.StatusBadRequest - case 4: // DEADLINE_EXCEEDED - return http.StatusGatewayTimeout - case 5: // NOT_FOUND - return http.StatusNotFound - case 6, 10: // ALREADY_EXISTS, ABORTED - return http.StatusConflict - case 7: // PERMISSION_DENIED - return http.StatusForbidden - case 8: // RESOURCE_EXHAUSTED - return http.StatusTooManyRequests - case 12: // UNIMPLEMENTED - return http.StatusNotImplemented - case 14: // UNAVAILABLE - return http.StatusServiceUnavailable - case 16: // UNAUTHENTICATED - return http.StatusUnauthorized - default: - return http.StatusInternalServerError + var te *envruntime.TransportError + if errors.As(err, &te) { + msg := cmp.Or(te.Op, te.Message) + return sdkerrors.NewSDKDefaultError(msg, te.StatusCode, te.Body, te.Resp) } -} - -func firstNonEmpty(values ...string) string { - for _, v := range values { - if v != "" { - return v - } - } - return "" + return err } func isNonRetryableControlPlaneError(err error) bool { - var sdkErr *sdkerrors.SDKDefaultError - if errors.As(err, &sdkErr) { - switch sdkErr.StatusCode { - case http.StatusBadRequest, http.StatusUnauthorized, http.StatusForbidden, http.StatusNotFound: - return true - } - } - return false -} - -func retryBackoff(base time.Duration, consecutive int) time.Duration { - if base <= 0 { - base = time.Second - } - var shift int - if consecutive > 1 { - shift = min(consecutive-1, 5) + if envruntime.IsNonRetryableError(err) { + return true } - return min(base*time.Duration(1< 10*time.Minute { - delete(p.failures, id) - } - } - st, ok := p.failures[interactionID] - if !ok { - st = &failureState{} - p.failures[interactionID] = st - } - st.consecutive++ - shift := min(st.consecutive-1, 5) - st.cooldownEnd = now.Add(min(time.Duration(1<= 0; i-- { - eval = filepath.Join(eval, suffix[i]) - } - return filepath.Clean(eval) - } - parent := filepath.Dir(cur) - if parent == cur { - break - } - suffix = append(suffix, filepath.Base(cur)) - cur = parent - } - return filepath.Clean(absPath) -} - -type fsServer struct { - rootDir string - workDir string -} - -func newFSServer(rootDir, workDir string) *fsServer { - if abs, err := filepath.Abs(rootDir); err == nil { - rootDir = canonicalizeLongestPrefix(abs) - } - cleanRoot := filepath.Clean(rootDir) - cleanWork := cleanRoot - if strings.TrimSpace(workDir) != "" { - if absWork, err := filepath.Abs(workDir); err == nil { - cleanWork = canonicalizeLongestPrefix(absWork) - } else { - cleanWork = filepath.Clean(workDir) - } - if !isWithinRoot(cleanRoot, cleanWork) { - cleanWork = cleanRoot - } - } - return &fsServer{ - rootDir: cleanRoot, - workDir: cleanWork, - } -} - -func (s *fsServer) ResolvePath(relPath string) (string, error) { - if strings.Contains(relPath, "\x00") { - return "", fmt.Errorf("access denied: path contains NUL byte") - } - var candidate string - if filepath.IsAbs(relPath) { - cleaned := filepath.Clean(relPath) - if isWithinRoot(s.rootDir, cleaned) { - candidate = cleaned - } else { - trimmed := strings.TrimPrefix(cleaned, filepath.VolumeName(cleaned)) - trimmed = strings.TrimPrefix(trimmed, string(filepath.Separator)) - candidate = filepath.Join(s.rootDir, trimmed) - } - } else { - candidate = filepath.Join(s.workDir, relPath) - } - if !isWithinRoot(s.rootDir, candidate) { - return "", fmt.Errorf("access denied: path %q escapes root directory", relPath) - } - - cur := candidate - var suffix []string - for { - if eval, err := filepath.EvalSymlinks(cur); err == nil { - resolved := eval - for i := len(suffix) - 1; i >= 0; i-- { - resolved = filepath.Join(resolved, suffix[i]) - if fi, lerr := os.Lstat(resolved); lerr == nil && fi.Mode()&os.ModeSymlink != 0 { - return "", fmt.Errorf("access denied: dangling symlink %q is not permitted", relPath) - } - } - if !isWithinRoot(s.rootDir, resolved) { - return "", fmt.Errorf("access denied: symlink target for %q escapes root directory", relPath) - } - return resolved, nil - } - if fi, lerr := os.Lstat(cur); lerr == nil && fi.Mode()&os.ModeSymlink != 0 { - return "", fmt.Errorf("access denied: dangling symlink %q is not permitted", relPath) - } - parent := filepath.Dir(cur) - if parent == cur { - break - } - suffix = append(suffix, filepath.Base(cur)) - cur = parent - } - return candidate, nil -} - -func (s *fsServer) resolveRelInRoot(relPath string) (string, error) { - resolved, err := s.ResolvePath(relPath) - if err != nil { - return "", err - } - return filepath.Rel(s.rootDir, resolved) -} - -func (s *fsServer) openInRoot(relPath string) (*os.Root, string, error) { - rel, err := s.resolveRelInRoot(relPath) - if err != nil { - return nil, "", err - } - root, err := os.OpenRoot(s.rootDir) - if err != nil { - return nil, "", err - } - return root, rel, nil -} - -func mkdirAllInRoot(root *os.Root, relDir string) error { - cur := "" - for _, part := range strings.Split(filepath.Clean(relDir), string(filepath.Separator)) { - if cur == "" { - cur = part - } else { - cur = cur + string(filepath.Separator) + part - } - if fi, err := root.Stat(cur); err == nil && fi.IsDir() { - continue - } - if err := root.Mkdir(cur, 0o755); err != nil { - return err - } - } - return nil -} - -func (s *fsServer) ReadRegularFile(path string) ([]byte, error) { - root, rel, err := s.openInRoot(path) - if err != nil { - return nil, err - } - defer root.Close() - - f, err := root.Open(rel) - if err != nil { - return nil, err - } - defer f.Close() - - info, err := f.Stat() - if err != nil { - return nil, err - } - if info.IsDir() { - return nil, fmt.Errorf("cannot read directory %q as file", path) - } - if !info.Mode().IsRegular() { - return nil, fmt.Errorf("cannot read non-regular file %q", path) - } - data, exceeded, err := readLimitedAll(f, maxReadFileBytes) - if err != nil { - return nil, err - } - if exceeded { - return nil, fmt.Errorf("file %q exceeds maximum full-read size (%d bytes); specify start_line and end_line", path, maxReadFileBytes) - } - return data, nil -} - -func readLimitedAll(r io.Reader, maxBytes int) ([]byte, bool, error) { - data, err := io.ReadAll(io.LimitReader(r, int64(maxBytes)+1)) - if err != nil { - return nil, false, err - } - if len(data) > maxBytes { - return nil, true, nil - } - return data, false, nil -} - -func (s *fsServer) ReadFileRange(path string, startLine, endLine int) (string, error) { - root, rel, err := s.openInRoot(path) - if err != nil { - return "", err - } - defer root.Close() - - f, err := root.Open(rel) - if err != nil { - return "", err - } - defer f.Close() - - info, err := f.Stat() - if err != nil { - return "", err - } - if info.IsDir() { - return "", fmt.Errorf("cannot read directory %q as file", path) - } - if !info.Mode().IsRegular() { - return "", fmt.Errorf("cannot read non-regular file %q", path) - } - if startLine <= 0 { - startLine = 1 - } - - reader := bufio.NewReader(f) - var selected []string - totalBytes := 0 - for lineNum := 1; ; lineNum++ { - part, readErr := reader.ReadString('\n') - if readErr != nil && !errors.Is(readErr, io.EOF) { - return "", readErr - } - line := strings.TrimSuffix(strings.TrimSuffix(part, "\n"), "\r") - if lineNum >= startLine && (endLine <= 0 || lineNum <= endLine) { - totalBytes += len(line) + 1 - if totalBytes > maxReadFileBytes { - return "", fmt.Errorf("line range [%d, %d] of %q exceeds maximum read size (%d bytes); narrow start_line and end_line", startLine, endLine, path, maxReadFileBytes) - } - selected = append(selected, line) - } - if (endLine > 0 && lineNum >= endLine) || errors.Is(readErr, io.EOF) { - break - } - } - return strings.Join(selected, "\n"), nil -} - -func (s *fsServer) WriteRegularFile(path, content string, overwrite, appendMode bool) (int, error) { - root, rel, err := s.openInRoot(path) - if err != nil { - return 0, err - } - defer root.Close() - - if fi, err := root.Lstat(rel); err == nil && !fi.Mode().IsRegular() { - return 0, fmt.Errorf("cannot write non-regular file %q", path) - } - if err := mkdirAllInRoot(root, filepath.Dir(rel)); err != nil { - return 0, err - } - flags := os.O_CREATE | os.O_WRONLY - switch { - case appendMode: - flags |= os.O_APPEND - case overwrite: - flags |= os.O_TRUNC - default: - flags |= os.O_EXCL - } - f, err := root.OpenFile(rel, flags, 0o644) - if err != nil { - return 0, err - } - defer f.Close() - return f.WriteString(content) -} - -func (s *fsServer) ListEntries(path string) ([]any, error) { - root, rel, err := s.openInRoot(path) - if err != nil { - return nil, err - } - defer root.Close() - - dirFile, err := root.Open(rel) - if err != nil { - return nil, err - } - defer dirFile.Close() - - entries, err := dirFile.ReadDir(-1) - if err != nil { - return nil, err - } - out := make([]any, 0, len(entries)) - for _, entry := range entries { - var size int64 - if info, infoErr := entry.Info(); infoErr == nil { - size = info.Size() - } - out = append(out, map[string]any{ - "name": entry.Name(), - "is_dir": entry.IsDir(), - "size": int(size), - }) - } - return out, nil -} - -func (s *fsServer) Grep(searchPath, query string, caseInsensitive, isRegex bool) ([]any, bool, error) { - fullPath, err := s.ResolvePath(searchPath) - if err != nil { - return nil, false, err - } - if query == "" { - return nil, false, fmt.Errorf("query is required") - } - pattern := query - if !isRegex { - pattern = regexp.QuoteMeta(query) - } - if caseInsensitive { - pattern = "(?i)" + pattern - } - re, err := regexp.Compile(pattern) - if err != nil { - return nil, false, fmt.Errorf("invalid search pattern: %w", err) - } - root, err := os.OpenRoot(s.rootDir) - if err != nil { - return nil, false, err - } - defer root.Close() - - var matches []any - totalBytes := 0 - err = filepath.WalkDir(fullPath, func(p string, d fs.DirEntry, walkErr error) error { - if walkErr != nil || d.IsDir() { - return nil - } - fileRel, err := s.resolveRelInRoot(p) - if err != nil { - return nil - } - f, err := root.Open(fileRel) - if err != nil { - return nil - } - info, statErr := f.Stat() - if statErr != nil || info.IsDir() || !info.Mode().IsRegular() || info.Size() > maxReadFileBytes { - _ = f.Close() - return nil - } - data, readErr := io.ReadAll(io.LimitReader(f, maxReadFileBytes+1)) - _ = f.Close() - if readErr != nil || len(data) > maxReadFileBytes { - return nil - } - rel, err := filepath.Rel(fullPath, p) - if err != nil || rel == "." { - rel = filepath.Base(p) - } - slashRel := filepath.ToSlash(rel) - for idx, rawLine := range strings.Split(string(data), "\n") { - line := strings.TrimSuffix(rawLine, "\r") - if re.MatchString(line) { - entryBytes := len(slashRel) + len(line) + 32 - if len(matches) >= maxGrepMatches || (len(matches) > 0 && totalBytes+entryBytes > maxGrepOutputBytes) { - return errSearchTruncated - } - totalBytes += entryBytes - matches = append(matches, map[string]any{ - "file": slashRel, - "line_number": idx + 1, - "line": line, - }) - } - } - return nil - }) - if errors.Is(err, errSearchTruncated) { - return matches, true, nil - } - return matches, false, err -} - -func (s *fsServer) FindFiles(searchDir, pattern string, maxDepth int, extensions ...string) ([]any, bool, error) { - fullPath, err := s.ResolvePath(searchDir) - if err != nil { - return nil, false, err - } - var extAllow map[string]bool - for _, rawExt := range extensions { - clean := strings.TrimPrefix(strings.ToLower(strings.TrimSpace(rawExt)), ".") - if clean != "" { - if extAllow == nil { - extAllow = make(map[string]bool, len(extensions)) - } - extAllow[clean] = true - } - } - patternLower := strings.ToLower(pattern) - var files []any - err = filepath.WalkDir(fullPath, func(p string, d fs.DirEntry, walkErr error) error { - if walkErr != nil || p == fullPath { - return nil - } - if _, err := s.ResolvePath(p); err != nil { - return nil - } - rel, err := filepath.Rel(fullPath, p) - if err != nil { - return nil - } - slashRel := filepath.ToSlash(rel) - if maxDepth > 0 && strings.Count(slashRel, "/")+1 > maxDepth { - if d.IsDir() { - return filepath.SkipDir - } - return nil - } - name := d.Name() - if len(extAllow) > 0 { - if d.IsDir() { - return nil - } - fileExt := strings.TrimPrefix(strings.ToLower(filepath.Ext(name)), ".") - if !extAllow[fileExt] { - return nil - } - } - matched, _ := filepath.Match(pattern, name) - if matched || patternLower == "" || strings.Contains(strings.ToLower(name), patternLower) { - if len(files) >= maxFindFiles { - return errSearchTruncated - } - files = append(files, slashRel) - } - return nil - }) - if errors.Is(err, errSearchTruncated) { - return files, true, nil - } - return files, false, err -} - -type shellServer struct { - workDir string -} - -func newShellServer(workDir string) *shellServer { - return &shellServer{workDir: workDir} -} - -type cappedBuffer struct { - buf bytes.Buffer - maxBytes int - truncated bool -} - -func (c *cappedBuffer) Write(p []byte) (int, error) { - origLen := len(p) - remaining := c.maxBytes - c.buf.Len() - if remaining <= 0 { - c.truncated = true - return origLen, nil - } - if len(p) > remaining { - c.buf.Write(p[:remaining]) - c.truncated = true - return origLen, nil - } - c.buf.Write(p) - return origLen, nil -} - -func (c *cappedBuffer) String() string { - if c.truncated { - return c.buf.String() + "\n...[output truncated at 4 MiB]..." - } - return c.buf.String() -} - -var scrubbedCredentialEnvKeys = map[string]bool{ - "GEMINI_ENVIRONMENT_KEY": true, - "ENVIRONMENT_KEY": true, - "GEMINI_API_KEY": true, - "GOOGLE_API_KEY": true, - "GEMINI_ACCESS_TOKEN": true, - "ACCESS_TOKEN": true, - "GOOGLE_OAUTH_ACCESS_TOKEN": true, - "CLOUDSDK_AUTH_ACCESS_TOKEN": true, - "GEMINI_LEASE_TOKEN": true, - "LEASE_TOKEN": true, -} - -func filterScrubbedEnv(baseEnv []string) []string { - out := make([]string, 0, len(baseEnv)) - for _, kv := range baseEnv { - key, _, _ := strings.Cut(kv, "=") - if scrubbedCredentialEnvKeys[strings.ToUpper(strings.TrimSpace(key))] { - continue - } - out = append(out, kv) - } - return out -} - -var allowedWorkerBashEnvKeys = map[string]bool{ - "PATH": true, - "HOME": true, - "USER": true, - "LOGNAME": true, - "SHELL": true, - "LANG": true, - "TERM": true, - "TMPDIR": true, - "SYSTEMROOT": true, - "WINDIR": true, - "COMSPEC": true, - "PATHEXT": true, - "TEMP": true, - "TMP": true, - "USERPROFILE": true, - "APPDATA": true, - "LOCALAPPDATA": true, - "PROGRAMFILES": true, - "PROGRAMFILES(X86)": true, - "PROGRAMDATA": true, - "HOMEDRIVE": true, - "HOMEPATH": true, - "COMPUTERNAME": true, - "USERNAME": true, - "USERDOMAIN": true, - "OS": true, - "SYSTEMDRIVE": true, - "PROCESSOR_ARCHITECTURE": true, - "NUMBER_OF_PROCESSORS": true, - "PSMODULEPATH": true, -} - -func buildWorkerBashEnv(baseEnv []string) []string { - out := make([]string, 0, 12) - for _, kv := range baseEnv { - key, _, ok := strings.Cut(kv, "=") - if !ok { - continue - } - upper := strings.ToUpper(strings.TrimSpace(key)) - if allowedWorkerBashEnvKeys[upper] || strings.HasPrefix(upper, "LC_") { - out = append(out, kv) - } - } - return out -} - -func resolveShellCommand(goos, command string) (string, []string) { - if p, err := exec.LookPath("bash"); err == nil { - return p, []string{"-c", command} - } - if p, err := exec.LookPath("sh"); err == nil { - return p, []string{"-c", command} - } - if goos == "windows" { - return "cmd.exe", []string{"/C", command} - } - return "sh", []string{"-c", command} -} - -func (s *shellServer) ExecuteBash(ctx context.Context, command, workDir string, timeout time.Duration) (string, string, int, error) { - if timeout <= 0 { - timeout = 2 * time.Minute - } - ctx, cancel := context.WithTimeout(ctx, timeout) - defer cancel() - - dir := s.workDir - if workDir != "" { - dir = workDir - } - - shellBin, shellArgs := resolveShellCommand(runtime.GOOS, command) - cmd := exec.CommandContext(ctx, shellBin, shellArgs...) - cmd.Dir = dir - configureSubprocessGroup(cmd) - cmd.WaitDelay = 5 * time.Second - cmd.Env = buildWorkerBashEnv(os.Environ()) - - stdoutBuf := &cappedBuffer{maxBytes: maxBashOutputBytes} - stderrBuf := &cappedBuffer{maxBytes: maxBashOutputBytes} - cmd.Stdout = stdoutBuf - cmd.Stderr = stderrBuf - - err := cmd.Run() - stdout := stdoutBuf.String() - stderr := stderrBuf.String() - - if errors.Is(ctx.Err(), context.DeadlineExceeded) { - return stdout, stderr, -1, fmt.Errorf("command timed out after %v", timeout) - } - if err != nil { - var exitErr *exec.ExitError - if errors.As(err, &exitErr) { - return stdout, stderr, exitErr.ExitCode(), nil - } - return stdout, stderr, -1, err - } - return stdout, stderr, 0, nil -} diff --git a/internal/cli/custom/runtime_test.go b/internal/cli/custom/runtime_test.go index b66fe07..0114dc7 100644 --- a/internal/cli/custom/runtime_test.go +++ b/internal/cli/custom/runtime_test.go @@ -17,27 +17,17 @@ package custom import ( "bytes" "context" - "encoding/binary" "encoding/json" - "errors" - "fmt" - "io" - "log" - "math" "net" "net/http" - "net/http/httptest" "os" - "os/exec" "path/filepath" - "reflect" - "slices" "strings" - "sync" "sync/atomic" "testing" "time" + "github.com/google-gemini/gemini-api-cli/internal/cli/custom/envruntime" "github.com/google-gemini/gemini-api-cli/internal/clierrors" "github.com/google-gemini/gemini-api-cli/internal/config" "github.com/google-gemini/gemini-api-cli/internal/flagutil" @@ -94,4319 +84,803 @@ func startH2CTestServer(t *testing.T, handler http.Handler) string { return "http://" + ln.Addr().String() } -func buildServerStepToolCallFrame(seq int64, callID, fnName string, args map[string]any) []byte { - fnCall := appendStringField(nil, 1, fnName) - fnCall = appendBytesField(fnCall, 2, encodeProtoStruct(args)) - toolCall := appendStringField(nil, 1, callID) - toolCall = appendBytesField(toolCall, 3, fnCall) - step := appendBytesField(nil, 4, toolCall) - payload := appendBytesField(nil, 11, step) - ev := appendBytesField(nil, 3, payload) - ev = appendVarintField(ev, 10, uint64(seq)) - return appendBytesField(nil, 1, ev) -} - -func buildServerReconnectRequiredFrame() []byte { - return appendBytesField(nil, 9, nil) -} - -func buildServerYieldFrame(nextInputSeq int64) []byte { - return appendBytesField(nil, 10, appendVarintField(nil, 1, uint64(nextInputSeq))) -} - -type parsedHandshake struct { - ackSeq int64 - interactionID string - nextOutSeq int64 - hasNextOutSeq bool - background bool - clientID string -} - -func parseHandshakeFrame(t *testing.T, frame []byte) parsedHandshake { +func newRuntimeCLIRootForTest(t *testing.T, serverURL string) (*cobra.Command, *bytes.Buffer, *bytes.Buffer) { t.Helper() - var hs parsedHandshake - rest := frame - for len(rest) > 0 { - num, _, v, sub, n, err := consumeField(rest) - if err != nil { - t.Fatalf("consumeField(handshake) failed: %v", err) - } - switch num { - case 2: - hs.ackSeq = int64(v) - case 8: - inner := sub - for len(inner) > 0 { - inum, _, iv, ib, in, ierr := consumeField(inner) - if ierr != nil { - t.Fatalf("consumeField(ReconnectInteraction) failed: %v", ierr) - } - switch inum { - case 1: - hs.interactionID = string(ib) - case 2: - hs.nextOutSeq = int64(iv) - hs.hasNextOutSeq = true - case 3: - hs.background = iv != 0 - case 4: - hs.clientID = string(ib) - } - inner = inner[in:] - } - } - rest = rest[n:] - } - return hs -} - -func TestWireFrameBoundsAndErrors(t *testing.T) { - var compressedHdr [5]byte - compressedHdr[0] = 1 - - var oversizedHdr [5]byte - binary.BigEndian.PutUint32(oversizedHdr[1:5], maxGRPCFrameBytes+1) - - var truncatedHdr [5]byte - binary.BigEndian.PutUint32(truncatedHdr[1:5], 16) + root := &cobra.Command{Use: "gemini-api", SilenceErrors: true, SilenceUsage: true} + pf := root.PersistentFlags() + pf.StringP("output-format", "o", "pretty", "") + pf.String("color", "never", "") + pf.StringP("jq", "q", "", "") + pf.Bool("raw-output", true, "") + pf.String("server-url", serverURL, "") + pf.StringArrayP("header", "H", nil, "") + pf.Bool("include-headers", false, "") + pf.String("timeout", "", "") + pf.Bool("interactive", false, "") + pf.Bool("no-interactive", true, "") + pf.Bool("usage", false, "") + pf.Bool("dry-run", false, "") + pf.BoolP("debug", "d", false, "") + pf.Bool("agent-mode", false, "") + pf.Bool("no-retries", false, "") + pf.Bool("no-keyring", true, "") + pf.String("api-key", "", "") + pf.String("access-token", "", "") + pf.String("api-version", "v1beta", "") + pf.String("api-revision", "", "") + pf.String("user-project", "", "") - tests := []struct { - name string - raw []byte - }{ - {name: "compressed_frame_rejected", raw: compressedHdr[:]}, - {name: "oversized_frame_rejected", raw: oversizedHdr[:]}, - {name: "truncated_frame_rejected", raw: append(truncatedHdr[:], 1, 2, 3)}, - } - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - if got, err := grpcReadFrame(bytes.NewReader(tc.raw)); err == nil { - t.Errorf("grpcReadFrame(%x) = %x, want non-nil error", tc.raw, got) - } - }) - } + envs := &cobra.Command{Use: "environments"} + envs.AddCommand(newEnvironmentsConnectCmd(), newEnvironmentsPollCmd(), newEnvironmentsServeCmd()) + root.AddCommand(envs) + root.SetContext(t.Context()) - // Build a deeply nested google.protobuf.Struct exceeding maxProtoNestingDepth. - var deepStruct []byte - for i := 0; i < maxProtoNestingDepth+4; i++ { - val := appendBytesField(nil, 5, deepStruct) - entry := appendStringField(nil, 1, "k") - entry = appendBytesField(entry, 2, val) - deepStruct = appendBytesField(nil, 1, entry) - } - if got, err := decodeProtoStruct(deepStruct, 0); err == nil { - t.Errorf("decodeProtoStruct(depth>%d) = %v, want depth limit error", maxProtoNestingDepth, got) - } + var stdout, stderr bytes.Buffer + root.SetOut(&stdout) + root.SetErr(&stderr) + return root, &stdout, &stderr } -func TestOutboundWireFieldNumbers(t *testing.T) { - batch := &clientEventBatch{ - AckOutputSeqNum: 7, - ReconnectInteraction: &reconnectInteractionEvent{ - InteractionID: "int_123", - NextOutputSeqNum: 8, - ClientID: "environment", - }, - Disconnect: &disconnectEvent{}, - Events: []*clientEvent{{ - SeqNum: 1, - TimestampMicros: 1700000000, - Payload: &wirePayload{ - Step: &wireStep{ - ToolResult: &toolResultStep{ - CallID: "call_1", - FunctionResult: &functionResultStep{ - Name: "view_file", - IsError: false, - Result: map[string]any{"content": "ok"}, - }, - }, - }, - }, - }}, - } - var buf bytes.Buffer - if err := grpcWriteFrame(&buf, batch.marshalWire()); err != nil { - t.Fatalf("grpcWriteFrame() unexpected error: %v", err) - } - frame, err := grpcReadFrame(&buf) - if err != nil { - t.Fatalf("grpcReadFrame() unexpected error: %v", err) - } - - seenFields := make(map[int]bool) - rest := frame - for len(rest) > 0 { - num, _, _, _, n, err := consumeField(rest) - if err != nil { - t.Fatalf("consumeField() unexpected error: %v", err) - } - seenFields[num] = true - rest = rest[n:] - } - for _, wantField := range []int{1, 2, 8, 9} { - if !seenFields[wantField] { - t.Errorf("clientEventBatch.marshalWire() missing protobuf tag %d in %v", wantField, seenFields) +func TestRegisterMountsRuntimeCommandsAsDynamic(t *testing.T) { + noopRunE := func(*cobra.Command, []string) error { return nil } + buildSkeleton := func() (*cobra.Command, *cobra.Command) { + root := &cobra.Command{Use: "gemini-api"} + for _, customName := range []string{"tts", "analyze", "transcribe"} { + root.AddCommand(&cobra.Command{Use: customName}) } - } - - gotHS := parseHandshakeFrame(t, frame) - wantHS := parsedHandshake{ - ackSeq: 7, - interactionID: "int_123", - nextOutSeq: 8, - hasNextOutSeq: true, - background: false, - clientID: "environment", - } - if gotHS != wantHS { - t.Errorf("parseHandshakeFrame(NextOutputSeqNum=8) = %+v, want %+v", gotHS, wantHS) - } - - zeroSeqFrame := (&clientEventBatch{ - ReconnectInteraction: &reconnectInteractionEvent{ - InteractionID: "int_zero", - NextOutputSeqNum: 0, - Background: true, - ClientID: "environment", - }, - }).marshalWire() - gotZeroHS := parseHandshakeFrame(t, zeroSeqFrame) - wantZeroHS := parsedHandshake{ - interactionID: "int_zero", - hasNextOutSeq: false, - background: true, - clientID: "environment", - } - if gotZeroHS != wantZeroHS { - t.Errorf("parseHandshakeFrame(NextOutputSeqNum=0) = %+v, want %+v", gotZeroHS, wantZeroHS) - } - if gotNil := (*clientEventBatch)(nil).marshalWire(); gotNil != nil { - t.Errorf("(*clientEventBatch)(nil).marshalWire() = %x, want nil", gotNil) - } -} -func TestSandboxConfinementWorkdirAndEnvAllowlist(t *testing.T) { - base := t.TempDir() - rootDir := filepath.Join(base, "sandbox") - workDir := filepath.Join(rootDir, "project") - outsideDir := filepath.Join(base, "outside") - if err := os.MkdirAll(workDir, 0o755); err != nil { - t.Fatalf("MkdirAll(workDir) failed: %v", err) - } - if err := os.MkdirAll(outsideDir, 0o755); err != nil { - t.Fatalf("MkdirAll(outsideDir) failed: %v", err) - } - if err := os.WriteFile(filepath.Join(outsideDir, "secret.txt"), []byte("top-secret"), 0o600); err != nil { - t.Fatalf("WriteFile(secret.txt) failed: %v", err) - } - if err := os.WriteFile(filepath.Join(rootDir, "shared.txt"), []byte("from-root"), 0o600); err != nil { - t.Fatalf("WriteFile(shared.txt) failed: %v", err) - } - if err := os.WriteFile(filepath.Join(workDir, "main.go"), []byte("package main\n"), 0o600); err != nil { - t.Fatalf("WriteFile(main.go) failed: %v", err) - } - if err := os.Symlink(outsideDir, filepath.Join(workDir, "sym_out")); err != nil { - t.Fatalf("Symlink(sym_out) failed: %v", err) - } - if err := os.Symlink(filepath.Join(outsideDir, "missing_target"), filepath.Join(workDir, "dangling_link")); err != nil { - t.Fatalf("Symlink(dangling_link) failed: %v", err) - } + filesGroup := &cobra.Command{Use: "files"} + for _, leaf := range []string{"get", "delete"} { + c := &cobra.Command{Use: leaf, RunE: noopRunE} + c.Flags().String("file", "", "") + filesGroup.AddCommand(c) + } + root.AddCommand(filesGroup) - fsSrv := newFSServer(rootDir, workDir) + modelsGroup := &cobra.Command{Use: "models"} + modelsGet := &cobra.Command{Use: "get", RunE: noopRunE} + modelsGet.Flags().String("model", "", "") + modelsGroup.AddCommand(modelsGet) + root.AddCommand(modelsGroup) - // Relative paths resolve against workDir while allowing parent navigation up to rootDir. - if got, err := fsSrv.ReadRegularFile("main.go"); err != nil || string(got) != "package main\n" { - t.Errorf("ReadRegularFile(main.go) = (%q, %v), want %q", got, err, "package main\n") - } - if got, err := fsSrv.ReadRegularFile("../shared.txt"); err != nil || string(got) != "from-root" { - t.Errorf("ReadRegularFile(../shared.txt) = (%q, %v), want %q", got, err, "from-root") - } + envsGroup := &cobra.Command{Use: "environments"} + envFilesGroup := &cobra.Command{Use: "files"} + envFilesList := &cobra.Command{Use: "list", RunE: noopRunE} + envFilesList.Flags().String("environment", "", "") + envFilesList.Flags().String("path", "", "") + envFilesGroup.AddCommand(envFilesList) + envsGroup.AddCommand(envFilesGroup) + root.AddCommand(envsGroup) - escapePaths := []struct { - name string - path string - }{ - {name: "dotdot_escape_from_workdir", path: "../../outside/secret.txt"}, - {name: "symlink_existing_target", path: "sym_out/secret.txt"}, - {name: "symlink_nonexistent_leaf", path: "sym_out/new_subdir/pwn.txt"}, - {name: "dangling_symlink", path: "dangling_link"}, - {name: "nul_byte_path", path: "safe\x00evil.txt"}, - } - for _, tc := range escapePaths { - t.Run(tc.name, func(t *testing.T) { - if got, err := fsSrv.ResolvePath(tc.path); err == nil { - t.Errorf("ResolvePath(%q) = %q, want confinement error", tc.path, got) - } - }) + for _, intent := range []string{"generate", "image", "video", "music"} { + c := &cobra.Command{Use: intent, RunE: noopRunE} + c.Flags().String("body", "", "") + root.AddCommand(c) + } + agentGroup := &cobra.Command{Use: "agent"} + agentRun := &cobra.Command{Use: "run", RunE: noopRunE} + agentRun.Flags().String("body", "", "") + agentGroup.AddCommand(agentRun) + root.AddCommand(agentGroup) + return root, envsGroup } - if _, _, err := resolveSandboxDirs(rootDir, outsideDir, false); err == nil { - t.Errorf("resolveSandboxDirs(%q, %q) = nil error, want escape error", rootDir, outsideDir) - } + root, envsGroup := buildSkeleton() + Register(root) - workerEnv := buildWorkerBashEnv([]string{ - "PATH=/usr/bin", - "HOME=/home/user", - "USER=user", - "LANG=en_US.UTF-8", - "LC_ALL=C.UTF-8", - "TERM=xterm-256color", - "TMPDIR=/tmp", - "AWS_SECRET_ACCESS_KEY=aws_secret", - "GITHUB_TOKEN=gh_secret", - "GH_TOKEN=gh2_secret", - "OPENAI_API_KEY=sk_secret", - "ANTHROPIC_API_KEY=ant_secret", - "GOOGLE_APPLICATION_CREDENTIALS=/creds.json", - "GEMINI_API_KEY=gem_secret", - "GEMINI_ENVIRONMENT_KEY=ek_secret", - "LEASE_TOKEN=lease_secret", - "UNLISTED_VAR=drop_me", - }) - joinedWorkerEnv := strings.Join(workerEnv, "\n") - for _, secret := range []string{"aws_secret", "gh_secret", "gh2_secret", "sk_secret", "ant_secret", "/creds.json", "gem_secret", "ek_secret", "lease_secret", "drop_me"} { - if strings.Contains(joinedWorkerEnv, secret) { - t.Errorf("buildWorkerBashEnv() leaked %q in %q", secret, joinedWorkerEnv) + for _, subName := range []string{"connect", "poll", "serve"} { + sub := findChild(envsGroup, subName) + if sub == nil { + t.Fatalf("findChild(environments, %q) = nil, want registered subcommand", subName) } - } - for _, keep := range []string{"PATH=/usr/bin", "HOME=/home/user", "USER=user", "LANG=en_US.UTF-8", "LC_ALL=C.UTF-8", "TERM=xterm-256color", "TMPDIR=/tmp"} { - if !strings.Contains(joinedWorkerEnv, keep) { - t.Errorf("buildWorkerBashEnv() dropped baseline var %q in %q", keep, joinedWorkerEnv) + if got := sub.Annotations["speakeasy_usage_dynamic"]; got != "true" { + t.Errorf("%s Annotations[speakeasy_usage_dynamic] = %q, want \"true\"", subName, got) } } - t.Setenv("GITHUB_TOKEN", "ghp_live_secret_token") - shellSrv := newShellServer(workDir) - stdout, _, exitCode, err := shellSrv.ExecuteBash(t.Context(), "env", "", 5*time.Second) - if err != nil || exitCode != 0 { - t.Fatalf("ExecuteBash(env) = (%d, %v), want exit 0", exitCode, err) - } - if strings.Contains(stdout, "ghp_live_secret_token") { - t.Errorf("ExecuteBash(env) leaked GITHUB_TOKEN in stdout:\n%s", stdout) - } -} - -func TestWorkerAllSandboxToolsBoundsAndBashControls(t *testing.T) { - rootDir := t.TempDir() - w := newPullWorker(nil, newFSServer(rootDir, rootDir), newShellServer(rootDir), pullWorkerOptions{ - WorkerID: "worker-tools", - CommandTimeout: 400 * time.Millisecond, - Logger: log.New(io.Discard, "", 0), - }) - - if res, isErr := w.dispatchTool(t.Context(), "write_to_file", map[string]any{ - "TargetFile": "notes/hello.txt", - "CodeContent": "alpha\nbeta\ngamma\n", - }); isErr { - t.Fatalf("dispatchTool(write_to_file) error = %v, want success", res) - } - - if res, isErr := w.dispatchTool(t.Context(), "replace_file_content", map[string]any{ - "TargetFile": "notes/hello.txt", - "TargetContent": "beta", - "ReplacementContent": "BETA", - }); isErr { - t.Fatalf("dispatchTool(replace_file_content) error = %v, want success", res) - } - - viewRes, isErr := w.dispatchTool(t.Context(), "view_file", map[string]any{ - "AbsolutePath": "notes/hello.txt", - "StartLine": 2, - "EndLine": 2, - }) - if isErr || viewRes["content"] != "BETA" { - t.Errorf("dispatchTool(view_file) = (%v, %v), want content=\"BETA\"", viewRes, isErr) - } - - // Oversized file (>4 MiB) rejects full view_file asking for start_line/end_line, while ranged view_file succeeds. - bigPath := filepath.Join(rootDir, "huge.log") - var bigBuilder strings.Builder - bigBuilder.WriteString("header-line-1\nheader-line-2\n") - chunk := strings.Repeat("x", 1024) + "\n" - for bigBuilder.Len() <= maxReadFileBytes+4096 { - bigBuilder.WriteString(chunk) - } - if err := os.WriteFile(bigPath, []byte(bigBuilder.String()), 0o600); err != nil { - t.Fatalf("WriteFile(huge.log) failed: %v", err) - } - fullBigRes, fullBigErr := w.dispatchTool(t.Context(), "view_file", map[string]any{"path": "huge.log"}) - if !fullBigErr || !strings.Contains(fmt.Sprint(fullBigRes["error"]), "start_line") { - t.Errorf("view_file(huge.log without line range) = (%v, %v), want error requesting start_line/end_line", fullBigRes, fullBigErr) - } - rangeBigRes, rangeBigErr := w.dispatchTool(t.Context(), "view_file", map[string]any{ - "path": "huge.log", - "start_line": 2, - "end_line": 2, - }) - if rangeBigErr || rangeBigRes["content"] != "header-line-2" { - t.Errorf("view_file(huge.log lines 2..2) = (%v, %v), want \"header-line-2\"", rangeBigRes, rangeBigErr) - } - - // Grep truncates at maxGrepMatches (1,000) and reports truncated: true. - var manyLines strings.Builder - for i := 0; i < maxGrepMatches+50; i++ { - fmt.Fprintf(&manyLines, "match-token-%d\n", i) - } - if err := os.WriteFile(filepath.Join(rootDir, "many.txt"), []byte(manyLines.String()), 0o600); err != nil { - t.Fatalf("WriteFile(many.txt) failed: %v", err) - } - grepTruncRes, grepTruncErr := w.dispatchTool(t.Context(), "grep_search", map[string]any{"Query": "match-token-"}) - if grepTruncErr { - t.Fatalf("grep_search(many.txt) unexpected error: %v", grepTruncRes) - } - if gotMatches := len(grepTruncRes["matches"].([]any)); gotMatches != maxGrepMatches { - t.Errorf("grep_search match count = %d, want %d", gotMatches, maxGrepMatches) - } - if grepTruncRes["truncated"] != true { - t.Errorf("grep_search truncated = %v, want true", grepTruncRes["truncated"]) - } - - listRes, isErr := w.dispatchTool(t.Context(), "list_dir", map[string]any{"DirectoryPath": "notes"}) - if isErr || len(listRes["entries"].([]any)) != 1 { - t.Errorf("dispatchTool(list_dir) = (%v, %v), want 1 entry", listRes, isErr) + root2, envsGroup2 := buildSkeleton() + preExisting := &cobra.Command{Use: "connect", Short: "pre-existing"} + envsGroup2.AddCommand(preExisting) + Register(root2) + if got := findChild(envsGroup2, "connect"); got != preExisting { + t.Errorf("findChild(environments, \"connect\") = %p, want pre-existing %p", got, preExisting) } - - findRes, isErr := w.dispatchTool(t.Context(), "find_by_name", map[string]any{"Pattern": "hello.txt"}) - if isErr || len(findRes["files"].([]any)) != 1 { - t.Errorf("dispatchTool(find_by_name) = (%v, %v), want 1 file", findRes, isErr) + connectCount := 0 + for _, c := range envsGroup2.Commands() { + if c.Name() == "connect" { + connectCount++ + } } - - // WaitMsBeforeAsync must NOT kill a command early, while timeout_seconds > CommandTimeout is clamped. - bashRes, isErr := w.dispatchTool(t.Context(), "run_command", map[string]any{ - "CommandLine": "sleep 0.12 && printf 'ok'", - "WaitMsBeforeAsync": 20, - }) - if isErr || bashRes["stdout"] != "ok" { - t.Errorf("run_command with WaitMsBeforeAsync=20 = (%v, %v), want stdout=\"ok\"", bashRes, isErr) + if connectCount != 1 { + t.Errorf("environments connect command count = %d, want 1", connectCount) } - - start := time.Now() - _, timedOutErr := w.dispatchTool(t.Context(), "execute_bash", map[string]any{ - "command": "sleep 5", - "timeout_seconds": 30, - }) - if !timedOutErr || time.Since(start) > 3*time.Second { - t.Errorf("execute_bash with timeout_seconds=30 took %v (isErr=%v), want clamped to CommandTimeout=400ms", time.Since(start), timedOutErr) + if findChild(envsGroup2, "poll") == nil || findChild(envsGroup2, "serve") == nil { + t.Errorf("Register() did not mount remaining runtime subcommands") } - // DisableBash rejects execute_bash immediately without spawning a shell. - disabledWorker := newPullWorker(nil, newFSServer(rootDir, rootDir), newShellServer(rootDir), pullWorkerOptions{ - WorkerID: "worker-no-bash", - DisableBash: true, - Logger: log.New(io.Discard, "", 0), + t.Run("missing_environments_group_panics", func(t *testing.T) { + defer func() { + if r := recover(); r == nil { + t.Errorf("Register(incomplete root) did not panic, want panic") + } + }() + incompleteRoot, missingEnvs := buildSkeleton() + incompleteRoot.RemoveCommand(missingEnvs) + Register(incompleteRoot) }) - if res, disabledErr := disabledWorker.dispatchTool(t.Context(), "execute_bash", map[string]any{"command": "echo nope"}); !disabledErr { - t.Errorf("dispatchTool(execute_bash) with DisableBash=true = (%v, false), want error", res) - } } -func TestBidiStreamDebugReconnectRetryAndAckOrdering(t *testing.T) { - var connections atomic.Int32 - var inBandHeartbeats atomic.Int32 - var sawDisconnect atomic.Bool - rootDir := t.TempDir() - counterFile := filepath.Join(rootDir, "executions.txt") - - var mu sync.Mutex - var handshakes []parsedHandshake - - srvURL := startH2CTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != "/google.learning.gemini.api.interactions.v1beta.InteractionsService/CreateInteractionBidiStreamByEnvironmentKey" { - http.Error(w, "bad path: "+r.URL.Path, http.StatusNotFound) - return - } - if r.Header.Get("x-goog-lease-token") != "tok_live_1" { - http.Error(w, "missing lease token", http.StatusUnauthorized) - return - } - connNum := connections.Add(1) +func TestEnvironmentsRuntimeCommandsFlagsDryRunAndExecution(t *testing.T) { + t.Setenv("GEMINI_API_KEY", "") + t.Setenv("GEMINI_ACCESS_TOKEN", "") + t.Setenv("GEMINI_ENVIRONMENT_KEY", "") + t.Setenv("GEMINI_ENVIRONMENT_ID", "") + t.Setenv("GEMINI_WORKER_ID", "") + t.Setenv("HOME", t.TempDir()) + t.Setenv("USERPROFILE", t.TempDir()) - // Connection 1 simulates a transient 503 during connectStream so executeLease must retry connectStream. - if connNum == 1 { - http.Error(w, `{"error":{"code":503,"message":"unavailable"}}`, http.StatusServiceUnavailable) - return + t.Run("flag_defaults_and_usage_kdl", func(t *testing.T) { + connectCmd := newEnvironmentsConnectCmd() + if got := connectCmd.Flags().Lookup("command-timeout").DefValue; got != "2m" { + t.Errorf("connect --command-timeout DefValue = %q, want \"2m\"", got) } - - w.Header().Set("Content-Type", "application/grpc") - w.WriteHeader(http.StatusOK) - flusher, _ := w.(http.Flusher) - flusher.Flush() - - hsFrame, err := grpcReadFrame(r.Body) - if err != nil { - t.Errorf("conn %d handshake read failed: %v", connNum, err) - return + if connectCmd.Flags().Lookup("worker-id") == nil { + t.Errorf("connect missing --worker-id flag") } - hs := parseHandshakeFrame(t, hsFrame) - mu.Lock() - handshakes = append(handshakes, hs) - mu.Unlock() - - // Connection 2: send a slow side-effecting tool call so in-band heartbeats fire BEFORE completion, - // then abort the stream right before reading the response so Send/ack must retry cleanly! - if connNum == 2 { - _ = grpcWriteFrame(w, buildServerStepToolCallFrame(1, "call_once", "execute_bash", map[string]any{ - "command": "sleep 0.12 && printf 'x' >> " + counterFile + " && printf 'done'", - })) - flusher.Flush() - - for { - frame, err := grpcReadFrame(r.Body) - if err != nil { - return - } - hasEvent := false - rest := frame - for len(rest) > 0 { - num, _, _, _, n, err := consumeField(rest) - if err != nil { - return - } - if num == 1 { - hasEvent = true - } - rest = rest[n:] - } - if !hasEvent { - inBandHeartbeats.Add(1) - continue - } - break - } - - _ = grpcWriteFrame(w, buildServerReconnectRequiredFrame()) - flusher.Flush() - return + if connectCmd.Flags().Lookup("disable-bash") == nil { + t.Errorf("connect missing --disable-bash flag") } - - // Connection 3: replay call_once (must hit executedCalls cache without re-running bash!) and send Yield. - _ = grpcWriteFrame(w, buildServerStepToolCallFrame(2, "call_once", "execute_bash", map[string]any{ - "command": "printf 'y' >> " + counterFile, - })) - flusher.Flush() - if _, err := grpcReadFrame(r.Body); err != nil { - t.Errorf("conn 3 cached response read failed: %v", err) - return + pollCmd := newEnvironmentsPollCmd() + if got := pollCmd.Flags().Lookup("provision-timeout").DefValue; got != "2m" { + t.Errorf("poll --provision-timeout DefValue = %q, want \"2m\"", got) } - _ = grpcWriteFrame(w, buildServerYieldFrame(99)) - flusher.Flush() - for { - frame, err := grpcReadFrame(r.Body) - if err != nil { - break - } - rest := frame - for len(rest) > 0 { - num, _, _, _, n, err := consumeField(rest) - if err != nil { - break - } - if num == 9 { - sawDisconnect.Store(true) - } - rest = rest[n:] + for _, sub := range []string{"connect", "poll", "serve"} { + root, stdout, _ := newRuntimeCLIRootForTest(t, "") + root.SetArgs([]string{"environments", sub, "--usage", "ignored-positional"}) + if err := root.Execute(); err != nil { + t.Fatalf("Execute(environments %s --usage) unexpected error: %v", sub, err) } - if sawDisconnect.Load() { - break + if !strings.Contains(stdout.String(), `cmd "`+sub+`"`) { + t.Errorf("environments %s --usage output missing cmd header:\n%s", sub, stdout.String()) } } - })) - - cmd := newTestRuntimeCmd(t, srvURL) - var debugBuf bytes.Buffer - cmd.SetErr(&debugBuf) - // Enable --debug specifically to guard against DebugClient deadlocking on stream io.Pipe! - _ = cmd.Root().PersistentFlags().Set("debug", "true") - - rt, err := newRuntimeTransport(cmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport() failed: %v", err) - } - var workerLog bytes.Buffer - client := NewHTTPClient(HTTPClientOptions{Transport: rt}) - worker := newPullWorker(client, newFSServer(rootDir, rootDir), newShellServer(rootDir), pullWorkerOptions{ - WorkerID: "worker-test", - HeartbeatInterval: 35 * time.Millisecond, - CommandTimeout: 5 * time.Second, - Logger: log.New(&workerLog, "", 0), }) - if err := runStartupLeases(t.Context(), worker, []*LeasedInteraction{{ - InteractionID: "int_abc", - EnvironmentID: "env_1", - LeaseToken: "tok_live_1", - }}); err != nil { - t.Fatalf("runStartupLeases() with --debug unexpected error: %v", err) - } - - for _, wantLog := range []string{ - "[int_abc] Attached interaction session", - "[int_abc] Executing tool execute_bash", - "[int_abc] Completed interaction session", - } { - if !strings.Contains(workerLog.String(), wantLog) { - t.Errorf("worker log missing %q:\n%s", wantLog, workerLog.String()) - } - } - - if got := connections.Load(); got != 3 { - t.Errorf("connections = %d, want 3 (1 transient 503 + 2 stream sessions)", got) - } - wantHandshakes := []parsedHandshake{ - {ackSeq: 0, interactionID: "int_abc", nextOutSeq: 1, hasNextOutSeq: true, clientID: "environment"}, - {ackSeq: 1, interactionID: "int_abc", nextOutSeq: 2, hasNextOutSeq: true, clientID: "environment"}, - } - if !slices.Equal(handshakes, wantHandshakes) { - t.Errorf("handshakes = %+v, want %+v", handshakes, wantHandshakes) - } - if got := inBandHeartbeats.Load(); got < 1 { - t.Errorf("inBandHeartbeats = %d, want >= 1 during slow bash tool execution", got) - } - if !sawDisconnect.Load() { - t.Errorf("sawDisconnect = false, want Disconnect frame sent on server Yield") - } - execBytes, err := os.ReadFile(counterFile) - if err != nil || string(execBytes) != "x" { - t.Errorf("side-effect execution file = (%q, %v), want exactly 1 execution (\"x\")", execBytes, err) - } - if strings.Contains(debugBuf.String(), "tok_live_1") { - t.Errorf("--debug stderr leaked raw x-goog-lease-token:\n%s", debugBuf.String()) - } -} - -func TestStreamAckNotAdvancedWhenSendFails(t *testing.T) { - // Directly test consumeStream + sessionState when Send fails because the pipe reader closed mid-call. - rootDir := t.TempDir() - counterFile := filepath.Join(rootDir, "once.txt") - w := newPullWorker(nil, newFSServer(rootDir, rootDir), newShellServer(rootDir), pullWorkerOptions{ - WorkerID: "worker-ack-test", - HeartbeatInterval: time.Hour, - CommandTimeout: 5 * time.Second, - Logger: log.New(io.Discard, "", 0), - }) - - // Build a single server frame carrying seq=1 tool call followed by immediate EOF on server body, - // while the client->server pipe reader is closed so stream.Send returns an error! - var serverBuf bytes.Buffer - _ = grpcWriteFrame(&serverBuf, buildServerStepToolCallFrame(1, "call_send_fail", "execute_bash", map[string]any{ - "command": "printf '1' >> " + counterFile, - })) - pr, pw := io.Pipe() - _ = pr.Close() // Force stream.Send to fail with io.ErrClosedPipe! - - stream := &grpcInteractionStream{ - resp: &http.Response{Body: io.NopCloser(&serverBuf)}, - pw: pw, - } - state := &sessionState{ - nextClientSeq: 1, - executedCalls: make(map[string]*wirePayload), - } - - _, err := w.consumeStream(t.Context(), stream, state) - if err == nil { - t.Fatalf("consumeStream() = nil error when Send failed, want non-nil error") - } - if gotAck := state.ackSeq(); gotAck != 0 { - t.Errorf("state.ackSeq() after failed Send = %d, want 0 so reconnect requests seq 1 again", gotAck) - } - if state.nextClientSeq != 1 { - t.Errorf("state.nextClientSeq after failed Send = %d, want rolled back to 1", state.nextClientSeq) - } - if len(state.executedCalls) != 1 { - t.Errorf("len(state.executedCalls) = %d, want 1 cached result ready for replay", len(state.executedCalls)) - } -} - -func TestPollerControllerHandoffHeartbeatAndRedaction(t *testing.T) { - var heartbeats atomic.Int32 - var leasedOnce atomic.Bool - var mu sync.Mutex - var capturedStdin ProvisionInput - - srvURL := startH2CTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - var req leaseInteractionsRequest - _ = json.NewDecoder(r.Body).Decode(&req) - w.Header().Set("Content-Type", "application/json") - if len(req.ActiveLeases) > 0 { - heartbeats.Add(1) - _ = json.NewEncoder(w).Encode(leaseInteractionsResponse{ - Leases: []*LeasedInteraction{{ - InteractionID: req.ActiveLeases[0].InteractionID, - LeaseToken: req.ActiveLeases[0].LeaseToken, - }}, - }) - return - } - if !leasedOnce.Swap(true) { - w.WriteHeader(http.StatusOK) - _, _ = w.Write([]byte(`{ - "leases": [{ - "interaction_id": "int_handoff", - "environment_id": "env_alpha", - "lease_token": "secret_lease_tok_999", - "lease_generation": "2", - "lease_expiration_time": "2026-10-22T10:00:00Z", - "environment_config": { - "network": { - "allowlist": [{"domain": "api.example.com", "credential": "cred_net"}] - }, - "sources": [ - {"type": "GCS", "source": "gs://bucket/prefix", "target": "/workspace/data"} - ], - "env": { - "PLAIN_KEY": "plain_val", - "OBJ_KEY": {"value": "obj_val"}, - "SECRET_KEY": {"credential": "cred_secret"} - }, - "type": "self_hosted" - } - }], - "poll_backoff": "5s" - }`)) - return - } - _ = json.NewEncoder(w).Encode(leaseInteractionsResponse{PollBackoff: "5s"}) - })) - - outFile := filepath.Join(t.TempDir(), "provision.json") - cmd := newTestRuntimeCmd(t, srvURL) - var debugBuf bytes.Buffer - cmd.SetErr(&debugBuf) - _ = cmd.Root().PersistentFlags().Set("debug", "true") - - rt, err := newRuntimeTransport(cmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport() failed: %v", err) - } - client := NewHTTPClient(HTTPClientOptions{ - Transport: rt, - EnvironmentKey: "ek_test_secret", - }) - - ctx, cancel := context.WithCancel(t.Context()) - poller := NewPollerController(client, NewCommandWorkerProvisioner("sleep 0.12 && cat > "+outFile), PollerOptions{ - PollerID: "poller-1", - MaxInteractions: 2, - PollInterval: 30 * time.Millisecond, - ProvisionTimeout: 5 * time.Second, - HandoffHeartbeatInterval: 35 * time.Millisecond, - Logger: log.New(io.Discard, "", 0), - }) - - done := make(chan error, 1) - go func() { - done <- poller.Run(ctx) - }() - - deadline := time.Now().Add(3 * time.Second) - for time.Now().Before(deadline) { - if raw, err := os.ReadFile(outFile); err == nil && len(raw) > 0 { - mu.Lock() - _ = json.Unmarshal(raw, &capturedStdin) - mu.Unlock() - break - } - time.Sleep(15 * time.Millisecond) - } - cancel() - <-done - - if capturedStdin.InteractionID != "int_handoff" || capturedStdin.EnvironmentID != "env_alpha" || capturedStdin.SchemaVersion != "v1" { - t.Errorf("ProvisionInput = %+v, want v1 int_handoff / env_alpha", capturedStdin) - } - var forwardedCfg map[string]any - if err := json.Unmarshal(capturedStdin.EnvironmentConfig, &forwardedCfg); err != nil { - t.Fatalf("json.Unmarshal(ProvisionInput.EnvironmentConfig) failed: %v", err) - } - envMap, _ := forwardedCfg["env"].(map[string]any) - if envMap["PLAIN_KEY"] != "plain_val" { - t.Errorf("forwarded env.PLAIN_KEY = %v, want %q", envMap["PLAIN_KEY"], "plain_val") - } - if objKey, _ := envMap["OBJ_KEY"].(map[string]any); objKey["value"] != "obj_val" { - t.Errorf("forwarded env.OBJ_KEY = %v, want {\"value\":\"obj_val\"}", envMap["OBJ_KEY"]) - } - if netObj, _ := forwardedCfg["network"].(map[string]any); netObj["allowlist"] == nil { - t.Errorf("forwarded network = %v, want allowlist preserved", forwardedCfg["network"]) - } - if got := heartbeats.Load(); got < 1 { - t.Errorf("handoff heartbeats = %d, want >= 1 during slow provision command", got) - } - if strings.Contains(debugBuf.String(), "secret_lease_tok_999") { - t.Errorf("--debug stderr leaked raw lease_token:\n%s", debugBuf.String()) - } -} - -func TestLeaseInteractionsWireParsing(t *testing.T) { - tests := []struct { - name string - envID string - respBody string - wantGeneration int64 - wantSubstr string - }{ - { - name: "gaos_allowlist_sources_and_envvar_objects", - envID: "env_scoped", - respBody: `{ - "leases": [{ - "interaction_id": "int_1", - "environment_id": "env_scoped", - "lease_token": "tok_1", - "lease_generation": "7", - "lease_expiration_time": "2026-10-22T12:00:00Z", - "environment_config": { - "network": { - "allowlist": [ - {"domain": "example.com", "credential": "cred_a"}, - {"domain": "*.googleapis.com"} - ] - }, - "sources": [ - {"type": "GCS", "source": "gs://my-bucket/src", "target": "/workspace/src"} - ], - "env": { - "PLAIN": "hello", - "CRED": {"credential": "cred_env"} - } - } - }], - "poll_backoff": "5s" - }`, - wantGeneration: 7, - wantSubstr: `"CRED": {"credential": "cred_env"}`, - }, - { - name: "gaos_disabled_network_and_numeric_generation", - envID: "env_disabled", - respBody: `{ - "leases": [{ - "interaction_id": "int_2", - "environment_id": "env_disabled", - "lease_token": "tok_2", - "lease_generation": 3, - "environment_config": { - "network": "disabled" - } - }] - }`, - wantGeneration: 3, - wantSubstr: `"network": "disabled"`, - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - var gotRawReq map[string]any - srvURL := startH2CTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - _ = json.NewDecoder(r.Body).Decode(&gotRawReq) - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(tc.respBody)) - })) - - cmd := newTestRuntimeCmd(t, srvURL) - rt, err := newRuntimeTransport(cmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport() failed: %v", err) - } - client := NewHTTPClient(HTTPClientOptions{ - Transport: rt, - EnvironmentID: tc.envID, - EnvironmentKey: "ek_live_1", - }) - - leases, err := client.LeaseInteractions(t.Context(), "worker-1", 4) - if err != nil { - t.Fatalf("LeaseInteractions() unexpected error: %v", err) - } - if gotRawReq["environment_id"] != tc.envID { - t.Errorf("request environment_id = %v, want %q", gotRawReq["environment_id"], tc.envID) - } - if len(leases) != 1 { - t.Fatalf("len(leases) = %d, want 1", len(leases)) - } - got := leases[0] - if got.LeaseGeneration != tc.wantGeneration { - t.Errorf("LeaseGeneration = %d, want %d", got.LeaseGeneration, tc.wantGeneration) - } - if !strings.Contains(string(got.EnvironmentConfig), tc.wantSubstr) { - t.Errorf("EnvironmentConfig = %s, want substring %q", got.EnvironmentConfig, tc.wantSubstr) - } - }) - } -} - -func TestHeartbeatInteractionRequestShape(t *testing.T) { - var rawReq map[string]any - srvURL := startH2CTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - _ = json.NewDecoder(r.Body).Decode(&rawReq) - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"leases":[{"interaction_id":"int_hb","lease_token":"tok_hb"}]}`)) - })) - - cmd := newTestRuntimeCmd(t, srvURL) - rt, err := newRuntimeTransport(cmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport() failed: %v", err) - } - client := NewHTTPClient(HTTPClientOptions{ - Transport: rt, - EnvironmentKey: "ek_live_hb", - }) - - res, err := client.HeartbeatInteraction(t.Context(), "poller-hb", "int_hb", "env_hb", "tok_hb") - if err != nil { - t.Fatalf("HeartbeatInteraction() unexpected error: %v", err) - } - if !res.Renewed || res.ShouldCancel { - t.Errorf("HeartbeatInteraction() = %+v, want Renewed=true ShouldCancel=false", res) - } - if _, badField := rawReq["heartbeat_leases"]; badField { - t.Errorf("request sent obsolete field heartbeat_leases: %v", rawReq) - } - if _, ok := rawReq["active_leases"]; !ok { - t.Errorf("request missing active_leases: %v", rawReq) - } - maxNew, hasMaxNew := rawReq["max_new_interactions"] - if !hasMaxNew || maxNew != float64(0) { - t.Errorf("request max_new_interactions = (%v, %v), want (0, true)", maxNew, hasMaxNew) - } - if rawReq["environment_id"] != "env_hb" { - t.Errorf("request environment_id = %v, want %q", rawReq["environment_id"], "env_hb") - } -} - -func TestWindowsSandboxFixes(t *testing.T) { - rootDir := t.TempDir() - fs := newFSServer(rootDir, rootDir) - worker := &pullWorker{fs: fs} - - // 1. Write a CRLF file and verify ReadFileRange, Grep ($ anchor), and replace_file_content (LF target -> CRLF preserved). - crlfContent := "[server]\r\nhost = 127.0.0.1\r\nport = 8080\r\nmode = debug\r\n" - if _, err := fs.WriteRegularFile("src/settings.ini", crlfContent, true, false); err != nil { - t.Fatalf("WriteRegularFile(src/settings.ini) failed: %v", err) - } - if _, err := fs.WriteRegularFile("logs/app.log", "2026-10-22 ERROR bind failed\r\n", true, false); err != nil { - t.Fatalf("WriteRegularFile(logs/app.log) failed: %v", err) - } - if _, err := fs.WriteRegularFile("README.md", "# Service\r\n", true, false); err != nil { - t.Fatalf("WriteRegularFile(README.md) failed: %v", err) - } - - ranged, err := fs.ReadFileRange("src/settings.ini", 2, 3) - if err != nil { - t.Fatalf("ReadFileRange() failed: %v", err) - } - if ranged != "host = 127.0.0.1\nport = 8080" { - t.Errorf("ReadFileRange() = %q, want %q", ranged, "host = 127.0.0.1\nport = 8080") - } - - matches, _, err := fs.Grep(".", "port = 8080$", false, true) - if err != nil { - t.Fatalf("Grep() failed: %v", err) - } - if len(matches) != 1 { - t.Fatalf("Grep('port = 8080$') len(matches) = %d, want 1", len(matches)) - } - - res, isErr := worker.toolReplaceFileContent(map[string]any{ - "TargetFile": "src/settings.ini", - "TargetContent": "host = 127.0.0.1\nport = 8080", - "ReplacementContent": "host = 127.0.0.1\nport = 9090", - "StartLine": 2, - "EndLine": 3, - }) - if isErr { - t.Fatalf("toolReplaceFileContent() failed: %v", res) - } - updatedBytes, err := fs.ReadRegularFile("src/settings.ini") - if err != nil { - t.Fatalf("ReadRegularFile() failed: %v", err) - } - wantCRLF := "[server]\r\nhost = 127.0.0.1\r\nport = 9090\r\nmode = debug\r\n" - if string(updatedBytes) != wantCRLF { - t.Errorf("updated CRLF file = %q, want %q", string(updatedBytes), wantCRLF) - } - - // 2. Verify find_by_name Extensions filter excludes directories and non-matching extensions. - findRes, findErr := worker.toolFindByName(map[string]any{ - "SearchDirectory": ".", - "Pattern": "*", - "Extensions": []any{"ini", ".LOG"}, - }) - if findErr { - t.Fatalf("toolFindByName() failed: %v", findRes) - } - gotFiles, _ := findRes["files"].([]any) - if len(gotFiles) != 2 || gotFiles[0] != "logs/app.log" || gotFiles[1] != "src/settings.ini" { - t.Errorf("toolFindByName(Extensions=[ini, .LOG]) = %v, want [logs/app.log src/settings.ini]", gotFiles) - } - - // 3. Verify Windows environment allowlist preserves COMPUTERNAME, USERNAME, OS, SYSTEMDRIVE, PSMODULEPATH. - filteredEnv := buildWorkerBashEnv([]string{ - "COMPUTERNAME=WIN-BYOS", - "USERNAME=svc_user", - "OS=Windows_NT", - "SYSTEMDRIVE=C:", - "PSMODULEPATH=C:\\Program Files\\WindowsPowerShell\\Modules", - "GEMINI_API_KEY=secret", - "UNLISTED_VAR=drop_me", - }) - if len(filteredEnv) != 5 { - t.Errorf("buildWorkerBashEnv() = %v, want 5 allowed Windows env vars", filteredEnv) - } -} - -func newRuntimeCLIRootForTest(t *testing.T, serverURL string) (*cobra.Command, *bytes.Buffer, *bytes.Buffer) { - t.Helper() - root := &cobra.Command{Use: "gemini-api", SilenceErrors: true, SilenceUsage: true} - pf := root.PersistentFlags() - pf.StringP("output-format", "o", "pretty", "") - pf.String("color", "never", "") - pf.StringP("jq", "q", "", "") - pf.Bool("raw-output", true, "") - pf.String("server-url", serverURL, "") - pf.StringArrayP("header", "H", nil, "") - pf.Bool("include-headers", false, "") - pf.String("timeout", "", "") - pf.Bool("interactive", false, "") - pf.Bool("no-interactive", true, "") - pf.Bool("usage", false, "") - pf.Bool("dry-run", false, "") - pf.BoolP("debug", "d", false, "") - pf.Bool("agent-mode", false, "") - pf.Bool("no-retries", false, "") - pf.Bool("no-keyring", true, "") - pf.String("api-key", "", "") - pf.String("access-token", "", "") - pf.String("api-version", "v1beta", "") - pf.String("api-revision", "", "") - pf.String("user-project", "", "") - - envs := &cobra.Command{Use: "environments"} - envs.AddCommand(newEnvironmentsConnectCmd(), newEnvironmentsPollCmd(), newEnvironmentsServeCmd()) - root.AddCommand(envs) - root.SetContext(t.Context()) - - var stdout, stderr bytes.Buffer - root.SetOut(&stdout) - root.SetErr(&stderr) - return root, &stdout, &stderr -} - -func TestRegisterMountsRuntimeCommandsAsDynamic(t *testing.T) { - noopRunE := func(*cobra.Command, []string) error { return nil } - buildSkeleton := func() (*cobra.Command, *cobra.Command) { - root := &cobra.Command{Use: "gemini-api"} - for _, customName := range []string{"tts", "analyze", "transcribe"} { - root.AddCommand(&cobra.Command{Use: customName}) - } - - filesGroup := &cobra.Command{Use: "files"} - for _, leaf := range []string{"get", "delete"} { - c := &cobra.Command{Use: leaf, RunE: noopRunE} - c.Flags().String("file", "", "") - filesGroup.AddCommand(c) - } - root.AddCommand(filesGroup) - - modelsGroup := &cobra.Command{Use: "models"} - modelsGet := &cobra.Command{Use: "get", RunE: noopRunE} - modelsGet.Flags().String("model", "", "") - modelsGroup.AddCommand(modelsGet) - root.AddCommand(modelsGroup) - - envsGroup := &cobra.Command{Use: "environments"} - envFilesGroup := &cobra.Command{Use: "files"} - envFilesList := &cobra.Command{Use: "list", RunE: noopRunE} - envFilesList.Flags().String("environment", "", "") - envFilesList.Flags().String("path", "", "") - envFilesGroup.AddCommand(envFilesList) - envsGroup.AddCommand(envFilesGroup) - root.AddCommand(envsGroup) - - for _, intent := range []string{"generate", "image", "video", "music"} { - c := &cobra.Command{Use: intent, RunE: noopRunE} - c.Flags().String("body", "", "") - root.AddCommand(c) - } - agentGroup := &cobra.Command{Use: "agent"} - agentRun := &cobra.Command{Use: "run", RunE: noopRunE} - agentRun.Flags().String("body", "", "") - agentGroup.AddCommand(agentRun) - root.AddCommand(agentGroup) - return root, envsGroup - } - - root, envsGroup := buildSkeleton() - Register(root) - - for _, subName := range []string{"connect", "poll", "serve"} { - sub := findChild(envsGroup, subName) - if sub == nil { - t.Fatalf("findChild(environments, %q) = nil, want registered subcommand", subName) - } - if got := sub.Annotations["speakeasy_usage_dynamic"]; got != "true" { - t.Errorf("%s Annotations[speakeasy_usage_dynamic] = %q, want \"true\"", subName, got) - } - } - - // Pre-existing subcommands under environments must not be overwritten or duplicated. - root2, envsGroup2 := buildSkeleton() - preExisting := &cobra.Command{Use: "connect", Short: "pre-existing"} - envsGroup2.AddCommand(preExisting) - Register(root2) - if got := findChild(envsGroup2, "connect"); got != preExisting { - t.Errorf("findChild(environments, \"connect\") = %p, want pre-existing %p", got, preExisting) - } - connectCount := 0 - for _, c := range envsGroup2.Commands() { - if c.Name() == "connect" { - connectCount++ - } - } - if connectCount != 1 { - t.Errorf("environments connect command count = %d, want 1", connectCount) - } - if findChild(envsGroup2, "poll") == nil || findChild(envsGroup2, "serve") == nil { - t.Errorf("Register() did not mount remaining runtime subcommands") - } - - t.Run("missing_environments_group_panics", func(t *testing.T) { - defer func() { - if r := recover(); r == nil { - t.Errorf("Register(incomplete root) did not panic, want panic") - } - }() - incompleteRoot, missingEnvs := buildSkeleton() - incompleteRoot.RemoveCommand(missingEnvs) - Register(incompleteRoot) - }) -} - -func TestEnvironmentsRuntimeCommandsFlagsDryRunAndExecution(t *testing.T) { - t.Setenv("GEMINI_API_KEY", "") - t.Setenv("GEMINI_ACCESS_TOKEN", "") - t.Setenv("GEMINI_ENVIRONMENT_KEY", "") - t.Setenv("GEMINI_ENVIRONMENT_ID", "") - t.Setenv("GEMINI_WORKER_ID", "") - t.Setenv("HOME", t.TempDir()) - t.Setenv("USERPROFILE", t.TempDir()) - - t.Run("flag_defaults_and_usage_kdl", func(t *testing.T) { - connectCmd := newEnvironmentsConnectCmd() - if got := connectCmd.Flags().Lookup("command-timeout").DefValue; got != "2m" { - t.Errorf("connect --command-timeout DefValue = %q, want \"2m\"", got) - } - if connectCmd.Flags().Lookup("worker-id") == nil { - t.Errorf("connect missing --worker-id flag") - } - if connectCmd.Flags().Lookup("disable-bash") == nil { - t.Errorf("connect missing --disable-bash flag") - } - pollCmd := newEnvironmentsPollCmd() - if got := pollCmd.Flags().Lookup("provision-timeout").DefValue; got != "2m" { - t.Errorf("poll --provision-timeout DefValue = %q, want \"2m\"", got) - } - - for _, sub := range []string{"connect", "poll", "serve"} { - root, stdout, _ := newRuntimeCLIRootForTest(t, "") - root.SetArgs([]string{"environments", sub, "--usage", "ignored-positional"}) - if err := root.Execute(); err != nil { - t.Fatalf("Execute(environments %s --usage) unexpected error: %v", sub, err) - } - if !strings.Contains(stdout.String(), `cmd "`+sub+`"`) { - t.Errorf("environments %s --usage output missing cmd header:\n%s", sub, stdout.String()) - } - } - }) - - t.Run("validation_errors", func(t *testing.T) { - sandboxDir := t.TempDir() - outsideDir := t.TempDir() - cases := []struct { - name string - args []string - wantSub string - }{ - { - name: "connect_positional_arg", - args: []string{"environments", "connect", "extra"}, - wantSub: "accepts 0 arguments, received 1", - }, - { - name: "connect_workdir_escapes_root", - args: []string{"environments", "connect", "--root-dir", sandboxDir, "--workdir", outsideDir}, - wantSub: "escapes root-dir", - }, - { - name: "connect_zero_heartbeat", - args: []string{"environments", "connect", "--heartbeat-interval", "0s"}, - wantSub: "--heartbeat-interval must be > 0", - }, - { - name: "connect_zero_command_timeout", - args: []string{"environments", "connect", "--command-timeout", "0s"}, - wantSub: "--command-timeout must be > 0", - }, - { - name: "connect_control_char_env_key", - args: []string{"environments", "connect", "--environment-key", "bad\x00key"}, - wantSub: "--environment-key contains control characters", - }, - { - name: "connect_zero_max_concurrent", - args: []string{"environments", "connect", "--max-concurrent", "0"}, - wantSub: "--max-concurrent must be >= 1", - }, - { - name: "connect_zero_poll_interval", - args: []string{"environments", "connect", "--poll-interval", "0s"}, - wantSub: "--poll-interval must be > 0", - }, - { - name: "connect_missing_credentials", - args: []string{"environments", "connect"}, - wantSub: "missing credentials", - }, - { - name: "poll_zero_max_interactions", - args: []string{"environments", "poll", "--max-interactions", "0"}, - wantSub: "--max-interactions must be >= 1", - }, - { - name: "poll_zero_poll_interval", - args: []string{"environments", "poll", "--poll-interval", "0s"}, - wantSub: "--poll-interval must be > 0", - }, - { - name: "poll_zero_provision_timeout", - args: []string{"environments", "poll", "--provision-timeout", "0s"}, - wantSub: "--provision-timeout must be > 0", - }, - { - name: "poll_zero_handoff_heartbeat", - args: []string{"environments", "poll", "--handoff-heartbeat-interval", "0s"}, - wantSub: "--handoff-heartbeat-interval must be > 0", - }, - { - name: "poll_control_char_env_key", - args: []string{"environments", "poll", "--environment-key", "bad\nkey"}, - wantSub: "--environment-key contains control characters", - }, - { - name: "poll_invalid_environment_id", - args: []string{"environments", "poll", "--environment", "environments/a/b"}, - wantSub: `invalid environment id "environments/a/b"`, - }, - { - name: "poll_missing_provision_cmd", - args: []string{"environments", "poll", "--environment-key", "ek_1"}, - wantSub: "--provision-cmd is required", - }, - { - name: "poll_missing_credentials", - args: []string{"environments", "poll", "--provision-cmd", "true"}, - wantSub: "missing credentials", - }, - { - name: "serve_only_interaction_id", - args: []string{"environments", "serve", "--interaction-id", "int_1"}, - wantSub: "--interaction-id and --lease-token must be provided together", - }, - { - name: "serve_invalid_interaction_id", - args: []string{"environments", "serve", "--interaction-id", "interactions/a/b", "--lease-token", "tok_1"}, - wantSub: "invalid interaction id", - }, - { - name: "serve_control_char_lease_token", - args: []string{"environments", "serve", "--interaction-id", "int_1", "--lease-token", "tok\r1"}, - wantSub: "--lease-token contains control characters", - }, - { - name: "serve_missing_lease", - args: []string{"environments", "serve"}, - wantSub: "missing lease", - }, - } - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - root, _, _ := newRuntimeCLIRootForTest(t, "http://127.0.0.1:1") - root.SetArgs(tc.args) - rawErr := root.Execute() - if rawErr == nil { - t.Fatalf("Execute(%v) = nil, want error containing %q", tc.args, tc.wantSub) - } - err := output.CLIError(root, rawErr) - if !strings.Contains(err.Error(), tc.wantSub) { - t.Errorf("Execute(%v) error = %q, want substring %q", tc.args, err.Error(), tc.wantSub) - } - if got := clierrors.ExitCode(err); got != clierrors.ExitUsage { - t.Errorf("ExitCode(%v) = %d, want %d (ExitUsage)", tc.args, got, clierrors.ExitUsage) - } - }) - } - }) - - t.Run("dry_run_previews", func(t *testing.T) { - dryCases := []struct { - name string - args []string - wantPath string - wantBody string - }{ - { - name: "connect_dry_run", - args: []string{"-o", "json", "--dry-run", "environments", "connect", "--environment-key", "ek_c", "--worker-id", "w-custom", "--max-concurrent", "2"}, - wantPath: "/v1beta/environments:leaseInteractions", - wantBody: "w-custom", - }, - { - name: "poll_dry_run_with_scoped_environment", - args: []string{"-o", "json", "--dry-run", "environments", "poll", "--environment", "environments/env_77", "--environment-key", "ek_p", "--poller-id", "p-custom", "--max-interactions", "3"}, - wantPath: "/v1beta/environments:leaseInteractions", - wantBody: "env_77", - }, - { - name: "serve_dry_run_default_sample", - args: []string{"-o", "json", "--dry-run", "environments", "serve"}, - wantPath: "/google.learning.gemini.api.interactions.v1beta.InteractionsService/CreateInteractionBidiStreamByEnvironmentKey", - }, - { - name: "serve_dry_run_explicit_lease", - args: []string{"-o", "json", "--dry-run", "environments", "serve", "--interaction-id", "interactions/int_s1", "--lease-token", "tok_s1"}, - wantPath: "/google.learning.gemini.api.interactions.v1beta.InteractionsService/CreateInteractionBidiStreamByEnvironmentKey", - }, - } - for _, tc := range dryCases { - t.Run(tc.name, func(t *testing.T) { - root, stdout, _ := newRuntimeCLIRootForTest(t, "http://127.0.0.1:1") - output.PreparseRenderingFlags(root, []string{"-o", "json"}) - root.SetArgs(tc.args) - if err := root.Execute(); err != nil { - t.Fatalf("Execute(%v) unexpected error: %v", tc.args, err) - } - var preview map[string]any - if err := json.Unmarshal(stdout.Bytes(), &preview); err != nil { - t.Fatalf("json.Unmarshal(dry-run output %q) failed: %v", stdout.String(), err) - } - reqObj, _ := preview["request"].(map[string]any) - if gotURL, _ := reqObj["url"].(string); !strings.HasSuffix(gotURL, tc.wantPath) { - t.Errorf("dry-run url = %q, want suffix %q", gotURL, tc.wantPath) - } - if tc.wantBody != "" && !strings.Contains(stdout.String(), tc.wantBody) { - t.Errorf("dry-run output = %s, want substring %q", stdout.String(), tc.wantBody) - } - }) - } - }) - - t.Run("live_connect_poll_and_serve_execution", func(t *testing.T) { - var connectWorkerID atomic.Value - var pollEnvID atomic.Value - var serveCompleted atomic.Bool - connectDone := make(chan struct{}, 1) - provisionFile := filepath.Join(t.TempDir(), "polled.json") - - srvURL := startH2CTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case "/v1beta/environments:leaseInteractions": - var req leaseInteractionsRequest - _ = json.NewDecoder(r.Body).Decode(&req) - w.Header().Set("Content-Type", "application/json") - if strings.HasPrefix(req.WorkerID, "poller") { - pollEnvID.Store(req.EnvironmentID) - if len(req.ActiveLeases) == 0 && pollEnvID.Load() == "env_live_poll" { - _ = json.NewEncoder(w).Encode(leaseInteractionsResponse{ - Leases: []*LeasedInteraction{{ - InteractionID: "int_polled", - EnvironmentID: "env_live_poll", - LeaseToken: "tok_polled", - }}, - }) - return - } - _ = json.NewEncoder(w).Encode(leaseInteractionsResponse{}) - return - } - if prev := connectWorkerID.Swap(req.WorkerID); prev == nil { - _ = json.NewEncoder(w).Encode(leaseInteractionsResponse{ - Leases: []*LeasedInteraction{{ - InteractionID: "int_connected", - EnvironmentID: "env_conn", - LeaseToken: "tok_conn", - }}, - }) - return - } - _ = json.NewEncoder(w).Encode(leaseInteractionsResponse{}) - - case "/google.learning.gemini.api.interactions.v1beta.InteractionsService/CreateInteractionBidiStreamByEnvironmentKey": - w.Header().Set("Content-Type", "application/grpc") - w.WriteHeader(http.StatusOK) - flusher, _ := w.(http.Flusher) - flusher.Flush() - hsFrame, err := grpcReadFrame(r.Body) - if err != nil { - return - } - hs := parseHandshakeFrame(t, hsFrame) - if hs.interactionID == "int_serve_bash_off" { - // Verify --disable-bash rejects execute_bash over the live serve command! - _ = grpcWriteFrame(w, buildServerStepToolCallFrame(1, "c_bash", "execute_bash", map[string]any{"command": "echo nope"})) - flusher.Flush() - if _, err := grpcReadFrame(r.Body); err == nil { - serveCompleted.Store(true) - } - } - // Terminate session with InteractionCompleted (field 6). - _ = grpcWriteFrame(w, appendBytesField(nil, 6, nil)) - flusher.Flush() - if hs.interactionID == "int_connected" { - select { - case connectDone <- struct{}{}: - default: - } - } - default: - http.NotFound(w, r) - } - })) - - // 1. Live `environments serve --disable-bash` - serveWork := filepath.Join(t.TempDir(), "serve-work") - serveRoot, _, _ := newRuntimeCLIRootForTest(t, srvURL) - serveRoot.SetArgs([]string{ - "--no-retries", - "environments", "serve", - "--interaction-id", "interactions/int_serve_bash_off", - "--lease-token", "tok_serve", - "--root-dir", serveWork, - "--workdir", serveWork, - "--disable-bash", - "--quiet", - }) - if err := serveRoot.Execute(); err != nil { - t.Fatalf("Execute(environments serve) unexpected error: %v", err) - } - if !serveCompleted.Load() { - t.Errorf("serveCompleted = false, want tool error frame sent before InteractionCompleted") - } - - // 2. Live `environments connect` - connCtx, cancelConn := context.WithCancel(t.Context()) - connDir := filepath.Join(t.TempDir(), "conn-work") - connRoot, _, _ := newRuntimeCLIRootForTest(t, srvURL) - connRoot.SetContext(connCtx) - connRoot.SetArgs([]string{ - "environments", "connect", - "--environment-key", "ek_conn", - "--worker-id", "worker-cli-1", - "--root-dir", connDir, - "--poll-interval", "20ms", - "--quiet", - }) - connErrCh := make(chan error, 1) - go func() { - connErrCh <- connRoot.Execute() - }() - select { - case <-connectDone: - case <-time.After(3 * time.Second): - t.Errorf("timed out waiting for environments connect to complete leased session") - } - cancelConn() - if err := <-connErrCh; err != nil { - t.Errorf("Execute(environments connect) = %v, want nil on context cancel", err) - } - if got, _ := connectWorkerID.Load().(string); got != "worker-cli-1" { - t.Errorf("connect WorkerID = %q, want \"worker-cli-1\"", got) - } - - // 3. Live `environments poll` - pollCtx, cancelPoll := context.WithCancel(t.Context()) - pollRoot, _, _ := newRuntimeCLIRootForTest(t, srvURL) - pollRoot.SetContext(pollCtx) - pollRoot.SetArgs([]string{ - "environments", "poll", - "--environment", "environments/env_live_poll", - "--environment-key", "ek_poll", - "--poller-id", "poller-cli-1", - "--poll-interval", "20ms", - "--provision-cmd", "cat > " + provisionFile, - "--quiet", - }) - pollErrCh := make(chan error, 1) - go func() { - pollErrCh <- pollRoot.Execute() - }() - deadline := time.Now().Add(3 * time.Second) - for time.Now().Before(deadline) { - if raw, err := os.ReadFile(provisionFile); err == nil && len(raw) > 0 { - break - } - time.Sleep(15 * time.Millisecond) - } - cancelPoll() - if err := <-pollErrCh; err != nil { - t.Errorf("Execute(environments poll) = %v, want nil on context cancel", err) - } - if got, _ := pollEnvID.Load().(string); got != "env_live_poll" { - t.Errorf("poll EnvironmentID = %q, want \"env_live_poll\"", got) - } - }) - - t.Run("helpers_logger_and_no_retries", func(t *testing.T) { - output.ResetAgentMode() - output.ResetPreparsedRendering() - config.Reset() - t.Cleanup(func() { - output.ResetAgentMode() - output.ResetPreparsedRendering() - config.Reset() - }) - - root, _, stderr := newRuntimeCLIRootForTest(t, "") - logNormal := newRuntimeLogger(root, false) - logNormal.Print("visible-msg") - if !strings.Contains(stderr.String(), "visible-msg") { - t.Errorf("newRuntimeLogger(quiet=false) did not write to stderr: %q", stderr.String()) - } - stderr.Reset() - logQuiet := newRuntimeLogger(root, true) - logQuiet.Print("hidden-msg") - if stderr.Len() != 0 { - t.Errorf("newRuntimeLogger(quiet=true) wrote %q, want empty", stderr.String()) - } - _ = root.PersistentFlags().Set("output-format", "json") - logMachine := newRuntimeLogger(root, false) - logMachine.Print("machine-hidden-msg") - if stderr.Len() != 0 { - t.Errorf("newRuntimeLogger(output-format=json) wrote %q, want empty", stderr.String()) - } - - if got := resolveNoRetries(root); got { - t.Errorf("resolveNoRetries(default) = true, want false") - } - _ = root.PersistentFlags().Set("no-retries", "true") - if got := resolveNoRetries(root); !got { - t.Errorf("resolveNoRetries(--no-retries=true) = false, want true") - } - }) -} - -func TestRuntimeTransportCredentialsHeadersAndValidation(t *testing.T) { - t.Setenv("GEMINI_API_KEY", "") - t.Setenv("GEMINI_ACCESS_TOKEN", "") - t.Setenv("GEMINI_NO_KEYRING", "true") - t.Setenv("HOME", t.TempDir()) - t.Setenv("USERPROFILE", t.TempDir()) - config.Reset() - t.Cleanup(config.Reset) - - t.Run("credentials_missing_and_headers", func(t *testing.T) { - cases := []struct { - name string - explicitCred bool - apiKey string - accessToken string - userProject string - apiRevision string - header string - wantMissing bool - wantAPIKey string - wantBearer string - }{ - { - name: "no_creds_without_explicit_runtime_cred", - explicitCred: false, - wantMissing: true, - }, - { - name: "no_creds_with_explicit_runtime_cred", - explicitCred: true, - wantMissing: false, - }, - { - name: "api_key_credential", - explicitCred: false, - apiKey: "key-abc", - userProject: "proj-123", - apiRevision: "rev-2", - header: "X-Custom-Trace: trace-99", - wantMissing: false, - wantAPIKey: "key-abc", - }, - { - name: "access_token_credential", - explicitCred: false, - accessToken: "oauth-tok-xyz", - wantMissing: false, - wantBearer: "Bearer oauth-tok-xyz", - }, - { - name: "api_key_wins_same_rank_tie", - explicitCred: false, - apiKey: "tie-key-123", - accessToken: "tie-tok-456", - wantMissing: false, - wantAPIKey: "tie-key-123", - wantBearer: "", - }, - } - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - cmd := newTestRuntimeCmd(t, "https://example.googleapis.com/") - pf := cmd.Root().PersistentFlags() - if tc.apiKey != "" { - _ = pf.Set("api-key", tc.apiKey) - } - if tc.accessToken != "" { - _ = pf.Set("access-token", tc.accessToken) - } - if tc.userProject != "" { - _ = pf.Set("user-project", tc.userProject) - } - if tc.apiRevision != "" { - _ = pf.Set("api-revision", tc.apiRevision) - } - if tc.header != "" { - _ = pf.Set("header", tc.header) - } - _ = pf.Set("timeout", "3s") - - rt, err := newRuntimeTransport(cmd, tc.explicitCred) - if err != nil { - t.Fatalf("newRuntimeTransport() unexpected error: %v", err) - } - if got := flagutil.CredentialsMissing(cmd); got != tc.wantMissing { - t.Errorf("CredentialsMissing() = %v, want %v", got, tc.wantMissing) - } - if rt.timeout != 3*time.Second { - t.Errorf("rt.timeout = %v, want 3s", rt.timeout) - } - - req, err := rt.newRequest(t.Context(), http.MethodPost, "/v1beta/test", http.NoBody) - if err != nil { - t.Fatalf("rt.newRequest() unexpected error: %v", err) - } - if got := req.Header.Get("x-goog-api-key"); got != tc.wantAPIKey { - t.Errorf("x-goog-api-key = %q, want %q", got, tc.wantAPIKey) - } - if got := req.Header.Get("Authorization"); got != tc.wantBearer { - t.Errorf("Authorization = %q, want %q", got, tc.wantBearer) - } - if tc.userProject != "" && req.Header.Get("x-goog-user-project") != tc.userProject { - t.Errorf("x-goog-user-project = %q, want %q", req.Header.Get("x-goog-user-project"), tc.userProject) - } - if tc.apiRevision != "" && req.Header.Get("Api-Revision") != tc.apiRevision { - t.Errorf("Api-Revision = %q, want %q", req.Header.Get("Api-Revision"), tc.apiRevision) - } - if tc.header != "" && req.Header.Get("X-Custom-Trace") != "trace-99" { - t.Errorf("X-Custom-Trace = %q, want \"trace-99\"", req.Header.Get("X-Custom-Trace")) - } - }) - } - }) - - t.Run("transport_validation_errors", func(t *testing.T) { - badCases := []struct { - name string - serverURL string - apiVer string - header string - timeout string - wantSub string - }{ - {name: "ftp_scheme", serverURL: "ftp://example.com", wantSub: "invalid --server-url"}, - {name: "url_with_credentials", serverURL: "https://user:pass@example.com", wantSub: "without credentials, query, or fragment"}, - {name: "url_with_query", serverURL: "https://example.com?foo=bar", wantSub: "without credentials, query, or fragment"}, - {name: "url_with_fragment", serverURL: "https://example.com#sec", wantSub: "without credentials, query, or fragment"}, - {name: "invalid_api_version", apiVer: "..", wantSub: "--api-version"}, - {name: "header_missing_colon", header: "BadHeader", wantSub: "invalid header format"}, - {name: "header_empty_key", header: " : val", wantSub: "invalid header name"}, - {name: "header_space_in_key", header: "Bad Key: val", wantSub: "invalid header name"}, - {name: "header_control_char_val", header: "X-Key: bad\x7fval", wantSub: "--header contains control characters"}, - {name: "invalid_timeout", timeout: "not-a-duration", wantSub: "invalid --timeout"}, - } - for _, tc := range badCases { - t.Run(tc.name, func(t *testing.T) { - cmd := newTestRuntimeCmd(t, tc.serverURL) - pf := cmd.Root().PersistentFlags() - if tc.apiVer != "" { - _ = pf.Set("api-version", tc.apiVer) - } - if tc.header != "" { - _ = pf.Set("header", tc.header) - } - if tc.timeout != "" { - _ = pf.Set("timeout", tc.timeout) - } - if got, err := newRuntimeTransport(cmd, true); err == nil || !strings.Contains(err.Error(), tc.wantSub) { - t.Errorf("newRuntimeTransport() = (%v, %v), want error containing %q", got, err, tc.wantSub) - } - }) - } - }) - - t.Run("leased_interaction_unmarshal_edge_cases", func(t *testing.T) { - validInputs := []struct { - raw string - wantGen int64 - }{ - {raw: `{"interaction_id":"i1","lease_token":"t1","lease_generation":null}`, wantGen: 0}, - {raw: `{"interaction_id":"i2","lease_token":"t2","lease_generation":""}`, wantGen: 0}, - } - for _, tc := range validInputs { - var li LeasedInteraction - if err := json.Unmarshal([]byte(tc.raw), &li); err != nil { - t.Errorf("json.Unmarshal(%s) unexpected error: %v", tc.raw, err) - } else if li.LeaseGeneration != tc.wantGen { - t.Errorf("LeaseGeneration(%s) = %d, want %d", tc.raw, li.LeaseGeneration, tc.wantGen) - } - } - - for _, bad := range []string{ - `{not-json`, - `{"interaction_id":"i1","lease_generation":"not-an-int"}`, - `{"interaction_id":"i1","lease_generation":true}`, - } { - var li LeasedInteraction - if err := json.Unmarshal([]byte(bad), &li); err == nil { - t.Errorf("json.Unmarshal(%s) = %+v, want error", bad, li) - } - } - }) -} - -func TestGRPCStatusParsingHeartbeatsAndPlainHTTPProbe(t *testing.T) { - t.Run("decode_grpc_message", func(t *testing.T) { - cases := []struct { - in string - want string - }{ - {in: "", want: ""}, - {in: " ", want: ""}, - {in: "permission%20denied%3A%20expired", want: "permission denied: expired"}, - {in: "invalid%ZZpercent", want: "invalid%ZZpercent"}, - } - for _, tc := range cases { - if got := decodeGRPCMessage(tc.in); got != tc.want { - t.Errorf("decodeGRPCMessage(%q) = %q, want %q", tc.in, got, tc.want) - } - } - }) - - t.Run("parse_grpc_status_and_code_mapping", func(t *testing.T) { - for _, okCode := range []string{"", "0", " 0 "} { - if err := parseGRPCStatus(okCode, "ignored", nil); err != nil { - t.Errorf("parseGRPCStatus(%q) = %v, want nil", okCode, err) - } - } - - statusCases := []struct { - rawCode string - rawMsg string - wantHTTP int - wantMsgSub string - }{ - {rawCode: "3", rawMsg: "bad%20arg", wantHTTP: http.StatusBadRequest, wantMsgSub: "bad arg"}, - {rawCode: "4", rawMsg: "", wantHTTP: http.StatusGatewayTimeout, wantMsgSub: "gRPC status 4"}, - {rawCode: "5", rawMsg: "missing", wantHTTP: http.StatusNotFound, wantMsgSub: "missing"}, - {rawCode: "6", rawMsg: "exists", wantHTTP: http.StatusConflict, wantMsgSub: "exists"}, - {rawCode: "7", rawMsg: "denied", wantHTTP: http.StatusForbidden, wantMsgSub: "denied"}, - {rawCode: "8", rawMsg: "quota", wantHTTP: http.StatusTooManyRequests, wantMsgSub: "quota"}, - {rawCode: "9", rawMsg: "precondition", wantHTTP: http.StatusBadRequest, wantMsgSub: "precondition"}, - {rawCode: "10", rawMsg: "aborted", wantHTTP: http.StatusConflict, wantMsgSub: "aborted"}, - {rawCode: "11", rawMsg: "range", wantHTTP: http.StatusBadRequest, wantMsgSub: "range"}, - {rawCode: "12", rawMsg: "unimplemented", wantHTTP: http.StatusNotImplemented, wantMsgSub: "unimplemented"}, - {rawCode: "14", rawMsg: "unavailable", wantHTTP: http.StatusServiceUnavailable, wantMsgSub: "unavailable"}, - {rawCode: "16", rawMsg: "unauthenticated", wantHTTP: http.StatusUnauthorized, wantMsgSub: "unauthenticated"}, - {rawCode: "13", rawMsg: "internal", wantHTTP: http.StatusInternalServerError, wantMsgSub: "internal"}, - {rawCode: "not-a-number", rawMsg: "", wantHTTP: http.StatusInternalServerError, wantMsgSub: "gRPC status 2"}, - } - for _, tc := range statusCases { - err := parseGRPCStatus(tc.rawCode, tc.rawMsg, nil) - if err == nil { - t.Fatalf("parseGRPCStatus(%q, %q) = nil, want HTTP %d error", tc.rawCode, tc.rawMsg, tc.wantHTTP) - } - wantStatusText := fmt.Sprintf("Status %d", tc.wantHTTP) - if !strings.Contains(err.Error(), wantStatusText) || !strings.Contains(err.Error(), tc.wantMsgSub) { - t.Errorf("parseGRPCStatus(%q, %q) = %q, want substrings %q and %q", tc.rawCode, tc.rawMsg, err.Error(), wantStatusText, tc.wantMsgSub) - } - } - }) - - t.Run("grpc_headers_and_trailers_on_stream", func(t *testing.T) { - var callNum atomic.Int32 - srvURL := startH2CTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - n := callNum.Add(1) - if n == 1 { - // Trailers-only gRPC error in initial response headers. - w.Header().Set("Content-Type", "application/grpc") - w.Header().Set("Grpc-Status", "16") - w.Header().Set("Grpc-Message", "invalid%20lease%20token") - w.WriteHeader(http.StatusOK) - return - } - // Stream opens 200 OK and terminates with non-zero gRPC trailer status. - w.Header().Set("Content-Type", "application/grpc") - w.Header().Set("Trailer", "Grpc-Status, Grpc-Message") - w.WriteHeader(http.StatusOK) - flusher, _ := w.(http.Flusher) - flusher.Flush() - _, _ = grpcReadFrame(r.Body) - w.Header().Set("Grpc-Status", "14") - w.Header().Set("Grpc-Message", "backend%20restarting") - })) - - cmd := newTestRuntimeCmd(t, srvURL) - rt, err := newRuntimeTransport(cmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport() failed: %v", err) - } - client := NewHTTPClient(HTTPClientOptions{Transport: rt}) - - // Call 1: Grpc-Status in HTTP headers fails connectStream with non-retryable 401. - _, err = client.connectStream(t.Context(), connectStreamRequest{InteractionID: "int_1", LeaseToken: "tok_1", NextOutputSeqNum: 1}) - if err == nil || !strings.Contains(err.Error(), "invalid lease token") || !isNonRetryableControlPlaneError(err) { - t.Errorf("connectStream(header Grpc-Status=16) = %v, want non-retryable 401 error", err) - } - - // Call 2: Grpc-Status in HTTP trailers surfaces on stream.Recv() EOF as retryable 503. - stream, err := client.connectStream(t.Context(), connectStreamRequest{InteractionID: "int_2", LeaseToken: "tok_2", NextOutputSeqNum: 1}) - if err != nil { - t.Fatalf("connectStream(call 2) unexpected error: %v", err) - } - defer stream.Close() - _, recvErr := stream.Recv() - if recvErr == nil || !strings.Contains(recvErr.Error(), "backend restarting") || isNonRetryableControlPlaneError(recvErr) { - t.Errorf("stream.Recv(trailer Grpc-Status=14) = %v, want retryable 503 error", recvErr) - } - }) - - t.Run("plain_http1_server_probe_surfaces_http_error", func(t *testing.T) { - http1Srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - http.Error(w, `{"error":{"code":403,"message":"forbidden plain http"}}`, http.StatusForbidden) - })) - defer http1Srv.Close() - - cmd := newTestRuntimeCmd(t, http1Srv.URL) - rt, err := newRuntimeTransport(cmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport() failed: %v", err) - } - client := NewHTTPClient(HTTPClientOptions{Transport: rt}) - _, err = client.connectStream(t.Context(), connectStreamRequest{InteractionID: "int_h1", LeaseToken: "tok_h1", NextOutputSeqNum: 1}) - if err == nil || !strings.Contains(err.Error(), "Status 403") { - t.Errorf("connectStream(HTTP/1.1 403) = %v, want probed Status 403 error", err) - } - }) - - t.Run("heartbeat_cancel_and_error_branches_and_backoff", func(t *testing.T) { - var mode atomic.Int32 - srvURL := startH2CTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - switch mode.Load() { - case 1: - http.Error(w, `{"error":{"code":409,"message":"lease stolen"}}`, http.StatusConflict) - case 2: - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"leases":[]}`)) - default: - http.Error(w, `{"error":{"code":500,"message":"boom"}}`, http.StatusInternalServerError) - } - })) - - cmd := newTestRuntimeCmd(t, srvURL) - _ = cmd.Root().PersistentFlags().Set("timeout", "2s") - rt, err := newRuntimeTransport(cmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport() failed: %v", err) - } - client := NewHTTPClient(HTTPClientOptions{Transport: rt, EnvironmentKey: "ek_hb"}) - - mode.Store(1) - res409, err := client.HeartbeatInteraction(t.Context(), "p1", "int_1", "env_1", "tok_1") - if err != nil || !res409.ShouldCancel || res409.Renewed { - t.Errorf("HeartbeatInteraction(409) = (%+v, %v), want ShouldCancel=true", res409, err) - } - - mode.Store(2) - resEmpty, err := client.HeartbeatInteraction(t.Context(), "p1", "int_1", "env_1", "tok_1") - if err != nil || !resEmpty.ShouldCancel || resEmpty.Renewed { - t.Errorf("HeartbeatInteraction(empty leases) = (%+v, %v), want ShouldCancel=true", resEmpty, err) - } - - mode.Store(3) - if _, err := client.HeartbeatInteraction(t.Context(), "p1", "int_1", "env_1", "tok_1"); err == nil { - t.Errorf("HeartbeatInteraction(500) = nil error, want 500 error") - } - - if isNonRetryableControlPlaneError(nil) { - t.Errorf("isNonRetryableControlPlaneError(nil) = true, want false") - } - if isNonRetryableControlPlaneError(errors.New("network reset")) { - t.Errorf("isNonRetryableControlPlaneError(plain error) = true, want false") - } - if got := retryBackoff(0, 1); got != time.Second { - t.Errorf("retryBackoff(0, 1) = %v, want 1s", got) - } - if got := retryBackoff(2*time.Second, 10); got != 30*time.Second { - t.Errorf("retryBackoff(2s, 10) = %v, want capped 30s", got) - } - }) -} - -func TestProvisionAndDispatchPollerDefaultsAndSubprocessGroup(t *testing.T) { - t.Run("resolve_provision_and_worker_shells", func(t *testing.T) { - bin, args := resolveProvisionShell("linux", "echo hi") - if !filepath.IsAbs(bin) || !slices.Equal(args, []string{"-c", "echo hi"}) { - t.Errorf("resolveProvisionShell(linux, normal PATH) = (%q, %v), want abs sh path and [-c echo hi]", bin, args) - } - - shOnlyDir := t.TempDir() - fakeSh := filepath.Join(shOnlyDir, "sh") - if err := os.WriteFile(fakeSh, []byte("#!/bin/sh\n"), 0o755); err != nil { - t.Fatalf("WriteFile(fakeSh) failed: %v", err) - } - t.Setenv("PATH", shOnlyDir) - wBinSh, wArgsSh := resolveShellCommand("linux", "pwd") - if wBinSh != fakeSh || !slices.Equal(wArgsSh, []string{"-c", "pwd"}) { - t.Errorf("resolveShellCommand(sh-only PATH) = (%q, %v), want (%q, [-c pwd])", wBinSh, wArgsSh, fakeSh) - } - - t.Setenv("PATH", "") - winBin, winArgs := resolveProvisionShell("windows", "echo win") - if winBin != "cmd.exe" || !slices.Equal(winArgs, []string{"/C", "echo win"}) { - t.Errorf("resolveProvisionShell(windows, empty PATH) = (%q, %v), want (cmd.exe, [/C echo win])", winBin, winArgs) - } - linBin, linArgs := resolveProvisionShell("linux", "echo lin") - if linBin != "sh" || !slices.Equal(linArgs, []string{"-c", "echo lin"}) { - t.Errorf("resolveProvisionShell(linux, empty PATH) = (%q, %v), want (sh, [-c echo lin])", linBin, linArgs) - } - - wWinBin, wWinArgs := resolveShellCommand("windows", "dir") - if wWinBin != "cmd.exe" || !slices.Equal(wWinArgs, []string{"/C", "dir"}) { - t.Errorf("resolveShellCommand(windows, empty PATH) = (%q, %v), want (cmd.exe, [/C dir])", wWinBin, wWinArgs) - } - wLinBin, wLinArgs := resolveShellCommand("linux", "ls") - if wLinBin != "sh" || !slices.Equal(wLinArgs, []string{"-c", "ls"}) { - t.Errorf("resolveShellCommand(linux, empty PATH) = (%q, %v), want (sh, [-c ls])", wLinBin, wLinArgs) - } - }) - - t.Run("provision_env_scrubbing_redaction_and_process_group_kill", func(t *testing.T) { - emptyProv := NewCommandWorkerProvisioner(" ") - if err := emptyProv.ProvisionAndDispatch(t.Context(), ProvisionInput{InteractionID: "i1", EnvironmentID: "e1"}); err == nil { - t.Errorf("ProvisionAndDispatch(empty cmd) = nil, want error") - } - - // Verify GEMINI_API_KEY and GEMINI_ENVIRONMENT_KEY are scrubbed from --provision-cmd's environment - // while non-credential environment variables are preserved, and empty SchemaVersion defaults to "v1". - t.Setenv("GEMINI_API_KEY", "scrub_this_api_key") - t.Setenv("GEMINI_ENVIRONMENT_KEY", "scrub_this_env_key") - t.Setenv("LEASE_TOKEN", "scrub_this_lease_tok") - t.Setenv("CUSTOM_ORCHESTRATOR_VAR", "keep_orchestrator_val") - - envDump := filepath.Join(t.TempDir(), "prov.env") - stdinDump := filepath.Join(t.TempDir(), "prov.json") - t.Setenv("TEST_STDIN_DUMP", stdinDump) - t.Setenv("TEST_ENV_DUMP", envDump) - prov := NewCommandWorkerProvisioner(`cat > "$TEST_STDIN_DUMP" && env > "$TEST_ENV_DUMP"`) - if err := prov.ProvisionAndDispatch(t.Context(), ProvisionInput{ - InteractionID: "int_env_test", - EnvironmentID: "env_test", - LeaseToken: "lt_secret_123", - }); err != nil { - t.Fatalf("ProvisionAndDispatch() unexpected error: %v", err) - } - - envBytes, err := os.ReadFile(envDump) - if err != nil { - t.Fatalf("ReadFile(envDump) failed: %v", err) - } - envText := string(envBytes) - for _, leaked := range []string{"scrub_this_api_key", "scrub_this_env_key", "scrub_this_lease_tok"} { - if strings.Contains(envText, leaked) { - t.Errorf("ProvisionAndDispatch leaked credential %q into child env:\n%s", leaked, envText) - } - } - if !strings.Contains(envText, "CUSTOM_ORCHESTRATOR_VAR=keep_orchestrator_val") { - t.Errorf("ProvisionAndDispatch dropped non-credential var CUSTOM_ORCHESTRATOR_VAR:\n%s", envText) - } - - var gotInput ProvisionInput - stdinBytes, err := os.ReadFile(stdinDump) - if err != nil || json.Unmarshal(stdinBytes, &gotInput) != nil { - t.Fatalf("reading ProvisionInput dump failed: %v", err) - } - if gotInput.SchemaVersion != "v1" { - t.Errorf("ProvisionInput.SchemaVersion = %q, want default \"v1\"", gotInput.SchemaVersion) - } - - // Verify stderr redaction of LeaseToken and empty-stderr failure formatting. - redactProv := NewCommandWorkerProvisioner("printf 'failed with lt_secret_abc\\n' >&2; exit 2") - err = redactProv.ProvisionAndDispatch(t.Context(), ProvisionInput{ - InteractionID: "int_fail", - EnvironmentID: "env_fail", - LeaseToken: "lt_secret_abc", - }) - if err == nil || strings.Contains(err.Error(), "lt_secret_abc") || !strings.Contains(err.Error(), "[REDACTED]") { - t.Errorf("ProvisionAndDispatch(stderr leak) error = %v, want [REDACTED] without raw token", err) - } - - silentFailProv := NewCommandWorkerProvisioner("exit 3") - if err := silentFailProv.ProvisionAndDispatch(t.Context(), ProvisionInput{InteractionID: "i3"}); err == nil || - !strings.HasPrefix(err.Error(), "provision command failed: ") || strings.Contains(err.Error(), "()") { - t.Errorf("ProvisionAndDispatch(exit 3, empty stderr) = %v, want 'provision command failed: ...' without empty '()'", err) - } - - // Verify configureSubprocessGroup kills background grandchild processes holding stderr open on timeout: - // without process-group SIGKILL, cmd.Run would block ~4s waiting for the grandchild's stderr pipe to close! - timeoutCtx, cancel := context.WithTimeout(t.Context(), 80*time.Millisecond) - defer cancel() - start := time.Now() - bgProv := NewCommandWorkerProvisioner("(sleep 4; echo orphan >&2) & sleep 4") - if err := bgProv.ProvisionAndDispatch(timeoutCtx, ProvisionInput{InteractionID: "i_bg"}); err == nil { - t.Errorf("ProvisionAndDispatch(timeout) = nil, want deadline error") - } - if elapsed := time.Since(start); elapsed > 1500*time.Millisecond { - t.Errorf("ProvisionAndDispatch(background child) took %v after 80ms timeout, want immediate process-group kill", elapsed) - } - }) - - t.Run("subprocess_group_unstarted_cancel_and_signals", func(t *testing.T) { - sigs := runtimeShutdownSignals() - if len(sigs) == 0 { - t.Errorf("runtimeShutdownSignals() = empty, want OS shutdown signals") - } - unstarted := exec.Command("sh", "-c", "true") - configureSubprocessGroup(unstarted) - if unstarted.Cancel == nil { - t.Fatalf("configureSubprocessGroup() left cmd.Cancel == nil") - } - if err := unstarted.Cancel(); err != nil { - t.Errorf("unstarted cmd.Cancel() = %v, want nil when cmd.Process is nil", err) - } - }) - - t.Run("poller_defaults_cooldown_eviction_cancel_and_no_retries", func(t *testing.T) { - defaultPoller := NewPollerController(nil, nil, PollerOptions{}) - if defaultPoller.pollerID == "" || defaultPoller.maxInteractions != 10 || defaultPoller.pollInterval != 2*time.Second || - defaultPoller.provisionTimeout != 2*time.Minute || defaultPoller.handoffHeartbeatInterval != defaultHeartbeatInterval || defaultPoller.logger == nil { - t.Errorf("NewPollerController(zero opts) = %+v, want populated defaults", defaultPoller) - } - - // Verify stale failure entries (>10m past cooldownEnd) are evicted while recent entries (<10m) are retained. - defaultPoller.failures["stale_int"] = &failureState{ - consecutive: 2, - cooldownEnd: time.Now().Add(-15 * time.Minute), - } - defaultPoller.failures["recent_int"] = &failureState{ - consecutive: 1, - cooldownEnd: time.Now().Add(-5 * time.Minute), - } - defaultPoller.recordDispatchFailure("fresh_int") - if _, exists := defaultPoller.failures["stale_int"]; exists { - t.Errorf("recordDispatchFailure() kept stale_int past 10m eviction window") - } - if _, exists := defaultPoller.failures["recent_int"]; !exists { - t.Errorf("recordDispatchFailure() evicted recent_int within 10m eviction window") - } - if st := defaultPoller.failures["fresh_int"]; st == nil || st.consecutive != 1 { - t.Errorf("recordDispatchFailure(fresh_int) = %+v, want consecutive=1", st) - } - - // Verify exponential cooldown progression (500ms -> 1s -> 2s -> 4s -> 8s -> 15s cap -> 15s cap). - wantCooldowns := []time.Duration{ - 500 * time.Millisecond, - 1 * time.Second, - 2 * time.Second, - 4 * time.Second, - 8 * time.Second, - 15 * time.Second, - 15 * time.Second, - } - for i, want := range wantCooldowns { - before := time.Now() - defaultPoller.recordDispatchFailure("backoff_seq") - after := time.Now() - st := defaultPoller.failures["backoff_seq"] - if st == nil || st.consecutive != i+1 { - t.Fatalf("recordDispatchFailure(backoff_seq #%d) = %+v, want consecutive=%d", i+1, st, i+1) - } - minEnd := before.Add(want - 50*time.Millisecond) - maxEnd := after.Add(want + 50*time.Millisecond) - if st.cooldownEnd.Before(minEnd) || st.cooldownEnd.After(maxEnd) { - t.Errorf("recordDispatchFailure(backoff_seq #%d) cooldown ≈ %v, want %v", i+1, st.cooldownEnd.Sub(before), want) - } - } - - // Verify ShouldCancel heartbeat cancels an active slow ProvisionAndDispatch immediately, - // and NoRetries=true returns transient 503 from Run() without retrying. - var leaseCount atomic.Int32 - srvURL := startH2CTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - var req leaseInteractionsRequest - _ = json.NewDecoder(r.Body).Decode(&req) - if len(req.ActiveLeases) > 0 { - http.Error(w, `{"error":{"code":404,"message":"canceled"}}`, http.StatusNotFound) - return - } - if leaseCount.Add(1) == 1 { - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"leases":[{"interaction_id":"int_cancel_me","lease_token":"tok_c"}]}`)) - return - } - http.Error(w, `{"error":{"code":503,"message":"transient"}}`, http.StatusServiceUnavailable) - })) - - cmd := newTestRuntimeCmd(t, srvURL) - rt, err := newRuntimeTransport(cmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport() failed: %v", err) - } - client := NewHTTPClient(HTTPClientOptions{Transport: rt, EnvironmentKey: "ek_1"}) - poller := NewPollerController(client, NewCommandWorkerProvisioner("sleep 5"), PollerOptions{ - PollerID: "poller-cancel-test", - MaxInteractions: 1, - PollInterval: 20 * time.Millisecond, - ProvisionTimeout: 10 * time.Second, - HandoffHeartbeatInterval: 25 * time.Millisecond, - NoRetries: true, - Logger: log.New(io.Discard, "", 0), - }) - - start := time.Now() - runErr := poller.Run(t.Context()) - if runErr == nil || !strings.Contains(runErr.Error(), "Status 503") { - t.Errorf("PollerController.Run(NoRetries=true) = %v, want Status 503 error", runErr) - } - if elapsed := time.Since(start); elapsed > 2*time.Second { - t.Errorf("PollerController ShouldCancel handoff abort took %v, want < 2s", elapsed) - } - }) -} - -func TestSandboxPathsCappedBufferAndFileOperations(t *testing.T) { - t.Run("resolve_sandbox_dirs_symlinks_dangling_and_create", func(t *testing.T) { - base := t.TempDir() - realRoot := filepath.Join(base, "real_root") - if err := os.MkdirAll(filepath.Join(realRoot, "sub"), 0o755); err != nil { - t.Fatalf("MkdirAll(realRoot/sub) failed: %v", err) - } - evalRealRoot, err := filepath.EvalSymlinks(realRoot) - if err != nil { - t.Fatalf("EvalSymlinks(realRoot) failed: %v", err) - } - - symRoot := filepath.Join(base, "sym_root") - if err := os.Symlink(realRoot, symRoot); err != nil { - t.Fatalf("Symlink(symRoot) failed: %v", err) - } - - // Symlinked rootDir with workDir under symlinked rootDir must canonicalize both! - gotRoot, gotWork, err := resolveSandboxDirs(symRoot, filepath.Join(symRoot, "sub"), false) - if err != nil { - t.Fatalf("resolveSandboxDirs(symRoot, symRoot/sub, false) failed: %v", err) - } - if gotRoot != evalRealRoot { - t.Errorf("gotRoot = %q, want canonical %q", gotRoot, evalRealRoot) - } - wantWork := filepath.Join(evalRealRoot, "sub") - if gotWork != wantWork { - t.Errorf("gotWork = %q, want canonical %q", gotWork, wantWork) - } - - // Dangling symlink as rootDir must be rejected by rejectDanglingSymlinkPrefix! - danglingRoot := filepath.Join(base, "dangling_root") - if err := os.Symlink(filepath.Join(base, "does_not_exist"), danglingRoot); err != nil { - t.Fatalf("Symlink(danglingRoot) failed: %v", err) - } - if _, _, err := resolveSandboxDirs(filepath.Join(danglingRoot, "nested"), "", false); err == nil || !strings.Contains(err.Error(), "dangling symlink") { - t.Errorf("resolveSandboxDirs(dangling symlink) = %v, want dangling symlink error", err) - } - - // create=true creates missing rootDir and workDir. - newRoot := filepath.Join(evalRealRoot, "auto_created_root") - newWork := filepath.Join(newRoot, "nested_work") - createdRoot, createdWork, err := resolveSandboxDirs(newRoot, newWork, true) - if err != nil { - t.Fatalf("resolveSandboxDirs(create=true) failed: %v", err) - } - if createdRoot != newRoot || createdWork != newWork { - t.Errorf("resolveSandboxDirs(create=true) = (%q, %q), want (%q, %q)", createdRoot, createdWork, newRoot, newWork) - } - - // Empty rootDir ("") and relative workDir ("rel_sub") resolve against cwd. - t.Chdir(evalRealRoot) - cwdRoot, cwdWork, err := resolveSandboxDirs("", "rel_sub", true) - if err != nil { - t.Fatalf("resolveSandboxDirs(\"\", \"rel_sub\", true) failed: %v", err) - } - if cwdRoot != evalRealRoot || cwdWork != filepath.Join(evalRealRoot, "rel_sub") { - t.Errorf("resolveSandboxDirs(cwd) = (%q, %q), want (%q, %q)", cwdRoot, cwdWork, evalRealRoot, filepath.Join(evalRealRoot, "rel_sub")) - } - }) - - t.Run("virtual_absolute_paths_and_file_edge_cases", func(t *testing.T) { - rootDir, err := filepath.EvalSymlinks(t.TempDir()) - if err != nil { - t.Fatalf("EvalSymlinks(TempDir) failed: %v", err) - } - workDir := filepath.Join(rootDir, "app") - if err := os.MkdirAll(workDir, 0o755); err != nil { - t.Fatalf("MkdirAll(workDir) failed: %v", err) - } - // newFSServer falls back to rootDir when workDir escapes rootDir. - if fallbackFS := newFSServer(rootDir, t.TempDir()); fallbackFS.workDir != rootDir { - t.Errorf("newFSServer(outside workDir).workDir = %q, want %q", fallbackFS.workDir, rootDir) - } - fs := newFSServer(rootDir, workDir) - - // Dangling symlink under an evaluated directory suffix is rejected. - realTargetDir := filepath.Join(rootDir, "real_target") - if err := os.MkdirAll(realTargetDir, 0o755); err != nil { - t.Fatalf("MkdirAll(realTargetDir) failed: %v", err) - } - if err := os.Symlink(filepath.Join(rootDir, "missing_leaf"), filepath.Join(realTargetDir, "dangling_leaf")); err != nil { - t.Fatalf("Symlink(dangling_leaf) failed: %v", err) - } - if _, err := fs.ResolvePath(filepath.Join("..", "real_target", "dangling_leaf", "sub")); err == nil || !strings.Contains(err.Error(), "dangling symlink") { - t.Errorf("ResolvePath(dangling suffix) = %v, want dangling symlink error", err) - } - - // Virtual absolute path "/workspace/hello.txt" outside rootDir maps cleanly inside rootDir! - gotVirt, err := fs.ResolvePath("/workspace/hello.txt") - wantVirt := filepath.Join(rootDir, "workspace", "hello.txt") - if err != nil || gotVirt != wantVirt { - t.Errorf("ResolvePath(/workspace/hello.txt) = (%q, %v), want (%q, nil)", gotVirt, err, wantVirt) - } - - // Real absolute path already inside rootDir stays unchanged. - realInside := filepath.Join(rootDir, "app", "main.py") - if gotReal, err := fs.ResolvePath(realInside); err != nil || gotReal != realInside { - t.Errorf("ResolvePath(%q) = (%q, %v), want (%q, nil)", realInside, gotReal, err, realInside) - } - if gotDot, err := fs.ResolvePath("."); err != nil || gotDot != workDir { - t.Errorf("ResolvePath(.) = (%q, %v), want (%q, nil)", gotDot, err, workDir) - } - if gotRoot, err := fs.ResolvePath(rootDir); err != nil || gotRoot != rootDir { - t.Errorf("ResolvePath(rootDir) = (%q, %v), want (%q, nil)", gotRoot, err, rootDir) - } - - // WriteRegularFile modes: overwrite=false (O_EXCL) and appendMode=true (O_APPEND). - if _, err := fs.WriteRegularFile("log.txt", "line1\n", false, false); err != nil { - t.Fatalf("WriteRegularFile(create excl) failed: %v", err) - } - if _, err := fs.WriteRegularFile("log.txt", "dup\n", false, false); err == nil { - t.Errorf("WriteRegularFile(overwrite=false on existing file) = nil, want O_EXCL error") - } - if _, err := fs.WriteRegularFile("log.txt", "line2\n", false, true); err != nil { - t.Fatalf("WriteRegularFile(appendMode=true) failed: %v", err) - } - data, err := fs.ReadRegularFile("log.txt") - if err != nil || string(data) != "line1\nline2\n" { - t.Errorf("ReadRegularFile(log.txt) = (%q, %v), want \"line1\\nline2\\n\"", data, err) - } - - // Directory read errors and reversed ReadFileRange. - if _, err := fs.ReadRegularFile("."); err == nil || !strings.Contains(err.Error(), "directory") { - t.Errorf("ReadRegularFile(.) = %v, want directory error", err) - } - if _, err := fs.ReadFileRange(".", 1, 2); err == nil || !strings.Contains(err.Error(), "directory") { - t.Errorf("ReadFileRange(.) = %v, want directory error", err) - } - if emptySlice, err := fs.ReadFileRange("log.txt", 5, 2); err != nil || emptySlice != "" { - t.Errorf("ReadFileRange(start=5, end=2) = (%q, %v), want (\"\", nil)", emptySlice, err) - } - - // Grep literal + case-insensitive, empty query error, and invalid regex error. - if _, err := fs.WriteRegularFile("a/b/deep.txt", "Alpha+Beta\n", true, false); err != nil { - t.Fatalf("WriteRegularFile(a/b/deep.txt) failed: %v", err) - } - matches, _, err := fs.Grep(".", "alpha+beta", true, false) - if err != nil || len(matches) != 1 { - t.Errorf("Grep(literal caseInsensitive) = (%v, %v), want 1 match", matches, err) - } - if _, _, err := fs.Grep(".", "", false, false); err == nil { - t.Errorf("Grep(empty query) = nil, want error") - } - if _, _, err := fs.Grep(".", "[unclosed", false, true); err == nil { - t.Errorf("Grep(invalid regex) = nil, want error") - } - - // FindFiles maxDepth=1 skips deeper directories. - shallow, _, err := fs.FindFiles(".", "", 1) - if err != nil { - t.Fatalf("FindFiles(maxDepth=1) failed: %v", err) - } - for _, item := range shallow { - if strings.Contains(fmt.Sprint(item), "/") { - t.Errorf("FindFiles(maxDepth=1) returned nested entry %v in %v", item, shallow) - } - } - - // Escape, missing-file, non-regular-write, single-file Grep, and ListEntries error paths. - if firstLine, err := fs.ReadFileRange("log.txt", 0, 1); err != nil || firstLine != "line1" { - t.Errorf("ReadFileRange(start=0, end=1) = (%q, %v), want (\"line1\", nil)", firstLine, err) - } - if _, err := fs.ReadRegularFile("missing.txt"); err == nil { - t.Errorf("ReadRegularFile(missing) = nil, want error") - } - if _, err := fs.ReadRegularFile("../../escape"); err == nil { - t.Errorf("ReadRegularFile(escape) = nil, want error") - } - if _, err := fs.ReadFileRange("missing.txt", 1, 2); err == nil { - t.Errorf("ReadFileRange(missing) = nil, want error") - } - if _, err := fs.ReadFileRange("../../escape", 1, 2); err == nil { - t.Errorf("ReadFileRange(escape) = nil, want error") - } - if _, err := fs.WriteRegularFile("../../escape", "x", true, false); err == nil { - t.Errorf("WriteRegularFile(escape) = nil, want error") - } - if _, err := fs.WriteRegularFile("a", "x", true, false); err == nil { - t.Errorf("WriteRegularFile(existing dir) = nil, want error") - } - if _, err := fs.ListEntries("missing_dir"); err == nil { - t.Errorf("ListEntries(missing_dir) = nil, want error") - } - if _, err := fs.ListEntries("../../escape"); err == nil { - t.Errorf("ListEntries(escape) = nil, want error") - } - if _, err := fs.ListEntries("log.txt"); err == nil { - t.Errorf("ListEntries(regular file) = nil, want error") - } - if singleMatch, _, err := fs.Grep("log.txt", "line1", false, false); err != nil || len(singleMatch) != 1 { - t.Errorf("Grep(single file) = (%v, %v), want 1 match", singleMatch, err) - } - if _, _, err := fs.Grep("../../escape", "x", false, false); err == nil { - t.Errorf("Grep(escape) = nil, want error") - } - if _, _, err := fs.FindFiles("../../escape", "*", 0); err == nil { - t.Errorf("FindFiles(escape) = nil, want error") - } - }) - - t.Run("capped_buffer_partial_and_full_truncation", func(t *testing.T) { - cb := &cappedBuffer{maxBytes: 10} - if n, err := cb.Write([]byte("123456")); err != nil || n != 6 { - t.Fatalf("cappedBuffer.Write(6 bytes) = (%d, %v), want (6, nil)", n, err) - } - if cb.String() != "123456" { - t.Errorf("cappedBuffer.String() before cap = %q, want \"123456\"", cb.String()) - } - // Partial write overflow (remaining = 4, len(p) = 8). - if n, err := cb.Write([]byte("ABCDEFGH")); err != nil || n != 8 { - t.Fatalf("cappedBuffer.Write(partial overflow) = (%d, %v), want (8, nil)", n, err) - } - // Full write overflow when remaining == 0. - if n, err := cb.Write([]byte("XYZ")); err != nil || n != 3 { - t.Fatalf("cappedBuffer.Write(zero remaining) = (%d, %v), want (3, nil)", n, err) - } - want := "123456ABCD\n...[output truncated at 4 MiB]..." - if got := cb.String(); got != want { - t.Errorf("cappedBuffer.String() = %q, want %q", got, want) - } - }) -} - -func TestWirePrimitivesProtoValuesAndWorkerConsumption(t *testing.T) { - t.Run("frame_limits_and_empty_frame", func(t *testing.T) { - var buf bytes.Buffer - oversized := make([]byte, maxGRPCFrameBytes+1) - if err := grpcWriteFrame(&buf, oversized); err == nil { - t.Errorf("grpcWriteFrame(oversized) = nil, want limit error") - } - - if err := grpcWriteFrame(&buf, nil); err != nil { - t.Fatalf("grpcWriteFrame(empty) failed: %v", err) - } - gotEmpty, err := grpcReadFrame(&buf) - if err != nil || gotEmpty != nil { - t.Errorf("grpcReadFrame(empty) = (%x, %v), want (nil, nil)", gotEmpty, err) - } - }) - - t.Run("proto_value_round_trip_and_wire_primitives", func(t *testing.T) { - type customStringer struct{ v string } - inStruct := map[string]any{ - "nil_val": nil, - "bool_true": true, - "bool_false": false, - "int_val": int(12), - "int32_val": int32(34), - "int64_val": int64(56), - "float_val": float64(7.5), - "str_val": "hello", - "any_list": []any{"a", float64(2), nil}, - "str_list": []string{"x", "y"}, - "nested": map[string]any{"k": "v"}, - "fallback": customStringer{v: "custom"}, - } - decoded, err := decodeProtoStruct(encodeProtoStruct(inStruct), 0) - if err != nil { - t.Fatalf("decodeProtoStruct() failed: %v", err) - } - if decoded["nil_val"] != nil || decoded["bool_true"] != true || decoded["bool_false"] != false || - decoded["int_val"] != float64(12) || decoded["int32_val"] != float64(34) || decoded["int64_val"] != float64(56) || - decoded["float_val"] != float64(7.5) || decoded["str_val"] != "hello" || decoded["fallback"] != "{custom}" { - t.Errorf("decodeProtoStruct() scalar mismatch: %+v", decoded) - } - if !reflect.DeepEqual(decoded["str_list"], []any{"x", "y"}) { - t.Errorf("decoded str_list = %#v, want []any{\"x\", \"y\"}", decoded["str_list"]) - } - - if _, err := decodeProtoValue(nil, maxProtoNestingDepth+1); err == nil { - t.Errorf("decodeProtoValue(depth overflow) = nil, want error") - } - if _, err := decodeProtoList(nil, maxProtoNestingDepth+1); err == nil { - t.Errorf("decodeProtoList(depth overflow) = nil, want error") - } - - // Exercise consumeVarint and consumeField wire types (fixed64, fixed32, overflow, invalid field/type). - if _, _, err := consumeVarint(nil); !errors.Is(err, io.ErrUnexpectedEOF) { - t.Errorf("consumeVarint(nil) = %v, want ErrUnexpectedEOF", err) - } - overflowVarint := bytes.Repeat([]byte{0x80}, 11) - if _, _, err := consumeVarint(overflowVarint); err == nil { - t.Errorf("consumeVarint(11-byte overflow) = nil, want overflow error") - } - if _, _, _, _, _, err := consumeField([]byte{0x00}); err == nil { - t.Errorf("consumeField(field 0) = nil, want invalid field number error") - } - - fixed64Buf := binary.LittleEndian.AppendUint64(appendTag(nil, 2, wireFixed64), math.Float64bits(42.25)) - if num, wt, vv, _, _, err := consumeField(fixed64Buf); err != nil || num != 2 || wt != wireFixed64 || math.Float64frombits(vv) != 42.25 { - t.Errorf("consumeField(fixed64) = (%d, %d, %v, %v), want field 2 = 42.25", num, wt, math.Float64frombits(vv), err) - } - if _, _, _, _, _, err := consumeField(appendTag(nil, 2, wireFixed64)); !errors.Is(err, io.ErrUnexpectedEOF) { - t.Errorf("consumeField(truncated fixed64) = %v, want ErrUnexpectedEOF", err) - } - - fixed32Buf := binary.LittleEndian.AppendUint32(appendTag(nil, 3, wireFixed32), 0x12345678) - if num, wt, vv, _, _, err := consumeField(fixed32Buf); err != nil || num != 3 || wt != wireFixed32 || vv != 0x12345678 { - t.Errorf("consumeField(fixed32) = (%d, %d, %#x, %v), want field 3 = 0x12345678", num, wt, vv, err) - } - if _, _, _, _, _, err := consumeField(appendTag(nil, 3, wireFixed32)); !errors.Is(err, io.ErrUnexpectedEOF) { - t.Errorf("consumeField(truncated fixed32) = %v, want ErrUnexpectedEOF", err) - } - if _, _, _, _, _, err := consumeField(appendTag(nil, 4, 3)); err == nil { - t.Errorf("consumeField(wireType=3) = nil, want unsupported wire type error") - } - }) - - t.Run("worker_defaults_reconnected_and_completed_stream_events", func(t *testing.T) { - defWorker := newPullWorker(nil, nil, nil, pullWorkerOptions{}) - if defWorker.workerID == "" || !strings.HasPrefix(defWorker.workerID, "worker-") { - t.Errorf("default workerID = %q, want worker--", defWorker.workerID) - } - if defWorker.maxConcurrent != 4 { - t.Errorf("default maxConcurrent = %d, want 4", defWorker.maxConcurrent) - } - if defWorker.pollInterval != 2*time.Second { - t.Errorf("default pollInterval = %v, want 2s", defWorker.pollInterval) - } - if defWorker.heartbeatInterval != defaultHeartbeatInterval { - t.Errorf("default heartbeatInterval = %v, want %v", defWorker.heartbeatInterval, defaultHeartbeatInterval) - } - if defWorker.commandTimeout != 2*time.Minute { - t.Errorf("default commandTimeout = %v, want 2m", defWorker.commandTimeout) - } - if defWorker.logger == nil { - t.Errorf("default logger = nil, want non-nil") - } - - customWorker := newPullWorker(nil, nil, nil, pullWorkerOptions{ - WorkerID: "w-explicit", - MaxConcurrent: 2, - PollInterval: 100 * time.Millisecond, - HeartbeatInterval: 80 * time.Millisecond, - CommandTimeout: 5 * time.Second, - Logger: log.New(io.Discard, "", 0), - }) - if customWorker.maxConcurrent != 2 || customWorker.heartbeatInterval != 80*time.Millisecond { - t.Errorf("explicit pullWorkerOptions = %+v, want maxConcurrent=2 heartbeatInterval=80ms", customWorker) - } - - // Build a ServerEventBatch stream carrying: - // 1. InteractionReconnected(next_input_seq_num=5) + AckInputSeqNum=4 - // 2. InteractionCompleted (field 6) - var serverFrames bytes.Buffer - reconnectedBatch := appendVarintField(nil, 2, 4) - reconnectedBatch = appendBytesField(reconnectedBatch, 8, appendVarintField(nil, 2, 5)) - _ = grpcWriteFrame(&serverFrames, reconnectedBatch) - _ = grpcWriteFrame(&serverFrames, appendBytesField(nil, 6, nil)) - - pr, pw := io.Pipe() - go func() { - _, _ = io.Copy(io.Discard, pr) - }() - stream := &grpcInteractionStream{ - resp: &http.Response{Body: io.NopCloser(&serverFrames)}, - pw: pw, - } - state := &sessionState{ - interactionID: "int_comp", - nextClientSeq: 1, - executedCalls: make(map[string]*wirePayload), - } - - reconnect, err := customWorker.consumeStream(t.Context(), stream, state) - if err != nil || reconnect { - t.Errorf("consumeStream(InteractionReconnected + InteractionCompleted) = (%v, %v), want (false, nil)", reconnect, err) - } - if state.nextClientSeq != 5 { - t.Errorf("state.nextClientSeq after InteractionReconnected(5) = %d, want 5", state.nextClientSeq) - } - if !stream.closed { - t.Errorf("stream.closed = false, want CloseSend called on InteractionCompleted") - } - }) - - t.Run("tool_helpers_and_replace_in_text_edges", func(t *testing.T) { - rootDir := t.TempDir() - subDir := filepath.Join(rootDir, "sub") - if err := os.MkdirAll(subDir, 0o755); err != nil { - t.Fatalf("MkdirAll(subDir) failed: %v", err) - } - w := newPullWorker(nil, newFSServer(rootDir, rootDir), newShellServer(rootDir), pullWorkerOptions{ - CommandTimeout: 5 * time.Second, - Logger: log.New(io.Discard, "", 0), - }) - - if res, isErr := w.dispatchTool(t.Context(), "nonexistent_tool", nil); !isErr || !strings.Contains(fmt.Sprint(res["error"]), "unsupported sandbox tool") { - t.Errorf("dispatchTool(nonexistent_tool) = (%v, %v), want unsupported tool error", res, isErr) - } - if _, isErr := w.toolExecuteBash(t.Context(), map[string]any{"command": "pwd", "Cwd": "../../escape"}); !isErr { - t.Errorf("toolExecuteBash(escaping Cwd) = false, want error") - } - bashOK, isErr := w.toolExecuteBash(t.Context(), map[string]any{"command": "printf ok", "Cwd": "sub", "timeout_ms": int64(2000)}) - if isErr || bashOK["stdout"] != "ok" { - t.Errorf("toolExecuteBash(valid Cwd + timeout_ms) = (%v, %v), want stdout=\"ok\"", bashOK, isErr) - } - - if _, isErr := w.toolViewFile(nil); !isErr { - t.Errorf("toolViewFile(nil) = false, want missing path error") - } - if _, isErr := w.toolWriteFile(nil); !isErr { - t.Errorf("toolWriteFile(nil) = false, want missing path error") - } - if _, isErr := w.toolReplaceFileContent(nil); !isErr { - t.Errorf("toolReplaceFileContent(nil) = false, want missing TargetFile error") - } - if _, isErr := w.toolReplaceFileContent(map[string]any{"TargetFile": "a.txt"}); !isErr { - t.Errorf("toolReplaceFileContent(missing TargetContent) = false, want error") - } - - // replaceInText edge cases. - if _, _, err := replaceInText("hello", "missing", "x", 0, 0, false); err == nil { - t.Errorf("replaceInText(missing) = nil, want not found error") - } - if _, _, err := replaceInText("dup dup", "dup", "x", 0, 0, false); err == nil { - t.Errorf("replaceInText(multiple with allowMultiple=false) = nil, want error") - } - gotMulti, cnt, err := replaceInText("dup dup", "dup", "ok", 0, 0, true) - if err != nil || cnt != 2 || gotMulti != "ok ok" { - t.Errorf("replaceInText(allowMultiple=true) = (%q, %d, %v), want (\"ok ok\", 2, nil)", gotMulti, cnt, err) - } - if _, _, err := replaceInText("a\nb\nc", "b", "B", 5, 2, false); err == nil { - t.Errorf("replaceInText(invalid range 5..2) = nil, want invalid line range error") - } - if _, _, err := replaceInText("a\nb\nb\nc", "b", "B", 2, 3, false); err == nil { - t.Errorf("replaceInText(ranged multiple with allowMultiple=false) = nil, want error") - } - if _, _, err := replaceInText("a\nb\nc", "z", "Z", 2, 2, false); err == nil { - t.Errorf("replaceInText(ranged missing) = nil, want error") - } - - // Type conversion helpers: stringArg, intArg, boolArg, stringSliceArg. - args := map[string]any{ - "num_str": 42, - "i32": int32(11), - "i64": int64(22), - "f64": float64(33), - "bad_bool": "not-a-bool", - "str_slice": []string{"a", "b"}, - "any_slice": []any{"c", 7, nil}, - } - if got := stringArg(args, "missing", "num_str"); got != "42" { - t.Errorf("stringArg(num_str) = %q, want \"42\"", got) - } - if got := intArg(args, 9, "i32"); got != 11 { - t.Errorf("intArg(i32) = %d, want 11", got) - } - if got := intArg(args, 9, "i64"); got != 22 { - t.Errorf("intArg(i64) = %d, want 22", got) - } - if got := intArg(args, 9, "f64"); got != 33 { - t.Errorf("intArg(f64) = %d, want 33", got) - } - if got := boolArg(args, true, "bad_bool"); !got { - t.Errorf("boolArg(bad_bool, fallback=true) = false, want true") - } - if got := stringSliceArg(args, "str_slice"); !slices.Equal(got, []string{"a", "b"}) { - t.Errorf("stringSliceArg(str_slice) = %v, want [a b]", got) - } - if got := stringSliceArg(args, "any_slice"); !slices.Equal(got, []string{"c", "7"}) { - t.Errorf("stringSliceArg(any_slice) = %v, want [c 7]", got) - } - }) -} - -type fakeHTTPDoer func(*http.Request) (*http.Response, error) - -func (f fakeHTTPDoer) Do(req *http.Request) (*http.Response, error) { - return f(req) -} - -type trackingReadCloser struct { - io.Reader - closed bool -} - -func (r *trackingReadCloser) Close() error { - r.closed = true - return nil -} - -type decodedClientEvent struct { - seqNum int64 - hasSeqNum bool - timestampMicros int64 - hasTimestamp bool - hasPayload bool - hasStep bool - hasToolResult bool - callID string - hasFnResult bool - fnName string - isError bool - result any -} - -func decodeClientEventWire(t *testing.T, data []byte) decodedClientEvent { - t.Helper() - var out decodedClientEvent - for len(data) > 0 { - num, wt, vv, vb, n, err := consumeField(data) - if err != nil { - t.Fatalf("consumeField(clientEvent) failed: %v", err) - } - data = data[n:] - switch { - case num == 3 && wt == wireBytes: - out.hasPayload = true - pData := vb - for len(pData) > 0 { - pnum, pwt, _, pvb, pn, perr := consumeField(pData) - if perr != nil { - t.Fatalf("consumeField(wirePayload) failed: %v", perr) - } - pData = pData[pn:] - if pnum == 11 && pwt == wireBytes { - out.hasStep = true - sData := pvb - for len(sData) > 0 { - snum, swt, _, svb, sn, serr := consumeField(sData) - if serr != nil { - t.Fatalf("consumeField(wireStep) failed: %v", serr) - } - sData = sData[sn:] - if snum == 5 && swt == wireBytes { - out.hasToolResult = true - trData := svb - for len(trData) > 0 { - tnum, twt, _, tvb, tn, terr := consumeField(trData) - if terr != nil { - t.Fatalf("consumeField(toolResultStep) failed: %v", terr) - } - trData = trData[tn:] - switch { - case tnum == 1 && twt == wireBytes: - out.callID = string(tvb) - case tnum == 3 && twt == wireBytes: - out.hasFnResult = true - frData := tvb - for len(frData) > 0 { - fnum, fwt, fvv, fvb, fn, ferr := consumeField(frData) - if ferr != nil { - t.Fatalf("consumeField(functionResultStep) failed: %v", ferr) - } - frData = frData[fn:] - switch { - case fnum == 2 && fwt == wireBytes: - out.fnName = string(fvb) - case fnum == 3 && fwt == wireVarint: - out.isError = fvv != 0 - case fnum == 7 && fwt == wireBytes: - v, derr := decodeProtoValue(fvb, 0) - if derr != nil { - t.Fatalf("decodeProtoValue(Result) failed: %v", derr) - } - out.result = v - } - } - } - } - } - } - } - } - case num == 10 && wt == wireVarint: - out.seqNum = int64(vv) - out.hasSeqNum = true - case num == 11 && wt == wireVarint: - out.timestampMicros = int64(vv) - out.hasTimestamp = true - } - } - return out -} - -func TestRuntimeClientBranchingHeadersAndStatusBoundaries(t *testing.T) { - t.Setenv("GEMINI_API_KEY", "") - t.Setenv("GEMINI_ACCESS_TOKEN", "") - t.Setenv("GEMINI_NO_KEYRING", "true") - t.Setenv("HOME", t.TempDir()) - t.Setenv("USERPROFILE", t.TempDir()) - config.Reset() - t.Cleanup(config.Reset) - - t.Run("default_server_url_stream_client_timeout_and_credential_completeness", func(t *testing.T) { - t.Setenv("GEMINI_TIMEOUT", "7s") - config.Reset() - if err := config.Init("gemini-api-cli", "GEMINI"); err != nil { - t.Fatalf("config.Init() unexpected error: %v", err) - } - - // Default --server-url falls back to sdk.ServerList[0] ("https://generativelanguage.googleapis.com") - // and --debug=true without --dry-run keeps streamClient distinct from DebugClient-wrapped httpClient. - defaultCmd := newTestRuntimeCmd(t, "") - _ = defaultCmd.Root().PersistentFlags().Set("debug", "true") - rtDefault, err := newRuntimeTransport(defaultCmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport(default serverURL) unexpected error: %v", err) - } - if rtDefault.baseURL != "https://generativelanguage.googleapis.com" || rtDefault.scheme != "https" { - t.Errorf("default baseURL/scheme = (%q, %q), want (https://generativelanguage.googleapis.com, https)", rtDefault.baseURL, rtDefault.scheme) - } - if rtDefault.timeout != 7*time.Second { - t.Errorf("rtDefault.timeout from GEMINI_TIMEOUT = %v, want 7s", rtDefault.timeout) - } - if rtDefault.streamClient == nil || rtDefault.streamClient == rtDefault.httpClient { - t.Errorf("rtDefault.streamClient = %v, want non-nil unwrapped client distinct from debug httpClient", rtDefault.streamClient) - } - - // --dry-run=true uses httpClient as streamClient so dry-run preview intercepts stream requests. - dryCmd := newTestRuntimeCmd(t, "") - _ = dryCmd.Root().PersistentFlags().Set("dry-run", "true") - _ = dryCmd.Root().PersistentFlags().Set("timeout", "3s") - // Explicit empty --api-key "" must be treated as incomplete so --access-token wins. - _ = dryCmd.Root().PersistentFlags().Set("api-key", "") - _ = dryCmd.Root().PersistentFlags().Set("access-token", "tok-complete") - rtDry, err := newRuntimeTransport(dryCmd, false) - if err != nil { - t.Fatalf("newRuntimeTransport(dry-run) unexpected error: %v", err) - } - if rtDry.streamClient == nil || rtDry.streamClient != rtDry.httpClient { - t.Errorf("rtDry.streamClient = %v, want equal to rtDry.httpClient (%v)", rtDry.streamClient, rtDry.httpClient) - } - if rtDry.timeout != 3*time.Second { - t.Errorf("rtDry.timeout with explicit --timeout=3s = %v, want 3s", rtDry.timeout) - } - req, err := rtDry.newRequest(t.Context(), http.MethodPost, "/v1beta/check", http.NoBody) - if err != nil { - t.Fatalf("rtDry.newRequest() unexpected error: %v", err) - } - if gotUA := req.Header.Get("User-Agent"); gotUA != rawUserAgent || gotUA == "" { - t.Errorf("User-Agent = %q, want non-empty %q", gotUA, rawUserAgent) - } - if gotAuth := req.Header.Get("Authorization"); gotAuth != "Bearer tok-complete" { - t.Errorf("Authorization with empty --api-key = %q, want \"Bearer tok-complete\"", gotAuth) - } - }) - - t.Run("lease_and_heartbeat_headers_empty_body_and_status_boundaries", func(t *testing.T) { - var gotHeaders http.Header - var statusCode atomic.Int32 - var respPayload atomic.Value - statusCode.Store(http.StatusOK) - respPayload.Store("") - - srvURL := startH2CTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - gotHeaders = r.Header.Clone() - w.WriteHeader(int(statusCode.Load())) - if body, _ := respPayload.Load().(string); body != "" { - _, _ = w.Write([]byte(body)) - } - })) - - cmd := newTestRuntimeCmd(t, srvURL) - rt, err := newRuntimeTransport(cmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport() unexpected error: %v", err) - } - - // Empty EnvironmentID and EnvironmentKey must NOT emit blank header keys. - bareClient := NewHTTPClient(HTTPClientOptions{Transport: rt}) - leasesEmpty, err := bareClient.LeaseInteractions(t.Context(), "w1", 1) - if err != nil || leasesEmpty != nil { - t.Errorf("LeaseInteractions(200 empty body) = (%v, %v), want (nil, nil)", leasesEmpty, err) - } - if gotCT := gotHeaders.Get("Content-Type"); gotCT != "application/json" { - t.Errorf("LeaseInteractions Content-Type = %q, want \"application/json\"", gotCT) - } - for _, absentHdr := range []string{"X-Goog-Environment-Key", "X-Goog-Environment-Id", "X-Goog-Lease-Token"} { - if vals := gotHeaders.Values(absentHdr); len(vals) != 0 { - t.Errorf("header %s = %v when unset, want absent", absentHdr, vals) - } - } - - // Populated EnvironmentID and EnvironmentKey must set exact headers. - scopedClient := NewHTTPClient(HTTPClientOptions{ - Transport: rt, - EnvironmentID: "env_hdr", - EnvironmentKey: "ek_hdr", - }) - respPayload.Store(`{"leases":[{"interaction_id":"i299","lease_token":"t299"}]}`) - statusCode.Store(299) - leases299, err := scopedClient.LeaseInteractions(t.Context(), "w1", 1) - if err != nil || len(leases299) != 1 || leases299[0].InteractionID != "i299" { - t.Errorf("LeaseInteractions(status=299) = (%v, %v), want 1 lease i299", leases299, err) - } - if gotHeaders.Get("X-Goog-Environment-Key") != "ek_hdr" || gotHeaders.Get("X-Goog-Environment-Id") != "env_hdr" { - t.Errorf("scoped headers = %v, want x-goog-environment-key=ek_hdr and x-goog-environment-id=env_hdr", gotHeaders) - } - - statusCode.Store(300) - if got, err := scopedClient.LeaseInteractions(t.Context(), "w1", 1); err == nil { - t.Errorf("LeaseInteractions(status=300) = %v, want non-2xx error", got) - } - rt199 := &runtimeTransport{ - baseURL: "https://example.googleapis.com", - scheme: "https", - apiVersion: "v1beta", - headers: make(map[string]string), - httpClient: fakeHTTPDoer(func(*http.Request) (*http.Response, error) { - return &http.Response{StatusCode: 199, Body: io.NopCloser(strings.NewReader("info"))}, nil - }), - } - if got, err := NewHTTPClient(HTTPClientOptions{Transport: rt199}).LeaseInteractions(t.Context(), "w1", 1); err == nil { - t.Errorf("LeaseInteractions(status=199) = %v, want non-2xx error", got) - } - - // HeartbeatInteraction handles nil, non-matching, and empty-body (parsed == nil) responses cleanly. - statusCode.Store(http.StatusOK) - respPayload.Store(`{"leases":[null,{"interaction_id":"other_id"},{"interaction_id":"target_id"}]}`) - hbMatch, err := scopedClient.HeartbeatInteraction(t.Context(), "p1", "target_id", "", "lt_1") - if err != nil || !hbMatch.Renewed || hbMatch.ShouldCancel { - t.Errorf("HeartbeatInteraction(mixed leases with match) = (%+v, %v), want Renewed=true", hbMatch, err) - } - respPayload.Store(`{"leases":[null,{"interaction_id":"other_id"}]}`) - hbMiss, err := scopedClient.HeartbeatInteraction(t.Context(), "p1", "target_id", "", "lt_1") - if err != nil || hbMiss.Renewed || !hbMiss.ShouldCancel { - t.Errorf("HeartbeatInteraction(mixed leases without match) = (%+v, %v), want ShouldCancel=true", hbMiss, err) - } - respPayload.Store("") - hbEmpty, err := scopedClient.HeartbeatInteraction(t.Context(), "p1", "target_id", "", "lt_1") - if err != nil || hbEmpty.Renewed || !hbEmpty.ShouldCancel { - t.Errorf("HeartbeatInteraction(empty 200 body) = (%+v, %v), want ShouldCancel=true", hbEmpty, err) - } - - // Explicit --access-token flag (rank 0) outranks GEMINI_API_KEY env var (rank 1). - t.Setenv("GEMINI_API_KEY", "env-api-key-rank1") - rankCmd := newTestRuntimeCmd(t, srvURL) - _ = rankCmd.Root().PersistentFlags().Set("access-token", "flag-oauth-rank0") - rankRT, err := newRuntimeTransport(rankCmd, false) - if err != nil { - t.Fatalf("newRuntimeTransport(rank precedence) failed: %v", err) - } - rankReq, err := rankRT.newRequest(t.Context(), http.MethodPost, "/v1beta/test", http.NoBody) - if err != nil { - t.Fatalf("rankRT.newRequest() failed: %v", err) - } - if gotKey, gotAuth := rankReq.Header.Get("x-goog-api-key"), rankReq.Header.Get("Authorization"); gotKey != "" || gotAuth != "Bearer flag-oauth-rank0" { - t.Errorf("headers with GEMINI_API_KEY + --access-token = (api-key=%q, auth=%q), want (\"\", \"Bearer flag-oauth-rank0\")", gotKey, gotAuth) - } - }) - - t.Run("preview_connect_and_stream_close_headers_and_nil_safety", func(t *testing.T) { - var capturedReq *http.Request - bodyTracker := &trackingReadCloser{Reader: strings.NewReader("")} - rt := &runtimeTransport{ - baseURL: "https://example.googleapis.com", - scheme: "https", - apiVersion: "v1beta", - headers: make(map[string]string), - httpClient: fakeHTTPDoer(func(req *http.Request) (*http.Response, error) { - capturedReq = req - return &http.Response{StatusCode: http.StatusOK, Body: bodyTracker}, nil - }), - } - c := NewHTTPClient(HTTPClientOptions{ - Transport: rt, - EnvironmentID: "env_prev", - EnvironmentKey: "ek_prev", - }) - if err := c.PreviewConnect(t.Context(), "lt_prev"); err != nil { - t.Fatalf("PreviewConnect() unexpected error: %v", err) - } - if !bodyTracker.closed { - t.Errorf("PreviewConnect() did not close response body") - } - if capturedReq.Header.Get("Content-Type") != "application/grpc" || - capturedReq.Header.Get("TE") != "trailers" || - capturedReq.Header.Get("X-Goog-Environment-Id") != "env_prev" || - capturedReq.Header.Get("X-Goog-Environment-Key") != "ek_prev" || - capturedReq.Header.Get("X-Goog-Lease-Token") != "lt_prev" { - t.Errorf("PreviewConnect headers = %v, want grpc/trailers and environment+lease headers", capturedReq.Header) - } - - // Nil response and nil body branches in PreviewConnect. - rt.httpClient = fakeHTTPDoer(func(*http.Request) (*http.Response, error) { - return &http.Response{StatusCode: http.StatusOK, Body: nil}, nil - }) - if err := c.PreviewConnect(t.Context(), "lt_prev"); err != nil { - t.Errorf("PreviewConnect(Body=nil) = %v, want nil", err) - } - rt.httpClient = fakeHTTPDoer(func(*http.Request) (*http.Response, error) { - return nil, nil - }) - if err := c.PreviewConnect(t.Context(), "lt_prev"); err != nil { - t.Errorf("PreviewConnect(resp=nil) = %v, want nil", err) - } - - // grpcInteractionStream.Close closes non-nil Body and tolerates nil Body / nil resp. - _, pw1 := io.Pipe() - streamTracker := &trackingReadCloser{Reader: strings.NewReader("")} - s1 := &grpcInteractionStream{resp: &http.Response{Body: streamTracker}, pw: pw1} - if err := s1.Close(); err != nil || !streamTracker.closed { - t.Errorf("stream.Close(non-nil Body) = (%v, closed=%v), want (nil, true)", err, streamTracker.closed) - } - _, pw2 := io.Pipe() - if err := (&grpcInteractionStream{resp: &http.Response{Body: nil}, pw: pw2}).Close(); err != nil { - t.Errorf("stream.Close(Body=nil) = %v, want nil", err) - } - _, pw3 := io.Pipe() - if err := (&grpcInteractionStream{resp: nil, pw: pw3}).Close(); err != nil { - t.Errorf("stream.Close(resp=nil) = %v, want nil", err) - } - }) - - t.Run("connect_stream_headers_debug_status_bounds_and_probe_conditions", func(t *testing.T) { - var streamHeaders http.Header - var streamStatus atomic.Int32 - streamStatus.Store(http.StatusOK) - - srvURL := startH2CTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - streamHeaders = r.Header.Clone() - code := int(streamStatus.Load()) - if code < 200 || code >= 300 { - w.WriteHeader(code) - _, _ = w.Write([]byte("stream rejected")) - return - } - w.Header().Set("Content-Type", "application/grpc") - w.WriteHeader(code) - flusher, _ := w.(http.Flusher) - flusher.Flush() - _, _ = grpcReadFrame(r.Body) - })) - - // 1. --debug=true writes both stream request and stream response debug banners. - dbgCmd := newTestRuntimeCmd(t, srvURL) - var dbgErr bytes.Buffer - dbgCmd.SetErr(&dbgErr) - _ = dbgCmd.Root().PersistentFlags().Set("debug", "true") - rtDbg, err := newRuntimeTransport(dbgCmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport(debug) failed: %v", err) - } - dbgClient := NewHTTPClient(HTTPClientOptions{Transport: rtDbg, EnvironmentID: "env_s", EnvironmentKey: "ek_s"}) - st, err := dbgClient.connectStream(t.Context(), connectStreamRequest{InteractionID: "i1", LeaseToken: "lt_s", NextOutputSeqNum: 1}) - if err != nil { - t.Fatalf("connectStream(200) unexpected error: %v", err) - } - _ = st.Close() - if streamHeaders.Get("Content-Type") != "application/grpc" || streamHeaders.Get("TE") != "trailers" { - t.Errorf("connectStream headers = %v, want Content-Type=application/grpc and TE=trailers", streamHeaders) - } - if !strings.Contains(dbgErr.String(), "=== DEBUG STREAM ===") || !strings.Contains(dbgErr.String(), "=== DEBUG STREAM RESPONSE ===") { - t.Errorf("debug stderr = %q, want both DEBUG STREAM banners", dbgErr.String()) - } - - // 2. --debug=false writes neither banner, and debug=true with errOut=nil does not panic. - quietCmd := newTestRuntimeCmd(t, srvURL) - var quietErr bytes.Buffer - quietCmd.SetErr(&quietErr) - rtQuiet, err := newRuntimeTransport(quietCmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport(quiet) failed: %v", err) - } - rtQuiet.errOut = &quietErr - quietClient := NewHTTPClient(HTTPClientOptions{Transport: rtQuiet}) - streamStatus.Store(299) - st299, err := quietClient.connectStream(t.Context(), connectStreamRequest{InteractionID: "i299", LeaseToken: "lt_299", NextOutputSeqNum: 1}) - if err != nil { - t.Fatalf("connectStream(299) unexpected error: %v", err) - } - _ = st299.Close() - if strings.Contains(quietErr.String(), "DEBUG STREAM") { - t.Errorf("connectStream(debug=false) wrote debug output: %q", quietErr.String()) - } - rtNilOut := *rtQuiet - rtNilOut.debug = true - rtNilOut.errOut = nil - stNilOut, err := NewHTTPClient(HTTPClientOptions{Transport: &rtNilOut}).connectStream(t.Context(), connectStreamRequest{InteractionID: "inil", LeaseToken: "lt_nil", NextOutputSeqNum: 1}) - if err != nil { - t.Fatalf("connectStream(debug=true, errOut=nil) unexpected error: %v", err) - } - _ = stNilOut.Close() - - // 3. Status boundary 199 and 300 fail connectStream. - streamStatus.Store(300) - if _, err := quietClient.connectStream(t.Context(), connectStreamRequest{InteractionID: "ibad", LeaseToken: "lt_bad", NextOutputSeqNum: 1}); err == nil { - t.Errorf("connectStream(status=300) = nil error, want non-2xx error") - } - rtStream199 := &runtimeTransport{ - baseURL: "https://example.googleapis.com", - scheme: "https", - apiVersion: "v1beta", - headers: make(map[string]string), - streamClient: fakeHTTPDoer(func(*http.Request) (*http.Response, error) { - return &http.Response{StatusCode: 199, Body: io.NopCloser(strings.NewReader("info"))}, nil - }), - } - if _, err := NewHTTPClient(HTTPClientOptions{Transport: rtStream199}).connectStream(t.Context(), connectStreamRequest{InteractionID: "i199", LeaseToken: "lt_199", NextOutputSeqNum: 1}); err == nil { - t.Errorf("connectStream(status=199) = nil error, want non-2xx error") - } - - // 4. probePlainHTTPError headers, canceled-context skip, and https-scheme skip. - var probeHeaders http.Header - var http1Hits atomic.Int32 - http1Srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - http1Hits.Add(1) - probeHeaders = r.Header.Clone() - http.Error(w, "bad gateway", http.StatusBadGateway) - })) - defer http1Srv.Close() - - probeCmd := newTestRuntimeCmd(t, http1Srv.URL) - rtProbe, err := newRuntimeTransport(probeCmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport(http1) failed: %v", err) - } - probeClient := NewHTTPClient(HTTPClientOptions{ - Transport: rtProbe, - EnvironmentID: "env_probe", - EnvironmentKey: "ek_probe", - }) - if _, err := probeClient.connectStream(t.Context(), connectStreamRequest{InteractionID: "iprobe", LeaseToken: "lt_probe", NextOutputSeqNum: 1}); err == nil || !strings.Contains(err.Error(), "Status 502") { - t.Errorf("connectStream(HTTP/1.1 502) = %v, want Status 502", err) - } - if probeHeaders.Get("Content-Type") != "application/grpc" || - probeHeaders.Get("X-Goog-Environment-Id") != "env_probe" || - probeHeaders.Get("X-Goog-Environment-Key") != "ek_probe" || - probeHeaders.Get("X-Goog-Lease-Token") != "lt_probe" { - t.Errorf("probePlainHTTPError headers = %v, want grpc + environment/lease headers", probeHeaders) - } - - // When ctx is already canceled, connectStream on http:// must NOT fire probePlainHTTPError. - var canceledProbeHits atomic.Int32 - rtHTTPCanceled := &runtimeTransport{ - baseURL: "http://127.0.0.1:8080", - scheme: "http", - apiVersion: "v1beta", - headers: make(map[string]string), - httpClient: fakeHTTPDoer(func(*http.Request) (*http.Response, error) { - canceledProbeHits.Add(1) - return &http.Response{StatusCode: http.StatusForbidden, Body: io.NopCloser(strings.NewReader("probe hit"))}, nil - }), - streamClient: fakeHTTPDoer(func(*http.Request) (*http.Response, error) { - return nil, context.Canceled - }), - } - canceledCtx, cancelNow := context.WithCancel(t.Context()) - cancelNow() - _, cancelErr := NewHTTPClient(HTTPClientOptions{Transport: rtHTTPCanceled}).connectStream(canceledCtx, connectStreamRequest{InteractionID: "icancel", LeaseToken: "lt_c", NextOutputSeqNum: 1}) - if !errors.Is(cancelErr, context.Canceled) || canceledProbeHits.Load() != 0 { - t.Errorf("connectStream(canceled ctx) = (%v, probeHits=%d), want (context.Canceled, 0)", cancelErr, canceledProbeHits.Load()) - } - - // When scheme is https, streamClient failure must NOT call httpClient probe. - var httpsProbeHits atomic.Int32 - rtHTTPSFail := &runtimeTransport{ - baseURL: "https://example.googleapis.com", - scheme: "https", - apiVersion: "v1beta", - headers: make(map[string]string), - httpClient: fakeHTTPDoer(func(*http.Request) (*http.Response, error) { - httpsProbeHits.Add(1) - return &http.Response{StatusCode: http.StatusForbidden, Body: io.NopCloser(strings.NewReader("nope"))}, nil - }), - streamClient: fakeHTTPDoer(func(*http.Request) (*http.Response, error) { - return nil, errors.New("tls dial failed") - }), - } - _, _ = NewHTTPClient(HTTPClientOptions{Transport: rtHTTPSFail}).connectStream(t.Context(), connectStreamRequest{InteractionID: "ihttps", LeaseToken: "lt_h", NextOutputSeqNum: 1}) - if got := httpsProbeHits.Load(); got != 0 { - t.Errorf("https stream failure called probePlainHTTPError %d times, want 0", got) - } - }) - - t.Run("retry_backoff_shift_and_cap_progression", func(t *testing.T) { + t.Run("validation_errors", func(t *testing.T) { + sbxDir := t.TempDir() + outsideDir := t.TempDir() cases := []struct { - base time.Duration - consecutive int - want time.Duration + name string + args []string + wantSub string }{ - {base: 100 * time.Millisecond, consecutive: -1, want: 100 * time.Millisecond}, - {base: 100 * time.Millisecond, consecutive: 0, want: 100 * time.Millisecond}, - {base: 100 * time.Millisecond, consecutive: 1, want: 100 * time.Millisecond}, - {base: 100 * time.Millisecond, consecutive: 2, want: 200 * time.Millisecond}, - {base: 100 * time.Millisecond, consecutive: 3, want: 400 * time.Millisecond}, - {base: 100 * time.Millisecond, consecutive: 6, want: 3200 * time.Millisecond}, - {base: 100 * time.Millisecond, consecutive: 7, want: 3200 * time.Millisecond}, - {base: 2 * time.Second, consecutive: 6, want: 30 * time.Second}, - } - for _, tc := range cases { - if got := retryBackoff(tc.base, tc.consecutive); got != tc.want { - t.Errorf("retryBackoff(%v, %d) = %v, want %v", tc.base, tc.consecutive, got, tc.want) - } - } - }) -} - -func TestPollerLoopRetryBackoffDeduplicationAndRecovery(t *testing.T) { - t.Run("empty_lease_token_stderr_not_corrupted_by_redactor", func(t *testing.T) { - prov := NewCommandWorkerProvisioner("printf 'plain boom' >&2; exit 2") - err := prov.ProvisionAndDispatch(t.Context(), ProvisionInput{ - InteractionID: "int_empty_tok", - EnvironmentID: "env_1", - LeaseToken: "", - }) - if err == nil || !strings.Contains(err.Error(), "plain boom") || strings.Contains(err.Error(), "[REDACTED]") { - t.Errorf("ProvisionAndDispatch(empty LeaseToken) = %v, want 'plain boom' without '[REDACTED]'", err) - } - }) - - t.Run("non_retryable_401_aborts_even_when_retries_enabled", func(t *testing.T) { - var calls atomic.Int32 - srvURL := startH2CTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - calls.Add(1) - http.Error(w, `{"error":{"code":401,"message":"unauthorized"}}`, http.StatusUnauthorized) - })) - cmd := newTestRuntimeCmd(t, srvURL) - rt, err := newRuntimeTransport(cmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport() failed: %v", err) - } - poller := NewPollerController(NewHTTPClient(HTTPClientOptions{Transport: rt, EnvironmentKey: "ek_bad"}), NewCommandWorkerProvisioner("true"), PollerOptions{ - PollerID: "p-401", - PollInterval: 20 * time.Millisecond, - NoRetries: false, - Logger: log.New(io.Discard, "", 0), - }) - if err := poller.Run(t.Context()); err == nil || !strings.Contains(err.Error(), "Status 401") { - t.Errorf("PollerController.Run(401) = %v, want Status 401", err) - } - if got := calls.Load(); got != 1 { - t.Errorf("401 call count = %d, want 1 (no retries on 401)", got) - } - }) - - t.Run("retry_backoff_reset_single_start_log_and_idle_sleep", func(t *testing.T) { - ctx, cancel := context.WithCancel(t.Context()) - defer cancel() - - var step atomic.Int32 - srvURL := startH2CTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - switch step.Add(1) { - case 1, 2: - http.Error(w, "unavailable", http.StatusServiceUnavailable) - case 3, 4: - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"leases":[]}`)) - case 5: - http.Error(w, "unavailable again", http.StatusServiceUnavailable) - default: - cancel() - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"leases":[]}`)) - } - })) - - cmd := newTestRuntimeCmd(t, srvURL) - rt, err := newRuntimeTransport(cmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport() failed: %v", err) - } - var logBuf bytes.Buffer - poller := NewPollerController(NewHTTPClient(HTTPClientOptions{Transport: rt, EnvironmentKey: "ek_ok"}), NewCommandWorkerProvisioner("true"), PollerOptions{ - PollerID: "p-backoff", - MaxInteractions: 2, - PollInterval: 20 * time.Millisecond, - NoRetries: false, - Logger: log.New(&logBuf, "", 0), - }) - _ = poller.Run(ctx) - - logs := logBuf.String() - if gotStarts := strings.Count(logs, "PollerController started"); gotStarts != 1 { - t.Errorf("start log count = %d, want exactly 1 in logs:\n%s", gotStarts, logs) - } - if got20 := strings.Count(logs, "(retrying in 20ms)"); got20 != 2 { - t.Errorf("'(retrying in 20ms)' count = %d, want 2 (initial error + reset after success) in logs:\n%s", got20, logs) - } - if got40 := strings.Count(logs, "(retrying in 40ms)"); got40 != 1 { - t.Errorf("'(retrying in 40ms)' count = %d, want 1 (2nd consecutive error) in logs:\n%s", got40, logs) - } - - // Verify empty lease responses sleep PollInterval instead of busy-spinning. - var idlePolls atomic.Int32 - idleSrv := startH2CTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - idlePolls.Add(1) - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"leases":[]}`)) - })) - idleCmd := newTestRuntimeCmd(t, idleSrv) - idleRT, err := newRuntimeTransport(idleCmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport(idle) failed: %v", err) - } - idleCtx, idleCancel := context.WithTimeout(t.Context(), 45*time.Millisecond) - defer idleCancel() - idlePoller := NewPollerController(NewHTTPClient(HTTPClientOptions{Transport: idleRT, EnvironmentKey: "ek_idle"}), NewCommandWorkerProvisioner("true"), PollerOptions{ - PollerID: "p-idle", - PollInterval: 120 * time.Millisecond, - Logger: log.New(io.Discard, "", 0), - }) - _ = idlePoller.Run(idleCtx) - if got := idlePolls.Load(); got != 1 { - t.Errorf("idle poll count in 45ms with 120ms PollInterval = %d, want 1", got) - } - }) - - t.Run("in_flight_deduplication_and_failure_state_recovery", func(t *testing.T) { - ctx, cancel := context.WithCancel(t.Context()) - defer cancel() - - markerFile := filepath.Join(t.TempDir(), "attempt_marker") - countFile := filepath.Join(t.TempDir(), "dup_count") - round2File := filepath.Join(t.TempDir(), "round2_seen") - t.Setenv("TEST_MARKER_FILE", markerFile) - t.Setenv("TEST_DUP_COUNT_FILE", countFile) - t.Setenv("TEST_ROUND2_FILE", round2File) - - var poller *PollerController - var mu sync.Mutex - sawFailureRecorded := false - recoverySent := false - var pollRound atomic.Int32 - var round1MaxNew, round2MaxNew atomic.Int32 - - srvURL := startH2CTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - var req leaseInteractionsRequest - _ = json.NewDecoder(r.Body).Decode(&req) - w.Header().Set("Content-Type", "application/json") - if len(req.ActiveLeases) > 0 { - _ = json.NewEncoder(w).Encode(leaseInteractionsResponse{ - Leases: []*LeasedInteraction{{InteractionID: req.ActiveLeases[0].InteractionID, LeaseToken: req.ActiveLeases[0].LeaseToken}}, - }) - return - } - switch pollRound.Add(1) { - case 1: - if req.MaxNewInteractions != nil { - round1MaxNew.Store(int32(*req.MaxNewInteractions)) - } - // Duplicate dup_1 entries in the same response + nil/empty entries: must dispatch dup_1 only ONCE! - _ = json.NewEncoder(w).Encode(leaseInteractionsResponse{ - Leases: []*LeasedInteraction{ - nil, - {InteractionID: ""}, - {InteractionID: "dup_1", EnvironmentID: "env_d", LeaseToken: "tok_d"}, - {InteractionID: "dup_1", EnvironmentID: "env_d", LeaseToken: "tok_d"}, - }, - }) - case 2: - if req.MaxNewInteractions != nil { - round2MaxNew.Store(int32(*req.MaxNewInteractions)) - } - _ = os.WriteFile(round2File, []byte("ok"), 0o600) - // Return int_flaky which fails on its first attempt. - _ = json.NewEncoder(w).Encode(leaseInteractionsResponse{ - Leases: []*LeasedInteraction{{InteractionID: "int_flaky", EnvironmentID: "env_f", LeaseToken: "tok_f"}}, - }) - default: - poller.mu.Lock() - st := poller.failures["int_flaky"] - if st != nil { - mu.Lock() - sawFailureRecorded = true - mu.Unlock() - // Fast-forward cooldown so recovery attempt runs without waiting 500ms. - st.cooldownEnd = time.Now() - } - poller.mu.Unlock() - mu.Lock() - sendRecovery := sawFailureRecorded && !recoverySent - if sendRecovery { - recoverySent = true - } - mu.Unlock() - if sendRecovery { - _ = json.NewEncoder(w).Encode(leaseInteractionsResponse{ - Leases: []*LeasedInteraction{{InteractionID: "int_flaky", EnvironmentID: "env_f", LeaseToken: "tok_f"}}, - }) - return - } - _ = json.NewEncoder(w).Encode(leaseInteractionsResponse{}) - } - })) - - // Script: - // - For dup_1: immediately append "1" to TEST_DUP_COUNT_FILE and stay in-flight until round 2 arrives. - // - For int_flaky: fail on first invocation (creating TEST_MARKER_FILE), succeed on second invocation (creating TEST_MARKER_FILE.ok). - provScript := `input=$(cat); case "$input" in *dup_1*) printf '1' >> "$TEST_DUP_COUNT_FILE"; i=0; while [ ! -f "$TEST_ROUND2_FILE" ] && [ "$i" -lt 100 ]; do sleep 0.01; i=$((i+1)); done ;; *int_flaky*) if [ ! -f "$TEST_MARKER_FILE" ]; then touch "$TEST_MARKER_FILE"; exit 1; fi; touch "$TEST_MARKER_FILE.ok" ;; esac` - - cmd := newTestRuntimeCmd(t, srvURL) - rt, err := newRuntimeTransport(cmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport() failed: %v", err) - } - var logBuf bytes.Buffer - poller = NewPollerController(NewHTTPClient(HTTPClientOptions{Transport: rt, EnvironmentKey: "ek_1"}), NewCommandWorkerProvisioner(provScript), PollerOptions{ - PollerID: "p-dedup", - MaxInteractions: 4, - PollInterval: 20 * time.Millisecond, - Logger: log.New(&logBuf, "", 0), - }) - - done := make(chan error, 1) - go func() { - done <- poller.Run(ctx) - }() - - deadline := time.Now().Add(3 * time.Second) - for time.Now().Before(deadline) { - _, okErr := os.Stat(markerFile + ".ok") - _, dupErr := os.Stat(countFile) - if okErr == nil && dupErr == nil { - break - } - time.Sleep(15 * time.Millisecond) - } - // Give dispatchLease's post-provision cleanup a brief moment to clear failures and log success. - deadline2 := time.Now().Add(time.Second) - for time.Now().Before(deadline2) { - poller.mu.Lock() - _, stillFailed := poller.failures["int_flaky"] - poller.mu.Unlock() - if !stillFailed && strings.Contains(logBuf.String(), "[int_flaky] Dispatched interaction lease") { - break - } - time.Sleep(10 * time.Millisecond) - } - cancel() - <-done - - dupBytes, err := os.ReadFile(countFile) - if err != nil || string(dupBytes) != "1" { - t.Errorf("dup_1 execution count bytes = (%q, %v), want \"1\" (deduplicated)", dupBytes, err) - } - if r1, r2 := round1MaxNew.Load(), round2MaxNew.Load(); r1 != 4 || r2 != 3 { - t.Errorf("max_new_interactions across polls = (round1=%d, round2=%d), want (4, 3) with 1 in-flight slot", r1, r2) - } - mu.Lock() - gotSawFailure := sawFailureRecorded - mu.Unlock() - if !gotSawFailure { - t.Errorf("poller.failures[int_flaky] was never recorded on failed dispatch") - } - poller.mu.Lock() - _, residualFailure := poller.failures["int_flaky"] - poller.mu.Unlock() - if residualFailure { - t.Errorf("poller.failures[int_flaky] was not cleared after successful recovery dispatch") - } - if !strings.Contains(logBuf.String(), "[int_flaky] ProvisionAndDispatch failed") || !strings.Contains(logBuf.String(), "[int_flaky] Dispatched interaction lease") { - t.Errorf("poller logs missing failure or recovery log lines:\n%s", logBuf.String()) - } - }) -} - -func TestSandboxSymlinkBoundariesAndShellExecution(t *testing.T) { - t.Run("whitespace_root_symlink_sibling_escape_and_symlink_workdir", func(t *testing.T) { - evalBase, err := filepath.EvalSymlinks(t.TempDir()) - if err != nil { - t.Fatalf("EvalSymlinks(TempDir) failed: %v", err) - } - d1d2 := filepath.Join(evalBase, "d1", "d2") - siblingD2 := filepath.Join(evalBase, "d2") - subInside := filepath.Join(d1d2, "real_work") - for _, dir := range []string{d1d2, siblingD2, subInside} { - if err := os.MkdirAll(dir, 0o755); err != nil { - t.Fatalf("MkdirAll(%q) failed: %v", dir, err) - } - } - symRoot := filepath.Join(evalBase, "sym_d1d2") - if err := os.Symlink(d1d2, symRoot); err != nil { - t.Fatalf("Symlink(symRoot) failed: %v", err) - } - symWorkInside := filepath.Join(d1d2, "sym_work") - if err := os.Symlink(subInside, symWorkInside); err != nil { - t.Fatalf("Symlink(symWorkInside) failed: %v", err) - } - - // Sibling escape: filepath.Rel(evalBase/sym_d1d2, evalBase/d2) == "../d2", which if joined onto - // evalBase/d1/d2 without checking isWithinRoot(origAbsRoot, absWork) would falsely resolve to d1/d2! - if gotR, gotW, err := resolveSandboxDirs(symRoot, siblingD2, false); err == nil { - t.Errorf("resolveSandboxDirs(symRoot=%q, escaping sibling=%q) = (%q, %q, nil), want escape error", symRoot, siblingD2, gotR, gotW) - } - - // Symlinked workDir inside rootDir must canonicalize to its real target path. - gotRoot, gotWork, err := resolveSandboxDirs(d1d2, symWorkInside, false) - if err != nil || gotRoot != d1d2 || gotWork != subInside { - t.Errorf("resolveSandboxDirs(symlinked workDir) = (%q, %q, %v), want (%q, %q, nil)", gotRoot, gotWork, err, d1d2, subInside) - } - - // Whitespace-only rootDir (" ") and workDir (" ") normalize to cwd ("."). - t.Chdir(subInside) - wsRoot, wsWork, err := resolveSandboxDirs(" ", " ", false) - if err != nil || wsRoot != subInside || wsWork != subInside { - t.Errorf("resolveSandboxDirs(\" \", \" \") = (%q, %q, %v), want (%q, %q, nil)", wsRoot, wsWork, err, subInside, subInside) - } - - // newFSServer canonicalizes symlinked rootDir/workDir and treats empty/whitespace workDir as rootDir even when cwd is a child of rootDir. - fsSym := newFSServer(symRoot, filepath.Join(symRoot, "real_work")) - if fsSym.rootDir != d1d2 || fsSym.workDir != subInside { - t.Errorf("newFSServer(symRoot, symWork) = (%q, %q), want (%q, %q)", fsSym.rootDir, fsSym.workDir, d1d2, subInside) - } - for _, blankWork := range []string{"", " "} { - fsBlank := newFSServer(d1d2, blankWork) - if fsBlank.workDir != d1d2 { - t.Errorf("newFSServer(d1d2, %q).workDir = %q (with cwd=%q), want %q", blankWork, fsBlank.workDir, subInside, d1d2) - } - } - if gotEmptyRel, err := fsSym.ResolvePath(""); err != nil || gotEmptyRel != subInside { - t.Errorf("ResolvePath(\"\") = (%q, %v), want (%q, nil)", gotEmptyRel, err, subInside) - } - }) - - t.Run("exact_read_size_boundaries_and_list_entries_size", func(t *testing.T) { - rootDir, err := filepath.EvalSymlinks(t.TempDir()) - if err != nil { - t.Fatalf("EvalSymlinks(TempDir) failed: %v", err) - } - fs := newFSServer(rootDir, rootDir) - - // Exact maxReadFileBytes file must succeed without truncating the final byte. - exactPayload := bytes.Repeat([]byte("z"), maxReadFileBytes) - if err := os.WriteFile(filepath.Join(rootDir, "exact.bin"), exactPayload, 0o600); err != nil { - t.Fatalf("WriteFile(exact.bin) failed: %v", err) - } - gotExact, err := fs.ReadRegularFile("exact.bin") - if err != nil || len(gotExact) != maxReadFileBytes { - t.Errorf("ReadRegularFile(exact maxReadFileBytes) len=%d err=%v, want len=%d nil", len(gotExact), err, maxReadFileBytes) - } - - // Direct unit assertions on readLimitedAll for exact limit, +1 overflow, exact maxBytes+1 read cap, and read error. - if got, exceeded, err := readLimitedAll(strings.NewReader("1234"), 4); err != nil || exceeded || string(got) != "1234" { - t.Errorf("readLimitedAll(exact 4) = (%q, %v, %v), want (\"1234\", false, nil)", got, exceeded, err) - } - tr := &countingReader{r: strings.NewReader("1234567")} - if _, exceeded, err := readLimitedAll(tr, 4); err != nil || !exceeded || tr.bytesRead != 5 { - t.Errorf("readLimitedAll(7 with max=4) = (exceeded=%v, err=%v, bytesRead=%d), want (true, nil, 5)", exceeded, err, tr.bytesRead) - } - if _, _, err := readLimitedAll(errReader{err: errors.New("disk read failed")}, 4); err == nil { - t.Errorf("readLimitedAll(errReader) = nil, want error") - } - - // Exact ReadFileRange byte-budget boundary: - // 2 lines of length (maxReadFileBytes/2 - 1) each -> totalBytes == (maxReadFileBytes/2)*2 == maxReadFileBytes (succeeds). - // Adding 1 byte to line 2 -> totalBytes == maxReadFileBytes + 1 (fails). - exactHalfLine := strings.Repeat("b", maxReadFileBytes/2-1) - if err := os.WriteFile(filepath.Join(rootDir, "two_exact_halves.txt"), []byte(exactHalfLine+"\n"+exactHalfLine), 0o600); err != nil { - t.Fatalf("WriteFile(two_exact_halves.txt) failed: %v", err) - } - if gotRangeExact, err := fs.ReadFileRange("two_exact_halves.txt", 1, 2); err != nil || len(gotRangeExact) != maxReadFileBytes-1 { - t.Errorf("ReadFileRange(exact maxReadFileBytes budget) len=%d err=%v, want len=%d nil", len(gotRangeExact), err, maxReadFileBytes-1) - } - if err := os.WriteFile(filepath.Join(rootDir, "two_over_by_one.txt"), []byte(exactHalfLine+"\n"+exactHalfLine+"x"), 0o600); err != nil { - t.Fatalf("WriteFile(two_over_by_one.txt) failed: %v", err) - } - if _, err := fs.ReadFileRange("two_over_by_one.txt", 1, 2); err == nil { - t.Errorf("ReadFileRange(maxReadFileBytes+1 budget) = nil, want size error") - } - - // Calling ReadFileRange with startLine=0, endLine=0 normalizes startLine=1 and formats "line range [1, 0]". - halfLine := strings.Repeat("a", maxReadFileBytes/2) - if err := os.WriteFile(filepath.Join(rootDir, "two_halves.txt"), []byte(halfLine+"\n"+halfLine+"\n"), 0o600); err != nil { - t.Fatalf("WriteFile(two_halves.txt) failed: %v", err) - } - _, rangeErr := fs.ReadFileRange("two_halves.txt", 0, 0) - if rangeErr == nil || !strings.Contains(rangeErr.Error(), "line range [1, 0]") { - t.Errorf("ReadFileRange(two_halves.txt, 0, 0) = %v, want 'line range [1, 0]' size error", rangeErr) - } - - // Multi-line ReadFileRange stops at endLine > 0, and stops cleanly at io.EOF when endLine == 0. - if _, err := fs.WriteRegularFile("sub/five.txt", "1\n2\n3\n4\n5\n", true, false); err != nil { - t.Fatalf("WriteRegularFile(sub/five.txt) failed: %v", err) - } - gotTwo, err := fs.ReadFileRange("sub/five.txt", 1, 2) - if err != nil || gotTwo != "1\n2" { - t.Errorf("ReadFileRange(1, 2) = (%q, %v), want (\"1\\n2\", nil)", gotTwo, err) - } - gotTailEOF, err := fs.ReadFileRange("sub/five.txt", 4, 0) - if err != nil || gotTailEOF != "4\n5\n" { - t.Errorf("ReadFileRange(4, 0 to EOF) = (%q, %v), want (\"4\\n5\\n\", nil)", gotTailEOF, err) - } - - // ListEntries populates exact byte size for regular files. - entries, err := fs.ListEntries("sub") - if err != nil || len(entries) != 1 { - t.Fatalf("ListEntries(sub) = (%v, %v), want 1 entry", entries, err) - } - entryMap, _ := entries[0].(map[string]any) - if entryMap["name"] != "five.txt" || entryMap["is_dir"] != false || entryMap["size"] != 10 { - t.Errorf("ListEntries(sub)[0] = %v, want name=five.txt is_dir=false size=10", entryMap) - } - }) - - t.Run("capped_buffer_isolated_zero_remaining_and_partial_overflow", func(t *testing.T) { - // Case 1: Exact fill (no overflow), followed by write when remaining == 0. - cbZero := &cappedBuffer{maxBytes: 6} - _, _ = cbZero.Write([]byte("123456")) - if strings.Contains(cbZero.String(), "truncated") { - t.Errorf("exact-capacity cappedBuffer marked truncated prematurely: %q", cbZero.String()) - } - _, _ = cbZero.Write([]byte("!")) - if want := "123456\n...[output truncated at 4 MiB]..."; cbZero.String() != want { - t.Errorf("cbZero.String() = %q, want %q", cbZero.String(), want) - } - - // Case 2: Single partial overflow write (len(p) > remaining) with no subsequent write. - cbPartial := &cappedBuffer{maxBytes: 6} - _, _ = cbPartial.Write([]byte("123456789")) - if want := "123456\n...[output truncated at 4 MiB]..."; cbPartial.String() != want { - t.Errorf("cbPartial.String() = %q, want %q", cbPartial.String(), want) - } - }) - - t.Run("shell_server_workdir_stderr_and_process_group_kill", func(t *testing.T) { - outerDir, err := filepath.EvalSymlinks(t.TempDir()) - if err != nil { - t.Fatalf("EvalSymlinks(TempDir) failed: %v", err) - } - defaultWork := filepath.Join(outerDir, "default_work") - overrideWork := filepath.Join(outerDir, "override_work") - for _, d := range []string{defaultWork, overrideWork} { - if err := os.MkdirAll(d, 0o755); err != nil { - t.Fatalf("MkdirAll(%q) failed: %v", d, err) - } - } - t.Chdir(outerDir) - sh := newShellServer(defaultWork) - - outDefault, errOut, code, err := sh.ExecuteBash(t.Context(), "pwd && printf 'stderr-text' >&2", "", 5*time.Second) - if err != nil || code != 0 || strings.TrimSpace(outDefault) != defaultWork || errOut != "stderr-text" { - t.Errorf("ExecuteBash(default workDir) = (stdout=%q, stderr=%q, code=%d, err=%v), want (%q, \"stderr-text\", 0, nil)", outDefault, errOut, code, err, defaultWork) - } - - outOverride, _, code, err := sh.ExecuteBash(t.Context(), "pwd", overrideWork, 5*time.Second) - if err != nil || code != 0 || strings.TrimSpace(outOverride) != overrideWork { - t.Errorf("ExecuteBash(override workDir) = (stdout=%q, code=%d, err=%v), want (%q, 0, nil)", outOverride, code, err, overrideWork) - } - - start := time.Now() - _, _, timeoutCode, timeoutErr := sh.ExecuteBash(t.Context(), "(sleep 4; echo child >&2) & sleep 4", "", 80*time.Millisecond) - if timeoutErr == nil || timeoutCode != -1 || time.Since(start) > 1500*time.Millisecond { - t.Errorf("ExecuteBash(process group timeout) = (code=%d, err=%v, elapsed=%v), want (-1, timeoutErr, <1.5s)", timeoutCode, timeoutErr, time.Since(start)) - } - }) -} - -func TestWireEncodingFieldNumbersSortingAndDepthBoundaries(t *testing.T) { - t.Run("client_event_and_batch_field_fidelity", func(t *testing.T) { - zeroBatch := (&clientEventBatch{ - AckOutputSeqNum: 0, - ReconnectInteraction: &reconnectInteractionEvent{ - InteractionID: "int_0", + { + name: "connect_positional_arg", + args: []string{"environments", "connect", "extra"}, + wantSub: "accepts 0 arguments, received 1", }, - }).marshalWire() - for rest := zeroBatch; len(rest) > 0; { - num, _, _, _, n, err := consumeField(rest) - if err != nil { - t.Fatalf("consumeField(zeroBatch) failed: %v", err) - } - if num == 2 { - t.Errorf("clientEventBatch(AckOutputSeqNum=0) unexpectedly emitted tag 2") - } - rest = rest[n:] - } - - errEv := (&clientEvent{ - SeqNum: 9, - TimestampMicros: 1234567, - Payload: &wirePayload{ - Step: &wireStep{ - ToolResult: &toolResultStep{ - CallID: "call_err_1", - FunctionResult: &functionResultStep{ - Name: "execute_bash", - IsError: true, - Result: map[string]any{"error": "failed"}, - }, - }, - }, + { + name: "connect_workdir_escapes_root", + args: []string{"environments", "connect", "--root-dir", sbxDir, "--workdir", outsideDir}, + wantSub: "escapes root-dir", + }, + { + name: "connect_zero_heartbeat", + args: []string{"environments", "connect", "--heartbeat-interval", "0s"}, + wantSub: "--heartbeat-interval must be > 0", + }, + { + name: "connect_zero_command_timeout", + args: []string{"environments", "connect", "--command-timeout", "0s"}, + wantSub: "--command-timeout must be > 0", + }, + { + name: "connect_control_char_env_key", + args: []string{"environments", "connect", "--environment-key", "bad\x00key"}, + wantSub: "--environment-key contains control characters", + }, + { + name: "connect_zero_max_concurrent", + args: []string{"environments", "connect", "--max-concurrent", "0"}, + wantSub: "--max-concurrent must be >= 1", + }, + { + name: "connect_zero_poll_interval", + args: []string{"environments", "connect", "--poll-interval", "0s"}, + wantSub: "--poll-interval must be > 0", + }, + { + name: "connect_missing_credentials", + args: []string{"environments", "connect"}, + wantSub: "missing credentials", + }, + { + name: "poll_zero_max_interactions", + args: []string{"environments", "poll", "--max-interactions", "0"}, + wantSub: "--max-interactions must be >= 1", + }, + { + name: "poll_zero_poll_interval", + args: []string{"environments", "poll", "--poll-interval", "0s"}, + wantSub: "--poll-interval must be > 0", + }, + { + name: "poll_zero_provision_timeout", + args: []string{"environments", "poll", "--provision-timeout", "0s"}, + wantSub: "--provision-timeout must be > 0", + }, + { + name: "poll_zero_handoff_heartbeat", + args: []string{"environments", "poll", "--handoff-heartbeat-interval", "0s"}, + wantSub: "--handoff-heartbeat-interval must be > 0", + }, + { + name: "poll_control_char_env_key", + args: []string{"environments", "poll", "--environment-key", "bad\nkey"}, + wantSub: "--environment-key contains control characters", + }, + { + name: "poll_invalid_environment_id", + args: []string{"environments", "poll", "--environment", "environments/a/b"}, + wantSub: `invalid environment id "environments/a/b"`, + }, + { + name: "poll_missing_provision_cmd", + args: []string{"environments", "poll", "--environment-key", "ek_1"}, + wantSub: "--provision-cmd is required", + }, + { + name: "poll_missing_credentials", + args: []string{"environments", "poll", "--provision-cmd", "true"}, + wantSub: "missing credentials", + }, + { + name: "serve_only_interaction_id", + args: []string{"environments", "serve", "--interaction-id", "int_1"}, + wantSub: "--interaction-id and --lease-token must be provided together", + }, + { + name: "serve_invalid_interaction_id", + args: []string{"environments", "serve", "--interaction-id", "interactions/a/b", "--lease-token", "tok_1"}, + wantSub: "invalid interaction id", + }, + { + name: "serve_control_char_lease_token", + args: []string{"environments", "serve", "--interaction-id", "int_1", "--lease-token", "tok\r1"}, + wantSub: "--lease-token contains control characters", + }, + { + name: "serve_missing_lease", + args: []string{"environments", "serve"}, + wantSub: "missing lease", }, - }).marshalWire() - gotErrEv := decodeClientEventWire(t, errEv) - if !gotErrEv.hasSeqNum || gotErrEv.seqNum != 9 || - !gotErrEv.hasTimestamp || gotErrEv.timestampMicros != 1234567 || - !gotErrEv.hasPayload || !gotErrEv.hasStep || !gotErrEv.hasToolResult || - gotErrEv.callID != "call_err_1" || !gotErrEv.hasFnResult || - gotErrEv.fnName != "execute_bash" || !gotErrEv.isError || - !reflect.DeepEqual(gotErrEv.result, map[string]any{"error": "failed"}) { - t.Errorf("decodeClientEventWire(full error event) = %+v, want all fields populated", gotErrEv) - } - - // Empty Step (ToolResult == nil), nil Step, and nil clientEvent batch element omit tags cleanly without panicking. - emptyPayloadEv := (&clientEvent{ - SeqNum: 0, - TimestampMicros: 0, - Payload: &wirePayload{Step: &wireStep{ToolResult: nil}}, - }).marshalWire() - gotEmptyEv := decodeClientEventWire(t, emptyPayloadEv) - if gotEmptyEv.hasSeqNum || gotEmptyEv.hasTimestamp || !gotEmptyEv.hasPayload || gotEmptyEv.hasStep || gotEmptyEv.hasToolResult { - t.Errorf("decodeClientEventWire(empty step, zero seq/ts) = %+v, want only empty payload tag 3", gotEmptyEv) - } - - nilStepEv := (&clientEvent{Payload: &wirePayload{Step: nil}}).marshalWire() - gotNilStepEv := decodeClientEventWire(t, nilStepEv) - if !gotNilStepEv.hasPayload || gotNilStepEv.hasStep || gotNilStepEv.hasToolResult { - t.Errorf("decodeClientEventWire(nil Step) = %+v, want hasPayload=true, hasStep=false", gotNilStepEv) - } - if gotNilElem := (&clientEventBatch{Events: []*clientEvent{nil}}).marshalWire(); len(gotNilElem) != 0 { - t.Errorf("clientEventBatch{Events:[nil]}.marshalWire() = %v, want empty", gotNilElem) - } - if gotEmptyStr := appendStringField(nil, 1, ""); len(gotEmptyStr) != 0 { - t.Errorf("appendStringField(empty) = %v, want proto3 omission (empty slice)", gotEmptyStr) - } - if gotFalseBool := appendBoolField(nil, 1, false); len(gotFalseBool) != 0 { - t.Errorf("appendBoolField(false) = %v, want proto3 omission (empty slice)", gotFalseBool) - } - for _, numIn := range []any{int32(7), int64(8), float64(9)} { - decNum, err := decodeProtoValue(encodeProtoValue(numIn), 0) - if err != nil || decNum == nil { - t.Errorf("decodeProtoValue(encodeProtoValue(%T(%v))) = (%v, %v), want numeric float64", numIn, numIn, decNum, err) - } - } - decStrList, err := decodeProtoValue(encodeProtoValue([]string{"s1", "s2"}), 0) - if err != nil || !reflect.DeepEqual(decStrList, []any{"s1", "s2"}) { - t.Errorf("decodeProtoValue([]string) = (%#v, %v), want []any{\"s1\", \"s2\"}", decStrList, err) } - decFallbackStr, err := decodeProtoValue(encodeProtoValue(struct{ N int }{5}), 0) - if err != nil || decFallbackStr != "{5}" { - t.Errorf("decodeProtoValue(custom struct) = (%#v, %v), want \"{5}\"", decFallbackStr, err) + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + root, _, _ := newRuntimeCLIRootForTest(t, "http://127.0.0.1:1") + root.SetArgs(tc.args) + rawErr := root.Execute() + if rawErr == nil { + t.Fatalf("Execute(%v) = nil, want error containing %q", tc.args, tc.wantSub) + } + err := output.CLIError(root, rawErr) + if !strings.Contains(err.Error(), tc.wantSub) { + t.Errorf("Execute(%v) error = %q, want substring %q", tc.args, err.Error(), tc.wantSub) + } + if got := clierrors.ExitCode(err); got != clierrors.ExitUsage { + t.Errorf("ExitCode(%v) = %d, want %d (ExitUsage)", tc.args, got, clierrors.ExitUsage) + } + }) } }) - t.Run("sorted_struct_keys_proto_lists_fixed_consumed_and_depth_boundaries", func(t *testing.T) { - // Verify encodeProtoStruct sorts keys deterministically across repeated runs. - unsorted := map[string]any{ - "k9": 1, "k2": 1, "k7": 1, "k0": 1, "k5": 1, - "k8": 1, "k1": 1, "k4": 1, "k6": 1, "k3": 1, + t.Run("dry_run_previews", func(t *testing.T) { + t.Cleanup(output.ResetPreparsedRendering) + dryCases := []struct { + name string + args []string + wantPath string + wantBody string + }{ + { + name: "connect_dry_run", + args: []string{"-o", "json", "--dry-run", "environments", "connect", "--environment-key", "ek_c", "--worker-id", "w-custom", "--max-concurrent", "2"}, + wantPath: "/v1beta/environments:leaseInteractions", + wantBody: "w-custom", + }, + { + name: "poll_dry_run_with_scoped_environment", + args: []string{"-o", "json", "--dry-run", "environments", "poll", "--environment", "environments/env_77", "--environment-key", "ek_p", "--poller-id", "p-custom", "--max-interactions", "3"}, + wantPath: "/v1beta/environments:leaseInteractions", + wantBody: "env_77", + }, + { + name: "serve_dry_run_default_sample", + args: []string{"-o", "json", "--dry-run", "environments", "serve"}, + wantPath: "/google.learning.gemini.api.interactions.v1beta.InteractionsService/CreateInteractionBidiStreamByEnvironmentKey", + }, + { + name: "serve_dry_run_explicit_lease", + args: []string{"-o", "json", "--dry-run", "environments", "serve", "--interaction-id", "interactions/int_s1", "--lease-token", "tok_s1"}, + wantPath: "/google.learning.gemini.api.interactions.v1beta.InteractionsService/CreateInteractionBidiStreamByEnvironmentKey", + }, } - wantOrder := []string{"k0", "k1", "k2", "k3", "k4", "k5", "k6", "k7", "k8", "k9"} - for run := 0; run < 5; run++ { - encoded := encodeProtoStruct(unsorted) - var gotKeys []string - for rest := encoded; len(rest) > 0; { - num, wt, _, vb, n, err := consumeField(rest) - if err != nil || num != 1 || wt != wireBytes { - t.Fatalf("consumeField(struct entry) = (%d, %d, %v)", num, wt, err) + for _, tc := range dryCases { + t.Run(tc.name, func(t *testing.T) { + root, stdout, _ := newRuntimeCLIRootForTest(t, "http://127.0.0.1:1") + output.PreparseRenderingFlags(root, []string{"-o", "json"}) + root.SetArgs(tc.args) + if err := root.Execute(); err != nil { + t.Fatalf("Execute(%v) unexpected error: %v", tc.args, err) } - rest = rest[n:] - enum, ewt, _, evb, _, eerr := consumeField(vb) - if eerr != nil || enum != 1 || ewt != wireBytes { - t.Fatalf("consumeField(entry key) = (%d, %d, %v)", enum, ewt, eerr) + var preview map[string]any + if err := json.Unmarshal(stdout.Bytes(), &preview); err != nil { + t.Fatalf("json.Unmarshal(dry-run output %q) failed: %v", stdout.String(), err) } - gotKeys = append(gotKeys, string(evb)) - } - if !slices.Equal(gotKeys, wantOrder) { - t.Fatalf("encodeProtoStruct key order = %v, want sorted %v", gotKeys, wantOrder) - } - } - - // Verify []any elements, nested struct value, NullValue overwrite, and unknown/wrong-wire-type list tags. - roundTrip, err := decodeProtoStruct(encodeProtoStruct(map[string]any{ - "any_list": []any{"a", float64(2), nil}, - "nested": map[string]any{"k": "v"}, - }), 0) - if err != nil { - t.Fatalf("decodeProtoStruct() failed: %v", err) - } - if !reflect.DeepEqual(roundTrip["any_list"], []any{"a", float64(2), nil}) { - t.Errorf("decoded any_list = %#v, want []any{\"a\", 2.0, nil}", roundTrip["any_list"]) - } - if !reflect.DeepEqual(roundTrip["nested"], map[string]any{"k": "v"}) { - t.Errorf("decoded nested = %#v, want map[string]any{\"k\": \"v\"}", roundTrip["nested"]) - } - - overwrittenNull, err := decodeProtoValue(append(encodeProtoValue("stale"), encodeProtoValue(nil)...), 0) - if err != nil || overwrittenNull != nil { - t.Errorf("decodeProtoValue(string then NullValue) = (%v, %v), want (nil, nil)", overwrittenNull, err) - } - - var listWithExtraTag []byte - listWithExtraTag = appendBytesField(listWithExtraTag, 1, encodeProtoValue("only")) - listWithExtraTag = appendVarintField(listWithExtraTag, 1, 77) - listWithExtraTag = appendVarintField(listWithExtraTag, 2, 77) - listWithExtraTag = appendBytesField(listWithExtraTag, 3, []byte{0xff}) - gotList, err := decodeProtoList(listWithExtraTag, 0) - if err != nil || !reflect.DeepEqual(gotList, []any{"only"}) { - t.Errorf("decodeProtoList(with mismatched wireType + unknown tags) = (%#v, %v), want []any{\"only\"}", gotList, err) - } - - // grpcWriteFrame skips second Write call when len(msg) == 0, and propagates payload Write error when len(msg) > 0. - cwEmpty := &countingWriter{} - if err := grpcWriteFrame(cwEmpty, nil); err != nil || cwEmpty.writes != 1 || !bytes.Equal(cwEmpty.buf.Bytes(), []byte{0, 0, 0, 0, 0}) { - t.Errorf("grpcWriteFrame(empty) = (err=%v, writes=%d, bytes=%v), want (nil, 1, [0 0 0 0 0])", err, cwEmpty.writes, cwEmpty.buf.Bytes()) - } - cwFailPayload := &countingWriter{failAtWrite: 2} - if err := grpcWriteFrame(cwFailPayload, []byte{0x42}); err == nil || cwFailPayload.writes != 2 { - t.Errorf("grpcWriteFrame(fail payload write) = (err=%v, writes=%d), want (error, 2)", err, cwFailPayload.writes) - } - - // consumeField consumed byte count for wireFixed32 (1-byte tag + 4-byte uint32 = 5). - fixed32WithTrailing := append(binary.LittleEndian.AppendUint32(appendTag(nil, 3, wireFixed32), 0xabcdef01), appendVarintField(nil, 4, 9)...) - _, _, vv32, _, consumed32, err := consumeField(fixed32WithTrailing) - if err != nil || vv32 != 0xabcdef01 || consumed32 != 5 { - t.Errorf("consumeField(wireFixed32) = (vv=%#x, consumed=%d, err=%v), want (0xabcdef01, 5, nil)", vv32, consumed32, err) + reqObj, _ := preview["request"].(map[string]any) + if gotURL, _ := reqObj["url"].(string); !strings.HasSuffix(gotURL, tc.wantPath) { + t.Errorf("dry-run url = %q, want suffix %q", gotURL, tc.wantPath) + } + if tc.wantBody != "" && !strings.Contains(stdout.String(), tc.wantBody) { + t.Errorf("dry-run output = %s, want substring %q", stdout.String(), tc.wantBody) + } + }) } + }) - // Verify strict (num == X && wt == wireY) matching across every wire decoder when unknown tags or mismatched wire types trail valid fields. - var recWire []byte - recWire = appendVarintField(recWire, 2, 13) - recWire = appendBytesField(recWire, 2, []byte("wrong_wt")) - recWire = appendVarintField(recWire, 99, 88) - rec, err := unmarshalInteractionReconnected(recWire) - if err != nil || rec.NextInputSeqNum != 13 { - t.Errorf("unmarshalInteractionReconnected(strict tag/wt) = (%+v, %v), want NextInputSeqNum=13", rec, err) - } + t.Run("live_connect_poll_and_serve_execution", func(t *testing.T) { + var connectWorkerID atomic.Value + var pollEnvID atomic.Value + var serveToolRejected atomic.Bool + connectDone := make(chan struct{}, 1) + provisionFile := filepath.Join(t.TempDir(), "polled.json") - var yieldWire []byte - yieldWire = appendVarintField(yieldWire, 1, 42) - yieldWire = appendBytesField(yieldWire, 1, []byte("wrong_wt")) - yieldWire = appendVarintField(yieldWire, 99, 7) - yieldParsed, err := unmarshalYield(yieldWire) - if err != nil || yieldParsed.NextInputSeqNum != 42 { - t.Errorf("unmarshalYield(strict tag/wt) = (%+v, %v), want NextInputSeqNum=42", yieldParsed, err) - } + srvURL := startH2CTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/v1beta/environments:leaseInteractions": + var req struct { + WorkerID string `json:"worker_id"` + EnvironmentID string `json:"environment_id"` + ActiveLeases []envruntime.LeaseEntry `json:"active_leases"` + } + _ = json.NewDecoder(r.Body).Decode(&req) + w.Header().Set("Content-Type", "application/json") + if strings.HasPrefix(req.WorkerID, "poller") { + pollEnvID.Store(req.EnvironmentID) + if len(req.ActiveLeases) == 0 && pollEnvID.Load() == "env_live_poll" { + _ = json.NewEncoder(w).Encode(map[string]any{ + "leases": []LeasedInteraction{{ + InteractionID: "int_polled", + EnvironmentID: "env_live_poll", + LeaseToken: "tok_polled", + }}, + }) + return + } + _, _ = w.Write([]byte(`{}`)) + return + } + if prev := connectWorkerID.Swap(req.WorkerID); prev == nil { + _ = json.NewEncoder(w).Encode(map[string]any{ + "leases": []LeasedInteraction{{ + InteractionID: "int_connected", + EnvironmentID: "env_conn", + LeaseToken: "tok_conn", + }}, + }) + return + } + _, _ = w.Write([]byte(`{}`)) - var fnCallWire []byte - fnCallWire = appendStringField(fnCallWire, 1, "view_file") - fnCallWire = appendBytesField(fnCallWire, 2, encodeProtoStruct(map[string]any{"p": "x"})) - fnCallWire = appendVarintField(fnCallWire, 1, 9) - fnCallWire = appendVarintField(fnCallWire, 2, 9) - fnCallWire = appendBytesField(fnCallWire, 99, []byte{0xff}) - fnCallParsed, err := unmarshalFunctionCallStep(fnCallWire) - if err != nil || fnCallParsed.Name != "view_file" || fnCallParsed.Arguments["p"] != "x" { - t.Errorf("unmarshalFunctionCallStep(strict tag/wt) = (%+v, %v), want Name=view_file, p=x", fnCallParsed, err) - } + case "/google.learning.gemini.api.interactions.v1beta.InteractionsService/CreateInteractionBidiStreamByEnvironmentKey": + w.Header().Set("Content-Type", "application/grpc") + w.WriteHeader(http.StatusOK) + flusher, _ := w.(http.Flusher) + flusher.Flush() - var tcWire []byte - tcWire = appendStringField(tcWire, 1, "cid_strict") - tcWire = appendBytesField(tcWire, 3, fnCallWire) - tcWire = appendVarintField(tcWire, 1, 9) - tcWire = appendVarintField(tcWire, 3, 9) - tcWire = appendBytesField(tcWire, 99, []byte{0xff}) - tcParsed, err := unmarshalToolCallStep(tcWire) - if err != nil || tcParsed.ID != "cid_strict" || tcParsed.FunctionCall == nil || tcParsed.FunctionCall.Name != "view_file" { - t.Errorf("unmarshalToolCallStep(strict tag/wt) = (%+v, %v), want ID=cid_strict, Name=view_file", tcParsed, err) - } + hsRaw, err := envruntime.GRPCReadFrame(r.Body) + if err != nil { + return + } + hsBatch, err := envruntime.UnmarshalClientEventBatch(hsRaw) + if err != nil || hsBatch.ReconnectInteraction == nil { + return + } + if hsBatch.ReconnectInteraction.ClientID != envruntime.DefaultEnvironmentClientID { + t.Errorf("handshake ClientID = %q, want %q", hsBatch.ReconnectInteraction.ClientID, envruntime.DefaultEnvironmentClientID) + } + intID := hsBatch.ReconnectInteraction.InteractionID + if intID == "int_serve_bash_off" { + _ = envruntime.GRPCWriteFrame(w, envruntime.MarshalServerEventBatch(&envruntime.ServerEventBatch{ + Events: []*envruntime.ServerEvent{{ + SeqNum: 1, + Payload: &envruntime.EventPayload{ + Step: &envruntime.Step{ + ToolCall: &envruntime.ToolCallStep{ + ID: "c_bash", + FunctionCall: &envruntime.FunctionCallStep{ + Name: "execute_bash", + Arguments: map[string]any{"command": "echo nope"}, + }, + }, + }, + }, + }}, + })) + flusher.Flush() + for { + replyRaw, err := envruntime.GRPCReadFrame(r.Body) + if err != nil { + return + } + if len(replyRaw) == 0 { + continue + } + replyBatch, err := envruntime.UnmarshalClientEventBatch(replyRaw) + if err == nil && len(replyBatch.Events) == 1 && replyBatch.Events[0].Payload.Step.ToolResult.FunctionResult.IsError { + serveToolRejected.Store(true) + } + break + } + } + _ = envruntime.GRPCWriteFrame(w, envruntime.MarshalServerEventBatch(&envruntime.ServerEventBatch{ + InteractionCompleted: &envruntime.InteractionCompletedEvent{}, + })) + flusher.Flush() + if intID == "int_connected" { + select { + case connectDone <- struct{}{}: + default: + } + } + default: + http.NotFound(w, r) + } + })) - var stepWire []byte - stepWire = appendBytesField(stepWire, 4, tcWire) - stepWire = appendVarintField(stepWire, 4, 9) - stepWire = appendBytesField(stepWire, 99, []byte{0xff}) - stepParsed, err := unmarshalStep(stepWire) - if err != nil || stepParsed.ToolCall == nil || stepParsed.ToolCall.ID != "cid_strict" { - t.Errorf("unmarshalStep(strict tag/wt) = (%+v, %v), want ToolCall.ID=cid_strict", stepParsed, err) + serveWork := filepath.Join(t.TempDir(), "serve-work") + serveRoot, _, _ := newRuntimeCLIRootForTest(t, srvURL) + serveRoot.SetArgs([]string{ + "--no-retries", + "environments", "serve", + "--interaction-id", "interactions/int_serve_bash_off", + "--lease-token", "tok_serve", + "--root-dir", serveWork, + "--workdir", serveWork, + "--disable-bash", + "--quiet", + }) + if err := serveRoot.Execute(); err != nil { + t.Fatalf("Execute(environments serve) unexpected error: %v", err) } - - var payloadWire []byte - payloadWire = appendBytesField(payloadWire, 11, stepWire) - payloadWire = appendVarintField(payloadWire, 11, 9) - payloadWire = appendBytesField(payloadWire, 99, []byte{0xff}) - payloadParsed, err := unmarshalPayload(payloadWire) - if err != nil || payloadParsed.Step == nil || payloadParsed.Step.ToolCall == nil || payloadParsed.Step.ToolCall.ID != "cid_strict" { - t.Errorf("unmarshalPayload(strict tag/wt) = (%+v, %v), want Step.ToolCall.ID=cid_strict", payloadParsed, err) + if !serveToolRejected.Load() { + t.Errorf("serveToolRejected = false, want --disable-bash tool error sent before InteractionCompleted") } - var sevWire []byte - sevWire = appendBytesField(sevWire, 3, payloadWire) - sevWire = appendVarintField(sevWire, 10, 29) - sevWire = appendVarintField(sevWire, 3, 9) - sevWire = appendBytesField(sevWire, 10, []byte("wrong_wt")) - sevWire = appendVarintField(sevWire, 99, 77) - sevWire = appendBytesField(sevWire, 98, []byte{0xff}) - sevParsed, err := unmarshalServerEvent(sevWire) - if err != nil || sevParsed.SeqNum != 29 || sevParsed.Payload == nil || sevParsed.Payload.Step == nil { - t.Errorf("unmarshalServerEvent(strict tag/wt) = (%+v, %v), want SeqNum=29 and Step!=nil", sevParsed, err) + connCtx, cancelConn := context.WithCancel(t.Context()) + connDir := filepath.Join(t.TempDir(), "conn-work") + connRoot, _, _ := newRuntimeCLIRootForTest(t, srvURL) + connRoot.SetContext(connCtx) + connRoot.SetArgs([]string{ + "environments", "connect", + "--environment-key", "ek_conn", + "--worker-id", "worker-cli-1", + "--root-dir", connDir, + "--poll-interval", "20ms", + "--quiet", + }) + connErrCh := make(chan error, 1) + go func() { + connErrCh <- connRoot.Execute() + }() + select { + case <-connectDone: + case <-time.After(3 * time.Second): + t.Errorf("timed out waiting for environments connect to complete leased session") } - - var batchWire []byte - batchWire = appendBytesField(batchWire, 1, sevWire) - batchWire = appendVarintField(batchWire, 2, 31) - batchWire = appendBytesField(batchWire, 8, recWire) - batchWire = appendBytesField(batchWire, 10, yieldWire) - batchWire = appendVarintField(batchWire, 1, 9) - batchWire = appendBytesField(batchWire, 2, []byte("wrong_wt")) - batchWire = appendVarintField(batchWire, 6, 9) - batchWire = appendVarintField(batchWire, 8, 9) - batchWire = appendVarintField(batchWire, 9, 9) - batchWire = appendVarintField(batchWire, 10, 9) - batchWire = appendVarintField(batchWire, 99, 77) - batchWire = appendBytesField(batchWire, 98, []byte{0xff}) - batch, err := unmarshalServerEventBatch(batchWire) - if err != nil || len(batch.Events) != 1 || batch.AckInputSeqNum != 31 || - batch.InteractionCompleted != nil || batch.ReconnectionRequired != nil || - batch.InteractionReconnected == nil || batch.InteractionReconnected.NextInputSeqNum != 13 || - batch.Yield == nil || batch.Yield.NextInputSeqNum != 42 { - t.Errorf("unmarshalServerEventBatch(strict tag/wt) = (%+v, %v), want exact fields preserved", batch, err) + cancelConn() + if err := <-connErrCh; err != nil { + t.Errorf("Execute(environments connect) = %v, want nil on context cancel", err) } - - var mapEntry []byte - mapEntry = appendStringField(mapEntry, 1, "k1") - mapEntry = appendBytesField(mapEntry, 2, encodeProtoValue("v1")) - mapEntry = appendVarintField(mapEntry, 1, 9) - mapEntry = appendVarintField(mapEntry, 2, 9) - mapEntry = appendBytesField(mapEntry, 99, []byte{0xff}) - var mapWire []byte - mapWire = appendBytesField(mapWire, 1, mapEntry) - mapWire = appendVarintField(mapWire, 1, 5) - mapWire = appendBytesField(mapWire, 99, []byte{0xff}) - strictStruct, err := decodeProtoStruct(mapWire, 0) - if err != nil || len(strictStruct) != 1 || strictStruct["k1"] != "v1" { - t.Errorf("decodeProtoStruct(strict tag/wt) = (%#v, %v), want map[k1:v1]", strictStruct, err) + if got, _ := connectWorkerID.Load().(string); got != "worker-cli-1" { + t.Errorf("connect WorkerID = %q, want \"worker-cli-1\"", got) } - var valWire []byte - valWire = append(valWire, encodeProtoValue("keep")...) - valWire = appendBytesField(valWire, 1, []byte("x")) - valWire = appendVarintField(valWire, 2, 1) - valWire = appendVarintField(valWire, 3, 1) - valWire = appendBytesField(valWire, 4, []byte("x")) - valWire = appendVarintField(valWire, 5, 1) - valWire = appendVarintField(valWire, 6, 1) - valWire = appendVarintField(valWire, 99, 1) - valWire = binary.LittleEndian.AppendUint64(appendTag(valWire, 98, wireFixed64), math.Float64bits(9.5)) - valWire = appendBytesField(valWire, 97, []byte{0xff}) - gotStrictVal, err := decodeProtoValue(valWire, 0) - if err != nil || gotStrictVal != "keep" { - t.Errorf("decodeProtoValue(strict tag/wt) = (%#v, %v), want \"keep\"", gotStrictVal, err) + pollCtx, cancelPoll := context.WithCancel(t.Context()) + pollRoot, _, _ := newRuntimeCLIRootForTest(t, srvURL) + pollRoot.SetContext(pollCtx) + pollRoot.SetArgs([]string{ + "environments", "poll", + "--environment", "environments/env_live_poll", + "--environment-key", "ek_poll", + "--poller-id", "poller-cli-1", + "--poll-interval", "20ms", + "--provision-cmd", "cat > " + provisionFile, + "--quiet", + }) + pollErrCh := make(chan error, 1) + go func() { + pollErrCh <- pollRoot.Execute() + }() + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + if raw, err := os.ReadFile(provisionFile); err == nil && len(raw) > 0 { + break + } + time.Sleep(15 * time.Millisecond) } - - // Exact depth+1 boundary assertions at maxProtoNestingDepth. - structWire := encodeProtoStruct(map[string]any{"k": "v"}) - if _, err := decodeProtoStruct(structWire, maxProtoNestingDepth-1); err != nil { - t.Errorf("decodeProtoStruct(depth=max-1) unexpected error: %v", err) + cancelPoll() + if err := <-pollErrCh; err != nil { + t.Errorf("Execute(environments poll) = %v, want nil on context cancel", err) } - if _, err := decodeProtoStruct(structWire, maxProtoNestingDepth); err == nil { - t.Errorf("decodeProtoStruct(depth=max) = nil, want depth+1 overflow error") + if got, _ := pollEnvID.Load().(string); got != "env_live_poll" { + t.Errorf("poll EnvironmentID = %q, want \"env_live_poll\"", got) } + }) - nestedStructVal := encodeProtoValue(map[string]any{"k": "v"}) - if _, err := decodeProtoValue(nestedStructVal, maxProtoNestingDepth-2); err != nil { - t.Errorf("decodeProtoValue(struct, depth=max-2) unexpected error: %v", err) - } - if _, err := decodeProtoValue(nestedStructVal, maxProtoNestingDepth-1); err == nil { - t.Errorf("decodeProtoValue(struct, depth=max-1) = nil, want inner decodeProtoStruct(depth=max) overflow error") - } - if _, err := decodeProtoValue(nestedStructVal, maxProtoNestingDepth); err == nil { - t.Errorf("decodeProtoValue(struct, depth=max) = nil, want depth+1 overflow error") - } + t.Run("helpers_logger_and_no_retries", func(t *testing.T) { + output.ResetAgentMode() + output.ResetPreparsedRendering() + config.Reset() + t.Cleanup(func() { + output.ResetAgentMode() + output.ResetPreparsedRendering() + config.Reset() + }) - nestedListVal := encodeProtoValue([]any{"item"}) - if _, err := decodeProtoValue(nestedListVal, maxProtoNestingDepth-2); err != nil { - t.Errorf("decodeProtoValue(list, depth=max-2) unexpected error: %v", err) + root, _, stderr := newRuntimeCLIRootForTest(t, "") + logNormal := newRuntimeLogger(root, false) + logNormal.Print("visible-msg") + if !strings.Contains(stderr.String(), "visible-msg") { + t.Errorf("newRuntimeLogger(quiet=false) did not write to stderr: %q", stderr.String()) } - if _, err := decodeProtoValue(nestedListVal, maxProtoNestingDepth-1); err == nil { - t.Errorf("decodeProtoValue(list, depth=max-1) = nil, want inner decodeProtoList(depth=max) overflow error") + stderr.Reset() + logQuiet := newRuntimeLogger(root, true) + logQuiet.Print("hidden-msg") + if stderr.Len() != 0 { + t.Errorf("newRuntimeLogger(quiet=true) wrote %q, want empty", stderr.String()) } - if _, err := decodeProtoValue(nestedListVal, maxProtoNestingDepth); err == nil { - t.Errorf("decodeProtoValue(list, depth=max) = nil, want depth+1 overflow error") + _ = root.PersistentFlags().Set("output-format", "json") + logMachine := newRuntimeLogger(root, false) + logMachine.Print("machine-hidden-msg") + if stderr.Len() != 0 { + t.Errorf("newRuntimeLogger(output-format=json) wrote %q, want empty", stderr.String()) } - rawListWire := appendBytesField(nil, 1, encodeProtoValue("item")) - if _, err := decodeProtoList(rawListWire, maxProtoNestingDepth-1); err != nil { - t.Errorf("decodeProtoList(depth=max-1) unexpected error: %v", err) + if got := resolveNoRetries(root); got { + t.Errorf("resolveNoRetries(default) = true, want false") } - if _, err := decodeProtoList(rawListWire, maxProtoNestingDepth); err == nil { - t.Errorf("decodeProtoList(depth=max) = nil, want depth+1 overflow error") + _ = root.PersistentFlags().Set("no-retries", "true") + if got := resolveNoRetries(root); !got { + t.Errorf("resolveNoRetries(--no-retries=true) = false, want true") } }) } -func TestWorkerStreamSequencingYieldAndToolBranching(t *testing.T) { - t.Run("no_retries_disables_reconnection_required_loop", func(t *testing.T) { - var conns atomic.Int32 - srvURL := startH2CTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - conns.Add(1) - w.Header().Set("Content-Type", "application/grpc") - w.WriteHeader(http.StatusOK) - flusher, _ := w.(http.Flusher) - flusher.Flush() - _, _ = grpcReadFrame(r.Body) - _ = grpcWriteFrame(w, buildServerReconnectRequiredFrame()) - flusher.Flush() - })) +func TestRuntimeTransportCredentialsHeadersValidationAndSDKErrorAdapter(t *testing.T) { + t.Setenv("GEMINI_API_KEY", "") + t.Setenv("GEMINI_ACCESS_TOKEN", "") + t.Setenv("GEMINI_NO_KEYRING", "true") + t.Setenv("HOME", t.TempDir()) + t.Setenv("USERPROFILE", t.TempDir()) + config.Reset() + t.Cleanup(config.Reset) - cmd := newTestRuntimeCmd(t, srvURL) - rt, err := newRuntimeTransport(cmd, true) - if err != nil { - t.Fatalf("newRuntimeTransport() failed: %v", err) - } - w := newPullWorker(NewHTTPClient(HTTPClientOptions{Transport: rt}), nil, nil, pullWorkerOptions{ - WorkerID: "w-no-reconnect", - NoRetries: true, - Logger: log.New(io.Discard, "", 0), - }) - err = w.executeLease(t.Context(), &LeasedInteraction{InteractionID: "int_nr", LeaseToken: "tok_nr"}) - if err == nil || !strings.Contains(err.Error(), "exceeded max stream reconnect attempts") { - t.Errorf("executeLease(NoRetries=true on ReconnectionRequired) = %v, want max reconnect error", err) - } - if got := conns.Load(); got != 1 { - t.Errorf("stream connection count with NoRetries=true = %d, want 1", got) + t.Run("credentials_missing_and_headers", func(t *testing.T) { + cases := []struct { + name string + explicitCred bool + apiKey string + accessToken string + userProject string + apiRevision string + header string + wantMissing bool + wantAPIKey string + wantBearer string + }{ + { + name: "no_creds_without_explicit_runtime_cred", + explicitCred: false, + wantMissing: true, + }, + { + name: "no_creds_with_explicit_runtime_cred", + explicitCred: true, + wantMissing: false, + }, + { + name: "api_key_credential", + explicitCred: false, + apiKey: "key-abc", + userProject: "proj-123", + apiRevision: "rev-2", + header: "X-Custom-Trace: trace-99", + wantMissing: false, + wantAPIKey: "key-abc", + }, + { + name: "access_token_credential", + explicitCred: false, + accessToken: "oauth-tok-xyz", + wantMissing: false, + wantBearer: "Bearer oauth-tok-xyz", + }, + { + name: "api_key_wins_same_rank_tie", + explicitCred: false, + apiKey: "tie-key-123", + accessToken: "tie-tok-456", + wantMissing: false, + wantAPIKey: "tie-key-123", + }, } - }) - - t.Run("ack_only_same_seq_tool_call_multi_seq_and_yield_next_input_seq_gate", func(t *testing.T) { - rootDir := t.TempDir() - w := newPullWorker(nil, newFSServer(rootDir, rootDir), newShellServer(rootDir), pullWorkerOptions{ - WorkerID: "w-seq", - HeartbeatInterval: time.Hour, - Logger: log.New(io.Discard, "", 0), - }) - - // Build 3 inbound ServerEventBatches: - // 1. SeqNum=5 with non-function Step (ToolCall without FunctionCall) -> ackUpdated=true, len(outEvents)==0 -> sends Ack=5. - // 2. SeqNum=5 (same seq -> ackUpdated=false!) with list_dir ToolCall (gets client SeqNum=1 -> nextSeq=2) + Yield(NextInputSeqNum=1 < 2 -> MUST NOT disconnect!). - // 3. SeqNum=6 with list_dir ToolCall (gets client SeqNum=2 -> nextSeq=3) + Yield(NextInputSeqNum=3 >= 3 -> MUST send Disconnect and exit!). - var serverBuf bytes.Buffer - nonFnStepPayload := appendBytesField(nil, 11, appendBytesField(nil, 4, appendStringField(nil, 1, "tc_no_fn"))) - batch1Ev := appendVarintField(appendBytesField(nil, 3, nonFnStepPayload), 10, 5) - _ = grpcWriteFrame(&serverBuf, appendBytesField(nil, 1, batch1Ev)) - - batch2 := append(buildServerStepToolCallFrame(5, "c_seq_1", "list_dir", nil), buildServerYieldFrame(1)...) - _ = grpcWriteFrame(&serverBuf, batch2) - - batch3 := append(buildServerStepToolCallFrame(6, "c_seq_2", "list_dir", nil), buildServerYieldFrame(3)...) - _ = grpcWriteFrame(&serverBuf, batch3) - - pr, pw := io.Pipe() - var outboundFrames [][]byte - readDone := make(chan struct{}) - go func() { - defer close(readDone) - for { - f, err := grpcReadFrame(pr) - if err != nil { - return + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + cmd := newTestRuntimeCmd(t, "https://example.googleapis.com/") + pf := cmd.Root().PersistentFlags() + if tc.apiKey != "" { + _ = pf.Set("api-key", tc.apiKey) } - outboundFrames = append(outboundFrames, f) - } - }() - - stream := &grpcInteractionStream{ - resp: &http.Response{Body: io.NopCloser(&serverBuf)}, - pw: pw, - } - state := &sessionState{ - interactionID: "int_seq", - nextClientSeq: 1, - executedCalls: make(map[string]*wirePayload), - } - reconnect, err := w.consumeStream(t.Context(), stream, state) - _ = pw.Close() - <-readDone - - if err != nil || reconnect { - t.Fatalf("consumeStream() = (%v, %v), want (false, nil)", reconnect, err) - } - // Expect 4 outbound frames: - // Frame 0: AckOutputSeqNum=5 (0 events) - // Frame 1: AckOutputSeqNum=5, Events=[SeqNum=1] - // Frame 2: AckOutputSeqNum=6, Events=[SeqNum=2] - // Frame 3: AckOutputSeqNum=6, Disconnect - if len(outboundFrames) != 4 { - t.Fatalf("len(outboundFrames) = %d, want 4", len(outboundFrames)) - } + if tc.accessToken != "" { + _ = pf.Set("access-token", tc.accessToken) + } + if tc.userProject != "" { + _ = pf.Set("user-project", tc.userProject) + } + if tc.apiRevision != "" { + _ = pf.Set("api-revision", tc.apiRevision) + } + if tc.header != "" { + _ = pf.Set("header", tc.header) + } + _ = pf.Set("timeout", "3s") - parseBatch := func(frame []byte) (ack int64, eventSeqs []int64, hasDisconnect bool) { - for rest := frame; len(rest) > 0; { - num, _, vv, vb, n, err := consumeField(rest) + rt, err := newRuntimeTransport(cmd, tc.explicitCred) if err != nil { - t.Fatalf("consumeField(outbound) failed: %v", err) + t.Fatalf("newRuntimeTransport() unexpected error: %v", err) } - rest = rest[n:] - switch num { - case 1: - ev := decodeClientEventWire(t, vb) - eventSeqs = append(eventSeqs, ev.seqNum) - case 2: - ack = int64(vv) - case 9: - hasDisconnect = true + if got := flagutil.CredentialsMissing(cmd); got != tc.wantMissing { + t.Errorf("CredentialsMissing() = %v, want %v", got, tc.wantMissing) + } + if rt.timeout != 3*time.Second { + t.Errorf("rt.timeout = %v, want 3s", rt.timeout) } - } - return ack, eventSeqs, hasDisconnect - } - - ack0, ev0, disc0 := parseBatch(outboundFrames[0]) - if ack0 != 5 || len(ev0) != 0 || disc0 { - t.Errorf("frame 0 = (ack=%d, ev=%v, disc=%v), want (5, [], false)", ack0, ev0, disc0) - } - ack1, ev1, disc1 := parseBatch(outboundFrames[1]) - if ack1 != 5 || !slices.Equal(ev1, []int64{1}) || disc1 { - t.Errorf("frame 1 = (ack=%d, ev=%v, disc=%v), want (5, [1], false)", ack1, ev1, disc1) - } - ack2, ev2, disc2 := parseBatch(outboundFrames[2]) - if ack2 != 6 || !slices.Equal(ev2, []int64{2}) || disc2 { - t.Errorf("frame 2 = (ack=%d, ev=%v, disc=%v), want (6, [2], false)", ack2, ev2, disc2) - } - ack3, ev3, disc3 := parseBatch(outboundFrames[3]) - if ack3 != 6 || len(ev3) != 0 || !disc3 { - t.Errorf("frame 3 = (ack=%d, ev=%v, disc=%v), want (6, [], true)", ack3, ev3, disc3) - } - // Also verify Yield with len(outEvents) == 0 disconnects even when NextInputSeqNum < state.nextClientSeq. - var yieldOnlyBuf bytes.Buffer - _ = grpcWriteFrame(&yieldOnlyBuf, buildServerYieldFrame(1)) - // Append a compressed-flag poison frame so failing to return on Yield immediately surfaces an error. - yieldOnlyBuf.Write([]byte{1, 0, 0, 0, 0}) - pr2, pw2 := io.Pipe() - var yieldOutbound [][]byte - readDone2 := make(chan struct{}) - go func() { - defer close(readDone2) - for { - f, err := grpcReadFrame(pr2) + req, err := rt.newRequest(t.Context(), http.MethodPost, "/v1beta/test", http.NoBody) if err != nil { - return + t.Fatalf("rt.newRequest() unexpected error: %v", err) } - yieldOutbound = append(yieldOutbound, f) - } - }() - stream2 := &grpcInteractionStream{ - resp: &http.Response{Body: io.NopCloser(&yieldOnlyBuf)}, - pw: pw2, - } - state2 := &sessionState{ - interactionID: "int_yield_empty", - nextClientSeq: 3, - lastAckSeq: 6, - executedCalls: make(map[string]*wirePayload), - } - reconnect2, err2 := w.consumeStream(t.Context(), stream2, state2) - _ = pw2.Close() - <-readDone2 - if err2 != nil || reconnect2 || len(yieldOutbound) != 1 { - t.Fatalf("consumeStream(empty outEvents with Yield=1 < nextSeq=3) = (reconnect=%v, err=%v, frames=%d), want (false, nil, 1)", reconnect2, err2, len(yieldOutbound)) - } - if ackY, evY, discY := parseBatch(yieldOutbound[0]); ackY != 6 || len(evY) != 0 || !discY { - t.Errorf("yield-only outbound frame = (ack=%d, ev=%v, disc=%v), want (6, [], true)", ackY, evY, discY) + if got := req.Header.Get("x-goog-api-key"); got != tc.wantAPIKey { + t.Errorf("x-goog-api-key = %q, want %q", got, tc.wantAPIKey) + } + if got := req.Header.Get("Authorization"); got != tc.wantBearer { + t.Errorf("Authorization = %q, want %q", got, tc.wantBearer) + } + if tc.userProject != "" && req.Header.Get("x-goog-user-project") != tc.userProject { + t.Errorf("x-goog-user-project = %q, want %q", req.Header.Get("x-goog-user-project"), tc.userProject) + } + if tc.apiRevision != "" && req.Header.Get("Api-Revision") != tc.apiRevision { + t.Errorf("Api-Revision = %q, want %q", req.Header.Get("Api-Revision"), tc.apiRevision) + } + if tc.header != "" && req.Header.Get("X-Custom-Trace") != "trace-99" { + t.Errorf("X-Custom-Trace = %q, want \"trace-99\"", req.Header.Get("X-Custom-Trace")) + } + }) } }) - t.Run("tool_cwd_zero_worker_timeout_view_endline_replace_and_arg_helpers", func(t *testing.T) { - rootDir, err := filepath.EvalSymlinks(t.TempDir()) - if err != nil { - t.Fatalf("EvalSymlinks(TempDir) failed: %v", err) - } - subDir := filepath.Join(rootDir, "nested") - if err := os.MkdirAll(subDir, 0o755); err != nil { - t.Fatalf("MkdirAll(subDir) failed: %v", err) - } - fs := newFSServer(rootDir, rootDir) - sh := newShellServer(rootDir) - - // Zero commandTimeout on pullWorker must honor per-call timeout_ms without clamping to 0 (2m). - zeroTimeoutWorker := &pullWorker{ - fs: fs, - shell: sh, - commandTimeout: 0, - logger: log.New(io.Discard, "", 0), - } - cwdRes, isErr := zeroTimeoutWorker.toolExecuteBash(t.Context(), map[string]any{"command": "pwd", "Cwd": "nested"}) - if isErr || strings.TrimSpace(fmt.Sprint(cwdRes["stdout"])) != subDir { - t.Errorf("toolExecuteBash(Cwd=nested, pwd) = (%v, %v), want stdout=%q", cwdRes, isErr, subDir) - } - - boundedCtx, cancel := context.WithTimeout(t.Context(), 500*time.Millisecond) - defer cancel() - timeoutRes, isErr := zeroTimeoutWorker.toolExecuteBash(boundedCtx, map[string]any{"command": "sleep 2", "timeout_ms": 60}) - if !isErr || !strings.Contains(fmt.Sprint(timeoutRes["error"]), "60ms") { - t.Errorf("toolExecuteBash(commandTimeout=0, timeout_ms=60) = (%v, %v), want '60ms' timeout error", timeoutRes, isErr) - } - - // Positive commandTimeout (5s) must NOT overwrite a smaller per-call timeout_ms (60ms). - cappedWorker := &pullWorker{ - fs: fs, - shell: sh, - commandTimeout: 5 * time.Second, - logger: log.New(io.Discard, "", 0), + t.Run("transport_validation_errors", func(t *testing.T) { + badCases := []struct { + name string + serverURL string + apiVer string + header string + timeout string + wantSub string + }{ + {name: "ftp_scheme", serverURL: "ftp://example.com", wantSub: "invalid --server-url"}, + {name: "url_with_credentials", serverURL: "https://user:pass@example.com", wantSub: "without credentials, query, or fragment"}, + {name: "url_with_query", serverURL: "https://example.com?foo=bar", wantSub: "without credentials, query, or fragment"}, + {name: "url_with_fragment", serverURL: "https://example.com#sec", wantSub: "without credentials, query, or fragment"}, + {name: "invalid_api_version", apiVer: "..", wantSub: "--api-version"}, + {name: "header_missing_colon", header: "BadHeader", wantSub: "invalid header format"}, + {name: "header_empty_key", header: " : val", wantSub: "invalid header name"}, + {name: "header_space_in_key", header: "Bad Key: val", wantSub: "invalid header name"}, + {name: "header_control_char_val", header: "X-Key: bad\x7fval", wantSub: "--header contains control characters"}, + {name: "invalid_timeout", timeout: "not-a-duration", wantSub: "invalid --timeout"}, } - boundedCtx2, cancel2 := context.WithTimeout(t.Context(), 500*time.Millisecond) - defer cancel2() - smallerTimeoutRes, isErr := cappedWorker.toolExecuteBash(boundedCtx2, map[string]any{"command": "sleep 2", "timeout_ms": 60}) - if !isErr || !strings.Contains(fmt.Sprint(smallerTimeoutRes["error"]), "60ms") { - t.Errorf("toolExecuteBash(commandTimeout=5s, timeout_ms=60) = (%v, %v), want '60ms' timeout preserved", smallerTimeoutRes, isErr) + for _, tc := range badCases { + t.Run(tc.name, func(t *testing.T) { + cmd := newTestRuntimeCmd(t, tc.serverURL) + pf := cmd.Root().PersistentFlags() + if tc.apiVer != "" { + _ = pf.Set("api-version", tc.apiVer) + } + if tc.header != "" { + _ = pf.Set("header", tc.header) + } + if tc.timeout != "" { + _ = pf.Set("timeout", tc.timeout) + } + if got, err := newRuntimeTransport(cmd, true); err == nil || !strings.Contains(err.Error(), tc.wantSub) { + t.Errorf("newRuntimeTransport() = (%v, %v), want error containing %q", got, err, tc.wantSub) + } + }) } + }) - // toolViewFile full read (StartLine==0 && EndLine==0) preserves exact CRLF bytes, plus EndLine-only and StartLine-only. - if _, err := fs.WriteRegularFile("lines.txt", "a\nb\nc\nd\n", true, false); err != nil { - t.Fatalf("WriteRegularFile(lines.txt) failed: %v", err) - } - viewFull, isErr := zeroTimeoutWorker.toolViewFile(map[string]any{"path": "lines.txt"}) - if isErr || viewFull["content"] != "a\nb\nc\nd\n" || viewFull["path"] != "lines.txt" { - t.Errorf("toolViewFile(full) = (%v, %v), want content=\"a\\nb\\nc\\nd\\n\"", viewFull, isErr) - } - if _, err := fs.WriteRegularFile("crlf_full.txt", "crlf1\r\ncrlf2\r\n", true, false); err != nil { - t.Fatalf("WriteRegularFile(crlf_full.txt) failed: %v", err) - } - viewCRLF, isErr := zeroTimeoutWorker.toolViewFile(map[string]any{"path": "crlf_full.txt"}) - if isErr || viewCRLF["content"] != "crlf1\r\ncrlf2\r\n" { - t.Errorf("toolViewFile(full CRLF) = (%v, %v), want exact \"crlf1\\r\\ncrlf2\\r\\n\" without range normalization", viewCRLF, isErr) - } - viewEndOnly, isErr := zeroTimeoutWorker.toolViewFile(map[string]any{"path": "lines.txt", "EndLine": 2}) - if isErr || viewEndOnly["content"] != "a\nb" { - t.Errorf("toolViewFile(StartLine=0, EndLine=2) = (%v, %v), want content=\"a\\nb\"", viewEndOnly, isErr) - } - viewStartOnly, isErr := zeroTimeoutWorker.toolViewFile(map[string]any{"path": "lines.txt", "StartLine": 3}) - if isErr || viewStartOnly["content"] != "c\nd\n" { - t.Errorf("toolViewFile(StartLine=3, EndLine=0) = (%v, %v), want content=\"c\\nd\\n\"", viewStartOnly, isErr) - } + t.Run("sdk_error_wrapping_and_debug_redaction", func(t *testing.T) { + srvURL := startH2CTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/v1beta/environments:leaseInteractions" { + http.Error(w, `{"error":{"message":"permission denied"}}`, http.StatusForbidden) + return + } + w.Header().Set("Content-Type", "application/grpc") + w.Header().Set("Trailer", "Grpc-Status, Grpc-Message") + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + w.Header().Set("Grpc-Status", "14") + w.Header().Set("Grpc-Message", "upstream%20down") + })) - // Explicit empty-path / empty-target error messages across file tools and default "." path for list_dir. - viewNoPath, isErr := zeroTimeoutWorker.toolViewFile(map[string]any{}) - if !isErr || viewNoPath["error"] != "path is required" { - t.Errorf("toolViewFile(missing path) = (%v, %v), want error=\"path is required\"", viewNoPath, isErr) - } - writeNoPath, isErr := zeroTimeoutWorker.toolWriteFile(map[string]any{}) - if !isErr || writeNoPath["error"] != "path is required" { - t.Errorf("toolWriteFile(missing path) = (%v, %v), want error=\"path is required\"", writeNoPath, isErr) - } - replaceNoPath, isErr := zeroTimeoutWorker.toolReplaceFileContent(map[string]any{"TargetContent": "a"}) - if !isErr || replaceNoPath["error"] != "TargetFile is required" { - t.Errorf("toolReplaceFileContent(missing TargetFile) = (%v, %v), want error=\"TargetFile is required\"", replaceNoPath, isErr) - } - replaceNoTarget, isErr := zeroTimeoutWorker.toolReplaceFileContent(map[string]any{"TargetFile": "lines.txt"}) - if !isErr || replaceNoTarget["error"] != "TargetContent is required" { - t.Errorf("toolReplaceFileContent(missing TargetContent) = (%v, %v), want error=\"TargetContent is required\"", replaceNoTarget, isErr) - } - openRoot, err := os.OpenRoot(rootDir) + cmd := newTestRuntimeCmd(t, srvURL) + var dbg bytes.Buffer + cmd.Root().SetErr(&dbg) + _ = cmd.Root().PersistentFlags().Set("debug", "true") + rt, err := newRuntimeTransport(cmd, true) if err != nil { - t.Fatalf("os.OpenRoot(rootDir) failed: %v", err) - } - defer openRoot.Close() - if err := mkdirAllInRoot(openRoot, "lines.txt"); err == nil { - t.Errorf("mkdirAllInRoot(regular file) = nil, want EEXIST error") - } - listDefault, isErr := zeroTimeoutWorker.toolListDir(map[string]any{}) - if isErr || listDefault["path"] != "." { - t.Errorf("toolListDir(empty args) = (%v, %v), want path=\".\"", listDefault, isErr) - } - - // replaceInText when CRLF content is matched with explicit CRLF target, and when StartLine==0 && EndLine>0. - gotCRLFTarget, cnt, err := replaceInText("a\r\nb\r\n", "a\r\nb", "x\r\ny", 0, 0, false) - if err != nil || cnt != 1 || gotCRLFTarget != "x\r\ny\r\n" { - t.Errorf("replaceInText(explicit CRLF target) = (%q, %d, %v), want (\"x\\r\\ny\\r\\n\", 1, nil)", gotCRLFTarget, cnt, err) - } - gotEndOnly, cnt, err := replaceInText("dup\nkeep\ndup", "dup", "NEW", 0, 2, false) - if err != nil || cnt != 1 || gotEndOnly != "NEW\nkeep\ndup" { - t.Errorf("replaceInText(StartLine=0, EndLine=2) = (%q, %d, %v), want (\"NEW\\nkeep\\ndup\", 1, nil)", gotEndOnly, cnt, err) - } - gotStartOnly, cnt, err := replaceInText("dup\nkeep\ndup", "dup", "NEW", 2, 0, false) - if err != nil || cnt != 1 || gotStartOnly != "dup\nkeep\nNEW" { - t.Errorf("replaceInText(StartLine=2, EndLine=0) = (%q, %d, %v), want (\"dup\\nkeep\\nNEW\", 1, nil)", gotStartOnly, cnt, err) - } - if invRange, err := fs.ReadFileRange("lines.txt", 4, 2); err != nil || invRange != "" { - t.Errorf("ReadFileRange(4, 2) = (%q, %v), want (\"\", nil)", invRange, err) - } - if rEmpty, wEmpty, err := resolveSandboxDirs(subDir, "", false); err != nil || rEmpty != subDir || wEmpty != subDir { - t.Errorf("resolveSandboxDirs(subDir, \"\") = (%q, %q, %v), want (%q, %q, nil)", rEmpty, wEmpty, err, subDir, subDir) - } - if rDot, wDot, err := resolveSandboxDirs(subDir, ".", false); err != nil || rDot != subDir || wDot != subDir { - t.Errorf("resolveSandboxDirs(subDir, \".\") = (%q, %q, %v), want (%q, %q, nil)", rDot, wDot, err, subDir, subDir) - } - if genZero, err := parseProtoInt64(json.RawMessage(`""`)); err != nil || genZero != 0 { - t.Errorf("parseProtoInt64(\"\") = (%d, %v), want (0, nil)", genZero, err) - } - schemaOutFile := filepath.Join(t.TempDir(), "schema.json") - t.Setenv("TEST_SCHEMA_OUT", schemaOutFile) - schemaProv := NewCommandWorkerProvisioner(`cat > "$TEST_SCHEMA_OUT"`) - if err := schemaProv.ProvisionAndDispatch(t.Context(), ProvisionInput{SchemaVersion: "v2", InteractionID: "i_v2"}); err != nil { - t.Fatalf("ProvisionAndDispatch(v2) failed: %v", err) - } - if rawV2, err := os.ReadFile(schemaOutFile); err != nil || !strings.Contains(string(rawV2), `"schema_version":"v2"`) { - t.Errorf("ProvisionAndDispatch(v2) payload = %q (err=%v), want schema_version=v2", rawV2, err) - } - - // stringArg skips nil values instead of formatting "", and boolArg returns explicit bool opposite fallback. - if got := stringArg(map[string]any{"first": nil, "second": "real"}, "first", "second"); got != "real" { - t.Errorf("stringArg(nil then real) = %q, want \"real\"", got) - } - if got := boolArg(map[string]any{"nil_v": nil, "explicit_false": false}, true, "nil_v", "explicit_false"); got { - t.Errorf("boolArg(explicit_false, fallback=true) = true, want false") - } - if got := boolArg(map[string]any{"explicit_true": true}, false, "explicit_true"); !got { - t.Errorf("boolArg(explicit_true, fallback=false) = false, want true") - } - if got := intArg(map[string]any{"nil_i": nil, "valid_i": 17}, 3, "nil_i", "valid_i"); got != 17 { - t.Errorf("intArg(nil then 17) = %d, want 17", got) + t.Fatalf("newRuntimeTransport(debug) failed: %v", err) } - if got := stringSliceArg(map[string]any{"nil_s": nil, "valid_s": []string{"x"}}, "nil_s", "valid_s"); !slices.Equal(got, []string{"x"}) { - t.Errorf("stringSliceArg(nil then [x]) = %v, want [x]", got) + hc := NewHTTPClient(HTTPClientOptions{Transport: rt, EnvironmentKey: "ek_err"}) + if hc.ServerURL() != srvURL { + t.Errorf("ServerURL() = %q, want %q", hc.ServerURL(), srvURL) } - // Dangling symlink root-dir rejection in resolveSandboxDirs (both direct and nested under dangling prefix). - danglingParent := t.TempDir() - danglingRoot := filepath.Join(danglingParent, "dangling_root") - if err := os.Symlink(filepath.Join(danglingParent, "nonexistent_target"), danglingRoot); err != nil { - t.Fatalf("os.Symlink(danglingRoot) failed: %v", err) - } - if _, _, err := resolveSandboxDirs(danglingRoot, "", false); err == nil || !strings.Contains(err.Error(), "dangling symlink") { - t.Errorf("resolveSandboxDirs(danglingRoot, create=false) = %v, want dangling symlink error", err) + _, leaseErr := hc.LeaseSimple(t.Context(), "w-err", 1) + if leaseErr == nil || !strings.Contains(leaseErr.Error(), "Status 403") { + t.Errorf("LeaseSimple(403) = %v, want SDKDefaultError containing \"Status 403\"", leaseErr) } - if _, _, err := resolveSandboxDirs(danglingRoot, "", true); err == nil || !strings.Contains(err.Error(), "dangling symlink") { - t.Errorf("resolveSandboxDirs(danglingRoot, create=true) = %v, want dangling symlink error", err) - } - if _, _, err := resolveSandboxDirs(filepath.Join(danglingRoot, "nested"), "", false); err == nil || !strings.Contains(err.Error(), "dangling symlink") { - t.Errorf("resolveSandboxDirs(danglingRoot/nested, create=false) = %v, want dangling symlink error", err) + if !isNonRetryableControlPlaneError(leaseErr) { + t.Errorf("isNonRetryableControlPlaneError(%v) = false, want true", leaseErr) } - // Nil-safe wire marshaling across clientEventBatch and clientEvent sub-structures. - if got := (*clientEventBatch)(nil).marshalWire(); len(got) != 0 { - t.Errorf("(*clientEventBatch)(nil).marshalWire() = %x, want empty", got) - } - nilPayloadWire := (&clientEvent{SeqNum: 7, TimestampMicros: 9, Payload: nil}).marshalWire() - if len(nilPayloadWire) == 0 || bytes.Contains(nilPayloadWire, []byte{0x1a}) { - t.Errorf("clientEvent{Payload:nil}.marshalWire() = %x, want SeqNum/Timestamp without field 3 tag 0x1a", nilPayloadWire) - } - emptyPayloadWire := (&clientEvent{Payload: &wirePayload{Step: nil}}).marshalWire() - if !bytes.Equal(emptyPayloadWire, []byte{0x1a, 0x00}) { - t.Errorf("clientEvent{Step:nil}.marshalWire() = %x, want 1a00", emptyPayloadWire) - } - emptyStepWire := (&clientEvent{Payload: &wirePayload{Step: &wireStep{ToolResult: nil}}}).marshalWire() - if !bytes.Equal(emptyStepWire, []byte{0x1a, 0x00}) { - t.Errorf("clientEvent{ToolResult:nil}.marshalWire() = %x, want 1a00", emptyStepWire) + stream, err := hc.ConnectInteraction(t.Context(), &envruntime.ConnectInteractionRequest{ + InteractionID: "int_tr", + LeaseToken: "tok_secret_dbg", + }) + if err != nil { + t.Fatalf("ConnectInteraction() failed: %v", err) } - nilFnResWire := (&clientEvent{Payload: &wirePayload{Step: &wireStep{ToolResult: &toolResultStep{CallID: "c1", FunctionResult: nil}}}}).marshalWire() - if len(nilFnResWire) == 0 { - t.Errorf("clientEvent{FunctionResult:nil}.marshalWire() = empty, want CallID encoded") + _, recvErr := stream.Recv(t.Context()) + _ = stream.Close() + if recvErr == nil || !strings.Contains(recvErr.Error(), "Status 503") { + t.Errorf("stream.Recv(Grpc-Status 14) = %v, want SDKDefaultError containing \"Status 503\"", recvErr) } - batchWithNilEv := (&clientEventBatch{ - AckOutputSeqNum: 4, - Events: []*clientEvent{nil, {SeqNum: 2}}, - ReconnectInteraction: &reconnectInteractionEvent{ - InteractionID: "i_rec", - NextOutputSeqNum: 3, - Background: true, - ClientID: "environment", - }, - Disconnect: &disconnectEvent{}, - }).marshalWire() - if len(batchWithNilEv) == 0 { - t.Errorf("clientEventBatch(with nil event + reconnect + disconnect) = empty") + if isNonRetryableControlPlaneError(recvErr) { + t.Errorf("isNonRetryableControlPlaneError(503) = true, want false") } - - // Nil-chain guards in handleServerEvent (Payload==nil, Step==nil, ToolCall==nil, FunctionCall==nil). - dummyState := &sessionState{interactionID: "i_nil", executedCalls: make(map[string]*wirePayload)} - for idx, sev := range []*serverEvent{ - {SeqNum: 1, Payload: nil}, - {SeqNum: 2, Payload: &wirePayload{Step: nil}}, - {SeqNum: 3, Payload: &wirePayload{Step: &wireStep{ToolCall: nil}}}, - {SeqNum: 4, Payload: &wirePayload{Step: &wireStep{ToolCall: &toolCallStep{ID: "tc_nil", FunctionCall: nil}}}}, - } { - if got := zeroTimeoutWorker.handleServerEvent(t.Context(), sev, dummyState); got != nil { - t.Errorf("handleServerEvent(case %d) = %+v, want nil", idx, got) - } + if strings.Contains(dbg.String(), "tok_secret_dbg") { + t.Errorf("debug stderr leaked lease token: %q", dbg.String()) } - // HeartbeatInteraction skips nil elements in parsed.Leases without panicking. - rtHB := &runtimeTransport{ - baseURL: "https://example.googleapis.com", - scheme: "https", - apiVersion: "v1beta", - headers: make(map[string]string), - httpClient: fakeHTTPDoer(func(*http.Request) (*http.Response, error) { - return &http.Response{ - StatusCode: http.StatusOK, - Body: io.NopCloser(strings.NewReader(`{"leases":[null,{"interaction_id":"int_hb_ok"}]}`)), - }, nil - }), - } - hbRes, hbErr := NewHTTPClient(HTTPClientOptions{Transport: rtHB}).HeartbeatInteraction(t.Context(), "p1", "int_hb_ok", "env1", "lt1") - if hbErr != nil || !hbRes.Renewed || hbRes.ShouldCancel { - t.Errorf("HeartbeatInteraction(with null lease entry) = (%+v, %v), want Renewed=true", hbRes, hbErr) + // Verify nil-Transport fallback constructor does not panic. + fallbackHC := NewHTTPClient(HTTPClientOptions{EnvironmentKey: "ek_fb"}) + if fallbackHC == nil { + t.Error("NewHTTPClient(nil Transport) = nil") } - // Empty gRPC frame read and unstarted subprocess group Cancel guard. - if emptyFrame, err := grpcReadFrame(bytes.NewReader([]byte{0, 0, 0, 0, 0})); err != nil || len(emptyFrame) != 0 { - t.Errorf("grpcReadFrame(empty frame) = (%x, %v), want (empty, nil)", emptyFrame, err) - } - unstarted := exec.Command("true") - configureSubprocessGroup(unstarted) - if err := unstarted.Cancel(); err != nil { - t.Errorf("unstarted.Cancel() = %v, want nil", err) - } - if pwdDefault, _, code, err := sh.ExecuteBash(t.Context(), "pwd", "", time.Second); err != nil || code != 0 || strings.TrimSpace(pwdDefault) != rootDir { - t.Errorf("ExecuteBash(workDir=\"\") = (%q, %d, %v), want (%q, 0, nil)", strings.TrimSpace(pwdDefault), code, err, rootDir) + // Verify SPEAKEASY_TEST_NAME keeps streamClient unwrapped even when --debug is enabled. + t.Setenv("SPEAKEASY_TEST_NAME", "contract_stream_case") + testCmd := newTestRuntimeCmd(t, srvURL) + _ = testCmd.Root().PersistentFlags().Set("debug", "true") + testRT, err := newRuntimeTransport(testCmd, true) + if err != nil { + t.Fatalf("newRuntimeTransport(SPEAKEASY_TEST_NAME + debug) failed: %v", err) } - if pwdSub, _, code, err := sh.ExecuteBash(t.Context(), "pwd", subDir, time.Second); err != nil || code != 0 || strings.TrimSpace(pwdSub) != subDir { - t.Errorf("ExecuteBash(workDir=subDir) = (%q, %d, %v), want (%q, 0, nil)", strings.TrimSpace(pwdSub), code, err, subDir) + if !testRT.hasTestClient || testRT.streamClient == testRT.httpClient { + t.Errorf("newRuntimeTransport(SPEAKEASY_TEST_NAME + debug): hasTestClient=%v, streamClient==httpClient=%v, want hasTestClient=true and unwrapped streamClient", testRT.hasTestClient, testRT.streamClient == testRT.httpClient) } }) } - -type countingReader struct { - r io.Reader - bytesRead int -} - -func (c *countingReader) Read(p []byte) (int, error) { - n, err := c.r.Read(p) - c.bytesRead += n - return n, err -} - -type errReader struct { - err error -} - -func (e errReader) Read([]byte) (int, error) { - return 0, e.err -} - -type countingWriter struct { - buf bytes.Buffer - writes int - failAtWrite int -} - -func (c *countingWriter) Write(p []byte) (int, error) { - c.writes++ - if c.failAtWrite > 0 && c.writes == c.failAtWrite { - return 0, errors.New("injected write failure") - } - return c.buf.Write(p) -} diff --git a/internal/cli/custom/runtime_wire.go b/internal/cli/custom/runtime_wire.go deleted file mode 100644 index c9c41b9..0000000 --- a/internal/cli/custom/runtime_wire.go +++ /dev/null @@ -1,658 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package custom - -import ( - "encoding/binary" - "fmt" - "io" - "math" - "net/url" - "sort" - "strings" -) - -const ( - wireVarint = 0 - wireFixed64 = 1 - wireBytes = 2 - wireFixed32 = 5 - - maxGRPCFrameBytes = 32 << 20 // 32 MiB - maxProtoNestingDepth = 64 -) - -// Client -> Server outbound wire messages (`ClientEventBatch`). - -type clientEventBatch struct { - Events []*clientEvent - AckOutputSeqNum int64 - ReconnectInteraction *reconnectInteractionEvent - Disconnect *disconnectEvent -} - -type reconnectInteractionEvent struct { - InteractionID string - NextOutputSeqNum int64 - Background bool - ClientID string -} - -type disconnectEvent struct{} - -type clientEvent struct { - Payload *wirePayload - SeqNum int64 - TimestampMicros int64 -} - -type wirePayload struct { - // Bidirectional step envelope (`Payload.step`, field 11). - Step *wireStep -} - -type wireStep struct { - ToolCall *toolCallStep - ToolResult *toolResultStep -} - -type toolCallStep struct { - ID string - FunctionCall *functionCallStep -} - -type functionCallStep struct { - Name string - Arguments map[string]any -} - -type toolResultStep struct { - CallID string - FunctionResult *functionResultStep -} - -type functionResultStep struct { - Name string - IsError bool - Result any -} - -// Server -> Client inbound wire messages (`ServerEventBatch`). - -type serverEventBatch struct { - Events []*serverEvent - AckInputSeqNum int64 - InteractionCompleted *interactionCompletedEvent - InteractionReconnected *interactionReconnectedEvent - ReconnectionRequired *reconnectionRequiredEvent - Yield *yieldEvent -} - -type serverEvent struct { - Payload *wirePayload - SeqNum int64 -} - -type interactionCompletedEvent struct{} - -type interactionReconnectedEvent struct { - NextInputSeqNum int64 -} - -type reconnectionRequiredEvent struct{} - -type yieldEvent struct { - NextInputSeqNum int64 -} - -// gRPC frame helpers (uncompressed 5-byte length-prefixed frames). - -func grpcWriteFrame(w io.Writer, msg []byte) error { - if len(msg) > maxGRPCFrameBytes { - return fmt.Errorf("outbound gRPC frame size %d exceeds limit %d", len(msg), maxGRPCFrameBytes) - } - var hdr [5]byte // hdr[0] == 0 denotes an uncompressed gRPC frame. - binary.BigEndian.PutUint32(hdr[1:5], uint32(len(msg))) - if _, err := w.Write(hdr[:]); err != nil { - return err - } - if len(msg) > 0 { - if _, err := w.Write(msg); err != nil { - return err - } - } - return nil -} - -func grpcReadFrame(r io.Reader) ([]byte, error) { - var hdr [5]byte - if _, err := io.ReadFull(r, hdr[:]); err != nil { - return nil, err - } - if hdr[0] != 0 { - return nil, fmt.Errorf("unsupported compressed gRPC frame flag %d", hdr[0]) - } - n := binary.BigEndian.Uint32(hdr[1:5]) - if n > maxGRPCFrameBytes { - return nil, fmt.Errorf("inbound gRPC frame size %d exceeds limit %d", n, maxGRPCFrameBytes) - } - if n == 0 { - return nil, nil - } - buf := make([]byte, n) - if _, err := io.ReadFull(r, buf); err != nil { - if err == io.EOF { - return nil, io.ErrUnexpectedEOF - } - return nil, err - } - return buf, nil -} - -func decodeGRPCMessage(raw string) string { - trimmed := strings.TrimSpace(raw) - if unescaped, err := url.PathUnescape(trimmed); err == nil { - return unescaped - } - return trimmed -} - -// Low-level protobuf binary encoding primitives. - -func appendTag(dst []byte, fieldNum int, wireType int) []byte { - return binary.AppendUvarint(dst, uint64(fieldNum<<3|wireType)) -} - -func appendVarintField(dst []byte, fieldNum int, v uint64) []byte { - dst = appendTag(dst, fieldNum, wireVarint) - return binary.AppendUvarint(dst, v) -} - -func appendBoolField(dst []byte, fieldNum int, v bool) []byte { - if !v { - return dst - } - return appendVarintField(dst, fieldNum, 1) -} - -func appendStringField(dst []byte, fieldNum int, s string) []byte { - if s == "" { - return dst - } - dst = appendTag(dst, fieldNum, wireBytes) - dst = binary.AppendUvarint(dst, uint64(len(s))) - return append(dst, s...) -} - -func appendBytesField(dst []byte, fieldNum int, b []byte) []byte { - dst = appendTag(dst, fieldNum, wireBytes) - dst = binary.AppendUvarint(dst, uint64(len(b))) - return append(dst, b...) -} - -// Outbound `ClientEventBatch` encoder. - -func (b *clientEventBatch) marshalWire() []byte { - if b == nil { - return nil - } - var out []byte - for _, ev := range b.Events { - if ev != nil { - out = appendBytesField(out, 1, ev.marshalWire()) - } - } - if b.AckOutputSeqNum != 0 { - out = appendVarintField(out, 2, uint64(b.AckOutputSeqNum)) - } - if r := b.ReconnectInteraction; r != nil { - var sub []byte - sub = appendStringField(sub, 1, r.InteractionID) - if r.NextOutputSeqNum != 0 { - sub = appendVarintField(sub, 2, uint64(r.NextOutputSeqNum)) - } - sub = appendBoolField(sub, 3, r.Background) - sub = appendStringField(sub, 4, r.ClientID) - out = appendBytesField(out, 8, sub) - } - if b.Disconnect != nil { - out = appendBytesField(out, 9, nil) - } - return out -} - -func (e *clientEvent) marshalWire() []byte { - var out []byte - if p := e.Payload; p != nil { - var sub []byte - if st := p.Step; st != nil && st.ToolResult != nil { - tr := st.ToolResult - var tb []byte - tb = appendStringField(tb, 1, tr.CallID) - if fnRes := tr.FunctionResult; fnRes != nil { - var frb []byte - frb = appendStringField(frb, 2, fnRes.Name) - frb = appendBoolField(frb, 3, fnRes.IsError) - frb = appendBytesField(frb, 7, encodeProtoValue(fnRes.Result)) - tb = appendBytesField(tb, 3, frb) - } - sub = appendBytesField(sub, 11, appendBytesField(nil, 5, tb)) - } - out = appendBytesField(out, 3, sub) - } - if e.SeqNum != 0 { - out = appendVarintField(out, 10, uint64(e.SeqNum)) - } - if e.TimestampMicros != 0 { - out = appendVarintField(out, 11, uint64(e.TimestampMicros)) - } - return out -} - -func encodeProtoValue(v any) []byte { - var out []byte - switch val := v.(type) { - case nil: - return appendVarintField(out, 1, 0) // NullValue.NULL_VALUE - case bool: - out = appendTag(out, 4, wireVarint) - if val { - return binary.AppendUvarint(out, 1) - } - return binary.AppendUvarint(out, 0) - case int: - return encodeNumberValue(float64(val)) - case int32: - return encodeNumberValue(float64(val)) - case int64: - return encodeNumberValue(float64(val)) - case float64: - return encodeNumberValue(val) - case string: - out = appendTag(out, 3, wireBytes) - out = binary.AppendUvarint(out, uint64(len(val))) - return append(out, val...) - case []any: - var listBuf []byte - for _, elem := range val { - listBuf = appendBytesField(listBuf, 1, encodeProtoValue(elem)) - } - return appendBytesField(out, 6, listBuf) - case []string: - var listBuf []byte - for _, elem := range val { - listBuf = appendBytesField(listBuf, 1, encodeProtoValue(elem)) - } - return appendBytesField(out, 6, listBuf) - case map[string]any: - return appendBytesField(out, 5, encodeProtoStruct(val)) - default: - s := fmt.Sprint(val) - out = appendTag(out, 3, wireBytes) - out = binary.AppendUvarint(out, uint64(len(s))) - return append(out, s...) - } -} - -func encodeNumberValue(f float64) []byte { - out := appendTag(nil, 2, wireFixed64) - return binary.LittleEndian.AppendUint64(out, math.Float64bits(f)) -} - -func encodeProtoStruct(m map[string]any) []byte { - keys := make([]string, 0, len(m)) - for k := range m { - keys = append(keys, k) - } - sort.Strings(keys) - var out []byte - for _, k := range keys { - var entry []byte - entry = appendTag(entry, 1, wireBytes) - entry = binary.AppendUvarint(entry, uint64(len(k))) - entry = append(entry, k...) - entry = appendBytesField(entry, 2, encodeProtoValue(m[k])) - out = appendBytesField(out, 1, entry) - } - return out -} - -// Low-level protobuf binary decoding primitives. - -func consumeVarint(b []byte) (uint64, int, error) { - v, n := binary.Uvarint(b) - if n == 0 { - return 0, 0, io.ErrUnexpectedEOF - } - if n < 0 { - return 0, 0, fmt.Errorf("protobuf varint overflows uint64") - } - return v, n, nil -} - -func consumeField(b []byte) (fieldNum int, wireType int, valVarint uint64, valBytes []byte, consumed int, err error) { - tag, n, err := consumeVarint(b) - if err != nil { - return 0, 0, 0, nil, 0, err - } - fieldNum = int(tag >> 3) - wireType = int(tag & 0x7) - if fieldNum <= 0 { - return 0, 0, 0, nil, 0, fmt.Errorf("invalid protobuf field number %d", fieldNum) - } - rest := b[n:] - switch wireType { - case wireVarint: - v, m, err := consumeVarint(rest) - if err != nil { - return 0, 0, 0, nil, 0, err - } - return fieldNum, wireType, v, nil, n + m, nil - case wireFixed64: - if len(rest) < 8 { - return 0, 0, 0, nil, 0, io.ErrUnexpectedEOF - } - return fieldNum, wireType, binary.LittleEndian.Uint64(rest[:8]), nil, n + 8, nil - case wireBytes: - l, m, err := consumeVarint(rest) - if err != nil { - return 0, 0, 0, nil, 0, err - } - after := rest[m:] - if l > uint64(len(after)) { - return 0, 0, 0, nil, 0, io.ErrUnexpectedEOF - } - return fieldNum, wireType, 0, after[:int(l)], n + m + int(l), nil - case wireFixed32: - if len(rest) < 4 { - return 0, 0, 0, nil, 0, io.ErrUnexpectedEOF - } - return fieldNum, wireType, uint64(binary.LittleEndian.Uint32(rest[:4])), nil, n + 4, nil - default: - return 0, 0, 0, nil, 0, fmt.Errorf("unsupported protobuf wire type %d", wireType) - } -} - -// Inbound `ServerEventBatch` decoder. - -func unmarshalServerEventBatch(data []byte) (*serverEventBatch, error) { - out := &serverEventBatch{} - for len(data) > 0 { - num, wt, vv, vb, n, err := consumeField(data) - if err != nil { - return nil, err - } - data = data[n:] - switch { - case num == 1 && wt == wireBytes: - ev, err := unmarshalServerEvent(vb) - if err != nil { - return nil, err - } - out.Events = append(out.Events, ev) - case num == 2 && wt == wireVarint: - out.AckInputSeqNum = int64(vv) - case num == 6 && wt == wireBytes: - out.InteractionCompleted = &interactionCompletedEvent{} - case num == 8 && wt == wireBytes: - rec, err := unmarshalInteractionReconnected(vb) - if err != nil { - return nil, err - } - out.InteractionReconnected = rec - case num == 9 && wt == wireBytes: - out.ReconnectionRequired = &reconnectionRequiredEvent{} - case num == 10 && wt == wireBytes: - y, err := unmarshalYield(vb) - if err != nil { - return nil, err - } - out.Yield = y - } - } - return out, nil -} - -func unmarshalInteractionReconnected(data []byte) (*interactionReconnectedEvent, error) { - out := &interactionReconnectedEvent{} - for len(data) > 0 { - num, wt, vv, _, n, err := consumeField(data) - if err != nil { - return nil, err - } - data = data[n:] - if num == 2 && wt == wireVarint { - out.NextInputSeqNum = int64(vv) - } - } - return out, nil -} - -func unmarshalYield(data []byte) (*yieldEvent, error) { - out := &yieldEvent{} - for len(data) > 0 { - num, wt, vv, _, n, err := consumeField(data) - if err != nil { - return nil, err - } - data = data[n:] - if num == 1 && wt == wireVarint { - out.NextInputSeqNum = int64(vv) - } - } - return out, nil -} - -func unmarshalServerEvent(data []byte) (*serverEvent, error) { - out := &serverEvent{} - for len(data) > 0 { - num, wt, vv, vb, n, err := consumeField(data) - if err != nil { - return nil, err - } - data = data[n:] - switch { - case num == 3 && wt == wireBytes: - p, err := unmarshalPayload(vb) - if err != nil { - return nil, err - } - out.Payload = p - case num == 10 && wt == wireVarint: - out.SeqNum = int64(vv) - } - } - return out, nil -} - -func unmarshalPayload(data []byte) (*wirePayload, error) { - out := &wirePayload{} - for len(data) > 0 { - num, wt, _, vb, n, err := consumeField(data) - if err != nil { - return nil, err - } - data = data[n:] - if num == 11 && wt == wireBytes { - st, err := unmarshalStep(vb) - if err != nil { - return nil, err - } - out.Step = st - } - } - return out, nil -} - -func unmarshalStep(data []byte) (*wireStep, error) { - out := &wireStep{} - for len(data) > 0 { - num, wt, _, vb, n, err := consumeField(data) - if err != nil { - return nil, err - } - data = data[n:] - if num == 4 && wt == wireBytes { - tc, err := unmarshalToolCallStep(vb) - if err != nil { - return nil, err - } - out.ToolCall = tc - } - } - return out, nil -} - -func unmarshalToolCallStep(data []byte) (*toolCallStep, error) { - out := &toolCallStep{} - for len(data) > 0 { - num, wt, _, vb, n, err := consumeField(data) - if err != nil { - return nil, err - } - data = data[n:] - switch { - case num == 1 && wt == wireBytes: - out.ID = string(vb) - case num == 3 && wt == wireBytes: - fc, err := unmarshalFunctionCallStep(vb) - if err != nil { - return nil, err - } - out.FunctionCall = fc - } - } - return out, nil -} - -func unmarshalFunctionCallStep(data []byte) (*functionCallStep, error) { - out := &functionCallStep{} - for len(data) > 0 { - num, wt, _, vb, n, err := consumeField(data) - if err != nil { - return nil, err - } - data = data[n:] - switch { - case num == 1 && wt == wireBytes: - out.Name = string(vb) - case num == 2 && wt == wireBytes: - m, err := decodeProtoStruct(vb, 0) - if err != nil { - return nil, err - } - out.Arguments = m - } - } - return out, nil -} - -func decodeProtoStruct(data []byte, depth int) (map[string]any, error) { - if depth > maxProtoNestingDepth { - return nil, fmt.Errorf("protobuf struct nesting exceeds maximum depth %d", maxProtoNestingDepth) - } - out := make(map[string]any) - for len(data) > 0 { - num, wt, _, vb, n, err := consumeField(data) - if err != nil { - return nil, err - } - data = data[n:] - if num != 1 || wt != wireBytes { - continue - } - var entryKey string - var entryVal any - entryData := vb - for len(entryData) > 0 { - enum, ewt, _, evb, en, err := consumeField(entryData) - if err != nil { - return nil, err - } - entryData = entryData[en:] - switch { - case enum == 1 && ewt == wireBytes: - entryKey = string(evb) - case enum == 2 && ewt == wireBytes: - v, err := decodeProtoValue(evb, depth+1) - if err != nil { - return nil, err - } - entryVal = v - } - } - out[entryKey] = entryVal - } - return out, nil -} - -func decodeProtoValue(data []byte, depth int) (any, error) { - if depth > maxProtoNestingDepth { - return nil, fmt.Errorf("protobuf value nesting exceeds maximum depth %d", maxProtoNestingDepth) - } - var res any - for len(data) > 0 { - num, wt, vv, vb, n, err := consumeField(data) - if err != nil { - return nil, err - } - data = data[n:] - switch { - case num == 1 && wt == wireVarint: - res = nil - case num == 2 && wt == wireFixed64: - res = math.Float64frombits(vv) - case num == 3 && wt == wireBytes: - res = string(vb) - case num == 4 && wt == wireVarint: - res = vv != 0 - case num == 5 && wt == wireBytes: - m, err := decodeProtoStruct(vb, depth+1) - if err != nil { - return nil, err - } - res = m - case num == 6 && wt == wireBytes: - items, err := decodeProtoList(vb, depth+1) - if err != nil { - return nil, err - } - res = items - } - } - return res, nil -} - -func decodeProtoList(data []byte, depth int) ([]any, error) { - if depth > maxProtoNestingDepth { - return nil, fmt.Errorf("protobuf list nesting exceeds maximum depth %d", maxProtoNestingDepth) - } - var out []any - for len(data) > 0 { - num, wt, _, vb, n, err := consumeField(data) - if err != nil { - return nil, err - } - data = data[n:] - if num == 1 && wt == wireBytes { - v, err := decodeProtoValue(vb, depth+1) - if err != nil { - return nil, err - } - out = append(out, v) - } - } - return out, nil -} diff --git a/internal/cli/custom/runtime_worker.go b/internal/cli/custom/runtime_worker.go deleted file mode 100644 index ae3c54f..0000000 --- a/internal/cli/custom/runtime_worker.go +++ /dev/null @@ -1,632 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package custom - -import ( - "context" - "errors" - "fmt" - "io" - "log" - "os" - "strings" - "sync" - "time" -) - -const ( - defaultMaxStreamReconnects = 3 - maxToolResultWireBytes = maxGRPCFrameBytes - 4096 -) - -type pullWorkerOptions struct { - WorkerID string - MaxConcurrent int - PollInterval time.Duration - HeartbeatInterval time.Duration - CommandTimeout time.Duration - DisableBash bool - NoRetries bool - Logger *log.Logger -} - -type pullWorker struct { - client *HTTPClient - fs *fsServer - shell *shellServer - workerID string - maxConcurrent int - pollInterval time.Duration - heartbeatInterval time.Duration - commandTimeout time.Duration - disableBash bool - noRetries bool - logger *log.Logger -} - -func newPullWorker(client *HTTPClient, fs *fsServer, shell *shellServer, opts pullWorkerOptions) *pullWorker { - if opts.WorkerID == "" { - opts.WorkerID = defaultRuntimeID("worker") - } - if opts.MaxConcurrent <= 0 { - opts.MaxConcurrent = 4 - } - if opts.PollInterval <= 0 { - opts.PollInterval = 2 * time.Second - } - if opts.HeartbeatInterval <= 0 { - opts.HeartbeatInterval = defaultHeartbeatInterval - } - if opts.CommandTimeout <= 0 { - opts.CommandTimeout = 2 * time.Minute - } - if opts.Logger == nil { - opts.Logger = log.New(os.Stderr, "", log.LstdFlags) - } - return &pullWorker{ - client: client, - fs: fs, - shell: shell, - workerID: opts.WorkerID, - maxConcurrent: opts.MaxConcurrent, - pollInterval: opts.PollInterval, - heartbeatInterval: opts.HeartbeatInterval, - commandTimeout: opts.CommandTimeout, - disableBash: opts.DisableBash, - noRetries: opts.NoRetries, - logger: opts.Logger, - } -} - -func (w *pullWorker) Run(ctx context.Context) error { - return runLeasePollLoop(ctx, leasePollLoopConfig{ - client: w.client, - callerID: w.workerID, - maxSlots: w.maxConcurrent, - pollInterval: w.pollInterval, - noRetries: w.noRetries, - logger: w.logger, - errLabel: "LeaseInteractions error", - startMsg: fmt.Sprintf("Worker %s polling (maxConcurrent=%d)", w.workerID, w.maxConcurrent), - }, func(ctx context.Context, item *LeasedInteraction) { - w.logger.Printf("[%s] Claimed interaction lease (env=%s)", item.InteractionID, firstNonEmpty(item.EnvironmentID, "default")) - if err := w.executeLease(ctx, item); err != nil && ctx.Err() == nil { - w.logger.Printf("[%s] session ended with error: %v", item.InteractionID, err) - return - } - if ctx.Err() == nil { - w.logger.Printf("[%s] Completed interaction session", item.InteractionID) - } - }) -} - -func runStartupLeases(ctx context.Context, w *pullWorker, leases []*LeasedInteraction) error { - for _, lease := range leases { - w.logger.Printf("[%s] Attached interaction session", lease.InteractionID) - if err := w.executeLease(ctx, lease); err != nil { - return err - } - w.logger.Printf("[%s] Completed interaction session", lease.InteractionID) - } - return nil -} - -type sessionState struct { - interactionID string - mu sync.Mutex - lastAckSeq int64 - nextClientSeq int64 - executedCalls map[string]*wirePayload -} - -func (s *sessionState) ackSeq() int64 { - s.mu.Lock() - defer s.mu.Unlock() - return s.lastAckSeq -} - -func (w *pullWorker) executeLease(ctx context.Context, lease *LeasedInteraction) error { - state := &sessionState{ - interactionID: lease.InteractionID, - nextClientSeq: 1, - executedCalls: make(map[string]*wirePayload), - } - maxReconnects := defaultMaxStreamReconnects - if w.noRetries { - maxReconnects = 0 - } - - for attempt := 0; attempt <= maxReconnects; attempt++ { - if err := ctx.Err(); err != nil { - return err - } - stream, err := w.client.connectStream(ctx, connectStreamRequest{ - InteractionID: lease.InteractionID, - EnvironmentID: lease.EnvironmentID, - LeaseToken: lease.LeaseToken, - NextOutputSeqNum: state.ackSeq() + 1, - AckOutputSeqNum: state.ackSeq(), - }) - if err != nil { - if isNonRetryableControlPlaneError(err) || attempt >= maxReconnects { - return err - } - select { - case <-ctx.Done(): - return ctx.Err() - case <-time.After(retryBackoff(200*time.Millisecond, attempt+1)): - continue - } - } - - reconnect, runErr := w.consumeStream(ctx, stream, state) - _ = stream.Close() - if runErr != nil { - if isNonRetryableControlPlaneError(runErr) || attempt >= maxReconnects { - return runErr - } - select { - case <-ctx.Done(): - return ctx.Err() - case <-time.After(retryBackoff(200*time.Millisecond, attempt+1)): - continue - } - } - if !reconnect { - return nil - } - } - return fmt.Errorf("exceeded max stream reconnect attempts for interaction %q", lease.InteractionID) -} - -func (w *pullWorker) consumeStream(ctx context.Context, stream *grpcInteractionStream, state *sessionState) (bool, error) { - // Renew the server lease in-band on the authenticated bidi stream every - // heartbeatInterval (default 5s = 15s server TTL / 3) so --lease-token-only - // sandbox workers renew leases without needing environment/API keys. - stopHB := make(chan struct{}) - hbDone := make(chan struct{}) - go func() { - defer close(hbDone) - ticker := time.NewTicker(w.heartbeatInterval) - defer ticker.Stop() - for { - select { - case <-ctx.Done(): - return - case <-stopHB: - return - case <-ticker.C: - _ = stream.Send(&clientEventBatch{AckOutputSeqNum: state.ackSeq()}) - } - } - }() - defer func() { - close(stopHB) - <-hbDone - }() - - for { - batch, err := stream.Recv() - if errors.Is(err, io.EOF) { - return false, nil - } - if err != nil { - return false, err - } - - if batch.InteractionReconnected != nil { - state.mu.Lock() - if batch.InteractionReconnected.NextInputSeqNum > state.nextClientSeq { - state.nextClientSeq = batch.InteractionReconnected.NextInputSeqNum - } - state.mu.Unlock() - } - - state.mu.Lock() - pendingAck := state.lastAckSeq - startClientSeq := state.nextClientSeq - state.mu.Unlock() - - var outEvents []*clientEvent - ackUpdated := false - for _, ev := range batch.Events { - if ev.SeqNum > pendingAck { - pendingAck = ev.SeqNum - ackUpdated = true - } - - if respPayload := w.handleServerEvent(ctx, ev, state); respPayload != nil { - state.mu.Lock() - seq := state.nextClientSeq - state.nextClientSeq++ - state.mu.Unlock() - outEvents = append(outEvents, &clientEvent{ - SeqNum: seq, - TimestampMicros: time.Now().UnixMicro(), - Payload: respPayload, - }) - } - } - - if len(outEvents) > 0 || ackUpdated { - if err := stream.Send(&clientEventBatch{ - AckOutputSeqNum: pendingAck, - Events: outEvents, - }); err != nil { - state.mu.Lock() - state.nextClientSeq = startClientSeq - state.mu.Unlock() - return false, err - } - state.mu.Lock() - state.lastAckSeq = pendingAck - state.mu.Unlock() - } - - if batch.InteractionCompleted != nil { - _ = stream.CloseSend() - return false, nil - } - if batch.ReconnectionRequired != nil { - return true, nil - } - if batch.Yield != nil { - state.mu.Lock() - nextSeq := state.nextClientSeq - lastAck := state.lastAckSeq - state.mu.Unlock() - if len(outEvents) == 0 || batch.Yield.NextInputSeqNum >= nextSeq { - _ = stream.Send(&clientEventBatch{ - AckOutputSeqNum: lastAck, - Disconnect: &disconnectEvent{}, - }) - _ = stream.CloseSend() - _, _ = stream.Recv() - return false, nil - } - } - } -} - -func (w *pullWorker) handleServerEvent(ctx context.Context, ev *serverEvent, state *sessionState) *wirePayload { - if ev.Payload == nil { - return nil - } - if ev.Payload.Step != nil && ev.Payload.Step.ToolCall != nil && ev.Payload.Step.ToolCall.FunctionCall != nil { - tc := ev.Payload.Step.ToolCall - fnCall := tc.FunctionCall - callKey := "step:" + firstNonEmpty(tc.ID, fmt.Sprintf("seq-%d", ev.SeqNum)) - state.mu.Lock() - cached, ok := state.executedCalls[callKey] - state.mu.Unlock() - if ok { - return cached - } - w.logger.Printf("[%s] Executing tool %s", state.interactionID, fnCall.Name) - resMap, isErr := w.dispatchTool(ctx, fnCall.Name, fnCall.Arguments) - out := &wirePayload{ - Step: &wireStep{ - ToolResult: &toolResultStep{ - CallID: tc.ID, - FunctionResult: &functionResultStep{ - Name: fnCall.Name, - IsError: isErr, - Result: resMap, - }, - }, - }, - } - if len((&clientEvent{SeqNum: 1, TimestampMicros: 1, Payload: out}).marshalWire()) > maxToolResultWireBytes { - out = &wirePayload{ - Step: &wireStep{ - ToolResult: &toolResultStep{ - CallID: tc.ID, - FunctionResult: &functionResultStep{ - Name: fnCall.Name, - IsError: true, - Result: map[string]any{ - "error": fmt.Sprintf("tool %q result exceeds maximum frame size (%d bytes); narrow the query or specify start_line and end_line", fnCall.Name, maxGRPCFrameBytes), - }, - }, - }, - }, - } - } - state.mu.Lock() - state.executedCalls[callKey] = out - state.mu.Unlock() - return out - } - - return nil -} - -// Dispatch across the 7 canonical sandbox client tools. -func (w *pullWorker) dispatchTool(ctx context.Context, toolName string, args map[string]any) (map[string]any, bool) { - normalized := strings.ToLower(strings.TrimSpace(toolName)) - switch normalized { - case "execute_bash", "run_command", "run_shell_command", "bash", "shell": - return w.toolExecuteBash(ctx, args) - case "view_file", "read_file": - return w.toolViewFile(args) - case "write_to_file", "write_file": - return w.toolWriteFile(args) - case "replace_file_content", "edit_file": - return w.toolReplaceFileContent(args) - case "list_dir", "read_dir", "list_directory": - return w.toolListDir(args) - case "grep_search": - return w.toolGrepSearch(args) - case "find_by_name", "find_file": - return w.toolFindByName(args) - default: - return map[string]any{"error": fmt.Sprintf("unsupported sandbox tool %q", toolName)}, true - } -} - -func (w *pullWorker) toolExecuteBash(ctx context.Context, args map[string]any) (map[string]any, bool) { - if w.disableBash { - return map[string]any{"error": "execute_bash is disabled (--disable-bash)"}, true - } - cmdLine := stringArg(args, "command", "CommandLine", "cmd") - resolvedWorkDir, err := w.fs.ResolvePath(stringArg(args, "Cwd", "cwd", "working_dir")) - if err != nil { - return map[string]any{"error": err.Error()}, true - } - timeout := w.commandTimeout - if timeoutMS := intArg(args, 0, "timeout_ms"); timeoutMS > 0 { - timeout = time.Duration(timeoutMS) * time.Millisecond - } else if timeoutSec := intArg(args, 0, "timeout_seconds", "timeout"); timeoutSec > 0 { - timeout = time.Duration(timeoutSec) * time.Second - } - if w.commandTimeout > 0 && timeout > w.commandTimeout { - timeout = w.commandTimeout - } - stdout, stderr, exitCode, err := w.shell.ExecuteBash(ctx, cmdLine, resolvedWorkDir, timeout) - if err != nil { - return map[string]any{ - "error": err.Error(), - "stdout": stdout, - "stderr": stderr, - "exit_code": exitCode, - }, true - } - return map[string]any{ - "stdout": stdout, - "stderr": stderr, - "exit_code": exitCode, - }, false -} - -func (w *pullWorker) toolViewFile(args map[string]any) (map[string]any, bool) { - path := stringArg(args, "AbsolutePath", "path", "file_path") - if path == "" { - return map[string]any{"error": "path is required"}, true - } - startLine := intArg(args, 0, "StartLine", "start_line") - endLine := intArg(args, 0, "EndLine", "end_line") - var content string - if startLine > 0 || endLine > 0 { - sliced, err := w.fs.ReadFileRange(path, startLine, endLine) - if err != nil { - return map[string]any{"error": err.Error()}, true - } - content = sliced - } else { - data, err := w.fs.ReadRegularFile(path) - if err != nil { - return map[string]any{"error": err.Error()}, true - } - content = string(data) - } - return map[string]any{ - "path": path, - "content": content, - }, false -} - -func (w *pullWorker) toolWriteFile(args map[string]any) (map[string]any, bool) { - path := stringArg(args, "TargetFile", "path", "file_path") - if path == "" { - return map[string]any{"error": "path is required"}, true - } - content := stringArg(args, "CodeContent", "content") - overwrite := boolArg(args, true, "Overwrite", "overwrite") - appendMode := boolArg(args, false, "Append", "append") - n, err := w.fs.WriteRegularFile(path, content, overwrite, appendMode) - if err != nil { - return map[string]any{"error": err.Error()}, true - } - return map[string]any{ - "status": "ok", - "path": path, - "bytes_written": n, - }, false -} - -func (w *pullWorker) toolReplaceFileContent(args map[string]any) (map[string]any, bool) { - path := stringArg(args, "TargetFile", "path", "file_path") - if path == "" { - return map[string]any{"error": "TargetFile is required"}, true - } - target := stringArg(args, "TargetContent", "old_string") - if target == "" { - return map[string]any{"error": "TargetContent is required"}, true - } - replacement := stringArg(args, "ReplacementContent", "new_string") - allowMultiple := boolArg(args, false, "AllowMultiple", "allow_multiple") - startLine := intArg(args, 0, "StartLine", "start_line") - endLine := intArg(args, 0, "EndLine", "end_line") - - raw, err := w.fs.ReadRegularFile(path) - if err != nil { - return map[string]any{"error": err.Error()}, true - } - updated, count, err := replaceInText(string(raw), target, replacement, startLine, endLine, allowMultiple) - if err != nil { - return map[string]any{"error": err.Error()}, true - } - if _, err := w.fs.WriteRegularFile(path, updated, true, false); err != nil { - return map[string]any{"error": err.Error()}, true - } - return map[string]any{ - "status": "ok", - "path": path, - "replacements": count, - }, false -} - -func (w *pullWorker) toolListDir(args map[string]any) (map[string]any, bool) { - path := stringArg(args, "DirectoryPath", "path", "dir_path") - if path == "" { - path = "." - } - entries, err := w.fs.ListEntries(path) - if err != nil { - return map[string]any{"error": err.Error()}, true - } - return map[string]any{ - "path": path, - "entries": entries, - }, false -} - -func (w *pullWorker) toolGrepSearch(args map[string]any) (map[string]any, bool) { - searchPath := firstNonEmpty(stringArg(args, "SearchPath", "path"), ".") - query := stringArg(args, "Query", "query") - caseInsensitive := boolArg(args, false, "CaseInsensitive", "case_insensitive") - isRegex := boolArg(args, false, "IsRegex", "is_regex") - matches, truncated, err := w.fs.Grep(searchPath, query, caseInsensitive, isRegex) - if err != nil { - return map[string]any{"error": err.Error()}, true - } - res := map[string]any{"matches": matches} - if truncated { - res["truncated"] = true - } - return res, false -} - -func (w *pullWorker) toolFindByName(args map[string]any) (map[string]any, bool) { - searchDir := firstNonEmpty(stringArg(args, "SearchDirectory", "path"), ".") - pattern := stringArg(args, "Pattern", "pattern") - maxDepth := intArg(args, 0, "MaxDepth", "max_depth") - extensions := stringSliceArg(args, "Extensions", "extensions") - files, truncated, err := w.fs.FindFiles(searchDir, pattern, maxDepth, extensions...) - if err != nil { - return map[string]any{"error": err.Error()}, true - } - res := map[string]any{"files": files} - if truncated { - res["truncated"] = true - } - return res, false -} - -func replaceInText(content, target, replacement string, startLine, endLine int, allowMultiple bool) (string, int, error) { - preserveCRLF := strings.Contains(content, "\r\n") && !strings.Contains(target, "\r\n") - if preserveCRLF { - content = strings.ReplaceAll(content, "\r\n", "\n") - replacement = strings.ReplaceAll(replacement, "\r\n", "\n") - } - restore := func(s string) string { - if preserveCRLF { - return strings.ReplaceAll(s, "\n", "\r\n") - } - return s - } - - lines := strings.Split(content, "\n") - if startLine <= 0 { - startLine = 1 - } - if endLine <= 0 || endLine > len(lines) { - endLine = len(lines) - } - if startLine > endLine { - return "", 0, fmt.Errorf("invalid line range [%d, %d]", startLine, endLine) - } - - region := strings.Join(lines[startLine-1:endLine], "\n") - count := strings.Count(region, target) - if count == 0 { - return "", 0, fmt.Errorf("target content not found in lines [%d, %d]", startLine, endLine) - } - if count > 1 && !allowMultiple { - return "", 0, fmt.Errorf("found %d occurrences in lines [%d, %d] when AllowMultiple is false", count, startLine, endLine) - } - - var parts []string - if startLine > 1 { - parts = append(parts, strings.Join(lines[:startLine-1], "\n")) - } - parts = append(parts, strings.ReplaceAll(region, target, replacement)) - if endLine < len(lines) { - parts = append(parts, strings.Join(lines[endLine:], "\n")) - } - return restore(strings.Join(parts, "\n")), count, nil -} - -func stringArg(args map[string]any, keys ...string) string { - for _, k := range keys { - if v := args[k]; v != nil { - return fmt.Sprint(v) - } - } - return "" -} - -func intArg(args map[string]any, fallback int, keys ...string) int { - for _, k := range keys { - switch val := args[k].(type) { - case int: - return val - case int32: - return int(val) - case int64: - return int(val) - case float64: - return int(val) - } - } - return fallback -} - -func boolArg(args map[string]any, fallback bool, keys ...string) bool { - for _, k := range keys { - if b, ok := args[k].(bool); ok { - return b - } - } - return fallback -} - -func stringSliceArg(args map[string]any, keys ...string) []string { - for _, k := range keys { - switch val := args[k].(type) { - case []string: - return val - case []any: - out := make([]string, 0, len(val)) - for _, item := range val { - if item != nil { - out = append(out, fmt.Sprint(item)) - } - } - return out - } - } - return nil -}