fix: external session ids in sse

Author: rhnvrm

kc/manager_test.go

@@ -493,6 +493,44 @@ func TestNewConfigConstructor(t *testing.T) {
 	})
 }
 
+// TestExternalSessionIDFromErrorLog tests the exact session ID from the error log
+func TestExternalSessionIDFromErrorLog(t *testing.T) {
+	manager, err := newTestManager("test_key", "test_secret")
+	if err != nil {
+		t.Fatalf("Expected no error creating manager, got: %v", err)
+	}
+
+	// This is the exact session ID from the error log that was failing
+	externalSessionID := "6f615000-2644-45a7-a27c-f579e20b5992"
+	
+	// Should be able to get or create session with external session ID
+	kiteSession, isNew, err := manager.GetOrCreateSession(externalSessionID)
+	if err != nil {
+		t.Errorf("Expected no error for external session ID from error log, got: %v", err)
+	}
+	if !isNew {
+		t.Error("Expected new session to be created for external session ID")
+	}
+	if kiteSession == nil {
+		t.Error("Expected non-nil Kite session data")
+	}
+	if kiteSession.Kite == nil {
+		t.Error("Expected Kite client to be initialized")
+	}
+
+	// Subsequent call should reuse existing session
+	kiteSession2, isNew2, err2 := manager.GetOrCreateSession(externalSessionID)
+	if err2 != nil {
+		t.Errorf("Expected no error on second call, got: %v", err2)
+	}
+	if isNew2 {
+		t.Error("Expected existing session to be reused")
+	}
+	if kiteSession2 != kiteSession {
+		t.Error("Expected same session instance to be returned")
+	}
+}
+
 // Helper function to check if string contains substring
 func managerContains(s, substr string) bool {
 	if len(substr) > len(s) {

kc/session.go

@@ -100,11 +100,19 @@ func (sm *SessionRegistry) GenerateWithData(data any) string {
 }
 
 // should be a valid uuid and start with the correct prefix.
+// checkSessionID validates the format of a MCP session ID
+// Accepts both internal format (kitemcp-<uuid>) and external format (plain uuid)
 func checkSessionID(sessionID string) error {
-	if !strings.HasPrefix(sessionID, mcpSessionPrefix) {
-		return fmt.Errorf("%s", errInvalidSessionIDFormat)
+	// Handle internal format with prefix
+	if strings.HasPrefix(sessionID, mcpSessionPrefix) {
+		if _, err := uuid.Parse(sessionID[len(mcpSessionPrefix):]); err != nil {
+			return fmt.Errorf("%s: %w", errInvalidSessionIDFormat, err)
+		}
+		return nil
 	}
-	if _, err := uuid.Parse(sessionID[len(mcpSessionPrefix):]); err != nil {
+	
+	// Handle external format (plain UUID from SSE/stdio modes)
+	if _, err := uuid.Parse(sessionID); err != nil {
 		return fmt.Errorf("%s: %w", errInvalidSessionIDFormat, err)
 	}
 	return nil
@@ -350,8 +358,20 @@ func (sm *SessionRegistry) GetOrCreateSessionData(sessionID string, createDataFn
 
 	session, exists := sm.sessions[sessionID]
 	if !exists {
-		sm.logger.Warn("Session not found for ID", "session_id", sessionID)
-		return nil, false, errors.New(errSessionNotFound)
+		// Create a new session for external session IDs (from SSE/stdio modes)
+		sm.logger.Info("Creating new session for external session ID", "session_id", sessionID)
+		now := time.Now()
+		expiresAt := now.Add(sm.sessionDuration)
+		
+		session = &MCPSession{
+			ID:         sessionID,
+			Terminated: false,
+			CreatedAt:  now,
+			ExpiresAt:  expiresAt,
+			Data:       nil,
+		}
+		sm.sessions[sessionID] = session
+		exists = true
 	}
 
 	now := time.Now()

kc/session_test.go

@@ -599,18 +599,75 @@ func TestSessionExpiration(t *testing.T) {
 		t.Errorf("Expected no error for new session, got: %v", err)
 	}
 	if isTerminated {
-		t.Error("Expected new session to not be terminated")
+		t.Error("Expected session to not be terminated initially")
 	}
 
-	// Wait for expiration
-	time.Sleep(2 * time.Millisecond)
+	// Wait for session to expire
+	time.Sleep(5 * time.Millisecond)
 
 	// Session should now be expired
 	isTerminated, err = manager.Validate(sessionID)
 	if err != nil {
 		t.Errorf("Expected no error for expired session validation, got: %v", err)
 	}
 	if !isTerminated {
-		t.Error("Expected expired session to be terminated")
+		t.Error("Expected session to be terminated after expiry")
+	}
+}
+
+func TestExternalSessionIDFormat(t *testing.T) {
+	manager := NewSessionRegistry(testLogger())
+
+	// Test external session ID (plain UUID format from SSE/stdio modes)
+	externalSessionID := "6f615000-2644-45a7-a27c-f579e20b5992"
+	
+	// Should be able to create session data with external session ID
+	testData := map[string]string{"test": "data"}
+	data, isNew, err := manager.GetOrCreateSessionData(externalSessionID, func() any {
+		return testData
+	})
+	
+	if err != nil {
+		t.Errorf("Expected no error for external session ID, got: %v", err)
+	}
+	if !isNew {
+		t.Error("Expected new session to be created")
+	}
+	retrievedData, ok := data.(map[string]string)
+	if !ok {
+		t.Errorf("Expected data to be map[string]string, got: %T", data)
+	}
+	if retrievedData["test"] != "data" {
+		t.Errorf("Expected data['test'] to be 'data', got: %v", retrievedData["test"])
+	}
+
+	// Should be able to validate external session ID
+	isTerminated, err := manager.Validate(externalSessionID)
+	if err != nil {
+		t.Errorf("Expected no error validating external session ID, got: %v", err)
+	}
+	if isTerminated {
+		t.Error("Expected external session to not be terminated")
+	}
+
+	// Test internal session ID format still works
+	internalSessionID := manager.Generate()
+	
+	data2, isNew2, err2 := manager.GetOrCreateSessionData(internalSessionID, func() any {
+		return testData
+	})
+	
+	if err2 != nil {
+		t.Errorf("Expected no error for internal session ID, got: %v", err2)
+	}
+	if !isNew2 {
+		t.Error("Expected new internal session to be created")
+	}
+	retrievedData2, ok2 := data2.(map[string]string)
+	if !ok2 {
+		t.Errorf("Expected internal data to be map[string]string, got: %T", data2)
+	}
+	if retrievedData2["test"] != "data" {
+		t.Errorf("Expected internal data['test'] to be 'data', got: %v", retrievedData2["test"])
 	}
 }