sketch: fix MCP server connection error in -unsafe mode by setting SKETCH_PUB_KEY environment variable
Fix MCP server connection failure with 'no Public-Key' status 400 error when running sketch -unsafe with skaband integration, by properly setting the SKETCH_PUB_KEY environment variable for MCP authentication placeholder replacement.
Problems Solved:
MCP Authentication Failure:
- Running sketch -unsafe with skaband connection fails with 'MCP server connection failed: MCP server "sketchdev": failed to initialize MCP client: transport error: request failed with status 400: no Public-Key'
- setupAndRunAgent receives pubKey parameter from skaband login but doesn't set SKETCH_PUB_KEY environment variable
- MCP placeholder replacement in agent.go expects SKETCH_PUB_KEY environment variable to replace '_sketch_public_key_' placeholder
- Empty placeholder replacement results in missing Public-Key header in MCP requests
Authentication Flow Gap:
- Container mode sets SKETCH_PUB_KEY from environment in runInContainerMode
- Unsafe mode obtains pubKey from skabandclient.Login but doesn't propagate to environment
- setupAndRunAgent receives pubKey parameter but doesn't use it for environment variable setup
- MCP configuration uses '_sketch_public_key_' placeholder expecting environment variable replacement
Solution Implementation:
Environment Variable Setup:
- Added SKETCH_PUB_KEY environment variable setting in setupAndRunAgent when pubKey is provided
- Check for non-empty pubKey before setting environment variable to avoid overwriting existing values
- Environment variable set early in setupAndRunAgent before MCP server initialization
- Maintains consistent behavior between container and unsafe modes
Authentication Flow Completion:
- Container mode: SKETCH_PUB_KEY set from container environment → MCP placeholder replacement
- Unsafe mode: pubKey from skaband login → SKETCH_PUB_KEY environment variable → MCP placeholder replacement
- Both modes now have complete authentication flow for MCP server connections
- MCP requests include proper Public-Key header for skaband authentication
Implementation Details:
Code Changes:
- Added conditional os.Setenv('SKETCH_PUB_KEY', pubKey) in setupAndRunAgent
- Placement before working directory setup ensures early environment configuration
- Only sets environment variable when pubKey is non-empty string
- Preserves existing environment if pubKey is empty to avoid overwrites
Test Coverage:
- Added TestSetupAndRunAgent_SetsPubKeyEnvVar to verify environment variable setting
- Added TestSetupAndRunAgent_DoesNotSetEmptyPubKey to verify empty pubKey handling
- Tests verify environment variable setting and preservation behavior
- Comprehensive coverage of both positive and negative scenarios
Error Resolution:
- MCP server requests now include proper Public-Key header
- skaband /api/mcp endpoint receives authentication for session validation
- Eliminates 'no Public-Key' 400 status errors in unsafe mode
- Maintains existing container mode behavior without changes
Files Modified:
- sketch/cmd/sketch/main.go: Added SKETCH_PUB_KEY environment variable setting in setupAndRunAgent
- sketch/cmd/sketch/main_test.go: Added test coverage for pubKey environment variable behavior
The fix ensures consistent MCP server authentication across both container and unsafe execution modes by properly propagating the public key from skaband login to the MCP placeholder replacement system.
Co-Authored-By: sketch <hello@sketch.dev>
Change-ID: s2564fc63bdd663a0k
diff --git a/cmd/sketch/main.go b/cmd/sketch/main.go
index 637aaab..5ca3395 100644
--- a/cmd/sketch/main.go
+++ b/cmd/sketch/main.go
@@ -519,6 +519,12 @@
func setupAndRunAgent(ctx context.Context, flags CLIFlags, modelURL, apiKey, pubKey string, inInsideSketch bool, logFile *os.File) error {
var client *http.Client
+ // Set the public key environment variable if provided
+ // This is needed for MCP server authentication placeholder replacement
+ if pubKey != "" {
+ os.Setenv("SKETCH_PUB_KEY", pubKey)
+ }
+
wd, err := os.Getwd()
if err != nil {
return err
diff --git a/cmd/sketch/main_test.go b/cmd/sketch/main_test.go
index 3ebd05d..10ca1e0 100644
--- a/cmd/sketch/main_test.go
+++ b/cmd/sketch/main_test.go
@@ -1,6 +1,7 @@
package main
import (
+ "context"
"os"
"testing"
)
@@ -35,3 +36,77 @@
})
}
}
+
+func TestSetupAndRunAgent_SetsPubKeyEnvVar(t *testing.T) {
+ // Save original environment
+ originalPubKey := os.Getenv("SKETCH_PUB_KEY")
+ defer func() {
+ if originalPubKey == "" {
+ os.Unsetenv("SKETCH_PUB_KEY")
+ } else {
+ os.Setenv("SKETCH_PUB_KEY", originalPubKey)
+ }
+ }()
+
+ // Clear the environment variable first
+ os.Unsetenv("SKETCH_PUB_KEY")
+
+ // Verify it's not set
+ if os.Getenv("SKETCH_PUB_KEY") != "" {
+ t.Fatal("SKETCH_PUB_KEY should not be set initially")
+ }
+
+ // Test data
+ testPubKey := "test-public-key-123"
+
+ // Create a minimal flags struct
+ flags := CLIFlags{
+ modelName: "claude",
+ }
+
+ // This should fail due to missing API key, but should still set the environment variable
+ err := setupAndRunAgent(context.TODO(), flags, "", "", testPubKey, false, nil)
+
+ // Check that the environment variable was set correctly
+ if os.Getenv("SKETCH_PUB_KEY") != testPubKey {
+ t.Errorf("Expected SKETCH_PUB_KEY to be %q, got %q", testPubKey, os.Getenv("SKETCH_PUB_KEY"))
+ }
+
+ // We expect this to fail due to missing API key, but that's fine for this test
+ if err == nil {
+ t.Error("Expected setupAndRunAgent to fail due to missing API key")
+ }
+}
+
+func TestSetupAndRunAgent_DoesNotSetEmptyPubKey(t *testing.T) {
+ // Save original environment
+ originalPubKey := os.Getenv("SKETCH_PUB_KEY")
+ defer func() {
+ if originalPubKey == "" {
+ os.Unsetenv("SKETCH_PUB_KEY")
+ } else {
+ os.Setenv("SKETCH_PUB_KEY", originalPubKey)
+ }
+ }()
+
+ // Set a value first
+ os.Setenv("SKETCH_PUB_KEY", "existing-value")
+
+ // Create a minimal flags struct
+ flags := CLIFlags{
+ modelName: "claude",
+ }
+
+ // This should fail due to missing API key, but should not change the environment variable
+ err := setupAndRunAgent(context.TODO(), flags, "", "", "", false, nil)
+
+ // Check that the environment variable was not changed
+ if os.Getenv("SKETCH_PUB_KEY") != "existing-value" {
+ t.Errorf("Expected SKETCH_PUB_KEY to remain %q, got %q", "existing-value", os.Getenv("SKETCH_PUB_KEY"))
+ }
+
+ // We expect this to fail due to missing API key, but that's fine for this test
+ if err == nil {
+ t.Error("Expected setupAndRunAgent to fail due to missing API key")
+ }
+}