jango-blockchained · pull · Aug 29, 2025 · Aug 29, 2025 · Aug 29, 2025 · Aug 29, 2025
diff --git a/cli/exp_scaletest.go b/cli/exp_scaletest.go
@@ -864,6 +864,7 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command {
 		tickInterval      time.Duration
 		bytesPerTick      int64
 		ssh               bool
+		disableDirect     bool
 		useHostLogin      bool
 		app               string
 		template          string
@@ -1023,15 +1024,16 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command {
 
 				// Setup our workspace agent connection.
 				config := workspacetraffic.Config{
-					AgentID:      agent.ID,
-					BytesPerTick: bytesPerTick,
-					Duration:     strategy.timeout,
-					TickInterval: tickInterval,
-					ReadMetrics:  metrics.ReadMetrics(ws.OwnerName, ws.Name, agent.Name),
-					WriteMetrics: metrics.WriteMetrics(ws.OwnerName, ws.Name, agent.Name),
-					SSH:          ssh,
-					Echo:         ssh,
-					App:          appConfig,
+					AgentID:       agent.ID,
+					BytesPerTick:  bytesPerTick,
+					Duration:      strategy.timeout,
+					TickInterval:  tickInterval,
+					ReadMetrics:   metrics.ReadMetrics(ws.OwnerName, ws.Name, agent.Name),
+					WriteMetrics:  metrics.WriteMetrics(ws.OwnerName, ws.Name, agent.Name),
+					SSH:           ssh,
+					DisableDirect: disableDirect,
+					Echo:          ssh,
+					App:           appConfig,
 				}
 
 				if webClient != nil {
@@ -1117,6 +1119,13 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command {
 			Description: "Send traffic over SSH, cannot be used with --app.",
 			Value:       serpent.BoolOf(&ssh),
 		},
+		{
+			Flag:        "disable-direct",
+			Env:         "CODER_SCALETEST_WORKSPACE_TRAFFIC_DISABLE_DIRECT_CONNECTIONS",
+			Default:     "false",
+			Description: "Disable direct connections for SSH traffic to workspaces. Does nothing if `--ssh` is not also set.",
+			Value:       serpent.BoolOf(&disableDirect),
+		},
 		{
 			Flag:        "app",
 			Env:         "CODER_SCALETEST_WORKSPACE_TRAFFIC_APP",

diff --git a/cli/exp_task_status_test.go b/cli/exp_task_status_test.go
@@ -243,13 +243,12 @@ STATE CHANGED  STATUS   STATE  MESSAGE
 				ctx    = testutil.Context(t, testutil.WaitShort)
 				now    = time.Now().UTC() // TODO: replace with quartz
 				srv    = httptest.NewServer(http.HandlerFunc(tc.hf(ctx, now)))
-				client = new(codersdk.Client)
+				client = codersdk.New(testutil.MustURL(t, srv.URL))
 				sb     = strings.Builder{}
 				args   = []string{"exp", "task", "status", "--watch-interval", testutil.IntervalFast.String()}
 			)
 
 			t.Cleanup(srv.Close)
-			client.URL = testutil.MustURL(t, srv.URL)
 			args = append(args, tc.args...)
 			inv, root := clitest.New(t, args...)
 			inv.Stdout = &sb

diff --git a/cli/exp_taskcreate_test.go b/cli/exp_taskcreate_test.go
@@ -5,14 +5,12 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
-	"net/url"
 	"strings"
 	"testing"
 	"time"
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 
 	"github.com/coder/coder/v2/cli/clitest"
 	"github.com/coder/coder/v2/cli/cliui"
@@ -236,17 +234,14 @@ func TestTaskCreate(t *testing.T) {
 			var (
 				ctx    = testutil.Context(t, testutil.WaitShort)
 				srv    = httptest.NewServer(tt.handler(t, ctx))
-				client = new(codersdk.Client)
+				client = codersdk.New(testutil.MustURL(t, srv.URL))
 				args   = []string{"exp", "task", "create"}
 				sb     strings.Builder
 				err    error
 			)
 
 			t.Cleanup(srv.Close)
 
-			client.URL, err = url.Parse(srv.URL)
-			require.NoError(t, err)
-
 			inv, root := clitest.New(t, append(args, tt.args...)...)
 			inv.Environ = serpent.ParseEnviron(tt.env, "")
 			inv.Stdout = &sb

diff --git a/cli/root.go b/cli/root.go
@@ -635,6 +635,9 @@ func (r *RootCmd) HeaderTransport(ctx context.Context, serverURL *url.URL) (*cod
 }
 
 func (r *RootCmd) configureClient(ctx context.Context, client *codersdk.Client, serverURL *url.URL, inv *serpent.Invocation) error {
+	if client.SessionTokenProvider == nil {
+		client.SessionTokenProvider = codersdk.FixedSessionTokenProvider{}
+	}
 	transport := http.DefaultTransport
 	transport = wrapTransportWithTelemetryHeader(transport, inv)
 	if !r.noVersionCheck {

diff --git a/coderd/coderdtest/oidctest/idp.go b/coderd/coderdtest/oidctest/idp.go
@@ -641,7 +641,7 @@ func (f *FakeIDP) LoginWithClient(t testing.TB, client *codersdk.Client, idToken
 
 // ExternalLogin does the oauth2 flow for external auth providers. This requires
 // an authenticated coder client.
-func (f *FakeIDP) ExternalLogin(t testing.TB, client *codersdk.Client, opts ...func(r *http.Request)) {
+func (f *FakeIDP) ExternalLogin(t testing.TB, client *codersdk.Client, opts ...codersdk.RequestOption) {
 	coderOauthURL, err := client.URL.Parse(fmt.Sprintf("/external-auth/%s/callback", f.externalProviderID))
 	require.NoError(t, err)
 	f.SetRedirect(t, coderOauthURL.String())
@@ -660,11 +660,7 @@ func (f *FakeIDP) ExternalLogin(t testing.TB, client *codersdk.Client, opts ...f
 	req, err := http.NewRequestWithContext(ctx, "GET", coderOauthURL.String(), nil)
 	require.NoError(t, err)
 	// External auth flow requires the user be authenticated.
-	headerName := client.SessionTokenHeader
-	if headerName == "" {
-		headerName = codersdk.SessionTokenHeader
-	}
-	req.Header.Set(headerName, client.SessionToken())
+	opts = append([]codersdk.RequestOption{client.SessionTokenProvider.AsRequestOption()}, opts...)
 	if cli.Jar == nil {
 		cli.Jar, err = cookiejar.New(nil)
 		require.NoError(t, err, "failed to create cookie jar")